kosiew commented on code in PR #20808: URL: https://github.com/apache/datafusion/pull/20808#discussion_r3014862064
########## datafusion/spark/src/function/datetime/quarter.rs: ########## @@ -0,0 +1,106 @@ +// 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. + +use arrow::array::{Array, ArrayRef}; +use arrow::compute::{CastOptions, DatePart, cast_with_options, date_part}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::logical_expr::{ + Coercion, ColumnarValue, Signature, TypeSignature, TypeSignatureClass, Volatility, +}; +use datafusion_common::types::{logical_date, logical_string}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, internal_err}; +use datafusion_expr::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_functions::utils::make_scalar_function; +use std::sync::Arc; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkQuarter { + signature: Signature, +} + +impl Default for SparkQuarter { + fn default() -> Self { + Self::new() + } +} + +impl SparkQuarter { + pub fn new() -> Self { + Self { + signature: Signature::one_of( + vec![ + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Timestamp, + )]), + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Native(logical_date()), + )]), + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Native(logical_string()), + )]), + ], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for SparkQuarter { + fn name(&self) -> &str { + "quarter" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { + Ok(Arc::new(Field::new( + self.name(), + DataType::Int32, + args.arg_fields[0].is_nullable(), + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + make_scalar_function(spark_quarter, vec![])(&args.args) + } +} + +fn spark_quarter(args: &[ArrayRef]) -> Result<ArrayRef> { + let [array] = take_function_args("quarter", args)?; + match array.data_type() { + DataType::Date32 | DataType::Timestamp(_, _) => { + let quarter = date_part(array, DatePart::Quarter)?; + Ok(quarter) + } + DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => { + let date_array = + cast_with_options(array, &DataType::Date32, &CastOptions::default())?; Review Comment: I am a bit concerned that the new string handling is narrower than the shared datetime coercion path. This currently forces every string through a `Date32` cast before calling `date_part`. That can reject valid timestamp-shaped strings that `date_part` already accepts elsewhere, for example `date_part('second', '2020-09-08T12:00:12.12345678+00:00')` in `datafusion/sqllogictest/test_files/datetime/date_part.slt`. Because this does not route through the existing `date_part('quarter', ...)` behavior, `quarter` can still diverge from the rest of the datetime coercion model for string inputs. Could we reuse the same coercion path here so the behavior stays aligned? ########## datafusion/spark/src/function/datetime/quarter.rs: ########## @@ -0,0 +1,106 @@ +// 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. + +use arrow::array::{Array, ArrayRef}; +use arrow::compute::{CastOptions, DatePart, cast_with_options, date_part}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::logical_expr::{ Review Comment: Nice addition overall. One thing to fix here is the direct import from `datafusion::logical_expr`. In `datafusion/spark/Cargo.toml`, `datafusion` is still an optional dependency behind the `core` feature, so this regresses a CI-critical configuration. `cargo check -p datafusion-spark --no-default-features` now fails with `use of unresolved module or unlinked crate 'datafusion'`. The rest of `datafusion-spark` pulls these expression types from `datafusion_expr`, which keeps the no-default-features build working. Could we switch this import to match the existing pattern? ########## datafusion/spark/src/function/datetime/quarter.rs: ########## @@ -0,0 +1,106 @@ +// 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. + +use arrow::array::{Array, ArrayRef}; +use arrow::compute::{CastOptions, DatePart, cast_with_options, date_part}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::logical_expr::{ + Coercion, ColumnarValue, Signature, TypeSignature, TypeSignatureClass, Volatility, +}; +use datafusion_common::types::{logical_date, logical_string}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, internal_err}; +use datafusion_expr::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_functions::utils::make_scalar_function; +use std::sync::Arc; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkQuarter { + signature: Signature, +} + +impl Default for SparkQuarter { + fn default() -> Self { + Self::new() + } +} + +impl SparkQuarter { + pub fn new() -> Self { + Self { + signature: Signature::one_of( + vec![ + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Timestamp, + )]), + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Native(logical_date()), + )]), + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Native(logical_string()), + )]), + ], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for SparkQuarter { + fn name(&self) -> &str { + "quarter" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { + Ok(Arc::new(Field::new( + self.name(), + DataType::Int32, + args.arg_fields[0].is_nullable(), Review Comment: I think the return-field nullability needs to be loosened here. Right now `return_field_from_args` mirrors the input field nullability, but the new string path can produce `NULL` even when the input is non-null. This patch adds cases like `quarter('abc'::string)` and `quarter(''::string)` returning `NULL`, so `quarter(non_null_utf8_col)` would still be advertised as `Int32 NOT NULL` even though execution can yield nulls. That looks like a schema contract bug. It also differs from existing Spark helpers like `next_day`, which force nullable output when invalid strings can map to `NULL`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
