geoffreyclaude commented on code in PR #26370:
URL: https://github.com/apache/datafusion/pull/26370#discussion_r4315163732
##########
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:
Added in a398f8768cad6b2df8db19552f179c2a5aecac47
--
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]
