Skip to main content

zeppelin_embed/planner/
prune.rs

1//! Conservative clustering-range pruning and schema validation.
2
3use crate::meta::{ColumnId, ColumnType, Predicate, PredicateValue, Schema, TIMESTAMP_COLUMN};
4use crate::segment::ClusteringKeyRange;
5
6/// A typed planning-time predicate rejection.
7#[derive(Clone, Debug, Eq, PartialEq)]
8pub enum PlanError {
9    /// The predicate named a column absent from the collection schema.
10    UnknownColumn(ColumnId),
11    /// A predicate value disagreed with its declared column type.
12    TypeMismatch {
13        /// Mismatched column.
14        column: ColumnId,
15        /// Declared type.
16        expected: ColumnType,
17        /// Supplied type.
18        actual: ColumnType,
19    },
20    /// A range was attached to a non-numeric column.
21    RangeRequiresNumericColumn(ColumnId),
22    /// Metadata and alive arrays described different row spaces.
23    RowCountMismatch {
24        /// Metadata row count.
25        columns: u32,
26        /// Alive row count.
27        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
57/// Validates every predicate node against the collection schema before execution.
58pub 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/// Returns `false` only when timestamp statistics prove no row can match.
130#[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}