feat(exec): Push dynamic filters through StreamingAggregation grouping keys (#17918)

Summary:
Pull Request resolved: #17918

Dynamic filters generated by a hash-join probe are pushed toward the
probe-side source by tracing each filtered column through the identity
projections of intervening operators. The trace stops at the first
operator that does not expose the filtered channel as an identity
projection.

`StreamingAggregation` did not register its grouping keys in
`identityProjections_`, unlike `HashAggregation`, so the trace broke at the
aggregation and the filter never reached the scan. A plan with a streaming
(segmented) aggregation between a hash-join probe and a table scan therefore
missed dynamic filtering on the aggregation's grouping keys.

Grouping keys are identity passthroughs to the output, so register them as
identity projections and let the existing Driver trace push dynamic filters
through to the source. No new operator support is needed.

Fixes #17827

Reviewed By: mbasmanova

Differential Revision: D109527298

fbshipit-source-id: 42ac6d6cd8ee358a3369926991b8d5c54975de8a

Author: kKPulla

velox/exec/StreamingAggregation.cpp

@@ -75,6 +75,14 @@ void StreamingAggregation::initialize() {
     groupingKeyTypes.push_back(inputType->childAt(channel));
   }
 
+  // Grouping keys pass through unchanged to the output: output channel 'i' is
+  // input channel 'groupingKeys_[i]'. Exposing them as identity projections
+  // lets dynamic filters on grouping keys be pushed down through this operator
+  // to the source.
+  for (column_index_t i = 0; i < groupingKeys_.size(); ++i) {
+    identityProjections_.emplace_back(groupingKeys_[i], i);
+  }
+
   std::shared_ptr<core::ExpressionEvaluator> expressionEvaluator;
   aggregates_ = toAggregateInfo(
       *aggregationNode_, *operatorCtx_, numKeys, expressionEvaluator, true);

velox/exec/tests/HashJoinTestExtra.cpp

@@ -1878,6 +1878,86 @@ TEST_P(HashJoinTest, dynamicFiltersPushDownThroughAgg) {
       .run();
 }
 
+TEST_P(HashJoinTest, dynamicFiltersPushDownThroughStreamingAgg) {
+  const int32_t numRowsProbe = 300;
+  const int32_t numRowsBuild = 100;
+
+  // Probe c0 is strictly increasing (row - 10), so the input is already
+  // clustered on the grouping key required by StreamingAggregation.
+  std::vector<RowVectorPtr> probeVectors{makeRowVector({
+      makeFlatVector<int32_t>(numRowsProbe, [&](auto row) { return row - 10; }),
+      makeFlatVector<int64_t>(numRowsProbe, folly::identity),
+  })};
+  std::shared_ptr<TempFilePath> probeFile = TempFilePath::create();
+  writeToFile(probeFile->getPath(), probeVectors);
+
+  // Create build data
+  std::vector<RowVectorPtr> buildVectors{makeRowVector(
+      {"u0"}, {makeFlatVector<int32_t>(numRowsBuild, [&](auto row) {
+        return 35 + 2 * (row + numRowsBuild / 5);
+      })})};
+
+  createDuckDbTable("t", probeVectors);
+  createDuckDbTable("u", buildVectors);
+
+  auto probeType = ROW({"c0", "c1"}, {INTEGER(), BIGINT()});
+
+  // For sum() the partial accumulator equals the final value, so the same
+  // reference query and verifier apply to both single and partial streaming
+  // aggregation steps.
+  for (const auto step :
+       {core::AggregationNode::Step::kSingle,
+        core::AggregationNode::Step::kPartial}) {
+    SCOPED_TRACE(fmt::format("step {}", core::mapAggregationStepToName(step)));
+    auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
+    auto buildSide =
+        PlanBuilder(planNodeIdGenerator).values(buildVectors).planNode();
+
+    // Inner join.
+    core::PlanNodeId scanNodeId;
+    core::PlanNodeId joinNodeId;
+    core::PlanNodeId aggNodeId;
+    auto op = PlanBuilder(planNodeIdGenerator, pool_.get())
+                  .tableScan(probeType)
+                  .capturePlanNodeId(scanNodeId)
+                  .streamingAggregation({"c0"}, {"sum(c1)"}, {}, step, false)
+                  .capturePlanNodeId(aggNodeId)
+                  .hashJoin(
+                      {"c0"},
+                      {"u0"},
+                      buildSide,
+                      "",
+                      {"c0", "a0"},
+                      core::JoinType::kInner)
+                  .capturePlanNodeId(joinNodeId)
+                  .planNode();
+
+    SplitPath splitPaths = {{scanNodeId, {probeFile->getPath()}}};
+    HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get())
+        .planNode(std::move(op))
+        .inputSplits(splitPaths)
+        .injectSpill(false)
+        .checkSpillStats(false)
+        .referenceQuery(
+            "SELECT c0, sum(c1) FROM t, u WHERE c0 = u0 group by c0")
+        .verifier([&](const std::shared_ptr<Task>& task, bool hasSpill) {
+          auto planStats = toPlanStats(task->taskStats());
+          auto dynamicFilterStats = planStats.at(scanNodeId).dynamicFilterStats;
+          ASSERT_EQ(
+              1, getFiltersProduced(task, getOperatorIndex(joinNodeId)).sum);
+          ASSERT_EQ(
+              1, getFiltersAccepted(task, getOperatorIndex(scanNodeId)).sum);
+          ASSERT_LT(
+              getInputPositions(task, getOperatorIndex(aggNodeId)),
+              numRowsProbe);
+          ASSERT_EQ(
+              dynamicFilterStats.producerNodeIds,
+              std::unordered_set({joinNodeId}));
+        })
+        .run();
+  }
+}
+
 TEST_P(HashJoinTest, noDynamicFiltersPushDownThroughRightJoin) {
   std::vector<RowVectorPtr> innerBuild = {makeRowVector(
       {"a"},