Fix crash in DWRF selective flat map reader fast path with all inline…

…d string keys (#8952) Summary: Pull Request resolved: #8952 Reviewed By: oerling Differential Revision: D54511991 fbshipit-source-id: 0db18e0471bb1b1cd6862f3dd1cd79b0ce804727
facebookincubator · Mar 5, 2024 · 1027b32 · 1027b32
1 parent f82112a
commit 1027b32
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git 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
@@ -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);
+}