From 47de46639a9e55b929d1bf6fd65c8c0657515bd8 Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Tue, 18 Aug 2026 22:22:56 +0200 Subject: [PATCH] fix: adapt batches with stricter nested nullability to aggregate plan schema (#24069) Enforces strict schema containment at the AggregateExec boundary via AdaptedInputRecordBatchStream and datafusion_common::nested_struct::adapt_batch_to_schema. When input batches contain stricter field nullability (e.g., non-nullable struct fields passed into an aggregate planned with nullable fields), the stream adapts incoming batches to conform to the declared self.input.schema() before feeding into aggregate streams. Fixes #24069 --- datafusion/common/src/nested_struct.rs | 194 +++++++++++++- datafusion/core/tests/sql/aggregates/mod.rs | 1 + .../sql/aggregates/nested_nullability.rs | 246 ++++++++++++++++++ .../src/aggregates/aggregate_stream.rs | 2 +- .../src/aggregates/grouped_hash_stream.rs | 2 +- .../src/aggregates/grouped_topk_stream.rs | 2 +- .../src/aggregates/hash_stream.rs | 4 +- .../physical-plan/src/aggregates/mod.rs | 53 +++- .../src/aggregates/ordered_final_stream.rs | 2 +- .../src/aggregates/ordered_partial_stream.rs | 2 +- .../src/aggregates/ordered_single_stream.rs | 2 +- .../src/aggregates/partial_reduce_stream.rs | 2 +- .../src/aggregates/single_stream.rs | 2 +- 13 files changed, 501 insertions(+), 13 deletions(-) create mode 100644 datafusion/core/tests/sql/aggregates/nested_nullability.rs diff --git a/datafusion/common/src/nested_struct.rs b/datafusion/common/src/nested_struct.rs index e915b91b911cc..52d0cb41a3e51 100644 --- a/datafusion/common/src/nested_struct.rs +++ b/datafusion/common/src/nested_struct.rs @@ -19,11 +19,12 @@ use crate::error::{_plan_err, Result}; use arrow::{ array::{ Array, ArrayRef, AsArray, DictionaryArray, FixedSizeListArray, GenericListArray, - GenericListViewArray, StructArray, downcast_integer, make_array, new_null_array, + GenericListViewArray, RecordBatch, StructArray, downcast_integer, make_array, + new_null_array, }, buffer::NullBuffer, compute::{CastOptions, can_cast_types, cast_with_options}, - datatypes::{DataType, DataType::Struct, Field, FieldRef}, + datatypes::{DataType, DataType::Struct, Field, FieldRef, SchemaRef}, }; use std::{collections::HashSet, sync::Arc}; @@ -1703,3 +1704,192 @@ mod tests { )); } } + +/// Adapts a [`RecordBatch`] to a target [`SchemaRef`]. +/// +/// If `batch` already has the target schema, it is returned immediately. +/// +/// If `batch` has a schema that is a legitimate subtype / stricter subset of +/// `target_schema` (as verified by [`arrow::datatypes::Schema::contains`]), +/// this function transforms the metadata/types of differing columns to match +/// `target_schema` without copying primitive buffer data. +/// +/// If `batch` does not conform to `target_schema` under schema containment, +/// an error is returned. +pub fn adapt_batch_to_schema( + batch: RecordBatch, + target_schema: &SchemaRef, +) -> Result { + if Arc::ptr_eq(batch.schema_ref(), target_schema) + || batch.schema().as_ref() == target_schema.as_ref() + { + return Ok(batch); + } + + if !target_schema.contains(batch.schema().as_ref()) { + return _plan_err!( + "Batch schema does not conform to expected schema. Expected: {target_schema}, got: {}", + batch.schema() + ); + } + + let mut columns = Vec::with_capacity(batch.num_columns()); + let mut needs_column_adaptation = false; + let cast_options = CastOptions::default(); + + for (target_field, col) in target_schema.fields().iter().zip(batch.columns()) { + if target_field.data_type() != col.data_type() { + needs_column_adaptation = true; + let adapted_col = cast_column(col, target_field.data_type(), &cast_options)?; + columns.push(adapted_col); + } else { + columns.push(Arc::clone(col)); + } + } + + if needs_column_adaptation { + Ok(RecordBatch::try_new(Arc::clone(target_schema), columns)?) + } else { + // Schema differs only in top-level metadata or field nullability, while + // column data types match exactly. Replace the schema on the batch. + Ok(RecordBatch::try_new( + Arc::clone(target_schema), + batch.columns().to_vec(), + )?) + } +} + +#[cfg(test)] +mod adapt_schema_tests { + use super::*; + use arrow::array::{BooleanArray, Int32Array, StringArray}; + use arrow::datatypes::{Field, Fields, Schema}; + + #[test] + fn test_adapt_batch_to_schema_identical() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec![Some("x"), None, Some("z")])), + ], + )?; + + let adapted = adapt_batch_to_schema(batch.clone(), &schema)?; + assert!(Arc::ptr_eq(batch.schema_ref(), adapted.schema_ref())); + assert_eq!(batch, adapted); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_stricter_nested_struct() -> Result<()> { + let declared_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "nested", + Struct(Fields::from(vec![Field::new( + "val", + DataType::Boolean, + true, + )])), + false, + ), + ])); + + let stricter_batch_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "nested", + Struct(Fields::from(vec![Field::new( + "val", + DataType::Boolean, + false, + )])), + false, + ), + ])); + + let struct_col = Arc::new(StructArray::new( + Fields::from(vec![Field::new("val", DataType::Boolean, false)]), + vec![Arc::new(BooleanArray::from(vec![true, false, true]))], + None, + )); + + let batch = RecordBatch::try_new( + stricter_batch_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3])), struct_col], + )?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + assert_eq!(adapted.num_rows(), 3); + assert_eq!(adapted.num_columns(), 2); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_top_level_nullability_only() -> Result<()> { + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + let stricter_batch_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + let batch = RecordBatch::try_new( + stricter_batch_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + assert_eq!(adapted.column(0).len(), 3); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_incompatible_rejected() { + let declared_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), // target is non-nullable + ])); + + let batch_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), // batch is nullable (not contained!) + ])); + + let batch = RecordBatch::try_new( + batch_schema, + vec![Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]))], + ) + .unwrap(); + + let res = adapt_batch_to_schema(batch, &declared_schema); + assert!(res.is_err()); + let err_msg = res.unwrap_err().to_string(); + assert!( + err_msg.contains("does not conform to expected schema"), + "unexpected error message: {err_msg}" + ); + } + + #[test] + fn test_adapt_batch_to_schema_incompatible_type_rejected() { + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)])); + + let batch_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + let batch = RecordBatch::try_new( + batch_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + + let res = adapt_batch_to_schema(batch, &declared_schema); + assert!(res.is_err()); + } +} diff --git a/datafusion/core/tests/sql/aggregates/mod.rs b/datafusion/core/tests/sql/aggregates/mod.rs index b209e91cc81e7..186297b639cbd 100644 --- a/datafusion/core/tests/sql/aggregates/mod.rs +++ b/datafusion/core/tests/sql/aggregates/mod.rs @@ -1021,3 +1021,4 @@ pub fn split_fuzz_timestamp_data_into_batches( pub mod basic; pub mod dict_nulls; +mod nested_nullability; diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs new file mode 100644 index 0000000000000..448759ad74c54 --- /dev/null +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -0,0 +1,246 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Regression tests for aggregating batches whose data types are *stricter* +//! than the table's declared schema. See +//! . +//! +//! Builds on the end-to-end reproducer from #24278 by @alamb. +//! +//! A `RecordBatch` is a valid instance of a schema that is a superset of its +//! own (see [`Schema::contains`] / `Field::contains`): most commonly the +//! schema declares a (possibly nested) field as nullable while the batch's +//! arrays mark it non-nullable. `MemTable::try_new` accepts such batches via +//! exactly that check, and engines embedding DataFusion (e.g. Comet) feed +//! such batches over FFI. Aggregations must therefore not fail when the +//! runtime arrays are stricter than the planned schema. +//! +//! [`Schema::contains`]: arrow::datatypes::Schema::contains + +use std::sync::Arc; + +use arrow::array::{BooleanArray, RecordBatch, StructArray, UInt32Array}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use datafusion::datasource::MemTable; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; +use datafusion::physical_plan::collect; +use datafusion::physical_plan::expressions::col; +use datafusion::prelude::*; +use datafusion_common::Result; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::FairSpillPool; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_functions_aggregate::array_agg::array_agg_udaf; + +/// Returns the fields of the struct column `b`: a single `colA Boolean`. +/// +/// `col_a_nullable` controls whether `colA` is declared nullable — the only +/// difference between the table's declared schema (`true`) and the actual +/// batches (`false`). +fn make_struct_fields(col_a_nullable: bool) -> Fields { + Fields::from(vec![Field::new("colA", DataType::Boolean, col_a_nullable)]) +} + +/// Returns the schema `(a UInt32 NOT NULL, b Struct("colA" Boolean) NOT NULL)` +/// with the nested field `b.colA` nullable per `col_a_nullable`. +/// +/// See [`make_struct_fields`]. +fn make_schema(col_a_nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Struct(make_struct_fields(col_a_nullable)), + false, + ), + ])) +} + +/// Runs a SQL aggregation over a table whose batches are stricter than its +/// declared schema. +/// +/// [`Self::run`] registers table `t(a UInt32, b Struct("colA" Boolean))` +/// where the declared schema marks the nested field `colA` as nullable, but +/// the batches carry a stricter, non-nullable `colA`, then runs the query +/// and returns the collected result. +struct AggregateBatchesTest { + /// Number of rows in the table. `a` is `0..num_rows` (so also the number + /// of groups for `GROUP BY a`) and `b.colA` alternates `true` / `false`. + num_rows: u32, + /// If set, the context uses a [`FairSpillPool`] of this size (and a small + /// batch size) so the aggregation is forced to spill. + memory_limit: Option, +} + +impl AggregateBatchesTest { + fn new() -> Self { + Self { + num_rows: 100, + memory_limit: None, + } + } + + fn with_num_rows(mut self, num_rows: u32) -> Self { + self.num_rows = num_rows; + self + } + + fn with_memory_limit(mut self, memory_limit: usize) -> Self { + self.memory_limit = Some(memory_limit); + self + } + + /// Runs `sql` against the table described above and asserts the result + /// has one output row per group (i.e. [`Self::num_rows`] rows in total). + async fn run(self, sql: &str) -> Result<()> { + // The table's declared schema: the nested field `b.colA` is + // nullable ... + let declared_schema = make_schema(true); + + // ... while the batches are stricter: `b.colA` is non-nullable. + // `MemTable::try_new` accepts this combination via + // `Schema::contains`. + let batch_struct_fields = make_struct_fields(false); + let batch = RecordBatch::try_new( + make_schema(false), + vec![ + Arc::new(UInt32Array::from_iter_values(0..self.num_rows)), + Arc::new(StructArray::new( + batch_struct_fields, + vec![Arc::new(BooleanArray::from_iter( + (0..self.num_rows).map(|i| Some(i % 2 == 0)), + ))], + None, + )), + ], + )?; + + let table = MemTable::try_new(declared_schema, vec![vec![batch]])?; + + let ctx = match self.memory_limit { + Some(limit) => { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(limit))) + .build_arc()?; + SessionContext::new_with_config_rt( + SessionConfig::new().with_batch_size(100), + runtime, + ) + } + None => SessionContext::new(), + }; + ctx.register_table("t", Arc::new(table))?; + + let result = ctx.sql(sql).await?.collect().await?; + + let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, self.num_rows as usize); + Ok(()) + } +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches() -> Result<()> { + AggregateBatchesTest::new() + .run("SELECT a, array_agg(b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches() -> Result<()> { + AggregateBatchesTest::new() + .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> { + AggregateBatchesTest::new() + .with_num_rows(10_000) + .with_memory_limit(4_000_000) + .run("SELECT a, array_agg(b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches_with_spilling() -> Result<()> { + AggregateBatchesTest::new() + .with_num_rows(10_000) + .with_memory_limit(4_000_000) + .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") + .await +} + +/// Direct unit test for `AggregateExec` boundary adaptation: +/// Feeds `AggregateExec` directly from a `MemorySourceConfig` whose batches carry +/// a stricter nested struct nullability than the plan schema without going +/// through `MemTable`. +#[tokio::test] +async fn test_aggregate_exec_direct_input_adaptation() -> Result<()> { + let declared_schema = make_schema(true); + let batch_struct_fields = make_struct_fields(false); + let num_rows = 100_u32; + let stricter_batch = RecordBatch::try_new( + make_schema(false), + vec![ + Arc::new(UInt32Array::from_iter_values(0..num_rows)), + Arc::new(StructArray::new( + batch_struct_fields, + vec![Arc::new(BooleanArray::from_iter( + (0..num_rows).map(|i| Some(i % 2 == 0)), + ))], + None, + )), + ], + )?; + + let input_plan: Arc = MemorySourceConfig::try_new_exec( + &[vec![stricter_batch]], + Arc::clone(&declared_schema), + None, + )?; + + let grouping_set = + PhysicalGroupBy::new_single(vec![(col("a", &declared_schema)?, "a".to_string())]); + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("b", &declared_schema)?]) + .schema(Arc::clone(&declared_schema)) + .alias("array_agg(b)") + .build()?, + )]; + + let agg_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + grouping_set, + aggregates, + vec![None], + input_plan, + Arc::clone(&declared_schema), + )?); + + let task_ctx = Arc::new(TaskContext::default()); + let results = collect(agg_exec, task_ctx).await?; + + let total_rows: usize = results.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, num_rows as usize); + Ok(()) +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index ac7727b459300..01bb7ef764c07 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -290,7 +290,7 @@ impl AggregateStream { let agg_filter_expr = Arc::clone(&agg.filter_expr); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?; let filter_expressions = match agg.mode.input_mode() { diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 99c101199459f..d6eb758c3eb83 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -389,7 +389,7 @@ impl GroupedHashAggregateStream { let agg_filter_expr = Arc::clone(&agg.filter_expr); let batch_size = context.session_config().batch_size(); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs index 193fdba4b0198..6ab122f6efa3c 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -68,7 +68,7 @@ impl GroupedTopKAggregateStream { ) -> Result { let agg_schema = Arc::clone(&aggr.schema); let group_by = Arc::clone(&aggr.group_by); - let input = aggr.input.execute(partition, Arc::clone(context))?; + let input = aggr.execute_input(partition, context)?; let baseline_metrics = BaselineMetrics::new(&aggr.metrics, partition); let group_by_metrics = GroupByMetrics::new(&aggr.metrics, partition); let aggregate_arguments = diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index f697e5a394f65..5f8efc2be1e37 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -434,7 +434,7 @@ impl PartialHashAggregateStream { debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); @@ -1020,7 +1020,7 @@ impl FinalHashAggregateStream { debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index f9dd90f6f98fe..a054921b8c94e 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -170,7 +170,7 @@ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, - InputOrderMode, SendableRecordBatchStream, Statistics, + InputOrderMode, RecordBatchStream, SendableRecordBatchStream, Statistics, }; use datafusion_common::config::ConfigOptions; use parking_lot::Mutex; @@ -876,6 +876,45 @@ pub struct AggregateExec { dynamic_filter: Option>, } +/// A stream wrapper that ensures every yielded batch matches the declared input schema. +pub(crate) struct AdaptedInputRecordBatchStream { + inner: SendableRecordBatchStream, + schema: SchemaRef, +} + +impl AdaptedInputRecordBatchStream { + pub(crate) fn new(inner: SendableRecordBatchStream, schema: SchemaRef) -> Self { + Self { inner, schema } + } +} + +impl futures::Stream for AdaptedInputRecordBatchStream { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + use futures::StreamExt; + match self.inner.poll_next_unpin(cx) { + std::task::Poll::Ready(Some(Ok(batch))) => { + let adapted = datafusion_common::nested_struct::adapt_batch_to_schema( + batch, + &self.schema, + ); + std::task::Poll::Ready(Some(adapted)) + } + other => other, + } + } +} + +impl RecordBatchStream for AdaptedInputRecordBatchStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + impl AggregateExec { /// Function used in `OptimizeAggregateOrder` optimizer rule, /// where we need parts of the new value, others cloned from the old one @@ -1154,6 +1193,18 @@ impl AggregateExec { /// Aggregation has multiple specialized implementations optimized for /// different workloads. This function picks the best available path. + pub(crate) fn execute_input( + &self, + partition: usize, + context: &Arc, + ) -> Result { + let input = self.input.execute(partition, Arc::clone(context))?; + Ok(Box::pin(AdaptedInputRecordBatchStream::new( + input, + self.input.schema(), + ))) + } + fn execute_typed( &self, partition: usize, diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 19deedc258c46..6ec0f337ec16e 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -270,7 +270,7 @@ impl OrderedFinalAggregateStream { )); debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; Self::new_with_input(agg, context, partition, input, &agg.input_order_mode) } diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 9e93a111a6466..afc2fa5b30ad1 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -129,7 +129,7 @@ impl OrderedPartialAggregateStream { debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index 2025bdce307d9..78700bb5ae8a1 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -331,7 +331,7 @@ impl OrderedSingleAggregateStream { debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 2f4535e66f4ef..06e8da1c7d76e 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -157,7 +157,7 @@ impl PartialReduceHashAggregateStream { debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index c6f25dc2cf28b..1360000de259f 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -322,7 +322,7 @@ impl SingleHashAggregateStream { debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); let schema = Arc::clone(&agg.schema); - let input = agg.input.execute(partition, Arc::clone(context))?; + let input = agg.execute_input(partition, context)?; let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);