You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/21 08:41:12 UTC
(datafusion) branch main updated: perf: extend fully ordered aggregation to more scalar types (#26370)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 6c1727fcd9 perf: extend fully ordered aggregation to more scalar types (#26370)
6c1727fcd9 is described below
commit 6c1727fcd9230dda428f8c3a2397fce235544bd5
Author: Geoffrey Claude <ge...@datadoghq.com>
AuthorDate: Wed Oct 21 06:47:51 2026 +0000
perf: extend fully ordered aggregation to more scalar types (#26370)
## Which issue does this PR close?
Closes #26369. Follow-up to #25150.
## Rationale for this change
Fully ordered multi-column `GROUP BY` still hashes temporal, decimal,
and fixed-size binary keys. These types can use the existing
adjacent-run grouping path because Arrow's adjacent equality agrees with
the grouping builders.
In our production usecases, we are especially interested in #26369's
optimization extending to `Timestamp`s (timeseries data) and
`FixedSizeBinary(16)` (for UUIDs).
## What changes are included in this PR?
- **Only production code change is 12 new types** added to the `match`
under `supports_schema`: the rest are tests and benches.
- Enable timestamps with all units and timezone metadata, dates, valid
time units, durations, intervals, Decimal32/64/128/256, and every valid
fixed-size binary width.
- Preserve type metadata and the existing selection of single-column
specializations.
- Add a `fully_ordered` group to the existing SQL aggregation Criterion
benchmark, with integer/string controls, timestamp run-length variants,
and a timestamp `SUM(v)` case. The first two commits contain only
benchmarks, followed by the type expansion and regression tests,
including the UUID coverage.
## What is the testing strategy for this PR?
Differential tests compare ordered and hash grouping across nulls,
units/timezones, decimal metadata, sliced batches, and group emission.
Dedicated tests cover zero-width binary and 16-byte binary keys that
share their first 15 bytes, including cross-batch continuation and keys
that differ only in the last byte. SQL regression tests cover
single-stage and partial/final aggregation, mixed sort directions, null
ordering, and aggregate filters.
The benchmarks check input types, row and group counts, and fully
clustered aggregation without sorting or repartitioning in the measured
query. The timestamp `SUM(v)` case also checks every group's sum before
timing.
<details>
<summary>SQL benchmark results</summary>
[Latest benchmark bot
report](https://github.com/apache/datafusion/pull/26370#issuecomment-6266981812),
comparing the benchmark-only baseline `dc6a10cf3c8` with `a398f8768ca`.
Both refs contain the same 17 benchmark cases. The table shows rounded
speedups from the report. The controls are roughly unchanged; the small
differences may be benchmark noise. Values below 1× mean slower.
| Second grouping key / query | Rows per group | Approx. speedup |
| --- | ---: | ---: |
| Int64 (control) | 31 | ~0.95× |
| Utf8View (control) | 31 | ~0.99× |
| Timestamp (ns) | 31 | ~7.0× |
| Timestamp (µs, UTC) | 31 | ~6.9× |
| Date32 | 31 | ~7.5× |
| Time32 (ms) | 31 | ~7.5× |
| Time64 (ns) | 31 | ~7.0× |
| Duration (ns) | 31 | ~7.0× |
| Interval (day/time) | 31 | ~6.4× |
| Interval (month/day/ns) | 31 | ~5.3× |
| Decimal128(30, 4) | 31 | ~6.1× |
| Decimal256(60, 4) | 31 | ~5.2× |
| FixedSizeBinary(16) | 31 | ~4.6× |
| FixedSizeBinary(7) | 31 | ~4.4× |
| Timestamp (ns), singleton groups | 1 | ~2.9× |
| Timestamp (ns), long groups | 512 | ~7.7× |
| Timestamp (ns), `SUM(v)` | 31 | ~4.0× |
The 16 key-only cases use the same two-key deduplication query:
```sql
SELECT k1, k2 FROM ordered_input GROUP BY k1, k2
```
The additional timestamp case includes accumulator work:
```sql
SELECT k1, k2, SUM(v) FROM ordered_input GROUP BY k1, k2
```
Here `v` is the nonconstant source row number. The timestamp `SUM(v)`
case improved from approximately 43.9 ms to 11.0 ms (~4.0×).
- **Input:** about 2.1 million rows. `k1` is an Int32 prefix that
advances every 32 groups; nullable `k2` cycles through 32 values,
including one null, using the type listed above. Both keys vary.
- **Group sizes:** normally 31 rows per key pair, so runs cross the
8,192-row batch boundaries. The additional timestamp cases use one or
512 rows per group to vary cardinality.
- **Before timing:** generate the data, sort by both keys, and
materialize it into one in-memory partition with its ordering declared
to the planner.
- **Measured work:** SQL planning, query execution, and collection of
all results. The key-only cases focus on grouping the keys; the
timestamp `SUM(v)` case also updates an accumulator.
- **Setup checks:** verify input type, row count, and output group
count; require fully clustered aggregation and fail if the measured plan
adds sorting or repartitioning. The `SUM(v)` case verifies every group's
result, including null timestamp keys and the final short group, before
timing.
</details>
## Are there any user-facing changes?
Already-sorted multi-column aggregations can use lower-overhead key
storage for these types. Query results stay the same.
Developed with AI assistance.
---
datafusion/core/benches/aggregate_query_sql.rs | 236 ++++++++++-
.../group_values/multi_group_by/clustered.rs | 469 ++++++++++++++++-----
.../test_files/ordered_aggregate_key_types.slt | 172 ++++++++
3 files changed, 770 insertions(+), 107 deletions(-)
diff --git a/datafusion/core/benches/aggregate_query_sql.rs b/datafusion/core/benches/aggregate_query_sql.rs
index d7e24aceba..2ea81cccc8 100644
--- a/datafusion/core/benches/aggregate_query_sql.rs
+++ b/datafusion/core/benches/aggregate_query_sql.rs
@@ -17,10 +17,17 @@
mod data_utils;
+use arrow::array::{Array, AsArray};
+use arrow::datatypes::{
+ DataType, Int32Type, Int64Type, IntervalUnit, TimeUnit, TimestampNanosecondType,
+};
use criterion::{Criterion, criterion_group, criterion_main};
use data_utils::create_table_provider;
+use datafusion::datasource::MemTable;
use datafusion::error::Result;
-use datafusion::execution::context::SessionContext;
+use datafusion::execution::context::{SessionConfig, SessionContext};
+use datafusion::physical_plan::{collect, displayable};
+use datafusion_expr::col;
use parking_lot::Mutex;
use std::hint::black_box;
use std::sync::Arc;
@@ -330,5 +337,230 @@ fn criterion_benchmark(c: &mut Criterion) {
});
}
-criterion_group!(benches, criterion_benchmark);
+/// Measure the multi-column fully ordered grouping path, with data generation
+/// and sorting outside timing. Both keys vary, and the nullable second key
+/// changes type while the shape of the input stays the same.
+fn fully_ordered(c: &mut Criterion) {
+ use DataType::*;
+ use IntervalUnit::{DayTime, MonthDayNano};
+ use TimeUnit::{Microsecond, Millisecond, Nanosecond};
+
+ let cases = [
+ ("int64", "signed_key", Int64, 31),
+ (
+ "utf8_view",
+ "arrow_cast(lpad(CAST(key_value AS VARCHAR), 4, '0'), 'Utf8View')",
+ Utf8View,
+ 31,
+ ),
+ (
+ "timestamp_ns",
+ "arrow_cast(signed_key, 'Timestamp(Nanosecond)')",
+ Timestamp(Nanosecond, None),
+ 31,
+ ),
+ (
+ "timestamp_us_utc",
+ "arrow_cast(signed_key, 'Timestamp(Microsecond, Some(\"UTC\"))')",
+ Timestamp(Microsecond, Some("UTC".into())),
+ 31,
+ ),
+ (
+ "date32",
+ "arrow_cast(CAST(signed_key AS INTEGER), 'Date32')",
+ Date32,
+ 31,
+ ),
+ (
+ "time32_ms",
+ "arrow_cast(CAST(key_value AS INTEGER), 'Time32(Millisecond)')",
+ Time32(Millisecond),
+ 31,
+ ),
+ (
+ "time64_ns",
+ "arrow_cast(key_value, 'Time64(Nanosecond)')",
+ Time64(Nanosecond),
+ 31,
+ ),
+ (
+ "duration_ns",
+ "arrow_cast(signed_key, 'Duration(Nanosecond)')",
+ Duration(Nanosecond),
+ 31,
+ ),
+ (
+ "interval_day_time",
+ "arrow_cast(concat(CAST(signed_key AS VARCHAR), ' days'), 'Interval(DayTime)')",
+ Interval(DayTime),
+ 31,
+ ),
+ (
+ "interval_month_day_nano",
+ "arrow_cast(arrow_cast(signed_key, 'Duration(Nanosecond)'), 'Interval(MonthDayNano)')",
+ Interval(MonthDayNano),
+ 31,
+ ),
+ (
+ "decimal128",
+ "arrow_cast(concat('123456789012345678901', lpad(CAST(key_value AS VARCHAR), 4, '0'), '.1250'), 'Decimal128(30, 4)')",
+ Decimal128(30, 4),
+ 31,
+ ),
+ (
+ "decimal256",
+ "arrow_cast(concat('123456789012345678901234567890123456789012345678901', lpad(CAST(key_value AS VARCHAR), 4, '0'), '.1250'), 'Decimal256(60, 4)')",
+ Decimal256(60, 4),
+ 31,
+ ),
+ (
+ "fixed_binary16",
+ "arrow_cast(decode(lpad(to_hex(key_value), 32, '0'), 'hex'), 'FixedSizeBinary(16)')",
+ FixedSizeBinary(16),
+ 31,
+ ),
+ (
+ "fixed_binary7",
+ "arrow_cast(decode(lpad(to_hex(key_value), 14, '0'), 'hex'), 'FixedSizeBinary(7)')",
+ FixedSizeBinary(7),
+ 31,
+ ),
+ (
+ "timestamp_ns_singleton",
+ "arrow_cast(signed_key, 'Timestamp(Nanosecond)')",
+ Timestamp(Nanosecond, None),
+ 1,
+ ),
+ (
+ "timestamp_ns_long_runs",
+ "arrow_cast(signed_key, 'Timestamp(Nanosecond)')",
+ Timestamp(Nanosecond, None),
+ 512,
+ ),
+ ];
+
+ let rt = Runtime::new().unwrap();
+ let mut group = c.benchmark_group("fully_ordered");
+ for (name, key_expr, key_type, run_length) in cases {
+ let ctx = rt
+ .block_on(fully_ordered_context(
+ key_expr, &key_type, run_length, false,
+ ))
+ .unwrap();
+ group.bench_function(name, |b| {
+ b.iter(|| query(ctx.clone(), &rt, FULLY_ORDERED_QUERY))
+ });
+ }
+ // Add accumulator work for one representative timestamp case without
+ // multiplying every key type and run length by the aggregate expression.
+ let ctx = rt
+ .block_on(fully_ordered_context(
+ "arrow_cast(signed_key, 'Timestamp(Nanosecond)')",
+ &Timestamp(Nanosecond, None),
+ 31,
+ true,
+ ))
+ .unwrap();
+ group.bench_function("timestamp_ns_sum", |b| {
+ b.iter(|| query(ctx.clone(), &rt, FULLY_ORDERED_SUM_QUERY))
+ });
+ group.finish();
+}
+
+const FULLY_ORDERED_QUERY: &str = "SELECT k1, k2 FROM ordered_input GROUP BY k1, k2";
+const FULLY_ORDERED_SUM_QUERY: &str =
+ "SELECT k1, k2, SUM(v) FROM ordered_input GROUP BY k1, k2";
+
+async fn fully_ordered_context(
+ key_expr: &str,
+ key_type: &DataType,
+ run_length: usize,
+ with_sum: bool,
+) -> Result<Arc<Mutex<SessionContext>>> {
+ const ROWS: usize = 2_097_152;
+ // Use source row numbers as nonconstant values for the aggregate case.
+ let (value_projection, input_value, sql) = if with_sum {
+ (", v", ", value AS v", FULLY_ORDERED_SUM_QUERY)
+ } else {
+ ("", "", FULLY_ORDERED_QUERY)
+ };
+ let config = SessionConfig::new()
+ .with_target_partitions(1)
+ .with_batch_size(8192)
+ .with_prefer_existing_sort(true);
+ let ctx = SessionContext::new_with_config(config);
+ // Runs of 31 cross batch boundaries; singleton and long-run cases vary
+ // cardinality without multiplying every type by every input shape.
+ let input = ctx
+ .sql(&format!(
+ "SELECT CAST(group_id / 32 AS INTEGER) AS k1, \
+ CASE WHEN key_value = 0 THEN NULL ELSE {key_expr} END AS k2{value_projection} \
+ FROM (SELECT value / {run_length} AS group_id, \
+ value / {run_length} % 32 AS key_value, \
+ value / {run_length} % 32 - 16 AS signed_key{input_value} \
+ FROM range({ROWS})) \
+ ORDER BY k1 ASC NULLS FIRST, k2 ASC NULLS FIRST"
+ ))
+ .await?
+ .create_physical_plan()
+ .await?;
+ let schema = input.schema();
+ assert_eq!(schema.field(1).data_type(), key_type);
+ let batches = collect(input, ctx.task_ctx()).await?;
+ assert_eq!(
+ batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
+ ROWS
+ );
+ // The setup query sorts both keys before materialization. Retain that
+ // known ordering explicitly when registering its results as a table.
+ let table = MemTable::try_new(schema, vec![batches])?.with_sort_order(vec![vec![
+ col("k1").sort(true, true),
+ col("k2").sort(true, true),
+ ]]);
+ ctx.register_table("ordered_input", Arc::new(table))?;
+
+ let plan = ctx.sql(sql).await?.create_physical_plan().await?;
+ let description = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ description.contains("group_clustering_mode=Full"),
+ "{description}"
+ );
+ // Input is already sorted in one partition. Fail if the planner adds
+ // sorting or repartitioning to the measured query.
+ assert!(!description.contains("SortExec"), "{description}");
+ assert!(!description.contains("RepartitionExec"), "{description}");
+ let result = collect(plan, ctx.task_ctx()).await?;
+ assert_eq!(
+ result.iter().map(|batch| batch.num_rows()).sum::<usize>(),
+ ROWS.div_ceil(run_length)
+ );
+ if with_sum {
+ let rows = i64::try_from(ROWS).unwrap();
+ let run_length = i64::try_from(run_length).unwrap();
+ for batch in &result {
+ let k1 = batch.column(0).as_primitive::<Int32Type>();
+ let k2 = batch.column(1).as_primitive::<TimestampNanosecondType>();
+ let sums = batch.column(2).as_primitive::<Int64Type>();
+ assert_eq!(sums.null_count(), 0);
+ for row in 0..batch.num_rows() {
+ // Each input group covers a contiguous range of source row
+ // numbers. Recover it from the keys, independently of output
+ // order, and check the arithmetic-series sum, including the
+ // final short group and groups with a NULL timestamp key.
+ let key_offset = if k2.is_null(row) {
+ 0
+ } else {
+ k2.value(row) + 16
+ };
+ let group_id = i64::from(k1.value(row)) * 32 + key_offset;
+ let start = group_id * run_length;
+ let end = (start + run_length).min(rows);
+ assert_eq!(sums.value(row), (start + end - 1) * (end - start) / 2);
+ }
+ }
+ }
+ Ok(Arc::new(Mutex::new(ctx)))
+}
+
+criterion_group!(benches, criterion_benchmark, fully_ordered);
criterion_main!(benches);
diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs
index 1fcd1fa0bf..d331c51016 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs
@@ -18,7 +18,7 @@
use std::mem;
use arrow::array::ArrayRef;
-use arrow::datatypes::{DataType, Schema, SchemaRef};
+use arrow::datatypes::{DataType, Schema, SchemaRef, TimeUnit};
use arrow_ord::partition::partition;
use datafusion_common::{Result, not_impl_err};
use datafusion_execution::memory_pool::proxy::VecAllocExt;
@@ -62,12 +62,24 @@ impl GroupValuesClustered {
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
+ | DataType::Timestamp(_, _)
+ | DataType::Date32
+ | DataType::Date64
+ | DataType::Time32(TimeUnit::Second | TimeUnit::Millisecond)
+ | DataType::Time64(TimeUnit::Microsecond | TimeUnit::Nanosecond)
+ | DataType::Duration(_)
+ | DataType::Interval(_)
+ | DataType::Decimal32(_, _)
+ | DataType::Decimal64(_, _)
+ | DataType::Decimal128(_, _)
+ | DataType::Decimal256(_, _)
| DataType::Utf8
| DataType::LargeUtf8
| DataType::Utf8View
| DataType::Binary
| DataType::LargeBinary
| DataType::BinaryView
+ | DataType::FixedSizeBinary(0..)
)
})
}
@@ -182,13 +194,157 @@ mod tests {
use std::sync::Arc;
use arrow::array::UInt32Array;
- use arrow::array::{Array, BooleanArray, Int32Array, Int64Array, StringArray};
+ use arrow::array::{
+ Array, AsArray, BooleanArray, Decimal32Array, Decimal64Array, Decimal128Array,
+ Decimal256Array, FixedSizeBinaryArray, Int32Array, Int64Array,
+ IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray,
+ StringArray,
+ };
+ use arrow::buffer::{Buffer, NullBuffer};
use arrow::compute::{cast, take};
- use arrow::datatypes::Field;
+ use arrow::datatypes::{
+ Field, IntervalDayTime, IntervalMonthDayNano, IntervalUnit, i256,
+ };
use datafusion_expr::GroupSelection;
use super::*;
+ // The values repeat in seven-row runs, with four keys per value of the
+ // first column. Construct logical types directly where casting would
+ // rescale decimals or cannot represent all interval components.
+ fn clustered_key_array(
+ key_type: &DataType,
+ rows: usize,
+ nullable: bool,
+ ) -> Result<ArrayRef> {
+ let keys = (0..rows)
+ .map(|row| {
+ let key = ((row / 7) % 4) as i32;
+ (!nullable || key != 0).then_some(key)
+ })
+ .collect::<Vec<_>>();
+ Ok(match key_type {
+ DataType::Boolean => Arc::new(BooleanArray::from_iter(
+ keys.iter().map(|key| key.map(|key| key >= 2)),
+ )),
+ DataType::Utf8
+ | DataType::LargeUtf8
+ | DataType::Utf8View
+ | DataType::Binary
+ | DataType::LargeBinary
+ | DataType::BinaryView => {
+ let strings = StringArray::from_iter(keys.iter().map(|key| {
+ key.map(|key| {
+ format!(
+ "key-{key}-{}",
+ "x".repeat(if key == 3 { 128 } else { 0 })
+ )
+ })
+ }));
+ cast(&strings, key_type)?
+ }
+ DataType::Decimal32(precision, scale) => Arc::new(
+ Decimal32Array::from_iter(keys.iter().map(|key| key.map(|key| key - 2)))
+ .with_precision_and_scale(*precision, *scale)?,
+ ),
+ DataType::Decimal64(precision, scale) => Arc::new(
+ Decimal64Array::from_iter(
+ keys.iter().map(|key| key.map(|key| i64::from(key - 2))),
+ )
+ .with_precision_and_scale(*precision, *scale)?,
+ ),
+ DataType::Decimal128(precision, scale) => {
+ // Equal low 64 bits must not hide differences in the high bits.
+ let values = [
+ -(1_i128 << 100),
+ 0,
+ 1_i128 << 100,
+ (1_i128 << 100) + (1_i128 << 64),
+ ];
+ let array = Decimal128Array::from_iter(
+ keys.iter().map(|key| key.map(|key| values[key as usize])),
+ );
+ array.validate_decimal_precision(*precision)?;
+ Arc::new(array.with_precision_and_scale(*precision, *scale)?)
+ }
+ DataType::Decimal256(precision, scale) => {
+ // These values share their low 128 bits but exceed i128's range.
+ let values = [
+ i256::from_parts(0, -(1_i128 << 100)),
+ i256::from_i128(0),
+ i256::from_parts(0, 1_i128 << 100),
+ i256::from_parts(0, (1_i128 << 100) + 1),
+ ];
+ let array = Decimal256Array::from_iter(
+ keys.iter().map(|key| key.map(|key| values[key as usize])),
+ );
+ array.validate_decimal_precision(*precision)?;
+ Arc::new(array.with_precision_and_scale(*precision, *scale)?)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ let values = [-13, 0, 1, 13];
+ Arc::new(IntervalYearMonthArray::from_iter(
+ keys.iter().map(|key| key.map(|key| values[key as usize])),
+ ))
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ let values = [
+ IntervalDayTime::new(-1, 0),
+ IntervalDayTime::new(0, 0),
+ IntervalDayTime::new(0, 1),
+ IntervalDayTime::new(1, 1),
+ ];
+ Arc::new(IntervalDayTimeArray::from_iter(
+ keys.iter().map(|key| key.map(|key| values[key as usize])),
+ ))
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ let values = [
+ IntervalMonthDayNano::new(0, 0, 0),
+ IntervalMonthDayNano::new(0, 0, 1),
+ IntervalMonthDayNano::new(0, 1, 1),
+ IntervalMonthDayNano::new(1, 1, 1),
+ ];
+ Arc::new(IntervalMonthDayNanoArray::from_iter(
+ keys.iter().map(|key| key.map(|key| values[key as usize])),
+ ))
+ }
+ DataType::FixedSizeBinary(width) => {
+ let bytes = keys
+ .iter()
+ .flat_map(|key| {
+ std::iter::repeat_n(
+ key.unwrap_or_default() as u8,
+ *width as usize,
+ )
+ })
+ .collect::<Vec<_>>();
+ let nulls = nullable.then(|| {
+ NullBuffer::from(keys.iter().map(Option::is_some).collect::<Vec<_>>())
+ });
+ Arc::new(FixedSizeBinaryArray::try_new_with_len(
+ *width,
+ Buffer::from(bytes),
+ nulls,
+ rows,
+ )?)
+ }
+ DataType::Date64 => {
+ let values = Int64Array::from_iter(
+ keys.iter()
+ .map(|key| key.map(|key| i64::from(key) * 86_400_000)),
+ );
+ cast(&values, key_type)?
+ }
+ DataType::Timestamp(_, _) | DataType::Time64(_) | DataType::Duration(_) => {
+ let values =
+ Int64Array::from_iter(keys.iter().map(|key| key.map(i64::from)));
+ cast(&values, key_type)?
+ }
+ _ => cast(&Int32Array::from(keys), key_type)?,
+ })
+ }
+
// Compare with the existing hash implementation, changing actual input
// boundaries and removing completed groups between batches. This checks the
// semantic contract rather than duplicating the adjacent-run algorithm.
@@ -198,7 +354,7 @@ mod tests {
fn clustered_keys_match_hash_grouping_across_batches_and_emits(
#[case] unsorted: bool,
) -> Result<()> {
- for key_type in [
+ let mut key_types = vec![
DataType::Boolean,
DataType::Int8,
DataType::Int16,
@@ -214,114 +370,135 @@ mod tests {
DataType::Binary,
DataType::LargeBinary,
DataType::BinaryView,
+ DataType::Date32,
+ DataType::Date64,
+ DataType::Time32(TimeUnit::Second),
+ DataType::Time32(TimeUnit::Millisecond),
+ DataType::Time64(TimeUnit::Microsecond),
+ DataType::Time64(TimeUnit::Nanosecond),
+ DataType::Interval(IntervalUnit::YearMonth),
+ DataType::Interval(IntervalUnit::DayTime),
+ DataType::Interval(IntervalUnit::MonthDayNano),
+ DataType::Decimal32(9, 2),
+ DataType::Decimal32(9, -2),
+ DataType::Decimal64(18, 2),
+ DataType::Decimal64(18, -2),
+ DataType::Decimal128(38, 2),
+ DataType::Decimal128(38, -2),
+ DataType::Decimal256(76, 2),
+ DataType::Decimal256(76, -2),
+ DataType::FixedSizeBinary(3),
+ DataType::FixedSizeBinary(16),
+ ];
+ for unit in [
+ TimeUnit::Second,
+ TimeUnit::Millisecond,
+ TimeUnit::Microsecond,
+ TimeUnit::Nanosecond,
] {
- let rows = 1025;
- let first: ArrayRef = Arc::new(Int32Array::from_iter((0..rows).map(|row| {
- let group = row / 7;
- let key = group / 4;
- // Swap neighboring key values while keeping each run contiguous.
- (group >= 4).then_some(if unsorted { key ^ 1 } else { key })
- })));
- let second: ArrayRef = match key_type {
- DataType::Boolean => {
- Arc::new(BooleanArray::from_iter((0..rows).map(|row| {
- let key = (row / 7) % 4;
- (key != 0).then_some(key >= 2)
- })))
- }
- DataType::Utf8
- | DataType::LargeUtf8
- | DataType::Utf8View
- | DataType::Binary
- | DataType::LargeBinary
- | DataType::BinaryView => {
- let strings = StringArray::from_iter((0..rows).map(|row| {
- let key = (row / 7) % 4;
- (key != 0).then(|| {
- format!(
- "key-{key}-{}",
- "x".repeat(if key == 3 { 128 } else { 0 })
- )
+ key_types.push(DataType::Duration(unit));
+ for timezone in [None, Some("UTC"), Some("Europe/Paris"), Some("+05:30")] {
+ key_types.push(DataType::Timestamp(unit, timezone.map(Into::into)));
+ }
+ }
+ for key_type in key_types {
+ for nullable in [false, true] {
+ let rows = 1025;
+ let first: ArrayRef =
+ Arc::new(Int32Array::from_iter((0..rows).map(|row| {
+ let group = (row / 7) as i32;
+ let key = group / 4;
+ // Swap neighboring key values while keeping each run contiguous.
+ (!nullable || group >= 4).then_some(if unsorted {
+ key ^ 1
+ } else {
+ key
})
- }));
- cast(&strings, &key_type)?
- }
- _ => {
- let values = Int32Array::from_iter((0..rows).map(|row| {
- let key = (row / 7) % 4;
- (key != 0).then_some(key)
- }));
- cast(&values, &key_type)?
- }
- };
- for descending in [false, true] {
- let mut input = vec![Arc::clone(&first), Arc::clone(&second)];
- if descending {
- let reverse = UInt32Array::from_iter_values((0..rows as u32).rev());
- input = input
- .iter()
- .map(|array| take(array, &reverse, None))
- .collect::<std::result::Result<_, _>>()?;
- }
- let schema = Arc::new(Schema::new(vec![
- Field::new("a", DataType::Int32, true),
- Field::new("b", key_type.clone(), true),
- ]));
- for batch_size in [1, 2, 7, 8, 63, 1024] {
- for emit_limit in [0, 1, 17, usize::MAX] {
- let mut clustered =
- GroupValuesClustered::try_new(Arc::clone(&schema))?;
- let mut hashed =
- GroupValuesColumn::<true>::try_new(Arc::clone(&schema))?;
- let mut actual = Vec::new();
- let mut expected = Vec::new();
- for offset in (0..rows as usize).step_by(batch_size) {
- let length = batch_size.min(rows as usize - offset);
- let batch = input
- .iter()
- .map(|array| array.slice(offset, length))
- .collect::<Vec<_>>();
- clustered.intern(&batch, &mut actual)?;
- hashed.intern(&batch, &mut expected)?;
- assert_eq!(
- actual, expected,
- "{key_type:?}, batch={batch_size}, offset={offset}, descending={descending}"
- );
- assert_eq!(clustered.len(), hashed.len());
- let selection = GroupSelection::all(clustered.len());
+ })));
+ let second = clustered_key_array(&key_type, rows, nullable)?;
+ for descending in [false, true] {
+ let mut input = vec![Arc::clone(&first), Arc::clone(&second)];
+ if descending {
+ let reverse =
+ UInt32Array::from_iter_values((0..rows as u32).rev());
+ input = input
+ .iter()
+ .map(|array| take(array, &reverse, None))
+ .collect::<std::result::Result<_, _>>()?;
+ }
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, nullable),
+ Field::new("b", key_type.clone(), nullable),
+ ]));
+ assert!(GroupValuesClustered::supports_schema(&schema), "{key_type}");
+ for batch_size in [1, 2, 7, 8, 63, 1024] {
+ for emit_limit in [0, 1, 17, usize::MAX] {
+ let mut clustered =
+ GroupValuesClustered::try_new(Arc::clone(&schema))?;
+ let mut hashed =
+ GroupValuesColumn::<true>::try_new(Arc::clone(&schema))?;
+ let mut actual = Vec::new();
+ let mut expected = Vec::new();
+ for offset in (0..rows).step_by(batch_size) {
+ let length = batch_size.min(rows - offset);
+ let batch = input
+ .iter()
+ .map(|array| array.slice(offset, length))
+ .collect::<Vec<_>>();
+ clustered.intern(&batch, &mut actual)?;
+ hashed.intern(&batch, &mut expected)?;
+ assert_eq!(
+ actual, expected,
+ "{key_type:?}, nullable={nullable}, unsorted={unsorted}, batch={batch_size}, offset={offset}, descending={descending}"
+ );
+ assert_eq!(clustered.len(), hashed.len());
+ let selection = GroupSelection::all(clustered.len());
+ let values = clustered.values_preserving(selection)?;
+ assert_eq!(values, hashed.values_preserving(selection)?);
+ assert_eq!(values[1].data_type(), &key_type);
+ // A zero-row batch must neither forget the boundary
+ // key nor leave stale group ids in the output buffer.
+ let empty = batch
+ .iter()
+ .map(|array| array.slice(0, 0))
+ .collect::<Vec<_>>();
+ clustered.intern(&empty, &mut actual)?;
+ assert!(actual.is_empty());
+ let emit =
+ emit_limit.min(clustered.len().saturating_sub(1));
+ let values = clustered.emit(EmitTo::First(emit))?;
+ assert_eq!(values, hashed.emit(EmitTo::First(emit))?);
+ assert_eq!(values[1].data_type(), &key_type);
+ }
+ let indices = [clustered.len() - 1, 0, clustered.len() - 1];
+ let selection = GroupSelection::try_from_indices(
+ &indices,
+ clustered.len(),
+ )?;
assert_eq!(
clustered.values_preserving(selection)?,
hashed.values_preserving(selection)?
);
- // A zero-row batch must neither forget the boundary
- // key nor leave stale group ids in the output buffer.
- let empty = batch
- .iter()
- .map(|array| array.slice(0, 0))
- .collect::<Vec<_>>();
- clustered.intern(&empty, &mut actual)?;
- assert!(actual.is_empty());
- let emit = emit_limit.min(clustered.len().saturating_sub(1));
- assert_eq!(
- clustered.emit(EmitTo::First(emit))?,
- hashed.emit(EmitTo::First(emit))?
- );
+ let values = clustered.emit(EmitTo::All)?;
+ assert_eq!(values, hashed.emit(EmitTo::All)?);
+ assert_eq!(values[1].data_type(), &key_type);
+ assert!(clustered.is_empty());
+ // All resets the clustered boundary state. The hash reference
+ // also needs its lookup table cleared after draining keys.
+ hashed.clear_shrink(0);
+ clustered.intern(&input, &mut actual)?;
+ hashed.intern(&input, &mut expected)?;
+ assert_eq!(actual, expected);
+ // Clearing populated storage must reset it as well.
+ clustered.clear_shrink(0);
+ hashed.clear_shrink(0);
+ clustered.intern(&input, &mut actual)?;
+ hashed.intern(&input, &mut expected)?;
+ assert_eq!(actual, expected);
+ let values = clustered.emit(EmitTo::All)?;
+ assert_eq!(values, hashed.emit(EmitTo::All)?);
+ assert_eq!(values[1].data_type(), &key_type);
}
- assert_eq!(
- clustered.emit(EmitTo::All)?,
- hashed.emit(EmitTo::All)?
- );
- assert!(clustered.is_empty());
- // All and clear_shrink reset the cross-batch state.
- clustered.clear_shrink(0);
- hashed.clear_shrink(0);
- clustered.intern(&input, &mut actual)?;
- hashed.intern(&input, &mut expected)?;
- assert_eq!(actual, expected);
- assert_eq!(
- clustered.emit(EmitTo::All)?,
- hashed.emit(EmitTo::All)?
- );
}
}
}
@@ -329,12 +506,93 @@ mod tests {
Ok(())
}
+ #[test]
+ fn fixed_size_binary_16_common_prefix() -> Result<()> {
+ let first = [0xa5_u8; 16];
+ let mut second = first;
+ second[15] = 0xa6;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::FixedSizeBinary(16), false),
+ ]));
+ let mut clustered = GroupValuesClustered::try_new(Arc::clone(&schema))?;
+ let mut hashed = GroupValuesColumn::<true>::try_new(schema)?;
+ let mut actual = Vec::new();
+ let mut expected = Vec::new();
+ // Continue the first group, then distinguish a key only by its last byte.
+ for (key, group_id) in [(first, 0), (first, 0), (second, 1)] {
+ let batch: [ArrayRef; 2] = [
+ Arc::new(Int32Array::from(vec![7])),
+ Arc::new(FixedSizeBinaryArray::try_from_iter([key].into_iter())?),
+ ];
+ clustered.intern(&batch, &mut actual)?;
+ hashed.intern(&batch, &mut expected)?;
+ assert_eq!(actual, [group_id]);
+ assert_eq!(actual, expected);
+ }
+ let values = clustered.emit(EmitTo::All)?;
+ assert_eq!(values, hashed.emit(EmitTo::All)?);
+ assert_eq!(
+ values[1].as_fixed_size_binary().value_data(),
+ [first, second].concat()
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn zero_width_binary_preserves_emitted_lengths_and_nulls() -> Result<()> {
+ for nullable in [false, true] {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::FixedSizeBinary(0), nullable),
+ ]));
+ let first: ArrayRef = Arc::new(Int32Array::from(vec![7; 6]));
+ let second: ArrayRef = Arc::new(FixedSizeBinaryArray::try_new_with_len(
+ 0,
+ Buffer::default(),
+ nullable.then(|| {
+ NullBuffer::from(vec![true, false, false, true, true, true])
+ }),
+ 6,
+ )?);
+ let input = [first.slice(1, 4), second.slice(1, 4)];
+ let mut clustered = GroupValuesClustered::try_new(schema)?;
+ let mut groups = Vec::new();
+ for offset in [0, 2] {
+ let batch = input
+ .iter()
+ .map(|array| array.slice(offset, 2))
+ .collect::<Vec<_>>();
+ clustered.intern(&batch, &mut groups)?;
+ assert_eq!(groups, vec![usize::from(nullable && offset == 2); 2]);
+ }
+ let values =
+ clustered.values_preserving(GroupSelection::all(clustered.len()))?;
+ assert_eq!(values[1].len(), if nullable { 2 } else { 1 });
+ assert_eq!(values[1].null_count(), usize::from(nullable));
+ let values = clustered.emit(EmitTo::First(1))?;
+ assert_eq!(values[1].len(), 1);
+ assert_eq!(values[1].null_count(), usize::from(nullable));
+ let values = clustered.emit(EmitTo::All)?;
+ assert_eq!(values[1].len(), usize::from(nullable));
+ assert_eq!(values[1].null_count(), 0);
+ assert_eq!(values[1].data_type(), &DataType::FixedSizeBinary(0));
+ }
+ Ok(())
+ }
+
#[test]
fn clustered_schema_gate_keeps_unvalidated_types_on_existing_paths() {
for data_type in [
+ DataType::Float16,
DataType::Float32,
DataType::Float64,
DataType::Null,
+ DataType::Time32(TimeUnit::Microsecond),
+ DataType::Time32(TimeUnit::Nanosecond),
+ DataType::Time64(TimeUnit::Second),
+ DataType::Time64(TimeUnit::Millisecond),
+ DataType::FixedSizeBinary(-1),
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
] {
@@ -343,6 +601,7 @@ mod tests {
Field::new("b", data_type, true),
]);
assert!(!GroupValuesClustered::supports_schema(&schema));
+ assert!(GroupValuesClustered::try_new(Arc::new(schema)).is_err());
}
assert!(!GroupValuesClustered::supports_schema(&Schema::new(vec![
Field::new("a", DataType::Int32, false)
diff --git a/datafusion/sqllogictest/test_files/ordered_aggregate_key_types.slt b/datafusion/sqllogictest/test_files/ordered_aggregate_key_types.slt
new file mode 100644
index 0000000000..cd47847420
--- /dev/null
+++ b/datafusion/sqllogictest/test_files/ordered_aggregate_key_types.slt
@@ -0,0 +1,172 @@
+# 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.
+
+# Fully ordered multi-column aggregation agrees with unordered aggregation for
+# parameterized key types. Two-row batches split both null and non-null groups.
+statement ok
+SET datafusion.execution.batch_size = 2;
+
+statement ok
+SET datafusion.execution.target_partitions = 1;
+
+statement ok
+SET datafusion.execution.enable_migration_aggregate = true;
+
+statement ok
+SET datafusion.optimizer.enable_round_robin_repartition = false;
+
+statement ok
+SET datafusion.optimizer.prefer_existing_sort = true;
+
+statement ok
+SET datafusion.explain.physical_plan_only = true;
+
+statement ok
+CREATE TABLE ordered_key_types AS
+SELECT
+ arrow_cast(t, 'Timestamp(Nanosecond, Some("Europe/Paris"))') AS ts_ns,
+ arrow_cast(t, 'Timestamp(Millisecond)') AS ts_ms,
+ arrow_cast(n, 'Decimal32(9, 2)') AS d32,
+ arrow_cast(n, 'Decimal64(18, 4)') AS d64,
+ arrow_cast(n, 'Decimal128(20, 3)') AS d128,
+ arrow_cast(n, 'Decimal256(50, 7)') AS d256,
+ arrow_cast(b, 'FixedSizeBinary(5)') AS b,
+ arrow_cast(CAST(t AS INT), 'Date32') AS date32,
+ arrow_cast(t * 86400000, 'Date64') AS date64,
+ arrow_cast(CAST(t AS INT), 'Time32(Millisecond)') AS time32,
+ arrow_cast(t, 'Time64(Microsecond)') AS time64,
+ arrow_cast(t, 'Duration(Nanosecond)') AS duration,
+ arrow_cast(CAST(t AS VARCHAR) || ' months', 'Interval(YearMonth)') AS iym,
+ arrow_cast(CAST(t AS VARCHAR) || ' days', 'Interval(DayTime)') AS idt,
+ arrow_cast(CAST(t AS VARCHAR) || ' nanoseconds', 'Interval(MonthDayNano)') AS imdn,
+ v
+FROM (VALUES
+ (0, -1.25, X'00000000ff', 10),
+ (0, -1.25, X'00000000ff', 0),
+ (0, -1.25, X'00000000ff', 20),
+ (0, -1.25, X'00000000ff', NULL),
+ (0, -1.25, X'0000000001', 4),
+ (1, NULL, NULL, 0),
+ (1, NULL, NULL, 8),
+ (NULL, 2.5, X'ff00000000', 12),
+ (NULL, 2.5, X'ff00000000', 4),
+ (NULL, NULL, NULL, 0)
+) AS input(t, n, b, v);
+
+# LIMIT preserves the subquery's ordering. It exceeds the fixture's row count.
+# Mixed directions and explicit null ordering must not change group equality.
+statement ok
+CREATE VIEW ordered_key_types_sorted AS
+SELECT * FROM ordered_key_types
+ORDER BY ts_ns DESC NULLS FIRST, ts_ms, d32 DESC NULLS LAST,
+ d64, d128, d256, b DESC NULLS FIRST, date32, date64,
+ time32, time64, duration, iym, idt, imdn
+LIMIT 1000;
+
+query TT
+EXPLAIN SELECT ts_ns, ts_ms, d32, d64, d128, d256, b,
+ date32, date64, time32, time64, duration, iym, idt, imdn, SUM(v)
+FROM ordered_key_types_sorted
+GROUP BY ts_ns, ts_ms, d32, d64, d128, d256, b,
+ date32, date64, time32, time64, duration, iym, idt, imdn;
+----
+physical_plan
+01)AggregateExec: mode=Single,<slt:ignore>group_clustering_mode=Full
+<slt:ignore>
+
+# Unordered reference. FILTER must also prevent division by zero.
+query IR?IIII rowsort ordered-key-types
+SELECT CAST(ts_ns AS BIGINT), d128, b, MIN(v), SUM(10 / v) FILTER (WHERE v <> 0),
+ COUNT(v) FILTER (WHERE v <> 0), COUNT(*)
+FROM ordered_key_types
+GROUP BY ts_ns, ts_ms, d32, d64, d128, d256, b,
+ date32, date64, time32, time64, duration, iym, idt, imdn;
+----
+0 -1.250 0000000001 4 2 1 1
+0 -1.250 00000000ff 0 1 2 4
+1 NULL NULL 0 1 1 2
+NULL 2.500 ff00000000 4 2 2 2
+NULL NULL NULL 0 NULL 0 1
+
+# The same keys and aggregates through fully ordered Single aggregation.
+query IR?IIII rowsort ordered-key-types
+SELECT CAST(ts_ns AS BIGINT), d128, b, MIN(v), SUM(10 / v) FILTER (WHERE v <> 0),
+ COUNT(v) FILTER (WHERE v <> 0), COUNT(*)
+FROM ordered_key_types_sorted
+GROUP BY ts_ns, ts_ms, d32, d64, d128, d256, b,
+ date32, date64, time32, time64, duration, iym, idt, imdn;
+----
+0 -1.250 0000000001 4 2 1 1
+0 -1.250 00000000ff 0 1 2 4
+1 NULL NULL 0 1 1 2
+NULL 2.500 ff00000000 4 2 2 2
+NULL NULL NULL 0 NULL 0 1
+
+statement ok
+SET datafusion.execution.target_partitions = 2;
+
+# Partial and final stages must both retain complete group-key ordering.
+query TT
+EXPLAIN SELECT ts_ns, ts_ms, d32, d64, d128, d256, b,
+ date32, date64, time32, time64, duration, iym, idt, imdn, SUM(v)
+FROM ordered_key_types_sorted
+GROUP BY ts_ns, ts_ms, d32, d64, d128, d256, b,
+ date32, date64, time32, time64, duration, iym, idt, imdn;
+----
+physical_plan
+01)AggregateExec: mode=FinalPartitioned,<slt:ignore>group_clustering_mode=Full
+02)--RepartitionExec:<slt:ignore>maintains_sort_order=true<slt:ignore>
+03)----AggregateExec: mode=Partial,<slt:ignore>group_clustering_mode=Full
+<slt:ignore>
+
+query IR?IIII rowsort ordered-key-types
+SELECT CAST(ts_ns AS BIGINT), d128, b, MIN(v), SUM(10 / v) FILTER (WHERE v <> 0),
+ COUNT(v) FILTER (WHERE v <> 0), COUNT(*)
+FROM ordered_key_types_sorted
+GROUP BY ts_ns, ts_ms, d32, d64, d128, d256, b,
+ date32, date64, time32, time64, duration, iym, idt, imdn;
+----
+0 -1.250 0000000001 4 2 1 1
+0 -1.250 00000000ff 0 1 2 4
+1 NULL NULL 0 1 1 2
+NULL 2.500 ff00000000 4 2 2 2
+NULL NULL NULL 0 NULL 0 1
+
+statement ok
+DROP VIEW ordered_key_types_sorted;
+
+statement ok
+DROP TABLE ordered_key_types;
+
+# Restore the sqllogictest runner's settings.
+statement ok
+RESET datafusion.execution.batch_size;
+
+statement ok
+SET datafusion.execution.target_partitions = 4;
+
+statement ok
+RESET datafusion.execution.enable_migration_aggregate;
+
+statement ok
+RESET datafusion.optimizer.enable_round_robin_repartition;
+
+statement ok
+RESET datafusion.optimizer.prefer_existing_sort;
+
+statement ok
+RESET datafusion.explain.physical_plan_only;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org