Skip to content

Commit fdbe9c4

Browse files
committed
support Date32/Date64 type in data generator.
1 parent 35525d5 commit fdbe9c4

File tree

3 files changed

+36
-2
lines changed

3 files changed

+36
-2
lines changed

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

+2
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,8 @@ fn baseline_config() -> DatasetGeneratorConfig {
164164
ColumnDescr::new("u16", DataType::UInt16),
165165
ColumnDescr::new("u32", DataType::UInt32),
166166
ColumnDescr::new("u64", DataType::UInt64),
167+
ColumnDescr::new("date32", DataType::Date32),
168+
ColumnDescr::new("date64", DataType::Date64),
167169
// TODO: date/time columns
168170
// todo decimal columns
169171
// begin string columns

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

+21-1
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
use std::sync::Arc;
1919

20-
use arrow::datatypes::{Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type};
20+
use arrow::datatypes::{Date32Type, Date64Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type};
2121
use arrow_array::{ArrayRef, RecordBatch};
2222
use arrow_schema::{DataType, Field, Schema};
2323
use datafusion_common::{arrow_datafusion_err, DataFusionError, Result};
@@ -392,6 +392,26 @@ impl RecordBatchGenerator {
392392
Float64Type
393393
)
394394
}
395+
DataType::Date32 => {
396+
generate_primitive_array!(
397+
self,
398+
num_rows,
399+
batch_gen_rng,
400+
array_gen_rng,
401+
i32,
402+
Date32Type
403+
)
404+
}
405+
DataType::Date64 => {
406+
generate_primitive_array!(
407+
self,
408+
num_rows,
409+
batch_gen_rng,
410+
array_gen_rng,
411+
i64,
412+
Date64Type
413+
)
414+
}
395415
DataType::Utf8 => {
396416
generate_string_array!(self, num_rows, batch_gen_rng, array_gen_rng, i32)
397417
}

test-utils/src/array_gen/primitive.rs

+13-1
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,19 @@ impl PrimitiveArrayGenerator {
5454
| DataType::UInt32
5555
| DataType::UInt64
5656
| DataType::Float32
57-
| DataType::Float64 => self.rng.gen::<N>(),
57+
| DataType::Float64
58+
| DataType::Date32 => self.rng.gen::<N>(),
59+
60+
DataType::Date64 => {
61+
// TODO: constrain this range to valid dates if necessary
62+
let date_value = self.rng.gen_range(i64::MIN..=i64::MAX);
63+
let millis_per_day: i64 = 86_400_000;
64+
let adjusted_value = date_value - (date_value % millis_per_day);
65+
// SAFETY: here we can convert i64 to N safely since we determine that the type N is i64
66+
unsafe {
67+
std::ptr::read(&adjusted_value as *const i64 as *const N)
68+
}
69+
}
5870

5971
_ => {
6072
let arrow_type = A::DATA_TYPE;

0 commit comments

Comments
 (0)