Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Return int32 for integer type date part #13466

Merged
merged 8 commits into from
Nov 22, 2024
+282 −243
Merged

Return int32 for integer type date part #13466

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
db289b3
return int for integar date part
jayzhan211 Nov 18, 2024
96342af
fix tpch test
jayzhan211 Nov 18, 2024
6be9d9f
type test
jayzhan211 Nov 19, 2024
0c17dbd
Update datafusion/functions/src/datetime/date_part.rs
jayzhan211 Nov 19, 2024
8a334d3
fix name
jayzhan211 Nov 20, 2024
7985684
use int for second
jayzhan211 Nov 20, 2024
68adb4e
rm dot
jayzhan211 Nov 21, 2024
c7afa29
Merge branch 'main' of https://github.com/apache/datafusion into date…
jayzhan211 Nov 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 54 additions & 25 deletions datafusion/functions/src/datetime/date_part.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,9 @@ use std::sync::{Arc, OnceLock};

use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::compute::kernels::cast_utils::IntervalUnit;
use arrow::compute::{binary, cast, date_part, DatePart};
use arrow::compute::{binary, date_part, DatePart};
use arrow::datatypes::DataType::{
Date32, Date64, Duration, Float64, Interval, Time32, Time64, Timestamp, Utf8,
Utf8View,
Date32, Date64, Duration, Interval, Time32, Time64, Timestamp, Utf8, Utf8View,
};
use arrow::datatypes::IntervalUnit::{DayTime, MonthDayNano, YearMonth};
use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
Expand All @@ -36,11 +35,12 @@ use datafusion_common::cast::{
as_timestamp_microsecond_array, as_timestamp_millisecond_array,
as_timestamp_nanosecond_array, as_timestamp_second_array,
};
use datafusion_common::{exec_err, Result, ScalarValue};
use datafusion_common::{exec_err, internal_err, ExprSchema, Result, ScalarValue};
use datafusion_expr::scalar_doc_sections::DOC_SECTION_DATETIME;
use datafusion_expr::TypeSignature::Exact;
use datafusion_expr::{
ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
ColumnarValue, Documentation, Expr, ScalarUDFImpl, Signature, Volatility,
TIMEZONE_WILDCARD,
};

#[derive(Debug)]
Expand Down Expand Up @@ -148,7 +148,21 @@ impl ScalarUDFImpl for DatePartFunc {
}

fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(Float64)
internal_err!("return_type_from_exprs shoud be called instead")
}

fn return_type_from_exprs(
&self,
args: &[Expr],
_schema: &dyn ExprSchema,
_arg_types: &[DataType],
) -> Result<DataType> {
match &args[0] {
Expr::Literal(ScalarValue::Utf8(Some(part))) if is_integar_part(part) => {
Ok(DataType::Int32)
}
_ => Ok(DataType::Float64),
}
}

fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
Expand All @@ -174,22 +188,18 @@ impl ScalarUDFImpl for DatePartFunc {
ColumnarValue::Scalar(scalar) => scalar.to_array()?,
};

// to remove quotes at most 2 characters
let part_trim = part.trim_matches(|c| c == '\'' || c == '\"');
if ![2, 0].contains(&(part.len() - part_trim.len())) {
return exec_err!("Date part '{part}' not supported");
}
let part_trim = part_normalization(part);

// using IntervalUnit here means we hand off all the work of supporting plurals (like "seconds")
// and synonyms ( like "ms,msec,msecond,millisecond") to Arrow
let arr = if let Ok(interval_unit) = IntervalUnit::from_str(part_trim) {
match interval_unit {
IntervalUnit::Year => date_part_f64(array.as_ref(), DatePart::Year)?,
IntervalUnit::Month => date_part_f64(array.as_ref(), DatePart::Month)?,
IntervalUnit::Week => date_part_f64(array.as_ref(), DatePart::Week)?,
IntervalUnit::Day => date_part_f64(array.as_ref(), DatePart::Day)?,
IntervalUnit::Hour => date_part_f64(array.as_ref(), DatePart::Hour)?,
IntervalUnit::Minute => date_part_f64(array.as_ref(), DatePart::Minute)?,
IntervalUnit::Year => date_part(array.as_ref(), DatePart::Year)?,
IntervalUnit::Month => date_part(array.as_ref(), DatePart::Month)?,
IntervalUnit::Week => date_part(array.as_ref(), DatePart::Week)?,
IntervalUnit::Day => date_part(array.as_ref(), DatePart::Day)?,
IntervalUnit::Hour => date_part(array.as_ref(), DatePart::Hour)?,
IntervalUnit::Minute => date_part(array.as_ref(), DatePart::Minute)?,
IntervalUnit::Second => seconds(array.as_ref(), Second)?,
IntervalUnit::Millisecond => seconds(array.as_ref(), Millisecond)?,
IntervalUnit::Microsecond => seconds(array.as_ref(), Microsecond)?,
Expand All @@ -200,9 +210,9 @@ impl ScalarUDFImpl for DatePartFunc {
} else {
// special cases that can be extracted (in postgres) but are not interval units
match part_trim.to_lowercase().as_str() {
"qtr" | "quarter" => date_part_f64(array.as_ref(), DatePart::Quarter)?,
"doy" => date_part_f64(array.as_ref(), DatePart::DayOfYear)?,
"dow" => date_part_f64(array.as_ref(), DatePart::DayOfWeekSunday0)?,
"qtr" | "quarter" => date_part(array.as_ref(), DatePart::Quarter)?,
"doy" => date_part(array.as_ref(), DatePart::DayOfYear)?,
"dow" => date_part(array.as_ref(), DatePart::DayOfWeekSunday0)?,
"epoch" => epoch(array.as_ref())?,
_ => return exec_err!("Date part '{part}' not supported"),
}
Expand All @@ -223,6 +233,30 @@ impl ScalarUDFImpl for DatePartFunc {
}
}

fn is_integar_part(part: &str) -> bool {
jayzhan211 marked this conversation as resolved.
Show resolved Hide resolved
let part = part_normalization(part);
matches!(
part.to_lowercase().as_str(),
"year"
| "month"
| "week"
| "day"
| "hour"
| "minute"
| "qtr"
| "quarter"
| "doy"
| "dow"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

second, millisecond, micro, nanos too?

Copy link
Contributor Author

@jayzhan211 jayzhan211 Nov 18, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

let r: Float64Array = secs
            .iter()
            .zip(subsecs)
            .map(|(secs, subsecs)| {
                secs.map(|secs| {
                    let subsecs = subsecs.unwrap_or(0);
                    (secs as f64 + ((subsecs % 1_000_000_000) as f64 / 1_000_000_000_f64))
                        * sf
                })
            })
            .collect();

Probably we still need f64 for seconds

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i didn't know extract(second .. returns seconds with fraction. In PostgreSQL it does so, in e.g. Trino and Snowflake it does not

postgres=# select extract(second from now());
  extract
-----------
 29.009019

trino> select extract(second from now());
 _col0
-------
    46

so yes, keep seconds with float for now

Copy link
Contributor

@comphead comphead Nov 18, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

extract for seconds is float correct, returning the fractions. however its not the case for DuckDB.

Copy link
Member

@findepi findepi Nov 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

filed #13482 for seconds

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jayzhan211 what about millisecond, micro, nano?

)
}

// Try to remove quote if exist, if the quote is invalid, return original string and let the downstream function handle the error
fn part_normalization(part: &str) -> &str {
part.strip_prefix(|c| c == '\'' || c == '\"')
.and_then(|s| s.strip_suffix(|c| c == '\'' || c == '\"'))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why would we want to support unmatched quotes like '..."?

the date_part argument shouldn't contain any quotes, why would we want to support quotes at all?

i see this is pre-existing, so okay to keep, but still wonder why we're doing this, doesn't feel right

.unwrap_or(part)
}

static DOCUMENTATION: OnceLock<Documentation> = OnceLock::new();

fn get_date_part_doc() -> &'static Documentation {
Expand Down Expand Up @@ -261,11 +295,6 @@ fn get_date_part_doc() -> &'static Documentation {
})
}

/// Invoke [`date_part`] and cast the result to Float64
fn date_part_f64(array: &dyn Array, part: DatePart) -> Result<ArrayRef> {
Ok(cast(date_part(array, part)?.as_ref(), &Float64)?)
}

/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
/// result to a total number of seconds, milliseconds, microseconds or
/// nanoseconds
Expand Down
2 changes: 1 addition & 1 deletion datafusion/sqllogictest/test_files/clickbench.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ SELECT "UserID", "SearchPhrase", COUNT(*) FROM hits GROUP BY "UserID", "SearchPh
519640690937130534 (empty) 2
7418527520126366595 (empty) 1

query IRTI rowsort
query IITI rowsort
SELECT "UserID", extract(minute FROM to_timestamp_seconds("EventTime")) AS m, "SearchPhrase", COUNT(*) FROM hits GROUP BY "UserID", m, "SearchPhrase" ORDER BY COUNT(*) DESC LIMIT 10;
----
-2461439046089301801 18 (empty) 1
Expand Down
Loading