refactor(nimble): Preserve MAP spec in flat map column readers

Summary:
X-link: facebookincubator/velox#15803

A production issue S602695 VeloxRuntimeError blocking feature injection pipeline exposed an issue in the Nimble flat map reader.

For context, earlier this half ROO feature injection was enabled in production and supported by both Velox map and flat map vector encodings. A particular pipeline was blocked this past weekend due to flat map reader crashing. A temporary mitigation was to use regular map instead of the flat map encoding as it appeared flat map reader was attempting to set values that were uninitialized by the reader.

The culprit was how scan spec is constructed for flat map. Rather than appending all features to keys and values children, we should follow map scan spec construction by preserving the keys and values and create child columns using the preserved values child spec.

Differential Revision: D89326920

Author: peterenescu

dwio/nimble/velox/FieldWriter.cpp

@@ -1060,10 +1060,6 @@ class FlatMapFieldWriter : public FieldWriter {
       return;
     }
 
-    // Only create keys on first call to write (with valid ranges). Subsequent
-    // calls must have the same set of keys, otherwise writer will throw.
-    bool populateMap = currentPassthroughFields_.empty();
-
     const auto& values = flatMapVector->mapValues();
     const auto& inMaps = flatMapVector->inMaps();
 
@@ -1082,8 +1078,13 @@ class FlatMapFieldWriter : public FieldWriter {
             "FlatMapVector keys are not distinct.");
         distinctKeySet.insert(key);
 
-        auto& writer = populateMap ? createPassthroughValueFieldWriter(key)
-                                   : findPassthroughValueFieldWriter(key);
+        // Only create keys on first call to write (with valid ranges). To
+        // account for key pruning, let's avoid failing on unfound subsequent
+        // keys.
+        auto existingPair = currentPassthroughFields_.find(key);
+        auto& writer = existingPair != currentPassthroughFields_.end()
+            ? *existingPair->second
+            : createPassthroughValueFieldWriter(key);
 
         if (inMaps[i]) {
           writer.write(values[i], inMaps[i], childRanges);

dwio/nimble/velox/selective/FlatMapColumnReader.cpp

@@ -66,9 +66,10 @@ std::vector<KeyNode<T>> makeKeyNodes(
 
   // Adjust the scan spec according to the output type.
   switch (outputType) {
-    // For a kMap output, just need a scan spec for map keys and one for map
-    // values.
-    case FlatMapOutput::kMap: {
+    // For a kMap and kFlatMap output, just need a scan spec for map keys and
+    // one for map values.
+    case FlatMapOutput::kMap:
+    case FlatMapOutput::kFlatMap: {
       keysSpec = scanSpec.getOrCreateChild(
           common::Subfield(common::ScanSpec::kMapKeysFieldName));
       valuesSpec = scanSpec.getOrCreateChild(
@@ -78,18 +79,6 @@ std::vector<KeyNode<T>> makeKeyNodes(
       valuesSpec->setProjectOut(true);
       break;
     }
-    // For a kFlatMap output, need to find the streams (distinct keys) to read
-    // from the file (nimbleType).
-    case FlatMapOutput::kFlatMap: {
-      for (int i = 0; i < childrenCount; ++i) {
-        auto key = parseKeyValue<T>(nimbleType.nameAt(i));
-        auto spec = scanSpec.getOrCreateChild(nimbleType.nameAt(i));
-        spec->setProjectOut(true);
-        spec->setChannel(i);
-        childSpecs[key] = spec;
-      }
-      break;
-    }
     // For a kStruct output, the streams to be read are part of the scan spec
     // already.
     case FlatMapOutput::kStruct: {
@@ -113,7 +102,7 @@ std::vector<KeyNode<T>> makeKeyNodes(
     if (auto it = childSpecs.find(node.key);
         it != childSpecs.end() && !it->second->isConstant()) {
       childSpec = it->second;
-    } else if (outputType != FlatMapOutput::kMap) {
+    } else if (outputType == FlatMapOutput::kStruct) {
       // Column not selected in 'scanSpec', skipping it.
       continue;
     } else {
@@ -243,6 +232,7 @@ class FlatMapColumnReader
     for (int i = 0; i < keyNodes_.size(); ++i) {
       keyNodes_[i].reader->scanSpec()->setSubscript(i);
       children_[i] = keyNodes_[i].reader.get();
+      children_[i]->setIsFlatMapValue(true);
 
       rawKeys[i] = keyNodes_[i].key.get();
     }
@@ -305,6 +295,60 @@ class FlatMapColumnReader
     // decoders.
   }
 
+  // Same as FlatMapAsMapColumnReader.
+  void read(int64_t offset, const RowSet& rows, const uint64_t* incomingNulls) {
+    numReads_ = scanSpec_->newRead();
+    prepareRead<char>(offset, rows, incomingNulls);
+    VELOX_DCHECK(!hasDeletion());
+    auto activeRows = rows;
+    auto* mapNulls =
+        nullsInReadRange_ ? nullsInReadRange_->as<uint64_t>() : nullptr;
+    if (scanSpec_->filter()) {
+      auto kind = scanSpec_->filter()->kind();
+      VELOX_CHECK(
+          kind == velox::common::FilterKind::kIsNull ||
+          kind == velox::common::FilterKind::kIsNotNull);
+      filterNulls<int32_t>(
+          rows, kind == velox::common::FilterKind::kIsNull, false);
+      if (outputRows_.empty()) {
+        for (auto* child : children_) {
+          child->addParentNulls(offset, mapNulls, rows);
+        }
+        readOffset_ = offset + rows.back() + 1;
+        return;
+      }
+      activeRows = outputRows_;
+    }
+    // Separate the loop to be cache friendly.
+    for (auto* child : children_) {
+      advanceFieldReader(child, offset);
+    }
+    for (auto* child : children_) {
+      child->read(offset, activeRows, mapNulls);
+      child->addParentNulls(offset, mapNulls, rows);
+    }
+    readOffset_ = offset + rows.back() + 1;
+  }
+
+  void getValues(const RowSet& rows, VectorPtr* result) override {
+    SelectiveFlatMapColumnReader::getValues(rows, result);
+
+    // After reading the flat map streams recursively, need to read the in map
+    // buffers.
+    if (result && *result) {
+      auto flatMapVector = (*result)->as<FlatMapVector>();
+      VELOX_CHECK(flatMapVector);
+
+      for (int i = 0; i < keyNodes_.size(); ++i) {
+        auto& nimbleData = children_[i]->formatData().template as<NimbleData>();
+        if (nimbleData.inMapBuffer()) {
+          auto inMapBuffer = nimbleData.inMapBuffer();
+          flatMapVector->inMapsAt(i, true) = inMapBuffer;
+        }
+      }
+    }
+  }
+
  private:
   std::vector<KeyNode<T>> keyNodes_;
 };

dwio/nimble/velox/selective/tests/SelectiveNimbleReaderTest.cpp

@@ -1523,7 +1523,6 @@ TEST_P(SelectiveNimbleReaderTest, nativeFlatMap) {
 
   // Dictionary wrapped keys
   {
-    LOG(INFO) << "Dictionary wrapped keys";
     testRoundtrip(constructFlatMap(
         BaseVector::wrapInDictionary(
             nullptr,
@@ -1564,6 +1563,44 @@ TEST_P(SelectiveNimbleReaderTest, nativeFlatMap) {
                 3, 0, makeFlatVector<int32_t>({1, 2, 3})))),
         "FlatMapVector keys are not distinct.");
   }
+
+  // Scan spec with map key filter
+  {
+    auto inputFlatMap = makeFlatMapVector<int64_t, int64_t>({
+        {{{1, 100}, {5, 500}, {10, 1000}, {15, 1500}, {20, 2000}}},
+        {{{2, 200}, {8, 800}, {12, 1200}}},
+        {{{3, 300}, {7, 700}, {18, 1800}, {25, 2500}}},
+    });
+    auto input = makeRowVector({inputFlatMap, inputFlatMap->toMapVector()});
+
+    VeloxWriterOptions writerOptions;
+    writerOptions.flatMapColumns = {"c0"};
+    auto fileContent =
+        test::createNimbleFile(*rootPool(), input, writerOptions, true);
+
+    // Test with map key filter [5, 12]
+    auto scanSpec = std::make_shared<common::ScanSpec>("root");
+    scanSpec->addAllChildFields(*input->type());
+    scanSpec->childByName("c0")
+        ->childByName(common::ScanSpec::kMapKeysFieldName)
+        ->setFilter(std::make_unique<common::BigintRange>(5, 12, false));
+    scanSpec->childByName("c1")
+        ->childByName(common::ScanSpec::kMapKeysFieldName)
+        ->setFilter(std::make_unique<common::BigintRange>(5, 12, false));
+    auto readers = makeReaders(input, fileContent, scanSpec, true);
+
+    // Expected output after filtering: only keys in [5, 12] remain
+    auto expectedFlatMap = makeFlatMapVector<int64_t, int64_t>({
+        {{{5, 500}, {10, 1000}}},
+        {{{8, 800}, {12, 1200}}},
+        {{{7, 700}}},
+    });
+    auto expected =
+        makeRowVector({expectedFlatMap->toMapVector(), expectedFlatMap});
+    validate(*expected, *readers.rowReader, inputFlatMap->size(), [](auto) {
+      return true;
+    });
+  }
 }
 
 TEST_P(SelectiveNimbleReaderTest, mapAsStruct) {