This is an automated email from the ASF dual-hosted git repository.
jayzhan 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 143ef97fc4 feat: Add GroupColumn `Decimal128Array` (#13564)
143ef97fc4 is described below
commit 143ef97fc46a4e10868b3cde86cdeaf629a753a5
Author: Jonathan Chen <[email protected]>
AuthorDate: Tue Dec 3 20:02:12 2024 -0500
feat: Add GroupColumn `Decimal128Array` (#13564)
* feat: Add GroupColumn `Decimal128Array
* fix clippy
* fix errors
* remove .with_data_type
* fix
* fmt
* remove adjust_output_array
* fix
* Add missing data type
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
.../src/aggregates/group_values/mod.rs | 5 ++-
.../aggregates/group_values/multi_group_by/mod.rs | 13 ++++++--
.../group_values/multi_group_by/primitive.rs | 8 ++---
.../group_values/single_group_by/primitive.rs | 1 +
datafusion/sqllogictest/test_files/group_by.slt | 39 ++++++++++++++++++++++
5 files changed, 59 insertions(+), 7 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index 58bc7bb90a..e4a7eb049e 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -19,7 +19,7 @@
use arrow::record_batch::RecordBatch;
use arrow_array::types::{
- Date32Type, Date64Type, Time32MillisecondType, Time32SecondType,
+ Date32Type, Date64Type, Decimal128Type, Time32MillisecondType,
Time32SecondType,
Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType,
TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
};
@@ -170,6 +170,9 @@ pub(crate) fn new_group_values(
TimeUnit::Microsecond =>
downcast_helper!(TimestampMicrosecondType, d),
TimeUnit::Nanosecond =>
downcast_helper!(TimestampNanosecondType, d),
},
+ DataType::Decimal128(_, _) => {
+ downcast_helper!(Decimal128Type, d);
+ }
DataType::Utf8 => {
return
Ok(Box::new(GroupValuesByes::<i32>::new(OutputType::Utf8)));
}
diff --git
a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
index aac186a6f6..034fb86d06 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
@@ -31,8 +31,8 @@ use crate::aggregates::group_values::GroupValues;
use ahash::RandomState;
use arrow::compute::cast;
use arrow::datatypes::{
- BinaryViewType, Date32Type, Date64Type, Float32Type, Float64Type,
Int16Type,
- Int32Type, Int64Type, Int8Type, StringViewType, Time32MillisecondType,
+ BinaryViewType, Date32Type, Date64Type, Decimal128Type, Float32Type,
Float64Type,
+ Int16Type, Int32Type, Int64Type, Int8Type, StringViewType,
Time32MillisecondType,
Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
TimestampMicrosecondType, TimestampMillisecondType,
TimestampNanosecondType,
TimestampSecondType, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
@@ -1008,6 +1008,14 @@ impl<const STREAMING: bool> GroupValues for
GroupValuesColumn<STREAMING> {
)
}
},
+ &DataType::Decimal128(_, _) => {
+ instantiate_primitive! {
+ v,
+ nullable,
+ Decimal128Type,
+ data_type
+ }
+ }
&DataType::Utf8 => {
let b =
ByteGroupValueBuilder::<i32>::new(OutputType::Utf8);
v.push(Box::new(b) as _)
@@ -1214,6 +1222,7 @@ fn supported_type(data_type: &DataType) -> bool {
| DataType::UInt64
| DataType::Float32
| DataType::Float64
+ | DataType::Decimal128(_, _)
| DataType::Utf8
| DataType::LargeUtf8
| DataType::Binary
diff --git
a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs
b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs
index 4686a78f24..4ceeb634ba 100644
---
a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs
+++
b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs
@@ -200,10 +200,10 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool>
GroupColumn
let first_n_nulls = if NULLABLE { self.nulls.take_n(n) } else { None };
- Arc::new(PrimitiveArray::<T>::new(
- ScalarBuffer::from(first_n),
- first_n_nulls,
- ))
+ Arc::new(
+ PrimitiveArray::<T>::new(ScalarBuffer::from(first_n),
first_n_nulls)
+ .with_data_type(self.data_type.clone()),
+ )
}
}
diff --git
a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
index 05214ec10d..85cd2e79b9 100644
---
a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
+++
b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
@@ -208,6 +208,7 @@ where
build_primitive(split, null_group)
}
};
+
Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))])
}
diff --git a/datafusion/sqllogictest/test_files/group_by.slt
b/datafusion/sqllogictest/test_files/group_by.slt
index 4acf519c5d..df7e21c2da 100644
--- a/datafusion/sqllogictest/test_files/group_by.slt
+++ b/datafusion/sqllogictest/test_files/group_by.slt
@@ -5499,3 +5499,42 @@ SELECT
GROUP BY ts, text
----
foo 2024-01-01T08:00:00+08:00
+
+# Test multi group by int + Decimal128
+statement ok
+create table source as values
+(1, '123.45'),
+(1, '123.45'),
+(2, '678.90'),
+(2, '1011.12'),
+(3, '1314.15'),
+(3, '1314.15'),
+(2, '1011.12'),
+(null, null),
+(null, '123.45'),
+(null, null),
+(null, '123.45'),
+(2, '678.90'),
+(2, '678.90'),
+(1, null)
+;
+
+statement ok
+create view t as select column1 as a, arrow_cast(column2, 'Decimal128(10, 2)')
as b from source;
+
+query IRI
+select a, b, count(*) from t group by a, b order by a, b;
+----
+1 123.45 2
+1 NULL 1
+2 678.9 3
+2 1011.12 2
+3 1314.15 2
+NULL 123.45 2
+NULL NULL 2
+
+statement ok
+drop view t
+
+statement ok
+drop table source;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]