support Date32/Date64 type in data generator.

apache · Oct 21, 2024 · fdbe9c4 · fdbe9c4
1 parent 35525d5
commit fdbe9c4
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 2 deletions.
diff --git a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs
@@ -164,6 +164,8 @@ fn baseline_config() -> DatasetGeneratorConfig {
         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),
         // TODO: date/time columns
         // todo decimal columns
         // begin string columns

diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/data_generator.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/data_generator.rs
@@ -17,7 +17,7 @@
 
 use std::sync::Arc;
 
-use arrow::datatypes::{Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type};
+use arrow::datatypes::{Date32Type, Date64Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type};
 use arrow_array::{ArrayRef, RecordBatch};
 use arrow_schema::{DataType, Field, Schema};
 use datafusion_common::{arrow_datafusion_err, DataFusionError, Result};
@@ -392,6 +392,26 @@ impl RecordBatchGenerator {
                     Float64Type
                 )
             }
+            DataType::Date32 => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    i32,
+                    Date32Type
+                )
+            }
+            DataType::Date64 => {
+                generate_primitive_array!(
+                    self,
+                    num_rows,
+                    batch_gen_rng,
+                    array_gen_rng,
+                    i64,
+                    Date64Type
+                )
+            }
             DataType::Utf8 => {
                 generate_string_array!(self, num_rows, batch_gen_rng, array_gen_rng, i32)
             }

diff --git a/test-utils/src/array_gen/primitive.rs b/test-utils/src/array_gen/primitive.rs
@@ -54,7 +54,19 @@ impl PrimitiveArrayGenerator {
                 | DataType::UInt32
                 | DataType::UInt64
                 | DataType::Float32
-                | DataType::Float64 => self.rng.gen::<N>(),
+                | DataType::Float64 
+                | DataType::Date32 => self.rng.gen::<N>(),
+
+                DataType::Date64 => {
+                    // TODO: constrain this range to valid dates if necessary
+                    let date_value = self.rng.gen_range(i64::MIN..=i64::MAX); 
+                    let millis_per_day: i64 = 86_400_000;
+                    let adjusted_value = date_value - (date_value % millis_per_day);
+                    // SAFETY: here we can convert i64 to N safely since we determine that the type N is i64
+                    unsafe {
+                        std::ptr::read(&adjusted_value as *const i64 as *const N)
+                    }
+                }
 
                 _ => {
                     let arrow_type = A::DATA_TYPE;