1use crate::meta::{ColumnId, ColumnType, Predicate, PredicateValue, Schema, TIMESTAMP_COLUMN};
4use crate::segment::ClusteringKeyRange;
5
6#[derive(Clone, Debug, Eq, PartialEq)]
8pub enum PlanError {
9 UnknownColumn(ColumnId),
11 TypeMismatch {
13 column: ColumnId,
15 expected: ColumnType,
17 actual: ColumnType,
19 },
20 RangeRequiresNumericColumn(ColumnId),
22 RowCountMismatch {
24 columns: u32,
26 alive: u32,
28 },
29}
30
31impl std::fmt::Display for PlanError {
32 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33 match self {
34 Self::UnknownColumn(column) => write!(formatter, "unknown column {}", column.get()),
35 Self::TypeMismatch {
36 column,
37 expected,
38 actual,
39 } => write!(
40 formatter,
41 "column {} expects {expected:?}, received {actual:?}",
42 column.get()
43 ),
44 Self::RangeRequiresNumericColumn(column) => {
45 write!(formatter, "column {} is not numeric", column.get())
46 }
47 Self::RowCountMismatch { columns, alive } => write!(
48 formatter,
49 "column row count {columns} differs from alive row count {alive}"
50 ),
51 }
52 }
53}
54
55impl std::error::Error for PlanError {}
56
57pub fn validate_predicate(predicate: &Predicate, schema: &Schema) -> Result<(), PlanError> {
59 match predicate {
60 Predicate::Eq { column, value } => validate_value(*column, value, schema),
61 Predicate::In { column, values } => {
62 require_column(*column, schema)?;
63 for value in values {
64 validate_value(*column, value, schema)?;
65 }
66 Ok(())
67 }
68 Predicate::Range(range) => {
69 let definition = require_column(range.column, schema)?;
70 if !matches!(
71 definition.column_type(),
72 ColumnType::U64 | ColumnType::I64 | ColumnType::F64
73 ) {
74 return Err(PlanError::RangeRequiresNumericColumn(range.column));
75 }
76 if let Some(lower) = &range.lower {
77 validate_value(range.column, &lower.value, schema)?;
78 }
79 if let Some(upper) = &range.upper {
80 validate_value(range.column, &upper.value, schema)?;
81 }
82 Ok(())
83 }
84 Predicate::Exists(column) | Predicate::IsNull(column) => {
85 require_column(*column, schema).map(|_| ())
86 }
87 Predicate::And(children) | Predicate::Or(children) => {
88 for child in children {
89 validate_predicate(child, schema)?;
90 }
91 Ok(())
92 }
93 Predicate::Not(child) => validate_predicate(child, schema),
94 }
95}
96
97fn require_column(
98 column: ColumnId,
99 schema: &Schema,
100) -> Result<&crate::meta::ColumnDefinition, PlanError> {
101 schema
102 .column(column)
103 .ok_or(PlanError::UnknownColumn(column))
104}
105
106fn validate_value(
107 column: ColumnId,
108 value: &PredicateValue,
109 schema: &Schema,
110) -> Result<(), PlanError> {
111 let expected = require_column(column, schema)?.column_type();
112 let actual = value.value_type();
113 let compatible = expected == actual
114 || matches!(
115 (expected, actual),
116 (ColumnType::DictionaryString, ColumnType::RawString)
117 );
118 if compatible {
119 Ok(())
120 } else {
121 Err(PlanError::TypeMismatch {
122 column,
123 expected,
124 actual,
125 })
126 }
127}
128
129#[must_use]
131pub fn segment_may_match(range: ClusteringKeyRange, predicate: &Predicate) -> bool {
132 match range {
133 ClusteringKeyRange::Unstamped => true,
134 ClusteringKeyRange::Empty => false,
135 ClusteringKeyRange::Bounded { min_ts, max_ts } => {
136 predicate_may_match_timestamp(predicate, min_ts, max_ts)
137 }
138 }
139}
140
141fn predicate_may_match_timestamp(predicate: &Predicate, min_ts: i64, max_ts: i64) -> bool {
142 match predicate {
143 Predicate::Eq { column, value } if *column == TIMESTAMP_COLUMN => {
144 matches!(value, PredicateValue::I64(value) if *value >= min_ts && *value <= max_ts)
145 }
146 Predicate::In { column, values } if *column == TIMESTAMP_COLUMN => values.iter().any(
147 |value| matches!(value, PredicateValue::I64(value) if *value >= min_ts && *value <= max_ts),
148 ),
149 Predicate::Range(range) if range.column == TIMESTAMP_COLUMN => {
150 let lower_overlaps = range.lower.as_ref().is_none_or(|bound| {
151 matches!(&bound.value, PredicateValue::I64(value) if if bound.inclusive {
152 max_ts >= *value
153 } else {
154 max_ts > *value
155 })
156 });
157 let upper_overlaps = range.upper.as_ref().is_none_or(|bound| {
158 matches!(&bound.value, PredicateValue::I64(value) if if bound.inclusive {
159 min_ts <= *value
160 } else {
161 min_ts < *value
162 })
163 });
164 lower_overlaps && upper_overlaps
165 }
166 Predicate::Exists(column) if *column == TIMESTAMP_COLUMN => true,
167 Predicate::IsNull(column) if *column == TIMESTAMP_COLUMN => false,
168 Predicate::And(children) => children
169 .iter()
170 .all(|child| predicate_may_match_timestamp(child, min_ts, max_ts)),
171 Predicate::Or(children) => children
172 .iter()
173 .any(|child| predicate_may_match_timestamp(child, min_ts, max_ts)),
174 Predicate::Not(_) => true,
175 Predicate::Eq { .. }
176 | Predicate::In { .. }
177 | Predicate::Range(_)
178 | Predicate::Exists(_)
179 | Predicate::IsNull(_) => true,
180 }
181}
182
183#[cfg(test)]
184#[allow(clippy::expect_used)]
185mod tests {
186 use rand::Rng;
187
188 use super::*;
189 use crate::meta::{RangeBound, RangePredicate};
190
191 #[test]
192 fn pruning_by_clustering_range_never_drops_a_matching_segment() {
193 let mut random = crate::test_support::seeded_rng(
194 "planner::prune::pruning_by_clustering_range_never_drops_a_matching_segment",
195 );
196 for _ in 0..2_048 {
197 let first = random.random_range(-100_i64..=100);
198 let second = random.random_range(-100_i64..=100);
199 let min_ts = first.min(second);
200 let max_ts = first.max(second);
201 let lower = random.random_range(-110_i64..=110);
202 let upper = random.random_range(-110_i64..=110);
203 let lower_inclusive = random.random();
204 let upper_inclusive = random.random();
205 let predicate = Predicate::Range(RangePredicate {
206 column: TIMESTAMP_COLUMN,
207 lower: Some(RangeBound {
208 value: PredicateValue::I64(lower),
209 inclusive: lower_inclusive,
210 }),
211 upper: Some(RangeBound {
212 value: PredicateValue::I64(upper),
213 inclusive: upper_inclusive,
214 }),
215 });
216 let row_truth = (min_ts..=max_ts).any(|value| {
217 (value > lower || (lower_inclusive && value == lower))
218 && (value < upper || (upper_inclusive && value == upper))
219 });
220 if row_truth {
221 assert!(segment_may_match(
222 ClusteringKeyRange::Bounded { min_ts, max_ts },
223 &predicate
224 ));
225 }
226 assert!(segment_may_match(ClusteringKeyRange::Unstamped, &predicate));
227 }
228 }
229}