Skip to content

Commit 4b0556a

Browse files
forsaken628forsaken628
committed
feat(query): propagate column statistics through UNION ALL
1 parent ec61b69 commit 4b0556a

19 files changed

Lines changed: 1254 additions & 288 deletions

File tree

‎src/common/statistics/src/stat_estimate.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -435,6 +435,18 @@ impl StatCount {
435435
}
436436
}
437437
}
438+
439+
pub fn sum(left: Self, right: Self) -> Self {
440+
match (left, right) {
441+
(StatCount::Exact(left), StatCount::Exact(right)) => {
442+
StatCount::exact(left.saturating_add(right))
443+
}
444+
_ => StatCount::estimate(
445+
left.expected() + right.expected(),
446+
left.upper() + right.upper(),
447+
),
448+
}
449+
}
438450
}
439451

440452
#[cfg(test)]

‎src/query/expression/src/stat_evaluator.rs‎

Lines changed: 138 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,18 @@ use std::collections::HashMap;
1818
use databend_common_exception::ErrorCode;
1919
use databend_common_exception::Result;
2020

21+
use super::Cast;
2122
use super::ColumnIndex;
23+
use super::ColumnRef;
24+
use super::ConstantFolder;
2225
use super::Expr;
2326
use super::FunctionCall;
2427
use super::FunctionContext;
2528
use super::FunctionDomain;
2629
use super::FunctionEval;
2730
use super::FunctionRegistry;
2831
use super::Scalar;
32+
use super::conversion::classify_conversion;
2933
use super::function_stat::DeriveStat;
3034
use super::stat_distribution::ArgStat;
3135
use super::stat_distribution::BorrowedDistribution;
@@ -93,10 +97,70 @@ impl<'a> StatEvaluator<'a> {
9397
Expr::FunctionCall(call) => Ok(self
9498
.eval_function_call(call, input_stats)?
9599
.map(CowStat::Owned)),
96-
Expr::Cast(_) | Expr::LambdaFunctionCall(_) => Ok(None),
100+
Expr::Cast(cast) => Ok(self.eval_cast(cast, input_stats)?.map(CowStat::Owned)),
101+
Expr::LambdaFunctionCall(_) => Ok(None),
97102
}
98103
}
99104

105+
fn eval_cast<'s, I: ColumnIndex>(
106+
&'a self,
107+
cast: &Cast<I>,
108+
input_stats: &'s HashMap<I, ArgStat<'_>>,
109+
) -> Result<Option<ReturnStat>> {
110+
let src_type = cast.expr.data_type();
111+
if cast.is_try || !classify_conversion(src_type, &cast.dest_type).is_lossless_injective() {
112+
return Ok(None);
113+
}
114+
115+
let Some(input) = self.eval(&cast.expr, input_stats)? else {
116+
return Ok(None);
117+
};
118+
let input = input.as_ref();
119+
120+
// Reuse the cast domain implementation without making statistics
121+
// evaluation depend on the physical value evaluator. The synthetic
122+
// column represents the already-derived statistics of the inner expr.
123+
let expr = Expr::Cast(Cast {
124+
span: cast.span,
125+
is_try: false,
126+
expr: Box::new(Expr::ColumnRef(ColumnRef {
127+
span: cast.span,
128+
id: 0,
129+
data_type: src_type.clone(),
130+
display_name: String::new(),
131+
})),
132+
dest_type: cast.dest_type.clone(),
133+
});
134+
let input_domains = HashMap::from([(0, input.domain.clone())]);
135+
let (_, Some(domain)) = ConstantFolder::fold_with_domain(
136+
&expr,
137+
&input_domains,
138+
self.func_ctx,
139+
self.fn_registry,
140+
) else {
141+
return Ok(None);
142+
};
143+
144+
let stat = ReturnStat {
145+
domain,
146+
// A lossless injective cast preserves distinct values and NULLs.
147+
// Histograms are typed, so their boundaries cannot be reused.
148+
ndv: input.ndv,
149+
null_count: input.null_count,
150+
distribution: OwnedDistribution::Unknown,
151+
};
152+
if let Err(msg) = stat.check_consistency_with_type(Some(&cast.dest_type)) {
153+
if cfg!(debug_assertions) {
154+
return Err(ErrorCode::Internal(format!(
155+
"Failed to derive statistics for cast: {msg}"
156+
)));
157+
}
158+
log::warn!(msg; "Derived invalid cast statistics");
159+
return Ok(None);
160+
}
161+
Ok(Some(stat))
162+
}
163+
100164
fn eval_function_call<'s, I: ColumnIndex>(
101165
&'a self,
102166
call: &FunctionCall<I>,
@@ -219,6 +283,8 @@ mod tests {
219283

220284
use super::*;
221285
use crate::types::DataType;
286+
use crate::types::NumberDataType;
287+
use crate::types::number::NumberScalar;
222288

223289
#[test]
224290
fn test_constant_null_uses_exact_input_cardinality() {
@@ -237,4 +303,75 @@ mod tests {
237303

238304
assert_eq!(stat.null_count, StatCount::exact(7));
239305
}
306+
307+
#[test]
308+
fn test_lossless_cast_preserves_basic_statistics() {
309+
let src_type = DataType::Number(NumberDataType::UInt8);
310+
let dest_type = src_type.clone().wrap_nullable();
311+
let expr = Expr::Cast(Cast {
312+
span: None,
313+
is_try: false,
314+
expr: Box::new(Expr::ColumnRef(ColumnRef {
315+
span: None,
316+
id: 0,
317+
data_type: src_type.clone(),
318+
display_name: "c0".to_string(),
319+
})),
320+
dest_type: dest_type.clone(),
321+
});
322+
let domain = crate::Domain::from_min_max(
323+
Scalar::Number(NumberScalar::UInt8(1)),
324+
Scalar::Number(NumberScalar::UInt8(3)),
325+
&src_type,
326+
);
327+
let input_stats = HashMap::from([(0, ArgStat {
328+
domain,
329+
ndv: NdvEstimate::exact(3.0),
330+
null_count: StatCount::exact(0),
331+
distribution: BorrowedDistribution::Unknown,
332+
})]);
333+
334+
let stat = StatEvaluator::run(
335+
&expr,
336+
&FunctionContext::default(),
337+
&FunctionRegistry::empty(),
338+
StatCardinality::exact(3),
339+
&input_stats,
340+
)
341+
.unwrap()
342+
.unwrap()
343+
.into_owned();
344+
345+
assert!(stat.domain.matches_data_type(&dest_type));
346+
assert_eq!(stat.ndv, NdvEstimate::exact(3.0));
347+
assert_eq!(stat.null_count, StatCount::exact(0));
348+
assert!(matches!(stat.distribution, OwnedDistribution::Unknown));
349+
}
350+
351+
#[test]
352+
fn test_lossy_cast_is_not_derived() {
353+
let expr = Expr::Cast(Cast {
354+
span: None,
355+
is_try: false,
356+
expr: Box::new(Expr::ColumnRef(ColumnRef {
357+
span: None,
358+
id: 0,
359+
data_type: DataType::Number(NumberDataType::Int64),
360+
display_name: "c0".to_string(),
361+
})),
362+
dest_type: DataType::Number(NumberDataType::UInt8),
363+
});
364+
365+
assert!(
366+
StatEvaluator::run(
367+
&expr,
368+
&FunctionContext::default(),
369+
&FunctionRegistry::empty(),
370+
StatCardinality::exact(3),
371+
&HashMap::new(),
372+
)
373+
.unwrap()
374+
.is_none()
375+
);
376+
}
240377
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,3 +52,4 @@ pub use stats::MAX_SELECTIVITY;
5252
pub use stats::SelectivityEstimator;
5353
pub use stats::TopNSet;
5454
pub use stats::UniformSampleSet;
55+
pub(crate) use stats::finite_range_ndv_upper;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ pub type ColumnStatSet = HashMap<Symbol, ColumnStat>;
3131
pub type TopNSet = HashMap<Symbol, ColumnTopN>;
3232
pub type CountMinSketchSet = HashMap<Symbol, ColumnCountMinSketch>;
3333

34-
#[derive(Debug, Clone)]
34+
#[derive(Debug, Clone, PartialEq)]
3535
/// Statistics information of a column
3636
pub struct ColumnStat {
3737
/// Min value of the column

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

Lines changed: 1 addition & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use databend_common_statistics::HistogramRangeBounds;
2121
use databend_common_statistics::NdvEstimate;
2222
use databend_common_statistics::StatCount;
2323

24+
use super::finite_range_ndv_upper;
2425
use crate::optimizer::ir::ColumnStat;
2526
use crate::plans::ComparisonOp;
2627

@@ -138,24 +139,6 @@ fn apply_range_bounds(column_stat: &mut ColumnStat, new_min: Datum, new_max: Dat
138139
Ok(())
139140
}
140141

141-
fn finite_range_ndv_upper(min: &Datum, max: &Datum) -> Option<f64> {
142-
if min == max {
143-
return Some(1.0);
144-
}
145-
match (min, max) {
146-
(Datum::Bool(false), Datum::Bool(true)) => Some(2.0),
147-
(Datum::Int(min), Datum::Int(max)) => max
148-
.checked_sub(*min)
149-
.and_then(|diff| diff.checked_add(1))
150-
.map(|value| value as f64),
151-
(Datum::UInt(min), Datum::UInt(max)) => max
152-
.checked_sub(*min)
153-
.and_then(|diff| diff.checked_add(1))
154-
.map(|value| value as f64),
155-
_ => None,
156-
}
157-
}
158-
159142
pub(super) fn clear_for_empty_result(column_stat: &mut ColumnStat) {
160143
column_stat.ndv = NdvEstimate::exact(0.0);
161144
column_stat.null_count = StatCount::exact(0);

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,28 @@ mod join;
1818
mod selectivity;
1919

2020
pub use column_stat::*;
21+
use databend_common_statistics::Datum;
2122
pub use databend_common_statistics::UniformSampleSet;
2223
pub(crate) use join::JoinConditionColumns;
2324
pub(crate) use join::JoinKeyStatUpdate;
2425
pub(crate) use join::JoinStatsEstimator;
2526
pub use selectivity::MAX_SELECTIVITY;
2627
pub use selectivity::SelectivityEstimator;
28+
29+
pub(crate) fn finite_range_ndv_upper(min: &Datum, max: &Datum) -> Option<f64> {
30+
if min == max {
31+
return Some(1.0);
32+
}
33+
match (min, max) {
34+
(Datum::Bool(false), Datum::Bool(true)) => Some(2.0),
35+
(Datum::Int(min), Datum::Int(max)) => max
36+
.checked_sub(*min)
37+
.and_then(|diff| diff.checked_add(1))
38+
.map(|value| value as f64),
39+
(Datum::UInt(min), Datum::UInt(max)) => max
40+
.checked_sub(*min)
41+
.and_then(|diff| diff.checked_add(1))
42+
.map(|value| value as f64),
43+
_ => None,
44+
}
45+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ impl EvalScalar {
8787
Ok(used_columns)
8888
}
8989

90-
fn derive_item_stat(
90+
pub(crate) fn derive_item_stat(
9191
scalar: &ScalarExpr,
9292
input_statistics: &Statistics,
9393
cardinality: StatCardinality,

0 commit comments

Comments
 (0)