Skip to content

Commit 1abe9d5

Browse files
waralexromwaralexrom
authored
feat: EXTRACT(EPOCH, ...) support (#136)
1 parent 3c85ef6 commit 1abe9d5

File tree

3 files changed

+75
-24
lines changed

3 files changed

+75
-24
lines changed

‎datafusion/core/tests/sql/expr.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -991,6 +991,19 @@ async fn test_extract_date_part() -> Result<()> {
991991
"date_part('doy', to_timestamp('2020-02-01T00:00:00+00:00'))",
992992
"32"
993993
);
994+
// EPOCH
995+
test_expression!(
996+
"EXTRACT(epoch FROM to_timestamp('2020-01-01T23:01:01.22+00:00'))",
997+
"1577919661.22"
998+
);
999+
test_expression!(
1000+
"date_part('epoch', to_timestamp('2020-01-01T23:01:01.22+00:00'))",
1001+
"1577919661.22"
1002+
);
1003+
test_expression!(
1004+
"date_part('epoch', CAST('2020-01-01' AS DATE))",
1005+
"1577836800"
1006+
);
9941007

9951008
// DOW
9961009
test_expression!(

‎datafusion/cube_ext/src/temporal.rs‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,9 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use arrow::array::{Array, Int32Array, Int32Builder, PrimitiveArray};
19-
use arrow::datatypes::{ArrowNumericType, ArrowTemporalType, DataType};
18+
use arrow::array::{Array, Float64Array, Int32Array, Int32Builder, PrimitiveArray};
19+
use arrow::compute::kernels::arity::unary;
20+
use arrow::datatypes::{ArrowNumericType, ArrowTemporalType, DataType, TimeUnit};
2021
use arrow::error::{ArrowError, Result};
2122

2223
use chrono::format::strftime::StrftimeItems;
@@ -144,6 +145,42 @@ where
144145
Ok(b.finish())
145146
}
146147

148+
pub fn epoch<T>(array: &PrimitiveArray<T>) -> Result<Float64Array>
149+
where
150+
T: ArrowTemporalType + ArrowNumericType,
151+
i64: From<T::Native>,
152+
{
153+
let b = match array.data_type() {
154+
DataType::Timestamp(tu, _) => {
155+
let scale = match tu {
156+
TimeUnit::Second => 1,
157+
TimeUnit::Millisecond => 1_000,
158+
TimeUnit::Microsecond => 1_000_000,
159+
TimeUnit::Nanosecond => 1_000_000_000,
160+
} as f64;
161+
unary(array, |n| {
162+
let n: i64 = n.into();
163+
n as f64 / scale
164+
})
165+
}
166+
DataType::Date32 => {
167+
let seconds_in_a_day = 86400_f64;
168+
unary(array, |n| {
169+
let n: i64 = n.into();
170+
n as f64 * seconds_in_a_day
171+
})
172+
}
173+
DataType::Date64 => unary(array, |n| {
174+
let n: i64 = n.into();
175+
n as f64 / 1_000_f64
176+
}),
177+
_ => {
178+
return_compute_error_with!("Can not convert {:?} to epoch", array.data_type())
179+
}
180+
};
181+
Ok(b)
182+
}
183+
147184
trait ChronoDateLikeExt {
148185
fn weekday_from_sunday(&self) -> i32;
149186
}

‎datafusion/physical-expr/src/datetime_expressions.rs‎

Lines changed: 23 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
//! DateTime expressions
1919
2020
use arrow::array::{Int64Array, IntervalDayTimeArray, IntervalYearMonthArray};
21+
use arrow::compute::cast;
2122
use arrow::{
2223
array::{Array, ArrayRef, GenericStringArray, PrimitiveArray, StringOffsetSizeTrait},
2324
compute::kernels::cast_utils::string_to_timestamp_nanos,
@@ -481,44 +482,44 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result<ColumnarValue> {
481482
}
482483

483484
macro_rules! extract_date_part {
484-
($ARRAY: expr, $FN:expr) => {
485+
($ARRAY: expr, $FN:expr, $RT: expr) => {
485486
match $ARRAY.data_type() {
486487
DataType::Date32 => {
487488
let array = $ARRAY.as_any().downcast_ref::<Date32Array>().unwrap();
488-
Ok($FN(array)?)
489+
Ok($FN(array).map(|v| cast(&(Arc::new(v) as ArrayRef), &$RT))?)
489490
}
490491
DataType::Date64 => {
491492
let array = $ARRAY.as_any().downcast_ref::<Date64Array>().unwrap();
492-
Ok($FN(array)?)
493+
Ok($FN(array).map(|v| cast(&(Arc::new(v) as ArrayRef), &$RT))?)
493494
}
494495
DataType::Timestamp(time_unit, None) => match time_unit {
495496
TimeUnit::Second => {
496497
let array = $ARRAY
497498
.as_any()
498499
.downcast_ref::<TimestampSecondArray>()
499500
.unwrap();
500-
Ok($FN(array)?)
501+
Ok($FN(array).map(|v| cast(&(Arc::new(v) as ArrayRef), &$RT))?)
501502
}
502503
TimeUnit::Millisecond => {
503504
let array = $ARRAY
504505
.as_any()
505506
.downcast_ref::<TimestampMillisecondArray>()
506507
.unwrap();
507-
Ok($FN(array)?)
508+
Ok($FN(array).map(|v| cast(&(Arc::new(v) as ArrayRef), &$RT))?)
508509
}
509510
TimeUnit::Microsecond => {
510511
let array = $ARRAY
511512
.as_any()
512513
.downcast_ref::<TimestampMicrosecondArray>()
513514
.unwrap();
514-
Ok($FN(array)?)
515+
Ok($FN(array).map(|v| cast(&(Arc::new(v) as ArrayRef), &$RT))?)
515516
}
516517
TimeUnit::Nanosecond => {
517518
let array = $ARRAY
518519
.as_any()
519520
.downcast_ref::<TimestampNanosecondArray>()
520521
.unwrap();
521-
Ok($FN(array)?)
522+
Ok($FN(array).map(|v| cast(&(Arc::new(v) as ArrayRef), &$RT))?)
522523
}
523524
},
524525
datatype => Err(DataFusionError::Internal(format!(
@@ -554,29 +555,29 @@ pub fn date_part(args: &[ColumnarValue]) -> Result<ColumnarValue> {
554555
};
555556

556557
let arr = match date_part.to_lowercase().as_str() {
557-
"doy" => extract_date_part!(array, cube_ext::temporal::doy),
558-
"dow" => extract_date_part!(array, cube_ext::temporal::dow),
559-
"year" => extract_date_part!(array, temporal::year),
560-
"quarter" => extract_date_part!(array, temporal::quarter),
561-
"month" => extract_date_part!(array, temporal::month),
562-
"week" => extract_date_part!(array, temporal::week),
563-
"day" => extract_date_part!(array, temporal::day),
564-
"hour" => extract_date_part!(array, temporal::hour),
565-
"minute" => extract_date_part!(array, temporal::minute),
566-
"second" => extract_date_part!(array, temporal::second),
558+
"doy" => extract_date_part!(array, cube_ext::temporal::doy, DataType::Int32),
559+
"dow" => extract_date_part!(array, cube_ext::temporal::dow, DataType::Int32),
560+
"year" => extract_date_part!(array, temporal::year, DataType::Int32),
561+
"quarter" => extract_date_part!(array, temporal::quarter, DataType::Int32),
562+
"month" => extract_date_part!(array, temporal::month, DataType::Int32),
563+
"week" => extract_date_part!(array, temporal::week, DataType::Int32),
564+
"day" => extract_date_part!(array, temporal::day, DataType::Int32),
565+
"hour" => extract_date_part!(array, temporal::hour, DataType::Int32),
566+
"minute" => extract_date_part!(array, temporal::minute, DataType::Int32),
567+
"second" => extract_date_part!(array, temporal::second, DataType::Int32),
568+
"epoch" => {
569+
extract_date_part!(array, cube_ext::temporal::epoch, DataType::Float64)
570+
}
567571
_ => Err(DataFusionError::Execution(format!(
568572
"Date part '{}' not supported",
569573
date_part
570574
))),
571575
}?;
572576

573577
Ok(if is_scalar {
574-
ColumnarValue::Scalar(ScalarValue::try_from_array(
575-
&(Arc::new(arr) as ArrayRef),
576-
0,
577-
)?)
578+
ColumnarValue::Scalar(ScalarValue::try_from_array(&arr?, 0)?)
578579
} else {
579-
ColumnarValue::Array(Arc::new(arr))
580+
ColumnarValue::Array(arr?)
580581
})
581582
}
582583

0 commit comments

Comments
 (0)