Skip to content

Commit

Permalink
Maintain time_zone in new_list
Browse files Browse the repository at this point in the history
  • Loading branch information
Dandandan committed Oct 21, 2023
1 parent 5bdc9af commit 8f81403
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 9 deletions.
38 changes: 31 additions & 7 deletions datafusion/common/src/scalar.rs
Original file line number Diff line number Diff line change
Expand Up @@ -620,26 +620,30 @@ macro_rules! build_timestamp_list {
TimestampSecondBuilder,
TimestampSecond,
values,
$SIZE
$SIZE,
$TIME_ZONE
)
}
TimeUnit::Millisecond => build_values_list_tz!(
TimestampMillisecondBuilder,
TimestampMillisecond,
values,
$SIZE
$SIZE,
$TIME_ZONE
),
TimeUnit::Microsecond => build_values_list_tz!(
TimestampMicrosecondBuilder,
TimestampMicrosecond,
values,
$SIZE
$SIZE,
$TIME_ZONE
),
TimeUnit::Nanosecond => build_values_list_tz!(
TimestampNanosecondBuilder,
TimestampNanosecond,
values,
$SIZE
$SIZE,
$TIME_ZONE
),
},
}
Expand Down Expand Up @@ -683,9 +687,10 @@ macro_rules! build_values_list {
}

macro_rules! build_values_list_tz {
($VALUE_BUILDER_TY:ident, $SCALAR_TY:ident, $VALUES:expr, $SIZE:expr) => {{
let mut builder =
ListBuilder::new($VALUE_BUILDER_TY::with_capacity($VALUES.len()));
($VALUE_BUILDER_TY:ident, $SCALAR_TY:ident, $VALUES:expr, $SIZE:expr, $TIME_ZONE:expr) => {{
let mut builder = ListBuilder::new(
$VALUE_BUILDER_TY::with_capacity($VALUES.len()).with_timezone_opt($TIME_ZONE),
);

for _ in 0..$SIZE {
for scalar_value in $VALUES {
Expand Down Expand Up @@ -5185,6 +5190,25 @@ mod tests {
assert_eq!(1, arr.len());
}

#[test]
fn test_newlist_timestamp_zone() {
let s: &'static str = "UTC";
let values = vec![ScalarValue::TimestampMillisecond(Some(1), Some(s.into()))];
let arr = ScalarValue::new_list(
&values,
&DataType::Timestamp(TimeUnit::Millisecond, Some(s.into())),
);
assert_eq!(1, arr.len());
assert_eq!(
arr.data_type(),
&DataType::List(Arc::new(Field::new(
"item",
DataType::Timestamp(TimeUnit::Millisecond, Some(s.into())),
true
)))
);
}

fn get_random_timestamps(sample_size: u64) -> Vec<ScalarValue> {
let vector_size = sample_size;
let mut timestamp = vec![];
Expand Down

0 comments on commit 8f81403

Please sign in to comment.