kosiew commented on code in PR #26370:
URL: https://github.com/apache/datafusion/pull/26370#discussion_r4313252118
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs:
##########
@@ -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, 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(
Review Comment:
Could we add a FixedSizeBinary(16) test where keys share the first 15 bytes
but differ in the last byte, with one group spanning batches? Comparing group
IDs and emitted bytes against hash grouping would strengthen UUID coverage,
since the current fixture uses repeated-byte keys.
##########
datafusion/core/benches/aggregate_query_sql.rs:
##########
@@ -330,5 +334,185 @@ 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))
+ .unwrap();
+ group.bench_function(name, |b| {
+ b.iter(|| query(ctx.clone(), &rt, FULLY_ORDERED_QUERY))
+ });
+ }
+ group.finish();
+}
+
+const FULLY_ORDERED_QUERY: &str = "SELECT k1, k2 FROM ordered_input GROUP BY
k1, k2";
Review Comment:
Could we also benchmark a timestamp grouping query with a value column and
SUM(v), alongside the existing key-only benchmark? This would help measure
accumulator work while keeping setup outside timing and verifying that the plan
still uses fully clustered aggregation without extra sorting or repartitioning.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
