From e1f89b69beef7fd923b8cbb1ae02bc83245ed06b Mon Sep 17 00:00:00 2001
From: Patrick Sullivan <patricksullivan@meta.com>
Date: Wed, 24 May 2023 14:26:20 -0700
Subject: [PATCH] Remove unnecessary member variable from FlatMapColumnReader
 (#4993)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/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
---
 .../dwio/dwrf/reader/FlatMapColumnReader.cpp  | 41 +++++++++----------
 velox/dwio/dwrf/reader/FlatMapColumnReader.h  |  4 --
 2 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/velox/dwio/dwrf/reader/FlatMapColumnReader.cpp b/velox/dwio/dwrf/reader/FlatMapColumnReader.cpp
index c7dbc3e0829a..c456e3a0d9a7 100644
--- a/velox/dwio/dwrf/reader/FlatMapColumnReader.cpp
+++ b/velox/dwio/dwrf/reader/FlatMapColumnReader.cpp
@@ -96,6 +96,16 @@ 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,
@@ -103,10 +113,12 @@ std::vector<std::unique_ptr<KeyNode<T>>> getKeyNodesFiltered(
     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,7 +223,7 @@ 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) {
@@ -227,8 +231,6 @@ FlatMapColumnReader<T>::FlatMapColumnReader(
   });
 
   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
index 9f0f446e274c..59ed0896bce6 100644
--- 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_;