Remove unnecessary member variable from FlatMapColumnReader (facebook…

…incubator#4993) Summary: Pull Request resolved: facebookincubator#4993 We don't need `keySelectionStats_` to be a member variable of `FlatMapColumnReader` and `FlatMapStructEncodingColumnReader`, we can simply create a stats in the constructors and use there instead. Reviewed By: xiaoxmeng, DanielMunozT, Magoja Differential Revision: D45977383 fbshipit-source-id: e5172eddd3587b8e3b16d9fdbfc2e93f27e0ce35
xiaoxmeng · May 24, 2023 · e1f89b6 · e1f89b6
1 parent eec9f42
commit e1f89b6
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 25 deletions.
diff --git a/velox/dwio/dwrf/reader/FlatMapColumnReader.cpp b/velox/dwio/dwrf/reader/FlatMapColumnReader.cpp
@@ -96,17 +96,29 @@ uint32_t visitUniqueStreamsOfNode(
   return streams;
 }
 
+void triggerKeySelectionNotification(
+    FlatMapContext& context,
+    facebook::velox::dwio::common::flatmap::FlatMapKeySelectionStats&
+        keySelectionStats) {
+  if (!context.keySelectionCallback) {
+    return;
+  }
+  context.keySelectionCallback(keySelectionStats);
+}
+
 template <typename T>
 std::vector<std::unique_ptr<KeyNode<T>>> getKeyNodesFiltered(
     const std::function<bool(const KeyValue<T>&)>& keyPredicate,
     const std::shared_ptr<const TypeWithId>& requestedType,
     const std::shared_ptr<const TypeWithId>& dataType,
     StripeStreams& stripe,
     memory::MemoryPool& memoryPool,
-    facebook::velox::dwio::common::flatmap::FlatMapKeySelectionStats&
-        keySelectionStats) {
+    FlatMapContext& flatMapContext) {
   std::vector<std::unique_ptr<KeyNode<T>>> keyNodes;
 
+  auto keySelectionStats =
+      facebook::velox::dwio::common::flatmap::FlatMapKeySelectionStats{};
+
   const auto requestedValueType = requestedType->childAt(1);
   const auto dataValueType = dataType->childAt(1);
   std::unordered_set<size_t> processed;
@@ -161,6 +173,8 @@ std::vector<std::unique_ptr<KeyNode<T>>> getKeyNodesFiltered(
   keySelectionStats.selectedKeys = keyNodes.size();
   keySelectionStats.totalKeys = processed.size();
 
+  triggerKeySelectionNotification(flatMapContext, keySelectionStats);
+
   VLOG(1) << "[Flat-Map] Initialized a flat-map column reader for node "
           << dataType->id << ", keys=" << keyNodes.size()
           << ", streams=" << streams;
@@ -190,16 +204,6 @@ std::vector<std::unique_ptr<KeyNode<T>>> rearrangeKeyNodesAsProjectedOrder(
 }
 } // namespace
 
-void triggerKeySelectionNotification(
-    FlatMapContext& context,
-    facebook::velox::dwio::common::flatmap::FlatMapKeySelectionStats&
-        keySelectionStats) {
-  if (!context.keySelectionCallback) {
-    return;
-  }
-  context.keySelectionCallback(keySelectionStats);
-}
-
 template <typename T>
 FlatMapColumnReader<T>::FlatMapColumnReader(
     const std::shared_ptr<const TypeWithId>& requestedType,
@@ -219,16 +223,14 @@ FlatMapColumnReader<T>::FlatMapColumnReader(
       dataType,
       stripe,
       memoryPool_,
-      keySelectionStats_);
+      flatMapContext_);
 
   // sort nodes by sequence id so order of keys is fixed
   std::sort(keyNodes_.begin(), keyNodes_.end(), [](auto& a, auto& b) {
     return a->getSequence() < b->getSequence();
   });
 
   initStringKeyBuffer();
-
-  triggerKeySelectionNotification(flatMapContext_, keySelectionStats_);
 }
 
 template <typename T>
@@ -539,8 +541,7 @@ std::vector<std::unique_ptr<KeyNode<T>>> getKeyNodesForStructEncoding(
     const std::shared_ptr<const TypeWithId>& dataType,
     StripeStreams& stripe,
     memory::MemoryPool& memoryPool,
-    facebook::velox::dwio::common::flatmap::FlatMapKeySelectionStats&
-        keySelectionStats) {
+    FlatMapContext& flatMapContext) {
   // `KeyNode` is ordered based on the projection. So if [3, 2, 1] is
   // projected, the vector of key node will be created [3, 2, 1].
   // If the key is not found in the stripe, the key node will be nullptr.
@@ -553,7 +554,7 @@ std::vector<std::unique_ptr<KeyNode<T>>> getKeyNodesForStructEncoding(
       dataType,
       stripe,
       memoryPool,
-      keySelectionStats);
+      flatMapContext);
 
   const auto& mapColumnIdAsStruct =
       stripe.getRowReaderOptions().getMapColumnIdAsStruct();
@@ -576,15 +577,13 @@ FlatMapStructEncodingColumnReader<T>::FlatMapStructEncodingColumnReader(
           dataType,
           stripe,
           memoryPool_,
-          keySelectionStats_)},
+          flatMapContext_)},
       nullColumnReader_{std::make_unique<NullColumnReader>(
           stripe,
           requestedType_->type->asMap().valueType())} {
   DWIO_ENSURE_EQ(nodeType_->id, dataType->id);
   DWIO_ENSURE(!keyNodes_.empty()); // "For struct encoding, keys to project must
                                    // be configured.";
-
-  triggerKeySelectionNotification(flatMapContext_, keySelectionStats_);
 }
 
 template <typename T>

diff --git a/velox/dwio/dwrf/reader/FlatMapColumnReader.h b/velox/dwio/dwrf/reader/FlatMapColumnReader.h
@@ -157,8 +157,6 @@ class FlatMapColumnReader : public ColumnReader {
 
  private:
   const std::shared_ptr<const dwio::common::TypeWithId> requestedType_;
-  facebook::velox::dwio::common::flatmap::FlatMapKeySelectionStats
-      keySelectionStats_;
   std::vector<std::unique_ptr<KeyNode<T>>> keyNodes_;
   std::unique_ptr<StringKeyBuffer> stringKeyBuffer_;
   bool returnFlatVector_;
@@ -187,8 +185,6 @@ class FlatMapStructEncodingColumnReader : public ColumnReader {
 
  private:
   const std::shared_ptr<const dwio::common::TypeWithId> requestedType_;
-  facebook::velox::dwio::common::flatmap::FlatMapKeySelectionStats
-      keySelectionStats_;
   std::vector<std::unique_ptr<KeyNode<T>>> keyNodes_;
   std::unique_ptr<NullColumnReader> nullColumnReader_;
   BufferPtr mergedNulls_;