Skip to content

Commit cedb05a

Browse files
compheadcomphead
authored
date_part support fractions of second (#4385)
* support seconds fraction in date_part * support seconds fraction in date_part * to be insync with pgsql * fix merge * fix bench * fix bench
1 parent 19cddf5 commit cedb05a

4 files changed

Lines changed: 105 additions & 26 deletions

File tree

‎benchmarks/src/tpch.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -192,18 +192,18 @@ pub fn get_answer_schema(n: usize) -> Schema {
192192
7 => Schema::new(vec![
193193
Field::new("supp_nation", DataType::Utf8, true),
194194
Field::new("cust_nation", DataType::Utf8, true),
195-
Field::new("l_year", DataType::Int32, true),
195+
Field::new("l_year", DataType::Float64, true),
196196
Field::new("revenue", DataType::Decimal128(15, 2), true),
197197
]),
198198

199199
8 => Schema::new(vec![
200-
Field::new("o_year", DataType::Int32, true),
200+
Field::new("o_year", DataType::Float64, true),
201201
Field::new("mkt_share", DataType::Decimal128(15, 2), true),
202202
]),
203203

204204
9 => Schema::new(vec![
205205
Field::new("nation", DataType::Utf8, true),
206-
Field::new("o_year", DataType::Int32, true),
206+
Field::new("o_year", DataType::Float64, true),
207207
Field::new("sum_profit", DataType::Decimal128(15, 2), true),
208208
]),
209209

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

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1285,12 +1285,36 @@ async fn test_extract_date_part() -> Result<()> {
12851285
"12"
12861286
);
12871287
test_expression!(
1288-
"EXTRACT(second FROM to_timestamp('2020-09-08T12:00:12+00:00'))",
1289-
"12"
1288+
"EXTRACT(second FROM to_timestamp('2020-09-08T12:00:12.12345678+00:00'))",
1289+
"12.12345678"
12901290
);
12911291
test_expression!(
1292-
"date_part('second', to_timestamp('2020-09-08T12:00:12+00:00'))",
1293-
"12"
1292+
"EXTRACT(millisecond FROM to_timestamp('2020-09-08T12:00:12.12345678+00:00'))",
1293+
"12123.45678"
1294+
);
1295+
test_expression!(
1296+
"EXTRACT(microsecond FROM to_timestamp('2020-09-08T12:00:12.12345678+00:00'))",
1297+
"12123456.78"
1298+
);
1299+
// test_expression!(
1300+
// "EXTRACT(nanosecond FROM to_timestamp('2020-09-08T12:00:12.12345678+00:00'))",
1301+
// "1212345678"
1302+
// );
1303+
test_expression!(
1304+
"date_part('second', to_timestamp('2020-09-08T12:00:12.12345678+00:00'))",
1305+
"12.12345678"
1306+
);
1307+
test_expression!(
1308+
"date_part('millisecond', to_timestamp('2020-09-08T12:00:12.12345678+00:00'))",
1309+
"12123.45678"
1310+
);
1311+
test_expression!(
1312+
"date_part('microsecond', to_timestamp('2020-09-08T12:00:12.12345678+00:00'))",
1313+
"12123456.78"
1314+
);
1315+
test_expression!(
1316+
"date_part('nanosecond', to_timestamp('2020-09-08T12:00:12.12345678+00:00'))",
1317+
"12123456780"
12941318
);
12951319
Ok(())
12961320
}

‎datafusion/expr/src/function.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ pub fn return_type(
133133
}
134134
BuiltinScalarFunction::Concat => Ok(DataType::Utf8),
135135
BuiltinScalarFunction::ConcatWithSeparator => Ok(DataType::Utf8),
136-
BuiltinScalarFunction::DatePart => Ok(DataType::Int32),
136+
BuiltinScalarFunction::DatePart => Ok(DataType::Float64),
137137
BuiltinScalarFunction::DateTrunc => {
138138
Ok(DataType::Timestamp(TimeUnit::Nanosecond, None))
139139
}

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

Lines changed: 73 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,18 @@
1717

1818
//! DateTime expressions
1919
20+
use arrow::compute::cast;
2021
use arrow::{
2122
array::TimestampNanosecondArray, compute::kernels::temporal, datatypes::TimeUnit,
2223
temporal_conversions::timestamp_ns_to_datetime,
2324
};
2425
use arrow::{
25-
array::{Array, ArrayRef, OffsetSizeTrait, PrimitiveArray},
26+
array::{Array, ArrayRef, Float64Array, OffsetSizeTrait, PrimitiveArray},
2627
compute::kernels::cast_utils::string_to_timestamp_nanos,
2728
datatypes::{
28-
ArrowPrimitiveType, DataType, IntervalDayTimeType, TimestampMicrosecondType,
29-
TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
29+
ArrowNumericType, ArrowPrimitiveType, ArrowTemporalType, DataType,
30+
IntervalDayTimeType, TimestampMicrosecondType, TimestampMillisecondType,
31+
TimestampNanosecondType, TimestampSecondType,
3032
},
3133
};
3234
use chrono::prelude::*;
@@ -401,30 +403,36 @@ pub fn date_bin(args: &[ColumnarValue]) -> Result<ColumnarValue> {
401403
macro_rules! extract_date_part {
402404
($ARRAY: expr, $FN:expr) => {
403405
match $ARRAY.data_type() {
404-
DataType::Date32 => match as_date32_array($ARRAY) {
405-
Ok(array) => Ok($FN(array)?),
406-
Err(e) => Err(e),
407-
},
406+
DataType::Date32 => {
407+
let array = as_date32_array($ARRAY)?;
408+
Ok($FN(array)
409+
.map(|v| cast(&(Arc::new(v) as ArrayRef), &DataType::Float64))?)
410+
}
408411
DataType::Date64 => {
409412
let array = as_date64_array($ARRAY)?;
410-
Ok($FN(array)?)
413+
Ok($FN(array)
414+
.map(|v| cast(&(Arc::new(v) as ArrayRef), &DataType::Float64))?)
411415
}
412416
DataType::Timestamp(time_unit, None) => match time_unit {
413417
TimeUnit::Second => {
414418
let array = as_timestamp_second_array($ARRAY)?;
415-
Ok($FN(array)?)
419+
Ok($FN(array)
420+
.map(|v| cast(&(Arc::new(v) as ArrayRef), &DataType::Float64))?)
416421
}
417422
TimeUnit::Millisecond => {
418423
let array = as_timestamp_millisecond_array($ARRAY)?;
419-
Ok($FN(array)?)
424+
Ok($FN(array)
425+
.map(|v| cast(&(Arc::new(v) as ArrayRef), &DataType::Float64))?)
420426
}
421427
TimeUnit::Microsecond => {
422428
let array = as_timestamp_microsecond_array($ARRAY)?;
423-
Ok($FN(array)?)
429+
Ok($FN(array)
430+
.map(|v| cast(&(Arc::new(v) as ArrayRef), &DataType::Float64))?)
424431
}
425432
TimeUnit::Nanosecond => {
426433
let array = as_timestamp_nanosecond_array($ARRAY)?;
427-
Ok($FN(array)?)
434+
Ok($FN(array)
435+
.map(|v| cast(&(Arc::new(v) as ArrayRef), &DataType::Float64))?)
428436
}
429437
},
430438
datatype => Err(DataFusionError::Internal(format!(
@@ -469,23 +477,70 @@ pub fn date_part(args: &[ColumnarValue]) -> Result<ColumnarValue> {
469477
"dow" => extract_date_part!(&array, temporal::num_days_from_sunday),
470478
"hour" => extract_date_part!(&array, temporal::hour),
471479
"minute" => extract_date_part!(&array, temporal::minute),
472-
"second" => extract_date_part!(&array, temporal::second),
480+
"second" => extract_date_part!(&array, seconds),
481+
"millisecond" => extract_date_part!(&array, millis),
482+
"microsecond" => extract_date_part!(&array, micros),
483+
"nanosecond" => extract_date_part!(&array, nanos),
473484
_ => Err(DataFusionError::Execution(format!(
474485
"Date part '{}' not supported",
475486
date_part
476487
))),
477488
}?;
478489

479490
Ok(if is_scalar {
480-
ColumnarValue::Scalar(ScalarValue::try_from_array(
481-
&(Arc::new(arr) as ArrayRef),
482-
0,
483-
)?)
491+
ColumnarValue::Scalar(ScalarValue::try_from_array(&arr?, 0)?)
484492
} else {
485-
ColumnarValue::Array(Arc::new(arr))
493+
ColumnarValue::Array(arr?)
486494
})
487495
}
488496

497+
fn to_ticks<T>(array: &PrimitiveArray<T>, frac: i32) -> Result<Float64Array>
498+
where
499+
T: ArrowTemporalType + ArrowNumericType,
500+
i64: From<T::Native>,
501+
{
502+
let zipped = temporal::second(array)?
503+
.values()
504+
.iter()
505+
.zip(temporal::nanosecond(array)?.values().iter())
506+
.map(|o| ((*o.0 as f64 + (*o.1 as f64) / 1_000_000_000.0) * (frac as f64)))
507+
.collect::<Vec<f64>>();
508+
509+
Ok(Float64Array::from(zipped))
510+
}
511+
512+
fn seconds<T>(array: &PrimitiveArray<T>) -> Result<Float64Array>
513+
where
514+
T: ArrowTemporalType + ArrowNumericType,
515+
i64: From<T::Native>,
516+
{
517+
to_ticks(array, 1)
518+
}
519+
520+
fn millis<T>(array: &PrimitiveArray<T>) -> Result<Float64Array>
521+
where
522+
T: ArrowTemporalType + ArrowNumericType,
523+
i64: From<T::Native>,
524+
{
525+
to_ticks(array, 1_000)
526+
}
527+
528+
fn micros<T>(array: &PrimitiveArray<T>) -> Result<Float64Array>
529+
where
530+
T: ArrowTemporalType + ArrowNumericType,
531+
i64: From<T::Native>,
532+
{
533+
to_ticks(array, 1_000_000)
534+
}
535+
536+
fn nanos<T>(array: &PrimitiveArray<T>) -> Result<Float64Array>
537+
where
538+
T: ArrowTemporalType + ArrowNumericType,
539+
i64: From<T::Native>,
540+
{
541+
to_ticks(array, 1_000_000_000)
542+
}
543+
489544
#[cfg(test)]
490545
mod tests {
491546
use std::sync::Arc;

0 commit comments

Comments
 (0)