Spark quarter function implementation#20808
Spark quarter function implementation#20808kazantsev-maksim wants to merge 15 commits intoapache:mainfrom
Conversation
kazantsev-maksim
changed the title
There was a problem hiding this comment.
Thanks for putting this together — the new quarter UDF is headed in a good direction, but I think there are a couple of compatibility and robustness issues we should address before merging.
The main blocker is that the current registration uses an exact Date32 signature, which looks narrower than Spark’s documented behavior for quarter. I also left a couple of smaller follow-ups around reusing the existing date_part("quarter", ...) path and avoiding the unwrap() panic path.
I’d also love to see coverage added for Spark’s documented string-literal example so we lock that compatibility in.
| impl SparkQuarter { | ||
| pub fn new() -> Self { | ||
| Self { | ||
| signature: Signature::exact(vec![DataType::Date32], Volatility::Immutable), |
There was a problem hiding this comment.
I think this is the main thing we should fix before merging. Right now the UDF is registered with an exact Date32 signature, which means we no longer preserve Spark’s documented call shape for quarter.
Spark’s SQL docs show SELECT quarter('2016-08-31'); returning 3, and this SLT file used to carry that example before it was replaced with explicit ::DATE casts. With the current signature, we only validate the casted form and could end up rejecting the plain string-literal case that Spark accepts.
Could we switch this to a coercible signature, or possibly just route through the existing date_part('quarter', ...) behavior, and add coverage for the uncasted query?
| } | ||
|
|
||
| fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { | ||
| let [arg] = take_function_args("quarter", args.args)?; |
There was a problem hiding this comment.
Small suggestion here: this seems to repeat the same scalar/array Date32 dispatching pattern we already have in other datetime helpers, while datafusion_functions::datetime::date_part() already supports "quarter".
Would it make sense to delegate to that existing implementation instead? It feels like that would help keep coercion rules, null handling, and any future date-part behavior aligned in one place.
| } | ||
|
|
||
| fn spark_quarter(days: i32) -> Result<i32> { | ||
| let quarter = Date32Type::to_naive_date_opt(days).unwrap().quarter(); |
There was a problem hiding this comment.
One thing that made me a little nervous here is the unwrap(). That introduces a panic path if a malformed Date32 value ever makes it down to this helper.
last_day handles the same kind of conversion by returning an explicit error instead, which feels a bit safer for a public UDF. Could we do the same here so bad inputs show up as query errors rather than aborting execution?
There was a problem hiding this comment.
I tried to rework it
|
Thanks for the review @kosiew. Could you please another see, when you have a time? |
2 participants
Which issue does this PR close?
N/A
Rationale for this change
Add new spark function: https://spark.apache.org/docs/latest/api/sql/index.html#quarter
What changes are included in this PR?
Are these changes tested?
Yes, tests added as part of this PR.
Are there any user-facing changes?
No, these are new function.