feat: Add Time/Interval/Decimal/Utf8View in aggregate fuzz testing

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -23,6 +23,10 @@ use arrow::datatypes::DataType;
 use arrow::record_batch::RecordBatch;
 use arrow::util::pretty::pretty_format_batches;
 use arrow_array::types::Int64Type;
+use arrow_schema::{
+    IntervalUnit, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE,
+    DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE,
+};
 use datafusion::common::Result;
 use datafusion::datasource::MemTable;
 use datafusion::physical_expr::aggregate::AggregateExprBuilder;
@@ -45,7 +49,7 @@ use crate::fuzz_cases::aggregation_fuzzer::{
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use hashbrown::HashMap;
 use rand::rngs::StdRng;
-use rand::{Rng, SeedableRng};
+use rand::{thread_rng, Rng, SeedableRng};
 use tokio::task::JoinSet;
 
 // ========================================================================
@@ -151,6 +155,7 @@ async fn test_count() {
 /// 1. Floating point numbers
 /// 1. structured types
 fn baseline_config() -> DatasetGeneratorConfig {
+    let mut rng = thread_rng();
     let columns = vec![
         ColumnDescr::new("i8", DataType::Int8),
         ColumnDescr::new("i16", DataType::Int16),
@@ -162,8 +167,41 @@ fn baseline_config() -> DatasetGeneratorConfig {
         ColumnDescr::new("u64", DataType::UInt64),
         ColumnDescr::new("date32", DataType::Date32),
         ColumnDescr::new("date64", DataType::Date64),
-        // TODO: date/time columns
-        // todo decimal columns
+        ColumnDescr::new("time32_s", DataType::Time32(TimeUnit::Second)),
+        ColumnDescr::new("time32_ms", DataType::Time32(TimeUnit::Millisecond)),
+        ColumnDescr::new("time64_us", DataType::Time64(TimeUnit::Microsecond)),
+        ColumnDescr::new("time64_ns", DataType::Time64(TimeUnit::Nanosecond)),
+        ColumnDescr::new(
+            "interval_year_month",
+            DataType::Interval(IntervalUnit::YearMonth),
+        ),
+        ColumnDescr::new(
+            "interval_day_time",
+            DataType::Interval(IntervalUnit::DayTime),
+        ),
+        ColumnDescr::new(
+            "interval_month_day_nano",
+            DataType::Interval(IntervalUnit::MonthDayNano),
+        ),
+        // begin decimal columns
+        ColumnDescr::new("decimal128", {
+            // Generate valid precision and scale for Decimal128 randomly.
+            let precision: u8 = rng.gen_range(0..=DECIMAL128_MAX_PRECISION);
+            // It's safe to cast `precision` to i8 type directly.
+            let scale: i8 = rng.gen_range(
+                i8::MIN..=std::cmp::min(precision as i8, DECIMAL128_MAX_SCALE),
+            );
+            DataType::Decimal128(precision, scale)
+        }),
+        ColumnDescr::new("decimal256", {
+            // Generate valid precision and scale for Decimal256 randomly.
+            let precision: u8 = rng.gen_range(0..=DECIMAL256_MAX_PRECISION);
+            // It's safe to cast `precision` to i8 type directly.
+            let scale: i8 = rng.gen_range(
+                i8::MIN..=std::cmp::min(precision as i8, DECIMAL256_MAX_SCALE),
+            );
+            DataType::Decimal256(precision, scale)
+        }),
         // begin string columns
         ColumnDescr::new("utf8", DataType::Utf8),
         ColumnDescr::new("largeutf8", DataType::LargeUtf8),

datafusion/core/tests/fuzz_cases/aggregation_fuzzer/data_generator.rs

@@ -19,10 +19,12 @@ use std::sync::Arc;
 
 use arrow::datatypes::{
     Date32Type, Date64Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type,
-    Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
+    Int8Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType,
+    Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
+    UInt16Type, UInt32Type, UInt64Type, UInt8Type,
 };
 use arrow_array::{ArrayRef, RecordBatch};
-use arrow_schema::{DataType, Field, Schema};
+use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
 use datafusion_common::{arrow_datafusion_err, DataFusionError, Result};
 use datafusion_physical_expr::{expressions::col, PhysicalSortExpr};
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
@@ -32,7 +34,7 @@ use rand::{
     thread_rng, Rng, SeedableRng,
 };
 use test_utils::{
-    array_gen::{PrimitiveArrayGenerator, StringArrayGenerator},
+    array_gen::{DecimalArrayGenerator, PrimitiveArrayGenerator, StringArrayGenerator},
     stagger_batch,
 };
 
@@ -236,21 +238,40 @@ macro_rules! generate_string_array {
     }};
 }
 
+macro_rules! generate_decimal_array {
+    ($SELF:ident, $NUM_ROWS:ident, $MAX_NUM_DISTINCT: expr, $BATCH_GEN_RNG:ident, $ARRAY_GEN_RNG:ident, $PRECISION: ident, $SCALE: ident, $NATIVE_TYPE: ident) => {{
+        let null_pct_idx = $BATCH_GEN_RNG.gen_range(0..$SELF.candidate_null_pcts.len());
+        let null_pct = $SELF.candidate_null_pcts[null_pct_idx];
+
+        let mut generator = DecimalArrayGenerator {
+            precision: $PRECISION,
+            scale: $SCALE,
+            num_decimals: $NUM_ROWS,
+            num_distinct_decimals: $MAX_NUM_DISTINCT,
+            null_pct,
+            rng: $ARRAY_GEN_RNG,
+        };
+
+        paste::paste! {
+            generator.[< gen_data_ $NATIVE_TYPE >]()
+        }
+    }};
+}
+
 macro_rules! generate_primitive_array {
-    ($SELF:ident, $NUM_ROWS:ident, $MAX_NUM_DISTINCT:expr, $BATCH_GEN_RNG:ident, $ARRAY_GEN_RNG:ident, $ARROW_TYPE:ident) => {
-        paste::paste! {{
-            let null_pct_idx = $BATCH_GEN_RNG.gen_range(0..$SELF.candidate_null_pcts.len());
-            let null_pct = $SELF.candidate_null_pcts[null_pct_idx];
-
-            let mut generator = PrimitiveArrayGenerator {
-                num_primitives: $NUM_ROWS,
-                num_distinct_primitives: $MAX_NUM_DISTINCT,
-                null_pct,
-                rng: $ARRAY_GEN_RNG,
-            };
-
-            generator.gen_data::<$ARROW_TYPE>()
-    }}}
+    ($SELF:ident, $NUM_ROWS:ident, $MAX_NUM_DISTINCT:expr, $BATCH_GEN_RNG:ident, $ARRAY_GEN_RNG:ident, $ARROW_TYPE:ident) => {{
+        let null_pct_idx = $BATCH_GEN_RNG.gen_range(0..$SELF.candidate_null_pcts.len());
+        let null_pct = $SELF.candidate_null_pcts[null_pct_idx];
+
+        let mut generator = PrimitiveArrayGenerator {
+            num_primitives: $NUM_ROWS,
+            num_distinct_primitives: $MAX_NUM_DISTINCT,
+            null_pct,
+            rng: $ARRAY_GEN_RNG,
+        };
+
+        generator.gen_data::<$ARROW_TYPE>()
+    }};
 }
 
 impl RecordBatchGenerator {
@@ -432,6 +453,100 @@ impl RecordBatchGenerator {
                     Date64Type
                 )
             }
+            DataType::Time32(TimeUnit::Second) => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    Time32SecondType
+                )
+            }
+            DataType::Time32(TimeUnit::Millisecond) => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    Time32MillisecondType
+                )
+            }
+            DataType::Time64(TimeUnit::Microsecond) => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    Time64MicrosecondType
+                )
+            }
+            DataType::Time64(TimeUnit::Nanosecond) => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    Time64NanosecondType
+                )
+            }
+            DataType::Interval(IntervalUnit::YearMonth) => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    IntervalYearMonthType
+                )
+            }
+            DataType::Interval(IntervalUnit::DayTime) => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    IntervalDayTimeType
+                )
+            }
+            DataType::Interval(IntervalUnit::MonthDayNano) => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    IntervalMonthDayNanoType
+                )
+            }
+            DataType::Decimal128(precision, scale) => {
+                generate_decimal_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    precision,
+                    scale,
+                    i128
+                )
+            }
+            DataType::Decimal256(precision, scale) => {
+                generate_decimal_array!(
+                    self,
+                    num_rows,
+                    max_num_distinct,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    precision,
+                    scale,
+                    i256
+                )
+            }
             DataType::Utf8 => {
                 generate_string_array!(
                     self,

test-utils/src/array_gen/decimal.rs

@@ -0,0 +1,92 @@
+// 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.
+
+use arrow::array::{ArrayRef, PrimitiveArray, PrimitiveBuilder, UInt32Array};
+use arrow::datatypes::{i256, Decimal128Type, Decimal256Type};
+use rand::rngs::StdRng;
+use rand::Rng;
+
+/// Randomly generate decimal arrays
+pub struct DecimalArrayGenerator {
+    /// The precision of the decimal type
+    pub precision: u8,
+    /// The scale of the decimal type
+    pub scale: i8,
+    /// The total number of decimals in the output
+    pub num_decimals: usize,
+    /// The number of distinct decimals in the columns
+    pub num_distinct_decimals: usize,
+    /// The percentage of nulls in the columns
+    pub null_pct: f64,
+    /// Random number generator
+    pub rng: StdRng,
+}
+
+macro_rules! impl_gen_decimal_data {
+    ($NATIVE_TYPE:ident, $DECIMAL_TYPE:ident) => {
+        paste::paste! {
+            /// Create a Decimal128Array / Decimal256Array with random values.
+            pub fn [< gen_data_ $NATIVE_TYPE >](&mut self) -> ArrayRef {
+                // table of decimals from which to draw
+                let distinct_decimals : PrimitiveArray<$DECIMAL_TYPE> = {
+                    let mut decimal_builder = PrimitiveBuilder::<$DECIMAL_TYPE>::with_capacity(self.num_distinct_decimals);
+                    for _ in 0..self.num_distinct_decimals {
+                        decimal_builder.append_option(Some(
+                            [< random_ $NATIVE_TYPE >](&mut self.rng)
+                        ));
+                    }
+
+                    let decimal_array = decimal_builder
+                        .finish()
+                        .with_precision_and_scale(self.precision, self.scale)
+                        .unwrap();
+                    decimal_array
+                };
+
+                // pick num_decimals randomly from the distinct decimal table
+                let indicies: UInt32Array = (0..self.num_decimals)
+                    .map(|_| {
+                        if self.rng.gen::<f64>() < self.null_pct {
+                            None
+                        } else if self.num_distinct_decimals > 1 {
+                            let range = 1..(self.num_distinct_decimals as u32);
+                            Some(self.rng.gen_range(range))
+                        } else {
+                            Some(0)
+                        }
+                    })
+                    .collect();
+
+                let options = None;
+                arrow::compute::take(&distinct_decimals, &indicies, options).unwrap()
+            }
+        }
+    };
+}
+
+impl DecimalArrayGenerator {
+    impl_gen_decimal_data!(i128, Decimal128Type);
+    impl_gen_decimal_data!(i256, Decimal256Type);
+}
+
+fn random_i128(rng: &mut StdRng) -> i128 {
+    rng.gen::<i128>()
+}
+
+fn random_i256(rng: &mut StdRng) -> i256 {
+    i256::from_parts(rng.gen::<u128>(), rng.gen::<i128>())
+}

test-utils/src/array_gen/mod.rs

@@ -15,8 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+mod decimal;
 mod primitive;
 mod string;
 
+pub use decimal::DecimalArrayGenerator;
 pub use primitive::PrimitiveArrayGenerator;
 pub use string::StringArrayGenerator;