fix: port GreptimeDB correctness fixes to DataFusion 54

datafusion/common/src/scalar/mod.rs

@@ -148,6 +148,30 @@ pub fn date_to_timestamp_multiplier(
     }
 }
 
+/// Returns the multiplier that converts the input timestamp representation into
+/// the desired timestamp unit, if the conversion requires a multiplication that
+/// can overflow an `i64`.
+pub fn timestamp_to_timestamp_multiplier(
+    source_type: &DataType,
+    target_type: &DataType,
+) -> Option<i64> {
+    let (DataType::Timestamp(source_unit, _), DataType::Timestamp(target_unit, _)) =
+        (source_type, target_type)
+    else {
+        return None;
+    };
+
+    match (source_unit, target_unit) {
+        (TimeUnit::Second, TimeUnit::Millisecond) => Some(1_000),
+        (TimeUnit::Second, TimeUnit::Microsecond) => Some(1_000_000),
+        (TimeUnit::Second, TimeUnit::Nanosecond) => Some(1_000_000_000),
+        (TimeUnit::Millisecond, TimeUnit::Microsecond) => Some(1_000),
+        (TimeUnit::Millisecond, TimeUnit::Nanosecond) => Some(1_000_000),
+        (TimeUnit::Microsecond, TimeUnit::Nanosecond) => Some(1_000),
+        _ => None,
+    }
+}
+
 /// Ensures the provided value can be represented as a timestamp with the given
 /// multiplier. Returns an [`DataFusionError::Execution`] when the converted
 /// value would overflow the timestamp range.
@@ -4224,10 +4248,21 @@ impl ScalarValue {
         cast_options: &CastOptions<'static>,
     ) -> Result<Self> {
         let source_type = self.data_type();
+
+        // Temporal casts that increase precision can overflow when the value is
+        // scaled to the target unit. For safe casts, return NULL for the
+        // overflowing scalar; otherwise preserve the regular cast error.
         if let Some(multiplier) = date_to_timestamp_multiplier(&source_type, target_type)
-            && let Some(value) = self.date_scalar_value_as_i64()
+            .or_else(|| timestamp_to_timestamp_multiplier(&source_type, target_type))
+            && let Some(value) = self.temporal_scalar_value_as_i64()
         {
-            ensure_timestamp_in_bounds(value, multiplier, &source_type, target_type)?;
+            if cast_options.safe {
+                if multiplier > 1 && value.checked_mul(multiplier).is_none() {
+                    return ScalarValue::try_new_null(target_type);
+                }
+            } else {
+                ensure_timestamp_in_bounds(value, multiplier, &source_type, target_type)?;
+            }
         }
 
         let scalar_array = self.to_array()?;
@@ -4247,10 +4282,14 @@ impl ScalarValue {
         ScalarValue::try_from_array(&cast_arr, 0)
     }
 
-    fn date_scalar_value_as_i64(&self) -> Option<i64> {
+    fn temporal_scalar_value_as_i64(&self) -> Option<i64> {
         match self {
             ScalarValue::Date32(Some(value)) => Some(i64::from(*value)),
             ScalarValue::Date64(Some(value)) => Some(*value),
+            ScalarValue::TimestampSecond(Some(value), _)
+            | ScalarValue::TimestampMillisecond(Some(value), _)
+            | ScalarValue::TimestampMicrosecond(Some(value), _)
+            | ScalarValue::TimestampNanosecond(Some(value), _) => Some(*value),
             _ => None,
         }
     }
@@ -9929,6 +9968,58 @@ mod tests {
         );
     }
 
+    #[test]
+    fn cast_timestamp_to_timestamp_overflow_returns_error() {
+        let scalar = ScalarValue::TimestampSecond(Some(i64::MAX), None);
+        let err = scalar
+            .cast_to(&DataType::Timestamp(TimeUnit::Nanosecond, None))
+            .expect_err("expected cast to fail");
+        assert!(
+            err.to_string()
+                .contains("converted value exceeds the representable i64 range"),
+            "unexpected error: {err}"
+        );
+    }
+
+    #[test]
+    fn cast_timestamp_to_finer_unit_safe_overflow_returns_typed_null() {
+        let timezone = Arc::<str>::from("UTC");
+        let scalar = ScalarValue::TimestampMillisecond(
+            Some(i64::MAX),
+            Some(Arc::clone(&timezone)),
+        );
+        let target_type =
+            DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::clone(&timezone)));
+
+        for cast_options in [
+            DEFAULT_CAST_OPTIONS,
+            CastOptions {
+                safe: false,
+                ..DEFAULT_CAST_OPTIONS
+            },
+        ] {
+            let err = scalar
+                .cast_to_with_options(&target_type, &cast_options)
+                .expect_err("unsafe cast should fail");
+            assert!(
+                err.to_string()
+                    .contains("converted value exceeds the representable i64 range"),
+                "unexpected error: {err}"
+            );
+        }
+
+        let cast_options = CastOptions {
+            safe: true,
+            ..DEFAULT_CAST_OPTIONS
+        };
+        assert_eq!(
+            scalar
+                .cast_to_with_options(&target_type, &cast_options)
+                .unwrap(),
+            ScalarValue::TimestampNanosecond(None, Some(timezone))
+        );
+    }
+
     #[test]
     fn null_dictionary_scalar_produces_null_dictionary_array() {
         let dictionary_scalar = ScalarValue::Dictionary(

datafusion/common/src/stats.rs

@@ -539,6 +539,9 @@ impl Statistics {
         skip: usize,
         n_partitions: usize,
     ) -> Result<Self> {
+        if fetch.is_none() && skip == 0 {
+            return Ok(self);
+        }
         let fetch_val = fetch.unwrap_or(usize::MAX);
 
         // Get the ratio of rows after / rows before on a per-partition basis
@@ -592,30 +595,30 @@ impl Statistics {
                 ..
             } => check_num_rows(fetch.and_then(|v| v.checked_mul(n_partitions)), false),
         };
-        let ratio: f64 = match (num_rows_before, self.num_rows) {
+        let ratio: Option<f64> = match (num_rows_before, self.num_rows) {
             (
                 Precision::Exact(nr_before) | Precision::Inexact(nr_before),
                 Precision::Exact(nr_after) | Precision::Inexact(nr_after),
             ) => {
                 if nr_before == 0 {
-                    0.0
+                    Some(0.0)
                 } else {
-                    nr_after as f64 / nr_before as f64
+                    Some(nr_after as f64 / nr_before as f64)
                 }
             }
-            _ => 0.0,
+            _ => None,
         };
         self.column_statistics = self
             .column_statistics
             .into_iter()
             .map(|cs| {
                 let mut cs = cs.to_inexact();
                 // Scale byte_size by the row ratio
-                cs.byte_size = match cs.byte_size {
-                    Precision::Exact(n) | Precision::Inexact(n) => {
+                cs.byte_size = match (cs.byte_size, ratio) {
+                    (Precision::Exact(n) | Precision::Inexact(n), Some(ratio)) => {
                         Precision::Inexact((n as f64 * ratio) as usize)
                     }
-                    Precision::Absent => Precision::Absent,
+                    _ => Precision::Absent,
                 };
                 // NDV can never exceed the number of rows
                 if let Some(&rows) = self.num_rows.get_value() {
@@ -637,11 +640,11 @@ impl Statistics {
             Some(sum) => Precision::Inexact(sum),
             None => {
                 // Fall back to scaling original total_byte_size if not all columns have byte_size
-                match &self.total_byte_size {
-                    Precision::Exact(n) | Precision::Inexact(n) => {
+                match (&self.total_byte_size, ratio) {
+                    (Precision::Exact(n) | Precision::Inexact(n), Some(ratio)) => {
                         Precision::Inexact((*n as f64 * ratio) as usize)
                     }
-                    Precision::Absent => Precision::Absent,
+                    _ => Precision::Absent,
                 }
             }
         };
@@ -2370,6 +2373,34 @@ mod tests {
         assert_eq!(result.total_byte_size, Precision::Exact(800));
     }
 
+    #[test]
+    fn test_with_fetch_no_limit_preserves_absent_num_rows() {
+        let original_stats = Statistics {
+            num_rows: Precision::Absent,
+            total_byte_size: Precision::Exact(800),
+            column_statistics: vec![col_stats_i64(10)],
+        };
+
+        let result = original_stats.clone().with_fetch(None, 0, 1).unwrap();
+
+        assert_eq!(result, original_stats);
+    }
+
+    #[test]
+    fn test_with_fetch_absent_num_rows_does_not_zero_byte_size() {
+        let original_stats = Statistics {
+            num_rows: Precision::Absent,
+            total_byte_size: Precision::Exact(800),
+            column_statistics: vec![col_stats_i64(10)],
+        };
+
+        let result = original_stats.with_fetch(Some(1), 0, 1).unwrap();
+
+        assert_eq!(result.num_rows, Precision::Inexact(1));
+        assert_eq!(result.total_byte_size, Precision::Absent);
+        assert_eq!(result.column_statistics[0].byte_size, Precision::Absent);
+    }
+
     #[test]
     fn test_with_fetch_with_skip() {
         // Test with both skip and fetch

datafusion/core/tests/physical_optimizer/filter_pushdown.rs

@@ -42,7 +42,9 @@ use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_expr::ScalarUDF;
 use datafusion_functions::math::random::RandomFunc;
 use datafusion_functions_aggregate::{count::count_udaf, min_max::min_udaf};
-use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col};
+use datafusion_physical_expr::{
+    LexOrdering, PhysicalSortExpr, expressions::col, utils::conjunction,
+};
 use datafusion_physical_expr::{
     Partitioning, ScalarFunctionExpr, aggregate::AggregateExprBuilder,
 };
@@ -738,6 +740,65 @@ fn test_pushdown_through_aggregates_on_grouping_columns() {
     );
 }
 
+#[test]
+fn test_pushdown_through_aggregates_preserves_parent_filter_order() {
+    // AggregateExec may push filters on grouping columns to its input, but must
+    // keep filters on aggregate outputs above itself. The parent-filter result
+    // order must match the incoming filter order, otherwise an unsupported
+    // aggregate-output filter can be reported as pushed down and removed.
+    let scan = TestScanBuilder::new(schema()).with_support(true).build();
+
+    let aggregate_expr = vec![
+        AggregateExprBuilder::new(count_udaf(), vec![col("a", &schema()).unwrap()])
+            .schema(schema())
+            .alias("cnt")
+            .build()
+            .map(Arc::new)
+            .unwrap(),
+    ];
+    let group_by = PhysicalGroupBy::new_single(vec![
+        (col("a", &schema()).unwrap(), "a".to_string()),
+        (col("b", &schema()).unwrap(), "b".to_string()),
+    ]);
+    let aggregate = Arc::new(
+        AggregateExec::try_new(
+            AggregateMode::Final,
+            group_by,
+            aggregate_expr,
+            vec![None],
+            scan,
+            schema(),
+        )
+        .unwrap(),
+    );
+
+    let aggregate_schema = aggregate.schema();
+    let aggregate_output_filter = col_lit_predicate(
+        "cnt",
+        ScalarValue::Int64(Some(1)),
+        aggregate_schema.as_ref(),
+    );
+    let grouping_key_filter = col_lit_predicate("b", "bar", aggregate_schema.as_ref());
+    let predicate = conjunction(vec![aggregate_output_filter, grouping_key_filter]);
+    let plan = Arc::new(FilterExec::try_new(predicate, aggregate).unwrap());
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, FilterPushdown::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - FilterExec: cnt@2 = 1 AND b@1 = bar
+        -   AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt]
+        -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
+      output:
+        Ok:
+          - FilterExec: cnt@2 = 1
+          -   AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=PartiallySorted([1])
+          -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar
+    "
+    );
+}
+
 /// Test various combinations of handling of child pushdown results
 /// in an ExecutionPlan in combination with support/not support in a DataSource.
 #[test]