This is an automated email from the ASF dual-hosted git repository.
alamb 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 99cdbccbe1 feat: Add random timezone to aggregate fuzz test (#13349)
99cdbccbe1 is described below
commit 99cdbccbe1f1ca51d1b7c1c856dfbdaafd3614c1
Author: Jonathan Chen <[email protected]>
AuthorDate: Wed Nov 13 08:59:08 2024 -0500
feat: Add random timezone to aggregate fuzz test (#13349)
* add random tz
* fixes
* add comment
---
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs | 2 ++
test-utils/Cargo.toml | 1 +
test-utils/src/array_gen/primitive.rs | 32 ++++++++++++++++++++--
3 files changed, 32 insertions(+), 3 deletions(-)
diff --git a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs
b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs
index 29e1d7bc22..09d0c8d5ca 100644
--- a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs
+++ b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs
@@ -172,6 +172,8 @@ fn baseline_config() -> DatasetGeneratorConfig {
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",
diff --git a/test-utils/Cargo.toml b/test-utils/Cargo.toml
index 414fa5569c..d2763f507f 100644
--- a/test-utils/Cargo.toml
+++ b/test-utils/Cargo.toml
@@ -27,6 +27,7 @@ workspace = true
[dependencies]
arrow = { workspace = true }
+chrono-tz = { version = "0.10.0", default-features = false }
datafusion-common = { workspace = true, default-features = true }
env_logger = { workspace = true }
paste = "1.0.15"
diff --git a/test-utils/src/array_gen/primitive.rs
b/test-utils/src/array_gen/primitive.rs
index 500a68143f..bfc151b43d 100644
--- a/test-utils/src/array_gen/primitive.rs
+++ b/test-utils/src/array_gen/primitive.rs
@@ -17,8 +17,9 @@
use arrow::array::{ArrayRef, ArrowPrimitiveType, PrimitiveArray, UInt32Array};
use arrow::datatypes::DataType;
-use rand::rngs::StdRng;
-use rand::Rng;
+use chrono_tz::{Tz, TZ_VARIANTS};
+use rand::{rngs::StdRng, seq::SliceRandom, thread_rng, Rng};
+use std::sync::Arc;
use super::random_data::RandomNativeData;
@@ -40,8 +41,16 @@ impl PrimitiveArrayGenerator {
where
A: ArrowPrimitiveType + RandomNativeData,
{
+ let data_type = match A::DATA_TYPE {
+ DataType::Timestamp(unit, _) => {
+ let timezone = Self::generate_timezone();
+ DataType::Timestamp(unit, timezone)
+ }
+ other => other,
+ };
+
// table of primitives from which to draw
- let distinct_primitives: PrimitiveArray<A> = match A::DATA_TYPE {
+ let distinct_primitives: PrimitiveArray<A> = match data_type {
DataType::Int8
| DataType::Int16
| DataType::Int32
@@ -86,4 +95,21 @@ impl PrimitiveArrayGenerator {
let options = None;
arrow::compute::take(&distinct_primitives, &indicies, options).unwrap()
}
+
+ // Generates a random timezone or returns `None`.
+ ///
+ /// Returns:
+ /// - `Some(Arc<String>)` containing the timezone name.
+ /// - `None` if no timezone is selected.
+ fn generate_timezone() -> Option<Arc<str>> {
+ let mut rng = thread_rng();
+
+ // Allows for timezones + None
+ let mut timezone_options: Vec<Option<&Tz>> = vec![None];
+ timezone_options.extend(TZ_VARIANTS.iter().map(Some));
+
+ let selected_option = timezone_options.choose(&mut
rng).cloned().flatten(); // random timezone/None
+
+ selected_option.map(|tz| Arc::from(tz.name()))
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
