Merge branch 'main' into implement_groups_accumulator_count_distinct_primitive_types

Author: coderfender

datafusion/expr-common/src/casts.rs

@@ -58,7 +58,28 @@ pub fn is_supported_type(data_type: &DataType) -> bool {
         || is_supported_binary_type(data_type)
 }
 
-/// Returns true if unwrap_cast_in_comparison support this numeric type
+fn is_date_type(data_type: &DataType) -> bool {
+    matches!(data_type, DataType::Date32 | DataType::Date64)
+}
+
+/// Returns true when unwrapping a date/timestamp cast could change comparison
+/// semantics.
+///
+/// A `Date` stores only a calendar day, while a `Timestamp` stores a specific
+/// instant or wall-clock time. `Timestamp -> Date` is lossy because it drops the
+/// time-of-day. `Date -> Timestamp` is also lossy in this optimizer context
+/// because there is no unique inverse: converting a date to a timestamp has to
+/// invent a time component such as midnight.
+///
+/// For example, `CAST(ts AS DATE) = DATE '2024-01-01'` means "any timestamp
+/// during that day", but unwrapping it to `ts = TIMESTAMP '2024-01-01
+/// 00:00:00'` matches only midnight.
+fn is_lossy_temporal_cast(from_type: &DataType, to_type: &DataType) -> bool {
+    (is_date_type(from_type) && to_type.is_temporal())
+        || (is_date_type(to_type) && from_type.is_temporal())
+}
+
+/// Returns true if unwrap_cast_in_comparison supports this numeric type
 fn is_supported_numeric_type(data_type: &DataType) -> bool {
     matches!(
         data_type,
@@ -70,6 +91,8 @@ fn is_supported_numeric_type(data_type: &DataType) -> bool {
             | DataType::Int16
             | DataType::Int32
             | DataType::Int64
+            | DataType::Date32
+            | DataType::Date64
             | DataType::Decimal32(_, _)
             | DataType::Decimal64(_, _)
             | DataType::Decimal128(_, _)
@@ -107,6 +130,10 @@ fn try_cast_numeric_literal(
         return None;
     }
 
+    if is_lossy_temporal_cast(&lit_data_type, target_type) {
+        return None;
+    }
+
     let mul = match target_type {
         DataType::UInt8
         | DataType::UInt16
@@ -115,7 +142,9 @@ fn try_cast_numeric_literal(
         | DataType::Int8
         | DataType::Int16
         | DataType::Int32
-        | DataType::Int64 => 1_i128,
+        | DataType::Int64
+        | DataType::Date32
+        | DataType::Date64 => 1_i128,
         DataType::Timestamp(_, _) => 1_i128,
         DataType::Decimal32(_, scale) => 10_i128.pow(*scale as u32),
         DataType::Decimal64(_, scale) => 10_i128.pow(*scale as u32),
@@ -129,8 +158,8 @@ fn try_cast_numeric_literal(
         DataType::UInt64 => (u64::MIN as i128, u64::MAX as i128),
         DataType::Int8 => (i8::MIN as i128, i8::MAX as i128),
         DataType::Int16 => (i16::MIN as i128, i16::MAX as i128),
-        DataType::Int32 => (i32::MIN as i128, i32::MAX as i128),
-        DataType::Int64 => (i64::MIN as i128, i64::MAX as i128),
+        DataType::Int32 | DataType::Date32 => (i32::MIN as i128, i32::MAX as i128),
+        DataType::Int64 | DataType::Date64 => (i64::MIN as i128, i64::MAX as i128),
         DataType::Timestamp(_, _) => (i64::MIN as i128, i64::MAX as i128),
         DataType::Decimal32(precision, _) => (
             // Different precision for decimal32 can store different range of value.
@@ -164,6 +193,8 @@ fn try_cast_numeric_literal(
         ScalarValue::UInt16(Some(v)) => (*v as i128).checked_mul(mul),
         ScalarValue::UInt32(Some(v)) => (*v as i128).checked_mul(mul),
         ScalarValue::UInt64(Some(v)) => (*v as i128).checked_mul(mul),
+        ScalarValue::Date32(Some(v)) => (*v as i128).checked_mul(mul),
+        ScalarValue::Date64(Some(v)) => (*v as i128).checked_mul(mul),
         ScalarValue::TimestampSecond(Some(v), _) => (*v as i128).checked_mul(mul),
         ScalarValue::TimestampMillisecond(Some(v), _) => (*v as i128).checked_mul(mul),
         ScalarValue::TimestampMicrosecond(Some(v), _) => (*v as i128).checked_mul(mul),
@@ -241,6 +272,8 @@ fn try_cast_numeric_literal(
                     DataType::Int16 => ScalarValue::Int16(Some(value as i16)),
                     DataType::Int32 => ScalarValue::Int32(Some(value as i32)),
                     DataType::Int64 => ScalarValue::Int64(Some(value as i64)),
+                    DataType::Date32 => ScalarValue::Date32(Some(value as i32)),
+                    DataType::Date64 => ScalarValue::Date64(Some(value as i64)),
                     DataType::UInt8 => ScalarValue::UInt8(Some(value as u8)),
                     DataType::UInt16 => ScalarValue::UInt16(Some(value as u16)),
                     DataType::UInt32 => ScalarValue::UInt32(Some(value as u32)),
@@ -700,6 +733,33 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_try_cast_to_type_date_timestamp_lossy_not_allowed() {
+        expect_cast(
+            ScalarValue::Date32(Some(1)),
+            DataType::Timestamp(TimeUnit::Second, None),
+            ExpectedCast::NoValue,
+        );
+
+        expect_cast(
+            ScalarValue::Date64(Some(86_400_000)),
+            DataType::Timestamp(TimeUnit::Millisecond, None),
+            ExpectedCast::NoValue,
+        );
+
+        expect_cast(
+            ScalarValue::TimestampSecond(Some(86_400), None),
+            DataType::Date32,
+            ExpectedCast::NoValue,
+        );
+
+        expect_cast(
+            ScalarValue::TimestampMillisecond(Some(86_400_000), None),
+            DataType::Date64,
+            ExpectedCast::NoValue,
+        );
+    }
+
     #[test]
     fn test_try_cast_to_type_unsupported() {
         // int64 to list

datafusion/expr-common/src/type_coercion/binary.rs

@@ -378,6 +378,16 @@ fn math_decimal_coercion(
             let (lhs_type, value_type) = math_decimal_coercion(lhs_type, value_type)?;
             Some((lhs_type, value_type))
         }
+        (RunEndEncoded(_, field), _) => {
+            let (value_type, rhs_type) =
+                math_decimal_coercion(field.data_type(), rhs_type)?;
+            Some((value_type, rhs_type))
+        }
+        (_, RunEndEncoded(_, field)) => {
+            let (lhs_type, value_type) =
+                math_decimal_coercion(lhs_type, field.data_type())?;
+            Some((lhs_type, value_type))
+        }
         (
             Null,
             Decimal32(_, _) | Decimal64(_, _) | Decimal128(_, _) | Decimal256(_, _),
@@ -1414,6 +1424,15 @@ fn mathematics_numerical_coercion(
         (_, Dictionary(_, value_type)) => {
             mathematics_numerical_coercion(lhs_type, value_type)
         }
+        (RunEndEncoded(_, lhs_field), RunEndEncoded(_, rhs_field)) => {
+            mathematics_numerical_coercion(lhs_field.data_type(), rhs_field.data_type())
+        }
+        (RunEndEncoded(_, field), _) => {
+            mathematics_numerical_coercion(field.data_type(), rhs_type)
+        }
+        (_, RunEndEncoded(_, field)) => {
+            mathematics_numerical_coercion(lhs_type, field.data_type())
+        }
         _ => numerical_coercion(lhs_type, rhs_type),
     }
 }
@@ -1493,6 +1512,15 @@ fn both_numeric_or_null_and_numeric(lhs_type: &DataType, rhs_type: &DataType) ->
         (_, Dictionary(_, value_type)) => {
             lhs_type.is_numeric() && value_type.is_numeric()
         }
+        (RunEndEncoded(_, lhs_field), RunEndEncoded(_, rhs_field)) => {
+            lhs_field.data_type().is_numeric() && rhs_field.data_type().is_numeric()
+        }
+        (RunEndEncoded(_, field), _) => {
+            field.data_type().is_numeric() && rhs_type.is_numeric()
+        }
+        (_, RunEndEncoded(_, field)) => {
+            lhs_type.is_numeric() && field.data_type().is_numeric()
+        }
         _ => lhs_type.is_numeric() && rhs_type.is_numeric(),
     }
 }

datafusion/expr-common/src/type_coercion/binary/tests/run_end_encoded.rs

@@ -16,11 +16,17 @@
 // under the License.
 
 use super::*;
+use DataType::*;
+
+fn ree(value_type: DataType) -> DataType {
+    RunEndEncoded(
+        Arc::new(Field::new("run_ends", Int32, false)),
+        Arc::new(Field::new("values", value_type, false)),
+    )
+}
 
 #[test]
 fn test_ree_type_coercion() {
-    use DataType::*;
-
     let lhs_type = RunEndEncoded(
         Arc::new(Field::new("run_ends", Int8, false)),
         Arc::new(Field::new("values", Int32, false)),
@@ -97,3 +103,29 @@ fn test_ree_type_coercion() {
         Some(rhs_type.clone())
     );
 }
+
+#[test]
+fn test_ree_arithmetic_coercion() -> Result<()> {
+    test_coercion_binary_rule!(ree(Int64), Int64, Operator::Plus, Int64);
+    test_coercion_binary_rule!(Int64, ree(Int64), Operator::Multiply, Int64);
+    test_coercion_binary_rule!(ree(Int32), ree(Int64), Operator::Plus, Int64);
+
+    // Decimal unwrapping through math_decimal_coercion
+    let (lhs, rhs) =
+        BinaryTypeCoercer::new(&ree(Decimal128(10, 2)), &Operator::Plus, &Int32)
+            .get_input_types()?;
+    assert_eq!(lhs, Decimal128(10, 2));
+    assert_eq!(rhs, Decimal128(10, 0));
+
+    let (lhs, rhs) =
+        BinaryTypeCoercer::new(&Int32, &Operator::Plus, &ree(Decimal128(10, 2)))
+            .get_input_types()?;
+    assert_eq!(lhs, Decimal128(10, 0));
+    assert_eq!(rhs, Decimal128(10, 2));
+
+    let result =
+        BinaryTypeCoercer::new(&ree(Utf8), &Operator::Plus, &Int32).get_input_types();
+    assert!(result.is_err());
+
+    Ok(())
+}

datafusion/functions/benches/left_right.rs

@@ -16,40 +16,40 @@
 // under the License.
 
 use std::hint::black_box;
+use std::ops::Range;
 use std::sync::Arc;
 
 use arrow::array::{ArrayRef, Int64Array};
 use arrow::datatypes::{DataType, Field};
 use arrow::util::bench_util::{
     create_string_array_with_len, create_string_view_array_with_len,
 };
-use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use criterion::{Criterion, criterion_group, criterion_main};
 use datafusion_common::config::ConfigOptions;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
 use datafusion_functions::unicode::{left, right};
 
+const BATCH_SIZE: usize = 8192;
+
 fn create_args(
-    size: usize,
     str_len: usize,
-    use_negative: bool,
+    n_range: Range<i64>,
     is_string_view: bool,
 ) -> Vec<ColumnarValue> {
     let string_arg = if is_string_view {
         ColumnarValue::Array(Arc::new(create_string_view_array_with_len(
-            size, 0.1, str_len, true,
+            BATCH_SIZE, 0.1, str_len, true,
         )))
     } else {
         ColumnarValue::Array(Arc::new(create_string_array_with_len::<i32>(
-            size, 0.1, str_len,
+            BATCH_SIZE, 0.1, str_len,
         )))
     };
 
-    // For negative n, we want to trigger the double-iteration code path
-    let n_values: Vec<i64> = if use_negative {
-        (0..size).map(|i| -((i % 10 + 1) as i64)).collect()
-    } else {
-        (0..size).map(|i| (i % 10 + 1) as i64).collect()
-    };
+    let n_span = (n_range.end - n_range.start) as usize;
+    let n_values: Vec<i64> = (0..BATCH_SIZE)
+        .map(|i| n_range.start + (i % n_span) as i64)
+        .collect();
     let n_array = Arc::new(Int64Array::from(n_values));
 
     vec![
@@ -59,68 +59,55 @@ fn create_args(
 }
 
 fn criterion_benchmark(c: &mut Criterion) {
-    let left_function = left();
-    let right_function = right();
+    // Short results (1-10 chars) produce inline StringView entries (≤12 bytes).
+    // Long results (20-29 chars) produce out-of-line entries.
+    let cases = [
+        ("short_result", 32, 1..11_i64),
+        ("long_result", 32, 20..30_i64),
+    ];
 
-    for function in [left_function, right_function] {
-        for is_string_view in [false, true] {
-            for is_negative in [false, true] {
-                for size in [1024, 4096] {
-                    let function_name = function.name();
-                    let mut group =
-                        c.benchmark_group(format!("{function_name} size={size}"));
+    for function in [left(), right()] {
+        let mut group = c.benchmark_group(function.name().to_string());
 
-                    let bench_name = format!(
-                        "{} {} n",
-                        if is_string_view {
-                            "string_view_array"
-                        } else {
-                            "string_array"
-                        },
-                        if is_negative { "negative" } else { "positive" },
-                    );
-                    let return_type = if is_string_view {
-                        DataType::Utf8View
-                    } else {
-                        DataType::Utf8
-                    };
-
-                    let args = create_args(size, 32, is_negative, is_string_view);
-                    group.bench_function(BenchmarkId::new(bench_name, size), |b| {
-                        let arg_fields = args
-                            .iter()
-                            .enumerate()
-                            .map(|(idx, arg)| {
-                                Field::new(format!("arg_{idx}"), arg.data_type(), true)
-                                    .into()
-                            })
-                            .collect::<Vec<_>>();
-                        let config_options = Arc::new(ConfigOptions::default());
+        for is_string_view in [false, true] {
+            let array_type = if is_string_view {
+                "string_view"
+            } else {
+                "string"
+            };
 
-                        b.iter(|| {
-                            black_box(
-                                function
-                                    .invoke_with_args(ScalarFunctionArgs {
-                                        args: args.clone(),
-                                        arg_fields: arg_fields.clone(),
-                                        number_rows: size,
-                                        return_field: Field::new(
-                                            "f",
-                                            return_type.clone(),
-                                            true,
-                                        )
-                                        .into(),
-                                        config_options: Arc::clone(&config_options),
-                                    })
-                                    .expect("should work"),
-                            )
-                        })
-                    });
+            for (case_name, str_len, n_range) in &cases {
+                let bench_name = format!("{array_type} {case_name}");
+                let args = create_args(*str_len, n_range.clone(), is_string_view);
+                let arg_fields: Vec<_> = args
+                    .iter()
+                    .enumerate()
+                    .map(|(idx, arg)| {
+                        Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
+                    })
+                    .collect();
+                let config_options = Arc::new(ConfigOptions::default());
+                let return_field = Field::new("f", DataType::Utf8View, true).into();
 
-                    group.finish();
-                }
+                group.bench_function(&bench_name, |b| {
+                    b.iter(|| {
+                        black_box(
+                            function
+                                .invoke_with_args(ScalarFunctionArgs {
+                                    args: args.clone(),
+                                    arg_fields: arg_fields.clone(),
+                                    number_rows: BATCH_SIZE,
+                                    return_field: Arc::clone(&return_field),
+                                    config_options: Arc::clone(&config_options),
+                                })
+                                .expect("should work"),
+                        )
+                    })
+                });
             }
         }
+
+        group.finish();
     }
 }