Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions datafusion/core/tests/dataframe/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -910,7 +910,7 @@ async fn test_aggregate_with_pk2() -> Result<()> {
assert_snapshot!(
physical_plan_to_string(&df).await,
@r"
AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[], ordering_mode=Sorted
AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[]
FilterExec: id@0 = 1 AND name@1 = a
DataSourceExec: partitions=1, partition_sizes=[1]
"
Expand Down Expand Up @@ -958,7 +958,7 @@ async fn test_aggregate_with_pk3() -> Result<()> {
assert_snapshot!(
physical_plan_to_string(&df).await,
@r"
AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[], ordering_mode=PartiallySorted([0])
AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[]
FilterExec: id@0 = 1
DataSourceExec: partitions=1, partition_sizes=[1]
"
Expand Down Expand Up @@ -1008,7 +1008,7 @@ async fn test_aggregate_with_pk4() -> Result<()> {
assert_snapshot!(
physical_plan_to_string(&df).await,
@r"
AggregateExec: mode=Single, gby=[id@0 as id], aggr=[], ordering_mode=Sorted
AggregateExec: mode=Single, gby=[id@0 as id], aggr=[]
FilterExec: id@0 = 1
DataSourceExec: partitions=1, partition_sizes=[1]
"
Expand Down
10 changes: 5 additions & 5 deletions datafusion/core/tests/physical_optimizer/filter_pushdown.rs
Original file line number Diff line number Diff line change
Expand Up @@ -740,12 +740,12 @@ fn test_pushdown_through_aggregates_on_grouping_columns() {
OptimizationTest:
input:
- FilterExec: b@1 = bar
- AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=PartiallySorted([0])
- AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt]
- FilterExec: a@0 = foo
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
output:
Ok:
- AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=Sorted
- AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt]
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=a@0 = foo AND b@1 = bar
"
);
Expand Down Expand Up @@ -804,7 +804,7 @@ fn test_pushdown_through_aggregates_preserves_parent_filter_order() {
output:
Ok:
- FilterExec: cnt@2 = 1
- AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=PartiallySorted([1])
- AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt]
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar
"
);
Expand Down Expand Up @@ -2099,7 +2099,7 @@ fn test_pushdown_filter_on_non_first_grouping_column() {
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
output:
Ok:
- AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=PartiallySorted([1])
- AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt]
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar
"
);
Expand Down Expand Up @@ -2435,7 +2435,7 @@ fn test_pushdown_through_aggregate_with_reordered_input_columns() {
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
output:
Ok:
- AggregateExec: mode=Final, gby=[a@1 as a, b@2 as b], aggr=[cnt], ordering_mode=PartiallySorted([1])
- AggregateExec: mode=Final, gby=[a@1 as a, b@2 as b], aggr=[cnt]
- ProjectionExec: expr=[c@2 as c, a@0 as a, b@1 as b]
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar
"
Expand Down
21 changes: 21 additions & 0 deletions datafusion/core/tests/physical_optimizer/sanity_checker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -321,6 +321,27 @@ async fn test_aggregate() -> Result<()> {
Ok(())
}

#[tokio::test]
async fn test_constant_grouping_expr_on_unbounded_input() -> Result<()> {
let case = QueryCase {
sql: "SELECT c1, c2, MIN(c4) FROM test WHERE c2 = 1 GROUP BY c1, c2".to_string(),
cases: vec![
Arc::new(UnaryTestCase {
source_type: SourceType::Bounded,
expect_fail: false,
}),
Arc::new(UnaryTestCase {
source_type: SourceType::Unbounded,
expect_fail: true,
}),
],
error_operator: "operator: AggregateExec".to_string(),
};

case.run().await?;
Ok(())
}

#[tokio::test]
async fn test_window_agg_hash_partition() -> Result<()> {
let test1 = UnaryTestCase {
Expand Down
105 changes: 100 additions & 5 deletions datafusion/physical-plan/src/aggregates/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1029,6 +1029,8 @@ impl AggregateExec {
let required_input_ordering =
LexRequirement::new(new_requirements).map(OrderingRequirements::new_soft);

// Constant expressions never change, so they cannot mark a completed group.
// Exclude them from both the ordering indices and the group expression count.
// If our aggregation has grouping sets then our base grouping exprs will
// be expanded based on the flags in `group_by.groups` where for each
// group we swap the grouping expr for `null` if the flag is `true`
Expand All @@ -1037,9 +1039,18 @@ impl AggregateExec {
let indices: Vec<usize> = indices
.into_iter()
.filter(|idx| group_by.groups.iter().all(|group| !group[*idx]))
.filter(|idx| {
input_eq_properties
.is_expr_constant(&groupby_exprs[*idx])
.is_none()
})
.collect();

let mut input_order_mode = if indices.len() == groupby_exprs.len()
let num_non_constant_groupby_exprs = groupby_exprs
.iter()
.filter(|expr| input_eq_properties.is_expr_constant(expr).is_none())
.count();
let mut input_order_mode = if indices.len() == num_non_constant_groupby_exprs
&& !indices.is_empty()
&& group_by.groups.len() == 1
{
Expand Down Expand Up @@ -3259,7 +3270,7 @@ mod tests {

use arrow::array::{
BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array,
Int64Array, NullArray, StructArray, UInt32Array, UInt64Array,
Int64Array, NullArray, StringArray, StructArray, UInt32Array, UInt64Array,
};
use arrow::compute::{SortOptions, concat_batches};
use arrow::datatypes::Int32Type;
Expand All @@ -3270,7 +3281,7 @@ mod tests {
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
use datafusion_expr::{
Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator,
Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, Operator,
Signature, Volatility,
};
use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf;
Expand All @@ -3284,7 +3295,7 @@ mod tests {
use datafusion_physical_expr::Partitioning;
use datafusion_physical_expr::PhysicalSortExpr;
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
use datafusion_physical_expr::expressions::{Literal, NotExpr};
use datafusion_physical_expr::expressions::{Literal, NotExpr, binary};

use crate::projection::ProjectionExec;
use crate::repartition::RepartitionExec;
Expand Down Expand Up @@ -4728,6 +4739,87 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn constant_grouping_expr_is_not_a_completion_boundary() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("key", DataType::Int32, true),
Field::new("market", DataType::Utf8, true),
Field::new("value", DataType::Int64, false),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int32Array::from(vec![None, Some(10), Some(10)])),
Arc::new(StringArray::from(vec![Some("US"), Some("US"), Some("US")])),
Arc::new(Int64Array::from(vec![3, 1, 2])),
],
)?;

let build_aggregate = |input: Arc<dyn ExecutionPlan>| -> Result<AggregateExec> {
let predicate =
binary(col("market", &schema)?, Operator::Eq, lit("US"), &schema)?;
let input = Arc::new(FilterExecBuilder::new(predicate, input).build()?);
AggregateExec::try_new(
AggregateMode::Single,
PhysicalGroupBy::new_single(vec![
(col("key", &schema)?, "key".to_string()),
(col("market", &schema)?, "market".to_string()),
]),
vec![Arc::new(
AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?])
.schema(Arc::clone(&schema))
.alias("COUNT(value)")
.build()?,
)],
vec![None],
input,
Arc::clone(&schema),
)
};

let unordered_input = TestMemoryExec::try_new_exec(
&[vec![batch.clone()]],
Arc::clone(&schema),
None,
)?;
let aggregate = build_aggregate(unordered_input)?;
assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear);
assert_eq!(
aggregate.schema().as_ref(),
&Schema::new(vec![
Field::new("key", DataType::Int32, true),
Field::new("market", DataType::Utf8, true),
Field::new("COUNT(value)", DataType::Int64, false),
])
);

let output =
collect(aggregate.execute(0, Arc::new(TaskContext::default()))?).await?;
assert_eq!(output.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
assert_snapshot!(batches_to_sort_string(&output), @r"
+-----+--------+--------------+
| key | market | COUNT(value) |
+-----+--------+--------------+
| | US | 1 |
| 10 | US | 2 |
+-----+--------+--------------+
");

let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new(
Column::new("key", 0),
))])
.unwrap();
let ordered_input =
TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)?
.try_with_sort_information(vec![ordering])?;
let ordered_input =
Arc::new(TestMemoryExec::update_cache(&Arc::new(ordered_input)));
let aggregate = build_aggregate(ordered_input)?;
assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Sorted);

Ok(())
}

fn partial_reduce_test_aggregate() -> Result<AggregateExec> {
partial_reduce_test_aggregate_with_batches(1)
}
Expand Down Expand Up @@ -6615,7 +6707,10 @@ mod tests {
");
}
}
Err(e) => assert!(matches!(e, DataFusionError::ResourcesExhausted(_))),
Err(e) => assert!(
matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)),
"unexpected error: {e}"
),
}

Ok(())
Expand Down
4 changes: 2 additions & 2 deletions datafusion/sqllogictest/test_files/aggregate.slt
Original file line number Diff line number Diff line change
Expand Up @@ -8751,9 +8751,9 @@ logical_plan
15)------------EmptyRelation: rows=1
physical_plan
01)ProjectionExec: expr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST]@1 as last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))@2 as sum(DISTINCT Int64(1))]
02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))], ordering_mode=Sorted
02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))]
03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=5
04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))], ordering_mode=Sorted
04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))]
05)--------UnionExec
06)----------ProjectionExec: expr=[1 as id, 2 as foo]
07)------------PlaceholderRowExec
Expand Down
4 changes: 2 additions & 2 deletions datafusion/sqllogictest/test_files/array_agg.slt
Original file line number Diff line number Diff line change
Expand Up @@ -574,9 +574,9 @@ logical_plan
15)------------EmptyRelation: rows=1
physical_plan
01)ProjectionExec: expr=[array_length(array_agg(DISTINCT a.foo)@1) as array_length(array_agg(DISTINCT a.foo)), sum(DISTINCT Int64(1))@2 as sum(DISTINCT Int64(1))]
02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted
02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))]
03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=5
04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted
04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))]
05)--------UnionExec
06)----------ProjectionExec: expr=[1 as id, CAST(2 AS Int64) as foo]
07)------------PlaceholderRowExec
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -512,9 +512,9 @@ query TT
EXPLAIN SELECT a, count(b) FROM agg_filter_pushdown WHERE a = 'x' GROUP BY a;
----
physical_plan
01)AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)], ordering_mode=Sorted
01)AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)]
02)--RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4
03)----AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)], ordering_mode=Sorted
03)----AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)]
04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_filter_pushdown.parquet]]}, projection=[a, b], file_type=parquet, predicate=a@0 = x, pruning_predicate=a_null_count@2 != row_count@3 AND a_min@0 <= x AND x <= a_max@1, required_guarantees=[a in (x)]

Expand Down
4 changes: 2 additions & 2 deletions datafusion/sqllogictest/test_files/union.slt
Original file line number Diff line number Diff line change
Expand Up @@ -783,7 +783,7 @@ logical_plan
physical_plan
01)UnionExec
02)--ProjectionExec: expr=[Int64(1)@0 as a]
03)----AggregateExec: mode=SinglePartitioned, gby=[1 as Int64(1)], aggr=[], ordering_mode=Sorted
03)----AggregateExec: mode=SinglePartitioned, gby=[1 as Int64(1)], aggr=[]
04)------PlaceholderRowExec
05)--ProjectionExec: expr=[2 as a]
06)----PlaceholderRowExec
Expand Down Expand Up @@ -811,7 +811,7 @@ logical_plan
physical_plan
01)UnionExec
02)--ProjectionExec: expr=[count(Int64(1))@1 as count, CAST(n@0 AS Int64) as n]
03)----AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))], ordering_mode=Sorted
03)----AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))]
04)------ProjectionExec: expr=[5 as n]
05)--------PlaceholderRowExec
06)--ProjectionExec: expr=[1 as count, max(Int64(10))@0 as n]
Expand Down
2 changes: 1 addition & 1 deletion datafusion/sqllogictest/test_files/window.slt
Original file line number Diff line number Diff line change
Expand Up @@ -276,7 +276,7 @@ physical_plan
03)----SortExec: expr=[b@0 ASC NULLS LAST], preserve_partitioning=[true]
04)------AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[max(d.a)]
05)--------RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=4
06)----------AggregateExec: mode=Partial, gby=[b@1 as b], aggr=[max(d.a)], ordering_mode=Sorted
06)----------AggregateExec: mode=Partial, gby=[b@1 as b], aggr=[max(d.a)]
07)------------UnionExec
08)--------------ProjectionExec: expr=[1 as a, aa as b]
09)----------------PlaceholderRowExec
Expand Down