| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, |
| // software distributed under the License is distributed on an |
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| // KIND, either express or implied. See the License for the |
| // specific language governing permissions and limitations |
| // under the License. |
| |
| //! temporal kernels |
| |
| use chrono::{DateTime, Datelike, Duration, NaiveDateTime, Timelike, Utc}; |
| |
| use std::sync::Arc; |
| |
| use arrow::{array::*, datatypes::DataType}; |
| use arrow_array::{ |
| downcast_dictionary_array, downcast_temporal_array, |
| temporal_conversions::*, |
| timezone::Tz, |
| types::{ArrowDictionaryKeyType, ArrowTemporalType, Date32Type, TimestampMicrosecondType}, |
| ArrowNumericType, |
| }; |
| |
| use arrow_schema::TimeUnit; |
| |
| use crate::SparkError; |
| |
| // Copied from arrow_arith/temporal.rs |
| macro_rules! return_compute_error_with { |
| ($msg:expr, $param:expr) => { |
| return { Err(SparkError::Internal(format!("{}: {:?}", $msg, $param))) } |
| }; |
| } |
| |
| // The number of days between the beginning of the proleptic gregorian calendar (0001-01-01) |
| // and the beginning of the Unix Epoch (1970-01-01) |
| const DAYS_TO_UNIX_EPOCH: i32 = 719_163; |
| |
| // Copied from arrow_arith/temporal.rs with modification to the output datatype |
| // Transforms a array of NaiveDate to an array of Date32 after applying an operation |
| fn as_datetime_with_op<A: ArrayAccessor<Item = T::Native>, T: ArrowTemporalType, F>( |
| iter: ArrayIter<A>, |
| mut builder: PrimitiveBuilder<Date32Type>, |
| op: F, |
| ) -> Date32Array |
| where |
| F: Fn(NaiveDateTime) -> i32, |
| i64: From<T::Native>, |
| { |
| iter.into_iter().for_each(|value| { |
| if let Some(value) = value { |
| match as_datetime::<T>(i64::from(value)) { |
| Some(dt) => builder.append_value(op(dt)), |
| None => builder.append_null(), |
| } |
| } else { |
| builder.append_null(); |
| } |
| }); |
| |
| builder.finish() |
| } |
| |
| #[inline] |
| fn as_datetime_with_op_single<F>( |
| value: Option<i32>, |
| builder: &mut PrimitiveBuilder<Date32Type>, |
| op: F, |
| ) where |
| F: Fn(NaiveDateTime) -> i32, |
| { |
| if let Some(value) = value { |
| match as_datetime::<Date32Type>(i64::from(value)) { |
| Some(dt) => builder.append_value(op(dt)), |
| None => builder.append_null(), |
| } |
| } else { |
| builder.append_null(); |
| } |
| } |
| |
| // Based on arrow_arith/temporal.rs:extract_component_from_datetime_array |
| // Transforms an array of DateTime<Tz> to an arrayOf TimeStampMicrosecond after applying an |
| // operation |
| fn as_timestamp_tz_with_op<A: ArrayAccessor<Item = T::Native>, T: ArrowTemporalType, F>( |
| iter: ArrayIter<A>, |
| mut builder: PrimitiveBuilder<TimestampMicrosecondType>, |
| tz: &str, |
| op: F, |
| ) -> Result<TimestampMicrosecondArray, SparkError> |
| where |
| F: Fn(DateTime<Tz>) -> i64, |
| i64: From<T::Native>, |
| { |
| let tz: Tz = tz.parse()?; |
| for value in iter { |
| match value { |
| Some(value) => match as_datetime_with_timezone::<T>(value.into(), tz) { |
| Some(time) => builder.append_value(op(time)), |
| _ => { |
| return Err(SparkError::Internal( |
| "Unable to read value as datetime".to_string(), |
| )); |
| } |
| }, |
| None => builder.append_null(), |
| } |
| } |
| Ok(builder.finish()) |
| } |
| |
| fn as_timestamp_tz_with_op_single<T: ArrowTemporalType, F>( |
| value: Option<T::Native>, |
| builder: &mut PrimitiveBuilder<TimestampMicrosecondType>, |
| tz: &Tz, |
| op: F, |
| ) -> Result<(), SparkError> |
| where |
| F: Fn(DateTime<Tz>) -> i64, |
| i64: From<T::Native>, |
| { |
| match value { |
| Some(value) => match as_datetime_with_timezone::<T>(value.into(), *tz) { |
| Some(time) => builder.append_value(op(time)), |
| _ => { |
| return Err(SparkError::Internal( |
| "Unable to read value as datetime".to_string(), |
| )); |
| } |
| }, |
| None => builder.append_null(), |
| } |
| Ok(()) |
| } |
| |
| #[inline] |
| fn as_days_from_unix_epoch(dt: Option<NaiveDateTime>) -> i32 { |
| dt.unwrap().num_days_from_ce() - DAYS_TO_UNIX_EPOCH |
| } |
| |
| // Apply the Tz to the Naive Date Time,,convert to UTC, and return as microseconds in Unix epoch |
| #[inline] |
| fn as_micros_from_unix_epoch_utc(dt: Option<DateTime<Tz>>) -> i64 { |
| dt.unwrap().with_timezone(&Utc).timestamp_micros() |
| } |
| |
| #[inline] |
| fn trunc_date_to_year<T: Datelike + Timelike>(dt: T) -> Option<T> { |
| Some(dt) |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| .and_then(|d| d.with_day0(0)) |
| .and_then(|d| d.with_month0(0)) |
| } |
| |
| /// returns the month of the beginning of the quarter |
| #[inline] |
| fn quarter_month<T: Datelike>(dt: &T) -> u32 { |
| 1 + 3 * ((dt.month() - 1) / 3) |
| } |
| |
| #[inline] |
| fn trunc_date_to_quarter<T: Datelike + Timelike>(dt: T) -> Option<T> { |
| Some(dt) |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| .and_then(|d| d.with_day0(0)) |
| .and_then(|d| d.with_month(quarter_month(&d))) |
| } |
| |
| #[inline] |
| fn trunc_date_to_month<T: Datelike + Timelike>(dt: T) -> Option<T> { |
| Some(dt) |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| .and_then(|d| d.with_day0(0)) |
| } |
| |
| #[inline] |
| fn trunc_date_to_week<T>(dt: T) -> Option<T> |
| where |
| T: Datelike + Timelike + std::ops::Sub<Duration, Output = T> + Copy, |
| { |
| Some(dt) |
| .map(|d| d - Duration::try_seconds(60 * 60 * 24 * d.weekday() as i64).unwrap()) |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| } |
| |
| #[inline] |
| fn trunc_date_to_day<T: Timelike>(dt: T) -> Option<T> { |
| Some(dt) |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| } |
| |
| #[inline] |
| fn trunc_date_to_hour<T: Timelike>(dt: T) -> Option<T> { |
| Some(dt) |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| } |
| |
| #[inline] |
| fn trunc_date_to_minute<T: Timelike>(dt: T) -> Option<T> { |
| Some(dt) |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| } |
| |
| #[inline] |
| fn trunc_date_to_second<T: Timelike>(dt: T) -> Option<T> { |
| Some(dt).and_then(|d| d.with_nanosecond(0)) |
| } |
| |
| #[inline] |
| fn trunc_date_to_ms<T: Timelike>(dt: T) -> Option<T> { |
| Some(dt).and_then(|d| d.with_nanosecond(1_000_000 * (d.nanosecond() / 1_000_000))) |
| } |
| |
| #[inline] |
| fn trunc_date_to_microsec<T: Timelike>(dt: T) -> Option<T> { |
| Some(dt).and_then(|d| d.with_nanosecond(1_000 * (d.nanosecond() / 1_000))) |
| } |
| |
| /// |
| /// Implements the spark [TRUNC](https://spark.apache.org/docs/latest/api/sql/index.html#trunc) |
| /// function where the specified format is a scalar value |
| /// |
| /// array is an array of Date32 values. The array may be a dictionary array. |
| /// |
| /// format is a scalar string specifying the format to apply to the timestamp value. |
| pub(crate) fn date_trunc_dyn(array: &dyn Array, format: String) -> Result<ArrayRef, SparkError> { |
| match array.data_type().clone() { |
| DataType::Dictionary(_, _) => { |
| downcast_dictionary_array!( |
| array => { |
| let truncated_values = date_trunc_dyn(array.values(), format)?; |
| Ok(Arc::new(array.with_values(truncated_values))) |
| } |
| dt => return_compute_error_with!("date_trunc does not support", dt), |
| ) |
| } |
| _ => { |
| downcast_temporal_array!( |
| array => { |
| date_trunc(array, format) |
| .map(|a| Arc::new(a) as ArrayRef) |
| } |
| dt => return_compute_error_with!("date_trunc does not support", dt), |
| ) |
| } |
| } |
| } |
| |
| pub(crate) fn date_trunc<T>( |
| array: &PrimitiveArray<T>, |
| format: String, |
| ) -> Result<Date32Array, SparkError> |
| where |
| T: ArrowTemporalType + ArrowNumericType, |
| i64: From<T::Native>, |
| { |
| let builder = Date32Builder::with_capacity(array.len()); |
| let iter = ArrayIter::new(array); |
| match array.data_type() { |
| DataType::Date32 => match format.to_uppercase().as_str() { |
| "YEAR" | "YYYY" | "YY" => Ok(as_datetime_with_op::<&PrimitiveArray<T>, T, _>( |
| iter, |
| builder, |
| |dt| as_days_from_unix_epoch(trunc_date_to_year(dt)), |
| )), |
| "QUARTER" => Ok(as_datetime_with_op::<&PrimitiveArray<T>, T, _>( |
| iter, |
| builder, |
| |dt| as_days_from_unix_epoch(trunc_date_to_quarter(dt)), |
| )), |
| "MONTH" | "MON" | "MM" => Ok(as_datetime_with_op::<&PrimitiveArray<T>, T, _>( |
| iter, |
| builder, |
| |dt| as_days_from_unix_epoch(trunc_date_to_month(dt)), |
| )), |
| "WEEK" => Ok(as_datetime_with_op::<&PrimitiveArray<T>, T, _>( |
| iter, |
| builder, |
| |dt| as_days_from_unix_epoch(trunc_date_to_week(dt)), |
| )), |
| _ => Err(SparkError::Internal(format!( |
| "Unsupported format: {:?} for function 'date_trunc'", |
| format |
| ))), |
| }, |
| dt => return_compute_error_with!( |
| "Unsupported input type '{:?}' for function 'date_trunc'", |
| dt |
| ), |
| } |
| } |
| |
| /// |
| /// Implements the spark [TRUNC](https://spark.apache.org/docs/latest/api/sql/index.html#trunc) |
| /// function where the specified format may be an array |
| /// |
| /// array is an array of Date32 values. The array may be a dictionary array. |
| /// |
| /// format is an array of strings specifying the format to apply to the corresponding date value. |
| /// The array may be a dictionary array. |
| pub(crate) fn date_trunc_array_fmt_dyn( |
| array: &dyn Array, |
| formats: &dyn Array, |
| ) -> Result<ArrayRef, SparkError> { |
| match (array.data_type().clone(), formats.data_type().clone()) { |
| (DataType::Dictionary(_, v), DataType::Dictionary(_, f)) => { |
| if !matches!(*v, DataType::Date32) { |
| return_compute_error_with!("date_trunc does not support", v) |
| } |
| if !matches!(*f, DataType::Utf8) { |
| return_compute_error_with!("date_trunc does not support format type ", f) |
| } |
| downcast_dictionary_array!( |
| formats => { |
| downcast_dictionary_array!( |
| array => { |
| date_trunc_array_fmt_dict_dict( |
| &array.downcast_dict::<Date32Array>().unwrap(), |
| &formats.downcast_dict::<StringArray>().unwrap()) |
| .map(|a| Arc::new(a) as ArrayRef) |
| } |
| dt => return_compute_error_with!("date_trunc does not support", dt) |
| ) |
| } |
| fmt => return_compute_error_with!("date_trunc does not support format type", fmt), |
| ) |
| } |
| (DataType::Dictionary(_, v), DataType::Utf8) => { |
| if !matches!(*v, DataType::Date32) { |
| return_compute_error_with!("date_trunc does not support", v) |
| } |
| downcast_dictionary_array!( |
| array => { |
| date_trunc_array_fmt_dict_plain( |
| &array.downcast_dict::<Date32Array>().unwrap(), |
| formats.as_any().downcast_ref::<StringArray>() |
| .expect("Unexpected value type in formats")) |
| .map(|a| Arc::new(a) as ArrayRef) |
| } |
| dt => return_compute_error_with!("date_trunc does not support", dt), |
| ) |
| } |
| (DataType::Date32, DataType::Dictionary(_, f)) => { |
| if !matches!(*f, DataType::Utf8) { |
| return_compute_error_with!("date_trunc does not support format type ", f) |
| } |
| downcast_dictionary_array!( |
| formats => { |
| downcast_temporal_array!(array => { |
| date_trunc_array_fmt_plain_dict( |
| array.as_any().downcast_ref::<Date32Array>() |
| .expect("Unexpected error in casting date array"), |
| &formats.downcast_dict::<StringArray>().unwrap()) |
| .map(|a| Arc::new(a) as ArrayRef) |
| } |
| dt => return_compute_error_with!("date_trunc does not support", dt), |
| ) |
| } |
| fmt => return_compute_error_with!("date_trunc does not support format type", fmt), |
| ) |
| } |
| (DataType::Date32, DataType::Utf8) => date_trunc_array_fmt_plain_plain( |
| array |
| .as_any() |
| .downcast_ref::<Date32Array>() |
| .expect("Unexpected error in casting date array"), |
| formats |
| .as_any() |
| .downcast_ref::<StringArray>() |
| .expect("Unexpected value type in formats"), |
| ) |
| .map(|a| Arc::new(a) as ArrayRef), |
| (dt, fmt) => Err(SparkError::Internal(format!( |
| "Unsupported datatype: {:}, format: {:?} for function 'date_trunc'", |
| dt, fmt |
| ))), |
| } |
| } |
| |
| macro_rules! date_trunc_array_fmt_helper { |
| ($array: ident, $formats: ident, $datatype: ident) => {{ |
| let mut builder = Date32Builder::with_capacity($array.len()); |
| let iter = $array.into_iter(); |
| match $datatype { |
| DataType::Date32 => { |
| for (index, val) in iter.enumerate() { |
| let op_result = match $formats.value(index).to_uppercase().as_str() { |
| "YEAR" | "YYYY" | "YY" => { |
| Ok(as_datetime_with_op_single(val, &mut builder, |dt| { |
| as_days_from_unix_epoch(trunc_date_to_year(dt)) |
| })) |
| } |
| "QUARTER" => Ok(as_datetime_with_op_single(val, &mut builder, |dt| { |
| as_days_from_unix_epoch(trunc_date_to_quarter(dt)) |
| })), |
| "MONTH" | "MON" | "MM" => { |
| Ok(as_datetime_with_op_single(val, &mut builder, |dt| { |
| as_days_from_unix_epoch(trunc_date_to_month(dt)) |
| })) |
| } |
| "WEEK" => Ok(as_datetime_with_op_single(val, &mut builder, |dt| { |
| as_days_from_unix_epoch(trunc_date_to_week(dt)) |
| })), |
| _ => Err(SparkError::Internal(format!( |
| "Unsupported format: {:?} for function 'date_trunc'", |
| $formats.value(index) |
| ))), |
| }; |
| op_result? |
| } |
| Ok(builder.finish()) |
| } |
| dt => return_compute_error_with!( |
| "Unsupported input type '{:?}' for function 'date_trunc'", |
| dt |
| ), |
| } |
| }}; |
| } |
| |
| fn date_trunc_array_fmt_plain_plain( |
| array: &Date32Array, |
| formats: &StringArray, |
| ) -> Result<Date32Array, SparkError> |
| where |
| { |
| let data_type = array.data_type(); |
| date_trunc_array_fmt_helper!(array, formats, data_type) |
| } |
| |
| fn date_trunc_array_fmt_plain_dict<K>( |
| array: &Date32Array, |
| formats: &TypedDictionaryArray<K, StringArray>, |
| ) -> Result<Date32Array, SparkError> |
| where |
| K: ArrowDictionaryKeyType, |
| { |
| let data_type = array.data_type(); |
| date_trunc_array_fmt_helper!(array, formats, data_type) |
| } |
| |
| fn date_trunc_array_fmt_dict_plain<K>( |
| array: &TypedDictionaryArray<K, Date32Array>, |
| formats: &StringArray, |
| ) -> Result<Date32Array, SparkError> |
| where |
| K: ArrowDictionaryKeyType, |
| { |
| let data_type = array.values().data_type(); |
| date_trunc_array_fmt_helper!(array, formats, data_type) |
| } |
| |
| fn date_trunc_array_fmt_dict_dict<K, F>( |
| array: &TypedDictionaryArray<K, Date32Array>, |
| formats: &TypedDictionaryArray<F, StringArray>, |
| ) -> Result<Date32Array, SparkError> |
| where |
| K: ArrowDictionaryKeyType, |
| F: ArrowDictionaryKeyType, |
| { |
| let data_type = array.values().data_type(); |
| date_trunc_array_fmt_helper!(array, formats, data_type) |
| } |
| |
| /// |
| /// Implements the spark [DATE_TRUNC](https://spark.apache.org/docs/latest/api/sql/index.html#date_trunc) |
| /// function where the specified format is a scalar value |
| /// |
| /// array is an array of Timestamp(Microsecond) values. Timestamp values must have a valid |
| /// timezone or no timezone. The array may be a dictionary array. |
| /// |
| /// format is a scalar string specifying the format to apply to the timestamp value. |
| pub(crate) fn timestamp_trunc_dyn( |
| array: &dyn Array, |
| format: String, |
| ) -> Result<ArrayRef, SparkError> { |
| match array.data_type().clone() { |
| DataType::Dictionary(_, _) => { |
| downcast_dictionary_array!( |
| array => { |
| let truncated_values = timestamp_trunc_dyn(array.values(), format)?; |
| Ok(Arc::new(array.with_values(truncated_values))) |
| } |
| dt => return_compute_error_with!("timestamp_trunc does not support", dt), |
| ) |
| } |
| _ => { |
| downcast_temporal_array!( |
| array => { |
| timestamp_trunc(array, format) |
| .map(|a| Arc::new(a) as ArrayRef) |
| } |
| dt => return_compute_error_with!("timestamp_trunc does not support", dt), |
| ) |
| } |
| } |
| } |
| |
| pub(crate) fn timestamp_trunc<T>( |
| array: &PrimitiveArray<T>, |
| format: String, |
| ) -> Result<TimestampMicrosecondArray, SparkError> |
| where |
| T: ArrowTemporalType + ArrowNumericType, |
| i64: From<T::Native>, |
| { |
| let builder = TimestampMicrosecondBuilder::with_capacity(array.len()); |
| let iter = ArrayIter::new(array); |
| match array.data_type() { |
| DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => { |
| match format.to_uppercase().as_str() { |
| "YEAR" | "YYYY" | "YY" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_year(dt)) |
| }) |
| } |
| "QUARTER" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_quarter(dt)) |
| }) |
| } |
| "MONTH" | "MON" | "MM" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_month(dt)) |
| }) |
| } |
| "WEEK" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_week(dt)) |
| }) |
| } |
| "DAY" | "DD" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_day(dt)) |
| }) |
| } |
| "HOUR" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_hour(dt)) |
| }) |
| } |
| "MINUTE" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_minute(dt)) |
| }) |
| } |
| "SECOND" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_second(dt)) |
| }) |
| } |
| "MILLISECOND" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_ms(dt)) |
| }) |
| } |
| "MICROSECOND" => { |
| as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter, builder, tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_microsec(dt)) |
| }) |
| } |
| _ => Err(SparkError::Internal(format!( |
| "Unsupported format: {:?} for function 'timestamp_trunc'", |
| format |
| ))), |
| } |
| } |
| dt => return_compute_error_with!( |
| "Unsupported input type '{:?}' for function 'timestamp_trunc'", |
| dt |
| ), |
| } |
| } |
| |
| /// |
| /// Implements the spark [DATE_TRUNC](https://spark.apache.org/docs/latest/api/sql/index.html#date_trunc) |
| /// function where the specified format may be an array |
| /// |
| /// array is an array of Timestamp(Microsecond) values. Timestamp values must have a valid |
| /// timezone or no timezone. The array may be a dictionary array. |
| /// |
| /// format is an array of strings specifying the format to apply to the corresponding timestamp |
| /// value. The array may be a dictionary array. |
| pub(crate) fn timestamp_trunc_array_fmt_dyn( |
| array: &dyn Array, |
| formats: &dyn Array, |
| ) -> Result<ArrayRef, SparkError> { |
| match (array.data_type().clone(), formats.data_type().clone()) { |
| (DataType::Dictionary(_, _), DataType::Dictionary(_, _)) => { |
| downcast_dictionary_array!( |
| formats => { |
| downcast_dictionary_array!( |
| array => { |
| timestamp_trunc_array_fmt_dict_dict( |
| &array.downcast_dict::<TimestampMicrosecondArray>().unwrap(), |
| &formats.downcast_dict::<StringArray>().unwrap()) |
| .map(|a| Arc::new(a) as ArrayRef) |
| } |
| dt => return_compute_error_with!("timestamp_trunc does not support", dt) |
| ) |
| } |
| fmt => return_compute_error_with!("timestamp_trunc does not support format type", fmt), |
| ) |
| } |
| (DataType::Dictionary(_, _), DataType::Utf8) => { |
| downcast_dictionary_array!( |
| array => { |
| timestamp_trunc_array_fmt_dict_plain( |
| &array.downcast_dict::<PrimitiveArray<TimestampMicrosecondType>>().unwrap(), |
| formats.as_any().downcast_ref::<StringArray>() |
| .expect("Unexpected value type in formats")) |
| .map(|a| Arc::new(a) as ArrayRef) |
| } |
| dt => return_compute_error_with!("timestamp_trunc does not support", dt), |
| ) |
| } |
| (DataType::Timestamp(TimeUnit::Microsecond, _), DataType::Dictionary(_, _)) => { |
| downcast_dictionary_array!( |
| formats => { |
| downcast_temporal_array!(array => { |
| timestamp_trunc_array_fmt_plain_dict( |
| array, |
| &formats.downcast_dict::<StringArray>().unwrap()) |
| .map(|a| Arc::new(a) as ArrayRef) |
| } |
| dt => return_compute_error_with!("timestamp_trunc does not support", dt), |
| ) |
| } |
| fmt => return_compute_error_with!("timestamp_trunc does not support format type", fmt), |
| ) |
| } |
| (DataType::Timestamp(TimeUnit::Microsecond, _), DataType::Utf8) => { |
| downcast_temporal_array!( |
| array => { |
| timestamp_trunc_array_fmt_plain_plain(array, |
| formats.as_any().downcast_ref::<StringArray>().expect("Unexpected value type in formats")) |
| .map(|a| Arc::new(a) as ArrayRef) |
| }, |
| dt => return_compute_error_with!("timestamp_trunc does not support", dt), |
| ) |
| } |
| (dt, fmt) => Err(SparkError::Internal(format!( |
| "Unsupported datatype: {:}, format: {:?} for function 'timestamp_trunc'", |
| dt, fmt |
| ))), |
| } |
| } |
| |
| macro_rules! timestamp_trunc_array_fmt_helper { |
| ($array: ident, $formats: ident, $datatype: ident) => {{ |
| let mut builder = TimestampMicrosecondBuilder::with_capacity($array.len()); |
| let iter = $array.into_iter(); |
| assert_eq!( |
| $array.len(), |
| $formats.len(), |
| "lengths of values array and format array must be the same" |
| ); |
| match $datatype { |
| DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => { |
| let tz: Tz = tz.parse()?; |
| for (index, val) in iter.enumerate() { |
| let op_result = match $formats.value(index).to_uppercase().as_str() { |
| "YEAR" | "YYYY" | "YY" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_year(dt)) |
| }) |
| } |
| "QUARTER" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_quarter(dt)) |
| }) |
| } |
| "MONTH" | "MON" | "MM" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_month(dt)) |
| }) |
| } |
| "WEEK" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_week(dt)) |
| }) |
| } |
| "DAY" | "DD" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_day(dt)) |
| }) |
| } |
| "HOUR" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_hour(dt)) |
| }) |
| } |
| "MINUTE" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_minute(dt)) |
| }) |
| } |
| "SECOND" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_second(dt)) |
| }) |
| } |
| "MILLISECOND" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_ms(dt)) |
| }) |
| } |
| "MICROSECOND" => { |
| as_timestamp_tz_with_op_single::<T, _>(val, &mut builder, &tz, |dt| { |
| as_micros_from_unix_epoch_utc(trunc_date_to_microsec(dt)) |
| }) |
| } |
| _ => Err(SparkError::Internal(format!( |
| "Unsupported format: {:?} for function 'timestamp_trunc'", |
| $formats.value(index) |
| ))), |
| }; |
| op_result? |
| } |
| Ok(builder.finish()) |
| } |
| dt => { |
| return_compute_error_with!( |
| "Unsupported input type '{:?}' for function 'timestamp_trunc'", |
| dt |
| ) |
| } |
| } |
| }}; |
| } |
| |
| fn timestamp_trunc_array_fmt_plain_plain<T>( |
| array: &PrimitiveArray<T>, |
| formats: &StringArray, |
| ) -> Result<TimestampMicrosecondArray, SparkError> |
| where |
| T: ArrowTemporalType + ArrowNumericType, |
| i64: From<T::Native>, |
| { |
| let data_type = array.data_type(); |
| timestamp_trunc_array_fmt_helper!(array, formats, data_type) |
| } |
| fn timestamp_trunc_array_fmt_plain_dict<T, K>( |
| array: &PrimitiveArray<T>, |
| formats: &TypedDictionaryArray<K, StringArray>, |
| ) -> Result<TimestampMicrosecondArray, SparkError> |
| where |
| T: ArrowTemporalType + ArrowNumericType, |
| i64: From<T::Native>, |
| K: ArrowDictionaryKeyType, |
| { |
| let data_type = array.data_type(); |
| timestamp_trunc_array_fmt_helper!(array, formats, data_type) |
| } |
| |
| fn timestamp_trunc_array_fmt_dict_plain<T, K>( |
| array: &TypedDictionaryArray<K, PrimitiveArray<T>>, |
| formats: &StringArray, |
| ) -> Result<TimestampMicrosecondArray, SparkError> |
| where |
| T: ArrowTemporalType + ArrowNumericType, |
| i64: From<T::Native>, |
| K: ArrowDictionaryKeyType, |
| { |
| let data_type = array.values().data_type(); |
| timestamp_trunc_array_fmt_helper!(array, formats, data_type) |
| } |
| |
| fn timestamp_trunc_array_fmt_dict_dict<T, K, F>( |
| array: &TypedDictionaryArray<K, PrimitiveArray<T>>, |
| formats: &TypedDictionaryArray<F, StringArray>, |
| ) -> Result<TimestampMicrosecondArray, SparkError> |
| where |
| T: ArrowTemporalType + ArrowNumericType, |
| i64: From<T::Native>, |
| K: ArrowDictionaryKeyType, |
| F: ArrowDictionaryKeyType, |
| { |
| let data_type = array.values().data_type(); |
| timestamp_trunc_array_fmt_helper!(array, formats, data_type) |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use crate::kernels::temporal::{ |
| date_trunc, date_trunc_array_fmt_dyn, timestamp_trunc, timestamp_trunc_array_fmt_dyn, |
| }; |
| use arrow_array::{ |
| builder::{PrimitiveDictionaryBuilder, StringDictionaryBuilder}, |
| iterator::ArrayIter, |
| types::{Date32Type, Int32Type, TimestampMicrosecondType}, |
| Array, Date32Array, PrimitiveArray, StringArray, TimestampMicrosecondArray, |
| }; |
| use std::sync::Arc; |
| |
| #[test] |
| #[cfg_attr(miri, ignore)] // test takes too long with miri |
| fn test_date_trunc() { |
| let size = 1000; |
| let mut vec: Vec<i32> = Vec::with_capacity(size); |
| for i in 0..size { |
| vec.push(i as i32); |
| } |
| let array = Date32Array::from(vec); |
| for fmt in [ |
| "YEAR", "YYYY", "YY", "QUARTER", "MONTH", "MON", "MM", "WEEK", |
| ] { |
| match date_trunc(&array, fmt.to_string()) { |
| Ok(a) => { |
| for i in 0..size { |
| assert!(array.values().get(i) >= a.values().get(i)) |
| } |
| } |
| _ => unreachable!(), |
| } |
| } |
| } |
| |
| #[test] |
| // This test only verifies that the various input array types work. Actually correctness to |
| // ensure this produces the same results as spark is verified in the JVM tests |
| fn test_date_trunc_array_fmt_dyn() { |
| let size = 10; |
| let formats = [ |
| "YEAR", "YYYY", "YY", "QUARTER", "MONTH", "MON", "MM", "WEEK", |
| ]; |
| let mut vec: Vec<i32> = Vec::with_capacity(size * formats.len()); |
| let mut fmt_vec: Vec<&str> = Vec::with_capacity(size * formats.len()); |
| for i in 0..size { |
| for fmt_value in &formats { |
| vec.push(i as i32 * 1_000_001); |
| fmt_vec.push(fmt_value); |
| } |
| } |
| |
| // timestamp array |
| let array = Date32Array::from(vec); |
| |
| // formats array |
| let fmt_array = StringArray::from(fmt_vec); |
| |
| // timestamp dictionary array |
| let mut date_dict_builder = PrimitiveDictionaryBuilder::<Int32Type, Date32Type>::new(); |
| for v in array.iter() { |
| date_dict_builder |
| .append(v.unwrap()) |
| .expect("Error in building timestamp array"); |
| } |
| let mut array_dict = date_dict_builder.finish(); |
| // apply timezone |
| array_dict = array_dict.with_values(Arc::new( |
| array_dict |
| .values() |
| .as_any() |
| .downcast_ref::<Date32Array>() |
| .unwrap() |
| .clone(), |
| )); |
| |
| // formats dictionary array |
| let mut formats_dict_builder = StringDictionaryBuilder::<Int32Type>::new(); |
| for v in fmt_array.iter() { |
| formats_dict_builder |
| .append(v.unwrap()) |
| .expect("Error in building formats array"); |
| } |
| let fmt_dict = formats_dict_builder.finish(); |
| |
| // verify input arrays |
| let iter = ArrayIter::new(&array); |
| let mut dict_iter = array_dict |
| .downcast_dict::<PrimitiveArray<Date32Type>>() |
| .unwrap() |
| .into_iter(); |
| for val in iter { |
| assert_eq!( |
| dict_iter |
| .next() |
| .expect("array and dictionary array do not match"), |
| val |
| ) |
| } |
| |
| // verify input format arrays |
| let fmt_iter = ArrayIter::new(&fmt_array); |
| let mut fmt_dict_iter = fmt_dict.downcast_dict::<StringArray>().unwrap().into_iter(); |
| for val in fmt_iter { |
| assert_eq!( |
| fmt_dict_iter |
| .next() |
| .expect("formats and dictionary formats do not match"), |
| val |
| ) |
| } |
| |
| // test cases |
| if let Ok(a) = date_trunc_array_fmt_dyn(&array, &fmt_array) { |
| for i in 0..array.len() { |
| assert!( |
| array.value(i) >= a.as_any().downcast_ref::<Date32Array>().unwrap().value(i) |
| ) |
| } |
| } else { |
| unreachable!() |
| } |
| if let Ok(a) = date_trunc_array_fmt_dyn(&array_dict, &fmt_array) { |
| for i in 0..array.len() { |
| assert!( |
| array.value(i) >= a.as_any().downcast_ref::<Date32Array>().unwrap().value(i) |
| ) |
| } |
| } else { |
| unreachable!() |
| } |
| if let Ok(a) = date_trunc_array_fmt_dyn(&array, &fmt_dict) { |
| for i in 0..array.len() { |
| assert!( |
| array.value(i) >= a.as_any().downcast_ref::<Date32Array>().unwrap().value(i) |
| ) |
| } |
| } else { |
| unreachable!() |
| } |
| if let Ok(a) = date_trunc_array_fmt_dyn(&array_dict, &fmt_dict) { |
| for i in 0..array.len() { |
| assert!( |
| array.value(i) >= a.as_any().downcast_ref::<Date32Array>().unwrap().value(i) |
| ) |
| } |
| } else { |
| unreachable!() |
| } |
| } |
| |
| #[test] |
| #[cfg_attr(miri, ignore)] // test takes too long with miri |
| fn test_timestamp_trunc() { |
| let size = 1000; |
| let mut vec: Vec<i64> = Vec::with_capacity(size); |
| for i in 0..size { |
| vec.push(i as i64); |
| } |
| let array = TimestampMicrosecondArray::from(vec).with_timezone_utc(); |
| for fmt in [ |
| "YEAR", |
| "YYYY", |
| "YY", |
| "QUARTER", |
| "MONTH", |
| "MON", |
| "MM", |
| "WEEK", |
| "DAY", |
| "DD", |
| "HOUR", |
| "MINUTE", |
| "SECOND", |
| "MILLISECOND", |
| "MICROSECOND", |
| ] { |
| match timestamp_trunc(&array, fmt.to_string()) { |
| Ok(a) => { |
| for i in 0..size { |
| assert!(array.values().get(i) >= a.values().get(i)) |
| } |
| } |
| _ => unreachable!(), |
| } |
| } |
| } |
| |
| #[test] |
| // test takes too long with miri |
| #[cfg_attr(miri, ignore)] |
| // This test only verifies that the various input array types work. Actually correctness to |
| // ensure this produces the same results as spark is verified in the JVM tests |
| fn test_timestamp_trunc_array_fmt_dyn() { |
| let size = 10; |
| let formats = [ |
| "YEAR", |
| "YYYY", |
| "YY", |
| "QUARTER", |
| "MONTH", |
| "MON", |
| "MM", |
| "WEEK", |
| "DAY", |
| "DD", |
| "HOUR", |
| "MINUTE", |
| "SECOND", |
| "MILLISECOND", |
| "MICROSECOND", |
| ]; |
| let mut vec: Vec<i64> = Vec::with_capacity(size * formats.len()); |
| let mut fmt_vec: Vec<&str> = Vec::with_capacity(size * formats.len()); |
| for i in 0..size { |
| for fmt_value in &formats { |
| vec.push(i as i64 * 1_000_000_001); |
| fmt_vec.push(fmt_value); |
| } |
| } |
| |
| // timestamp array |
| let array = TimestampMicrosecondArray::from(vec).with_timezone_utc(); |
| |
| // formats array |
| let fmt_array = StringArray::from(fmt_vec); |
| |
| // timestamp dictionary array |
| let mut timestamp_dict_builder = |
| PrimitiveDictionaryBuilder::<Int32Type, TimestampMicrosecondType>::new(); |
| for v in array.iter() { |
| timestamp_dict_builder |
| .append(v.unwrap()) |
| .expect("Error in building timestamp array"); |
| } |
| let mut array_dict = timestamp_dict_builder.finish(); |
| // apply timezone |
| array_dict = array_dict.with_values(Arc::new( |
| array_dict |
| .values() |
| .as_any() |
| .downcast_ref::<TimestampMicrosecondArray>() |
| .unwrap() |
| .clone() |
| .with_timezone_utc(), |
| )); |
| |
| // formats dictionary array |
| let mut formats_dict_builder = StringDictionaryBuilder::<Int32Type>::new(); |
| for v in fmt_array.iter() { |
| formats_dict_builder |
| .append(v.unwrap()) |
| .expect("Error in building formats array"); |
| } |
| let fmt_dict = formats_dict_builder.finish(); |
| |
| // verify input arrays |
| let iter = ArrayIter::new(&array); |
| let mut dict_iter = array_dict |
| .downcast_dict::<PrimitiveArray<TimestampMicrosecondType>>() |
| .unwrap() |
| .into_iter(); |
| for val in iter { |
| assert_eq!( |
| dict_iter |
| .next() |
| .expect("array and dictionary array do not match"), |
| val |
| ) |
| } |
| |
| // verify input format arrays |
| let fmt_iter = ArrayIter::new(&fmt_array); |
| let mut fmt_dict_iter = fmt_dict.downcast_dict::<StringArray>().unwrap().into_iter(); |
| for val in fmt_iter { |
| assert_eq!( |
| fmt_dict_iter |
| .next() |
| .expect("formats and dictionary formats do not match"), |
| val |
| ) |
| } |
| |
| // test cases |
| if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_array) { |
| for i in 0..array.len() { |
| assert!( |
| array.value(i) |
| >= a.as_any() |
| .downcast_ref::<TimestampMicrosecondArray>() |
| .unwrap() |
| .value(i) |
| ) |
| } |
| } else { |
| unreachable!() |
| } |
| if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_array) { |
| for i in 0..array.len() { |
| assert!( |
| array.value(i) |
| >= a.as_any() |
| .downcast_ref::<TimestampMicrosecondArray>() |
| .unwrap() |
| .value(i) |
| ) |
| } |
| } else { |
| unreachable!() |
| } |
| if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_dict) { |
| for i in 0..array.len() { |
| assert!( |
| array.value(i) |
| >= a.as_any() |
| .downcast_ref::<TimestampMicrosecondArray>() |
| .unwrap() |
| .value(i) |
| ) |
| } |
| } else { |
| unreachable!() |
| } |
| if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_dict) { |
| for i in 0..array.len() { |
| assert!( |
| array.value(i) |
| >= a.as_any() |
| .downcast_ref::<TimestampMicrosecondArray>() |
| .unwrap() |
| .value(i) |
| ) |
| } |
| } else { |
| unreachable!() |
| } |
| } |
| } |
