Merge branch 'main' into alamb/clickbenchmaxxing_2

Author: alamb

datafusion/core/src/execution/context/mod.rs

@@ -2218,7 +2218,9 @@ mod tests {
     use crate::test;
     use crate::test_util::{plan_and_collect, populate_csv_partitions};
     use arrow::datatypes::{DataType, TimeUnit};
+    use arrow_schema::FieldRef;
     use datafusion_common::DataFusionError;
+    use datafusion_common::datatype::DataTypeExt;
     use std::error::Error;
     use std::path::PathBuf;
 
@@ -2735,7 +2737,7 @@ mod tests {
     struct MyTypePlanner {}
 
     impl TypePlanner for MyTypePlanner {
-        fn plan_type(&self, sql_type: &ast::DataType) -> Result<Option<DataType>> {
+        fn plan_type_field(&self, sql_type: &ast::DataType) -> Result<Option<FieldRef>> {
             match sql_type {
                 ast::DataType::Datetime(precision) => {
                     let precision = match precision {
@@ -2745,7 +2747,9 @@ mod tests {
                         None | Some(9) => TimeUnit::Nanosecond,
                         _ => unreachable!(),
                     };
-                    Ok(Some(DataType::Timestamp(precision, None)))
+                    Ok(Some(
+                        DataType::Timestamp(precision, None).into_nullable_field_ref(),
+                    ))
                 }
                 _ => Ok(None),
             }

datafusion/expr/src/planner.rs

@@ -434,10 +434,28 @@ pub trait TypePlanner: Debug + Send + Sync {
     /// Plan SQL [`sqlparser::ast::DataType`] to DataFusion [`DataType`]
     ///
     /// Returns None if not possible
+    #[deprecated(since = "53.0.0", note = "Use plan_type_field()")]
     fn plan_type(
         &self,
         _sql_type: &sqlparser::ast::DataType,
     ) -> Result<Option<DataType>> {
         Ok(None)
     }
+
+    /// Plan SQL [`sqlparser::ast::DataType`] to DataFusion [`FieldRef`]
+    ///
+    /// Returns None if not possible. Unlike [`Self::plan_type`], `plan_type_field()`
+    /// makes it possible to express extension types (e.g., `arrow.uuid`) or otherwise
+    /// insert metadata into the DataFusion type representation. The default implementation
+    /// falls back on [`Self::plan_type`] for backward compatibility and wraps the result
+    /// in a nullable field reference.
+    fn plan_type_field(
+        &self,
+        sql_type: &sqlparser::ast::DataType,
+    ) -> Result<Option<FieldRef>> {
+        #[expect(deprecated)]
+        Ok(self
+            .plan_type(sql_type)?
+            .map(|data_type| data_type.into_nullable_field_ref()))
+    }
 }

datafusion/functions/src/datetime/date_part.rs

@@ -28,10 +28,11 @@ use arrow::datatypes::DataType::{
 };
 use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
 use arrow::datatypes::{
-    DataType, Date32Type, Date64Type, Field, FieldRef, IntervalUnit as ArrowIntervalUnit,
-    TimeUnit,
+    ArrowTimestampType, DataType, Date32Type, Date64Type, Field, FieldRef,
+    IntervalUnit as ArrowIntervalUnit, TimeUnit, TimestampMicrosecondType,
+    TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
 };
-use chrono::{Datelike, NaiveDate, TimeZone, Utc};
+use chrono::{Datelike, NaiveDate};
 use datafusion_common::types::{NativeType, logical_date};
 
 use datafusion_common::{
@@ -337,37 +338,32 @@ fn date_to_scalar(date: NaiveDate, target_type: &DataType) -> Option<ScalarValue
 
         Timestamp(unit, tz_opt) => {
             let naive_midnight = date.and_hms_opt(0, 0, 0)?;
-
-            let utc_dt = if let Some(tz_str) = tz_opt {
-                let tz: Tz = tz_str.parse().ok()?;
-
-                let local = tz.from_local_datetime(&naive_midnight);
-
-                let local_dt = match local {
-                    chrono::offset::LocalResult::Single(dt) => dt,
-                    chrono::offset::LocalResult::Ambiguous(dt1, _dt2) => dt1,
-                    chrono::offset::LocalResult::None => local.earliest()?,
-                };
-
-                local_dt.with_timezone(&Utc)
-            } else {
-                Utc.from_utc_datetime(&naive_midnight)
-            };
+            let tz: Option<Tz> = tz_opt.clone().and_then(|s| s.parse().ok());
 
             match unit {
-                Second => {
-                    ScalarValue::TimestampSecond(Some(utc_dt.timestamp()), tz_opt.clone())
-                }
+                Second => ScalarValue::TimestampSecond(
+                    TimestampSecondType::from_naive_datetime(naive_midnight, tz.as_ref()),
+                    tz_opt.clone(),
+                ),
                 Millisecond => ScalarValue::TimestampMillisecond(
-                    Some(utc_dt.timestamp_millis()),
+                    TimestampMillisecondType::from_naive_datetime(
+                        naive_midnight,
+                        tz.as_ref(),
+                    ),
                     tz_opt.clone(),
                 ),
                 Microsecond => ScalarValue::TimestampMicrosecond(
-                    Some(utc_dt.timestamp_micros()),
+                    TimestampMicrosecondType::from_naive_datetime(
+                        naive_midnight,
+                        tz.as_ref(),
+                    ),
                     tz_opt.clone(),
                 ),
                 Nanosecond => ScalarValue::TimestampNanosecond(
-                    Some(utc_dt.timestamp_nanos_opt()?),
+                    TimestampNanosecondType::from_naive_datetime(
+                        naive_midnight,
+                        tz.as_ref(),
+                    ),
                     tz_opt.clone(),
                 ),
             }

datafusion/functions/src/string/common.rs

@@ -22,7 +22,7 @@ use std::sync::Arc;
 use crate::strings::make_and_append_view;
 use arrow::array::{
     Array, ArrayRef, GenericStringArray, GenericStringBuilder, NullBufferBuilder,
-    OffsetSizeTrait, StringBuilder, StringViewArray, new_null_array,
+    OffsetSizeTrait, StringViewArray, StringViewBuilder, new_null_array,
 };
 use arrow::buffer::{Buffer, ScalarBuffer};
 use arrow::datatypes::DataType;
@@ -358,10 +358,8 @@ where
             >(array, op)?)),
             DataType::Utf8View => {
                 let string_array = as_string_view_array(array)?;
-                let mut string_builder = StringBuilder::with_capacity(
-                    string_array.len(),
-                    string_array.get_array_memory_size(),
-                );
+                let mut string_builder =
+                    StringViewBuilder::with_capacity(string_array.len());
 
                 for str in string_array.iter() {
                     if let Some(str) = str {
@@ -386,7 +384,7 @@ where
             }
             ScalarValue::Utf8View(a) => {
                 let result = a.as_ref().map(|x| op(x));
-                Ok(ColumnarValue::Scalar(ScalarValue::Utf8(result)))
+                Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(result)))
             }
             other => exec_err!("Unsupported data type {other:?} for function {name}"),
         },

datafusion/functions/src/string/lower.rs

@@ -19,7 +19,6 @@ use arrow::datatypes::DataType;
 use std::any::Any;
 
 use crate::string::common::to_lower;
-use crate::utils::utf8_to_str_type;
 use datafusion_common::Result;
 use datafusion_common::types::logical_string;
 use datafusion_expr::{
@@ -82,7 +81,7 @@ impl ScalarUDFImpl for LowerFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        utf8_to_str_type(&arg_types[0], "lower")
+        Ok(arg_types[0].clone())
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
@@ -97,8 +96,7 @@ impl ScalarUDFImpl for LowerFunc {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::{Array, ArrayRef, StringArray};
-    use arrow::datatypes::DataType::Utf8;
+    use arrow::array::{Array, ArrayRef, StringArray, StringViewArray};
     use arrow::datatypes::Field;
     use datafusion_common::config::ConfigOptions;
     use std::sync::Arc;
@@ -111,7 +109,7 @@ mod tests {
             number_rows: input.len(),
             args: vec![ColumnarValue::Array(input)],
             arg_fields,
-            return_field: Field::new("f", Utf8, true).into(),
+            return_field: Field::new("f", expected.data_type().clone(), true).into(),
             config_options: Arc::new(ConfigOptions::default()),
         };
 
@@ -197,4 +195,21 @@ mod tests {
 
         to_lower(input, expected)
     }
+
+    #[test]
+    fn lower_utf8view() -> Result<()> {
+        let input = Arc::new(StringViewArray::from(vec![
+            Some("ARROW"),
+            None,
+            Some("TSCHÜSS"),
+        ])) as ArrayRef;
+
+        let expected = Arc::new(StringViewArray::from(vec![
+            Some("arrow"),
+            None,
+            Some("tschüss"),
+        ])) as ArrayRef;
+
+        to_lower(input, expected)
+    }
 }

datafusion/functions/src/string/upper.rs

@@ -16,7 +16,6 @@
 // under the License.
 
 use crate::string::common::to_upper;
-use crate::utils::utf8_to_str_type;
 use arrow::datatypes::DataType;
 use datafusion_common::Result;
 use datafusion_common::types::logical_string;
@@ -81,7 +80,7 @@ impl ScalarUDFImpl for UpperFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        utf8_to_str_type(&arg_types[0], "upper")
+        Ok(arg_types[0].clone())
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
@@ -96,8 +95,7 @@ impl ScalarUDFImpl for UpperFunc {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::{Array, ArrayRef, StringArray};
-    use arrow::datatypes::DataType::Utf8;
+    use arrow::array::{Array, ArrayRef, StringArray, StringViewArray};
     use arrow::datatypes::Field;
     use datafusion_common::config::ConfigOptions;
     use std::sync::Arc;
@@ -110,7 +108,7 @@ mod tests {
             number_rows: input.len(),
             args: vec![ColumnarValue::Array(input)],
             arg_fields: vec![arg_field],
-            return_field: Field::new("f", Utf8, true).into(),
+            return_field: Field::new("f", expected.data_type().clone(), true).into(),
             config_options: Arc::new(ConfigOptions::default()),
         };
 
@@ -196,4 +194,21 @@ mod tests {
 
         to_upper(input, expected)
     }
+
+    #[test]
+    fn upper_utf8view() -> Result<()> {
+        let input = Arc::new(StringViewArray::from(vec![
+            Some("arrow"),
+            None,
+            Some("tschüß"),
+        ])) as ArrayRef;
+
+        let expected = Arc::new(StringViewArray::from(vec![
+            Some("ARROW"),
+            None,
+            Some("TSCHÜSS"),
+        ])) as ArrayRef;
+
+        to_upper(input, expected)
+    }
 }