From 684f6ef544f3ff65b8800ef4963ee98c136916ed Mon Sep 17 00:00:00 2001
From: Jimmy Lu <jimmylu@meta.com>
Date: Mon, 4 Mar 2024 14:51:39 -0800
Subject: [PATCH] Fix DWRF selective flat map reader fast path crash with all
 inlined string keys (#8952)

Summary: Pull Request resolved: https://github.com/facebookincubator/velox/pull/8952

Differential Revision: D54511991
---
 .../reader/SelectiveFlatMapColumnReader.cpp   |  2 +-
 velox/dwio/dwrf/test/ReaderTest.cpp           | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/velox/dwio/dwrf/reader/SelectiveFlatMapColumnReader.cpp b/velox/dwio/dwrf/reader/SelectiveFlatMapColumnReader.cpp
index 44dad5b9aaad..c25c6cae4beb 100644
--- a/velox/dwio/dwrf/reader/SelectiveFlatMapColumnReader.cpp
+++ b/velox/dwio/dwrf/reader/SelectiveFlatMapColumnReader.cpp
@@ -617,7 +617,7 @@ class SelectiveFlatMapReader : public SelectiveStructColumnReaderBase {
         keyBuffer->template asMutable<T>());
 
     std::vector<BufferPtr> keyStrings;
-    if constexpr (std::is_same_v<T, StringView>) {
+    if (std::is_same_v<T, StringView> && keyStrings_) {
       keyStrings.push_back(keyStrings_);
     }
     auto keyVector = std::make_shared<FlatVector<T>>(
diff --git a/velox/dwio/dwrf/test/ReaderTest.cpp b/velox/dwio/dwrf/test/ReaderTest.cpp
index 7d28696b74b5..a10c4680a1ae 100644
--- a/velox/dwio/dwrf/test/ReaderTest.cpp
+++ b/velox/dwio/dwrf/test/ReaderTest.cpp
@@ -3104,3 +3104,22 @@ TEST_F(TestReader, selectiveStringDirect) {
   });
   assertEqualVectors(expected, actual);
 }
+
+TEST_F(TestReader, selectiveFlatMapFastPathAllInlinedStringKeys) {
+  auto maps = makeMapVector<std::string, int64_t>(
+      {{{"a", 0}, {"b", 0}}, {{"a", 1}, {"b", 1}}});
+  auto row = makeRowVector({"c0"}, {maps});
+  auto config = std::make_shared<dwrf::Config>();
+  config->set(dwrf::Config::FLATTEN_MAP, true);
+  config->set(dwrf::Config::MAP_FLAT_COLS, {0});
+  auto [writer, reader] = createWriterReader({row}, pool(), config);
+  auto schema = asRowType(row->type());
+  auto spec = std::make_shared<common::ScanSpec>("<root>");
+  spec->addAllChildFields(*schema);
+  RowReaderOptions rowReaderOpts;
+  rowReaderOpts.setScanSpec(spec);
+  auto rowReader = reader->createRowReader(rowReaderOpts);
+  VectorPtr batch = BaseVector::create(schema, 0, pool());
+  ASSERT_EQ(rowReader->next(10, batch), 2);
+  assertEqualVectors(batch, row);
+}