feat(query): propagate column statistics through UNION ALL

src/common/statistics/src/stat_estimate.rs

@@ -435,6 +435,18 @@ impl StatCount {
             }
         }
     }
+
+    pub fn sum(left: Self, right: Self) -> Self {
+        match (left, right) {
+            (StatCount::Exact(left), StatCount::Exact(right)) => {
+                StatCount::exact(left.saturating_add(right))
+            }
+            _ => StatCount::estimate(
+                left.expected() + right.expected(),
+                left.upper() + right.upper(),
+            ),
+        }
+    }
 }
 
 #[cfg(test)]

src/query/expression/src/stat_evaluator.rs

@@ -18,14 +18,18 @@ use std::collections::HashMap;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 
+use super::Cast;
 use super::ColumnIndex;
+use super::ColumnRef;
+use super::ConstantFolder;
 use super::Expr;
 use super::FunctionCall;
 use super::FunctionContext;
 use super::FunctionDomain;
 use super::FunctionEval;
 use super::FunctionRegistry;
 use super::Scalar;
+use super::conversion::classify_conversion;
 use super::function_stat::DeriveStat;
 use super::stat_distribution::ArgStat;
 use super::stat_distribution::BorrowedDistribution;
@@ -93,10 +97,70 @@ impl<'a> StatEvaluator<'a> {
             Expr::FunctionCall(call) => Ok(self
                 .eval_function_call(call, input_stats)?
                 .map(CowStat::Owned)),
-            Expr::Cast(_) | Expr::LambdaFunctionCall(_) => Ok(None),
+            Expr::Cast(cast) => Ok(self.eval_cast(cast, input_stats)?.map(CowStat::Owned)),
+            Expr::LambdaFunctionCall(_) => Ok(None),
         }
     }
 
+    fn eval_cast<'s, I: ColumnIndex>(
+        &'a self,
+        cast: &Cast<I>,
+        input_stats: &'s HashMap<I, ArgStat<'_>>,
+    ) -> Result<Option<ReturnStat>> {
+        let src_type = cast.expr.data_type();
+        if cast.is_try || !classify_conversion(src_type, &cast.dest_type).is_lossless_injective() {
+            return Ok(None);
+        }
+
+        let Some(input) = self.eval(&cast.expr, input_stats)? else {
+            return Ok(None);
+        };
+        let input = input.as_ref();
+
+        // Reuse the cast domain implementation without making statistics
+        // evaluation depend on the physical value evaluator. The synthetic
+        // column represents the already-derived statistics of the inner expr.
+        let expr = Expr::Cast(Cast {
+            span: cast.span,
+            is_try: false,
+            expr: Box::new(Expr::ColumnRef(ColumnRef {
+                span: cast.span,
+                id: 0,
+                data_type: src_type.clone(),
+                display_name: String::new(),
+            })),
+            dest_type: cast.dest_type.clone(),
+        });
+        let input_domains = HashMap::from([(0, input.domain.clone())]);
+        let (_, Some(domain)) = ConstantFolder::fold_with_domain(
+            &expr,
+            &input_domains,
+            self.func_ctx,
+            self.fn_registry,
+        ) else {
+            return Ok(None);
+        };
+
+        let stat = ReturnStat {
+            domain,
+            // A lossless injective cast preserves distinct values and NULLs.
+            // Histograms are typed, so their boundaries cannot be reused.
+            ndv: input.ndv,
+            null_count: input.null_count,
+            distribution: OwnedDistribution::Unknown,
+        };
+        if let Err(msg) = stat.check_consistency_with_type(Some(&cast.dest_type)) {
+            if cfg!(debug_assertions) {
+                return Err(ErrorCode::Internal(format!(
+                    "Failed to derive statistics for cast: {msg}"
+                )));
+            }
+            log::warn!(msg; "Derived invalid cast statistics");
+            return Ok(None);
+        }
+        Ok(Some(stat))
+    }
+
     fn eval_function_call<'s, I: ColumnIndex>(
         &'a self,
         call: &FunctionCall<I>,
@@ -219,6 +283,8 @@ mod tests {
 
     use super::*;
     use crate::types::DataType;
+    use crate::types::NumberDataType;
+    use crate::types::number::NumberScalar;
 
     #[test]
     fn test_constant_null_uses_exact_input_cardinality() {
@@ -237,4 +303,75 @@ mod tests {
 
         assert_eq!(stat.null_count, StatCount::exact(7));
     }
+
+    #[test]
+    fn test_lossless_cast_preserves_basic_statistics() {
+        let src_type = DataType::Number(NumberDataType::UInt8);
+        let dest_type = src_type.clone().wrap_nullable();
+        let expr = Expr::Cast(Cast {
+            span: None,
+            is_try: false,
+            expr: Box::new(Expr::ColumnRef(ColumnRef {
+                span: None,
+                id: 0,
+                data_type: src_type.clone(),
+                display_name: "c0".to_string(),
+            })),
+            dest_type: dest_type.clone(),
+        });
+        let domain = crate::Domain::from_min_max(
+            Scalar::Number(NumberScalar::UInt8(1)),
+            Scalar::Number(NumberScalar::UInt8(3)),
+            &src_type,
+        );
+        let input_stats = HashMap::from([(0, ArgStat {
+            domain,
+            ndv: NdvEstimate::exact(3.0),
+            null_count: StatCount::exact(0),
+            distribution: BorrowedDistribution::Unknown,
+        })]);
+
+        let stat = StatEvaluator::run(
+            &expr,
+            &FunctionContext::default(),
+            &FunctionRegistry::empty(),
+            StatCardinality::exact(3),
+            &input_stats,
+        )
+        .unwrap()
+        .unwrap()
+        .into_owned();
+
+        assert!(stat.domain.matches_data_type(&dest_type));
+        assert_eq!(stat.ndv, NdvEstimate::exact(3.0));
+        assert_eq!(stat.null_count, StatCount::exact(0));
+        assert!(matches!(stat.distribution, OwnedDistribution::Unknown));
+    }
+
+    #[test]
+    fn test_lossy_cast_is_not_derived() {
+        let expr = Expr::Cast(Cast {
+            span: None,
+            is_try: false,
+            expr: Box::new(Expr::ColumnRef(ColumnRef {
+                span: None,
+                id: 0,
+                data_type: DataType::Number(NumberDataType::Int64),
+                display_name: "c0".to_string(),
+            })),
+            dest_type: DataType::Number(NumberDataType::UInt8),
+        });
+
+        assert!(
+            StatEvaluator::run(
+                &expr,
+                &FunctionContext::default(),
+                &FunctionRegistry::empty(),
+                StatCardinality::exact(3),
+                &HashMap::new(),
+            )
+            .unwrap()
+            .is_none()
+        );
+    }
 }

src/query/sql/src/planner/optimizer/ir/mod.rs

@@ -52,3 +52,4 @@ pub use stats::MAX_SELECTIVITY;
 pub use stats::SelectivityEstimator;
 pub use stats::TopNSet;
 pub use stats::UniformSampleSet;
+pub(crate) use stats::finite_range_ndv_upper;

src/query/sql/src/planner/optimizer/ir/stats/column_stat.rs

@@ -31,7 +31,7 @@ pub type ColumnStatSet = HashMap<Symbol, ColumnStat>;
 pub type TopNSet = HashMap<Symbol, ColumnTopN>;
 pub type CountMinSketchSet = HashMap<Symbol, ColumnCountMinSketch>;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 /// Statistics information of a column
 pub struct ColumnStat {
     /// Min value of the column

src/query/sql/src/planner/optimizer/ir/stats/constraint.rs

@@ -21,6 +21,7 @@ use databend_common_statistics::HistogramRangeBounds;
 use databend_common_statistics::NdvEstimate;
 use databend_common_statistics::StatCount;
 
+use super::finite_range_ndv_upper;
 use crate::optimizer::ir::ColumnStat;
 use crate::plans::ComparisonOp;
 
@@ -138,24 +139,6 @@ fn apply_range_bounds(column_stat: &mut ColumnStat, new_min: Datum, new_max: Dat
     Ok(())
 }
 
-fn finite_range_ndv_upper(min: &Datum, max: &Datum) -> Option<f64> {
-    if min == max {
-        return Some(1.0);
-    }
-    match (min, max) {
-        (Datum::Bool(false), Datum::Bool(true)) => Some(2.0),
-        (Datum::Int(min), Datum::Int(max)) => max
-            .checked_sub(*min)
-            .and_then(|diff| diff.checked_add(1))
-            .map(|value| value as f64),
-        (Datum::UInt(min), Datum::UInt(max)) => max
-            .checked_sub(*min)
-            .and_then(|diff| diff.checked_add(1))
-            .map(|value| value as f64),
-        _ => None,
-    }
-}
-
 pub(super) fn clear_for_empty_result(column_stat: &mut ColumnStat) {
     column_stat.ndv = NdvEstimate::exact(0.0);
     column_stat.null_count = StatCount::exact(0);

src/query/sql/src/planner/optimizer/ir/stats/mod.rs

@@ -18,9 +18,28 @@ mod join;
 mod selectivity;
 
 pub use column_stat::*;
+use databend_common_statistics::Datum;
 pub use databend_common_statistics::UniformSampleSet;
 pub(crate) use join::JoinConditionColumns;
 pub(crate) use join::JoinKeyStatUpdate;
 pub(crate) use join::JoinStatsEstimator;
 pub use selectivity::MAX_SELECTIVITY;
 pub use selectivity::SelectivityEstimator;
+
+pub(crate) fn finite_range_ndv_upper(min: &Datum, max: &Datum) -> Option<f64> {
+    if min == max {
+        return Some(1.0);
+    }
+    match (min, max) {
+        (Datum::Bool(false), Datum::Bool(true)) => Some(2.0),
+        (Datum::Int(min), Datum::Int(max)) => max
+            .checked_sub(*min)
+            .and_then(|diff| diff.checked_add(1))
+            .map(|value| value as f64),
+        (Datum::UInt(min), Datum::UInt(max)) => max
+            .checked_sub(*min)
+            .and_then(|diff| diff.checked_add(1))
+            .map(|value| value as f64),
+        _ => None,
+    }
+}

src/query/sql/src/planner/plans/eval_scalar.rs

@@ -87,7 +87,7 @@ impl EvalScalar {
         Ok(used_columns)
     }
 
-    fn derive_item_stat(
+    pub(crate) fn derive_item_stat(
         scalar: &ScalarExpr,
         input_statistics: &Statistics,
         cardinality: StatCardinality,