Implement sort fuzzer

Author: 2010YOUY01

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -18,15 +18,12 @@
 use std::sync::Arc;
 
 use crate::fuzz_cases::aggregation_fuzzer::{
-    AggregationFuzzerBuilder, ColumnDescr, DatasetGeneratorConfig, QueryBuilder,
+    AggregationFuzzerBuilder, DatasetGeneratorConfig, QueryBuilder,
 };
 
 use arrow::array::{types::Int64Type, Array, ArrayRef, AsArray, Int64Array, RecordBatch};
 use arrow::compute::{concat_batches, SortOptions};
-use arrow::datatypes::{
-    DataType, IntervalUnit, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE,
-    DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE,
-};
+use arrow::datatypes::DataType;
 use arrow::util::pretty::pretty_format_batches;
 use datafusion::common::Result;
 use datafusion::datasource::memory::MemorySourceConfig;
@@ -51,6 +48,8 @@ use test_utils::{add_empty_batches, StringBatchGenerator};
 use rand::rngs::StdRng;
 use rand::{thread_rng, Rng, SeedableRng};
 
+use super::record_batch_generator::get_supported_types_columns;
+
 // ========================================================================
 //  The new aggregation fuzz tests based on [`AggregationFuzzer`]
 // ========================================================================
@@ -201,81 +200,7 @@ async fn test_median() {
 /// 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),
-        ColumnDescr::new("i32", DataType::Int32),
-        ColumnDescr::new("i64", DataType::Int64),
-        ColumnDescr::new("u8", DataType::UInt8),
-        ColumnDescr::new("u16", DataType::UInt16),
-        ColumnDescr::new("u32", DataType::UInt32),
-        ColumnDescr::new("u64", DataType::UInt64),
-        ColumnDescr::new("date32", DataType::Date32),
-        ColumnDescr::new("date64", DataType::Date64),
-        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)),
-        // `None` is passed in here however when generating the array, it will generate
-        // random timezones.
-        ColumnDescr::new("timestamp_s", DataType::Timestamp(TimeUnit::Second, None)),
-        ColumnDescr::new(
-            "timestamp_ms",
-            DataType::Timestamp(TimeUnit::Millisecond, None),
-        ),
-        ColumnDescr::new(
-            "timestamp_us",
-            DataType::Timestamp(TimeUnit::Microsecond, None),
-        ),
-        ColumnDescr::new(
-            "timestamp_ns",
-            DataType::Timestamp(TimeUnit::Nanosecond, None),
-        ),
-        ColumnDescr::new("float32", DataType::Float32),
-        ColumnDescr::new("float64", DataType::Float64),
-        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(1..=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(1..=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),
-        ColumnDescr::new("utf8view", DataType::Utf8View),
-        // low cardinality columns
-        ColumnDescr::new("u8_low", DataType::UInt8).with_max_num_distinct(10),
-        ColumnDescr::new("utf8_low", DataType::Utf8).with_max_num_distinct(10),
-        ColumnDescr::new("bool", DataType::Boolean),
-        ColumnDescr::new("binary", DataType::Binary),
-        ColumnDescr::new("large_binary", DataType::LargeBinary),
-        ColumnDescr::new("binaryview", DataType::BinaryView),
-    ];
+    let columns = get_supported_types_columns(rng.gen());
 
     let min_num_rows = 512;
     let max_num_rows = 1024;