nodedb_sql/optimizer/
point_get.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Detect equality on the primary key → convert Scan to PointGet.
4
5use nodedb_types::DatabaseId;
6
7use crate::catalog::SqlCatalog;
8use crate::planner::cp_projection::to_cp_computed;
9use crate::types::*;
10
11/// If a Scan has a single equality filter on the collection's primary key,
12/// convert it to a PointGet.
13///
14/// The primary key is resolved from the catalog rather than assumed to be the
15/// conventional `id` / `document_id` / `key`: a collection created without an
16/// explicit `PRIMARY KEY` gets an auto-generated `_rowid` key, so a filter on a
17/// regular `id` column is NOT a point lookup and must stay a scan (routing it
18/// to PointGet would resolve a surrogate for the wrong key and return zero
19/// rows). A declared key of any name is a point lookup: `WHERE sku = 'p1'`
20/// reads the one row keyed `p1`, as `WHERE id = 'p1'` does. When the catalog
21/// cannot resolve a primary key (unknown collection), the conventional names
22/// apply. A catalog error fails the plan: `RetryableSchemaChanged` must reach
23/// the caller.
24///
25/// A point lookup returns the whole stored row and runs no expression. A
26/// computed SELECT item over it becomes [`Projection::CpComputed`]: the
27/// Control Plane evaluates it once on the returned row. `SELECT to_jsonb(*)
28/// ... WHERE <key> = $1` stays a point lookup that way.
29pub fn optimize(plan: SqlPlan, catalog: &dyn SqlCatalog) -> crate::Result<SqlPlan> {
30    match plan {
31        SqlPlan::Scan {
32            ref collection,
33            ref alias,
34            ref engine,
35            ref filters,
36            ref projection,
37            ref temporal,
38            ..
39        } if filters.len() == 1 && !temporal.is_temporal() && has_point_lookup(engine) => {
40            let pk = catalog
41                .get_collection(DatabaseId::DEFAULT, collection)?
42                .and_then(|info| info.primary_key);
43            if let Some((key_col, key_val)) = extract_pk_equality(&filters[0], pk.as_deref()) {
44                return Ok(SqlPlan::PointGet {
45                    collection: collection.clone(),
46                    alias: alias.clone(),
47                    engine: *engine,
48                    key_column: key_col,
49                    key_value: key_val,
50                    projection: projection.iter().cloned().map(to_cp_computed).collect(),
51                });
52            }
53            Ok(plan)
54        }
55        _ => Ok(plan),
56    }
57}
58
59/// Whether the engine answers a key equality with a point lookup.
60///
61/// Timeseries and array refuse point lookups in their engine rules. A key
62/// equality on them stays a scan.
63fn has_point_lookup(engine: &EngineType) -> bool {
64    match engine {
65        EngineType::DocumentSchemaless
66        | EngineType::DocumentStrict
67        | EngineType::KeyValue
68        | EngineType::Columnar
69        | EngineType::Spatial => true,
70        EngineType::Timeseries | EngineType::Array | EngineType::ContinuousAggregate => false,
71    }
72}
73
74/// Extract a simple equality filter eligible for the point-get rewrite.
75///
76/// When the catalog resolves the primary key, the candidate column is that
77/// key, whatever its name. A regular `id` column on a collection keyed by
78/// another column stays a scan: a point lookup would resolve a surrogate for
79/// the wrong key and return zero rows. When the catalog resolves no key, the
80/// conventional document-key names (`id` / `document_id` / `key`) are the
81/// candidates.
82fn extract_pk_equality(filter: &Filter, pk: Option<&str>) -> Option<(String, SqlValue)> {
83    let (column, value) = match &filter.expr {
84        FilterExpr::Comparison {
85            field,
86            op: CompareOp::Eq,
87            value,
88        } => (field.as_str(), value),
89        FilterExpr::Expr(SqlExpr::BinaryOp {
90            left,
91            op: BinaryOp::Eq,
92            right,
93        }) => match (left.as_ref(), right.as_ref()) {
94            (SqlExpr::Column { name, .. }, SqlExpr::Literal(value)) => (name.as_str(), value),
95            _ => return None,
96        },
97        _ => return None,
98    };
99    key_column_named(column, pk).map(|key| (key, value.clone()))
100}
101
102/// The key column `column` names, spelled as the catalog declares it, or
103/// `None` when `column` is not the key.
104fn key_column_named(column: &str, pk: Option<&str>) -> Option<String> {
105    match pk {
106        Some(pk) => column.eq_ignore_ascii_case(pk).then(|| pk.to_string()),
107        None => {
108            let column = column.to_lowercase();
109            matches!(column.as_str(), "id" | "document_id" | "key").then_some(column)
110        }
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use crate::catalog::SqlCatalogError;
118    use crate::temporal::TemporalScope;
119
120    struct ChangingCatalog;
121
122    impl SqlCatalog for ChangingCatalog {
123        fn get_collection(
124            &self,
125            _: DatabaseId,
126            _: &str,
127        ) -> std::result::Result<Option<CollectionInfo>, SqlCatalogError> {
128            Err(SqlCatalogError::RetryableSchemaChanged {
129                descriptor: "collection users".into(),
130            })
131        }
132    }
133
134    fn id_scan() -> SqlPlan {
135        SqlPlan::Scan {
136            collection: "users".into(),
137            alias: None,
138            engine: EngineType::DocumentSchemaless,
139            filters: vec![Filter {
140                expr: FilterExpr::Comparison {
141                    field: "id".into(),
142                    op: CompareOp::Eq,
143                    value: SqlValue::String("u1".into()),
144                },
145            }],
146            projection: Vec::new(),
147            sort_keys: Vec::new(),
148            limit: None,
149            offset: 0,
150            distinct: false,
151            window_functions: Vec::new(),
152            temporal: TemporalScope::default(),
153        }
154    }
155
156    /// A schemaless `users` collection keyed by `id`.
157    struct UsersCatalog;
158
159    impl SqlCatalog for UsersCatalog {
160        fn get_collection(
161            &self,
162            _: DatabaseId,
163            name: &str,
164        ) -> std::result::Result<Option<CollectionInfo>, SqlCatalogError> {
165            Ok((name == "users").then(|| CollectionInfo {
166                name: "users".to_string(),
167                engine: EngineType::DocumentSchemaless,
168                columns: Vec::new(),
169                primary_key: Some("id".to_string()),
170                has_auto_tier: false,
171                indexes: Vec::new(),
172                bitemporal: false,
173                primary: nodedb_types::PrimaryEngine::Document,
174                vector_primary: None,
175                partition_strategy: nodedb_types::PartitionStrategy::CollectionHomed,
176                open_schema: CollectionInfo::open_schema_for(EngineType::DocumentSchemaless),
177            }))
178        }
179    }
180
181    #[test]
182    fn a_computed_item_over_a_key_lookup_is_control_plane_computed() {
183        let SqlPlan::Scan {
184            collection,
185            alias,
186            engine,
187            filters,
188            sort_keys,
189            limit,
190            offset,
191            distinct,
192            window_functions,
193            temporal,
194            ..
195        } = id_scan()
196        else {
197            panic!("id_scan is a scan");
198        };
199        let scan = SqlPlan::Scan {
200            collection,
201            alias,
202            engine,
203            filters,
204            projection: vec![Projection::Computed {
205                expr: SqlExpr::Function {
206                    name: "to_jsonb".into(),
207                    args: vec![SqlExpr::Wildcard],
208                    distinct: false,
209                },
210                alias: "document".into(),
211            }],
212            sort_keys,
213            limit,
214            offset,
215            distinct,
216            window_functions,
217            temporal,
218        };
219        match optimize(scan, &UsersCatalog).unwrap() {
220            SqlPlan::PointGet { projection, .. } => assert!(matches!(
221                projection.as_slice(),
222                [Projection::CpComputed {
223                    expr: SqlExpr::Function { name, args, .. },
224                    alias,
225                }] if name == "to_jsonb"
226                    && matches!(args.as_slice(), [SqlExpr::Wildcard])
227                    && alias == "document"
228            )),
229            other => panic!("expected a point lookup, got {other:?}"),
230        }
231    }
232
233    #[test]
234    fn a_catalog_error_fails_the_point_get_rewrite() {
235        let error = optimize(id_scan(), &ChangingCatalog).unwrap_err();
236        assert_eq!(
237            error,
238            crate::SqlError::from(SqlCatalogError::RetryableSchemaChanged {
239                descriptor: "collection users".into(),
240            })
241        );
242    }
243
244    /// A collection keyed by the declared column `Sku`.
245    struct SkuCatalog;
246
247    impl SqlCatalog for SkuCatalog {
248        fn get_collection(
249            &self,
250            _: DatabaseId,
251            name: &str,
252        ) -> std::result::Result<Option<CollectionInfo>, SqlCatalogError> {
253            Ok(Some(CollectionInfo {
254                name: name.to_string(),
255                engine: EngineType::DocumentSchemaless,
256                columns: Vec::new(),
257                primary_key: Some("Sku".to_string()),
258                has_auto_tier: false,
259                indexes: Vec::new(),
260                bitemporal: false,
261                primary: nodedb_types::PrimaryEngine::Document,
262                vector_primary: None,
263                partition_strategy: nodedb_types::PartitionStrategy::CollectionHomed,
264                open_schema: CollectionInfo::open_schema_for(EngineType::DocumentSchemaless),
265            }))
266        }
267    }
268
269    fn equality_scan(field: &str, engine: EngineType) -> SqlPlan {
270        let SqlPlan::Scan {
271            collection,
272            alias,
273            projection,
274            sort_keys,
275            limit,
276            offset,
277            distinct,
278            window_functions,
279            temporal,
280            ..
281        } = id_scan()
282        else {
283            panic!("id_scan is a scan");
284        };
285        SqlPlan::Scan {
286            collection,
287            alias,
288            engine,
289            filters: vec![Filter {
290                expr: FilterExpr::Comparison {
291                    field: field.into(),
292                    op: CompareOp::Eq,
293                    value: SqlValue::String("p1".into()),
294                },
295            }],
296            projection,
297            sort_keys,
298            limit,
299            offset,
300            distinct,
301            window_functions,
302            temporal,
303        }
304    }
305
306    #[test]
307    fn a_declared_key_equality_is_a_point_lookup_on_that_key() {
308        match optimize(
309            equality_scan("sku", EngineType::DocumentSchemaless),
310            &SkuCatalog,
311        )
312        .unwrap()
313        {
314            SqlPlan::PointGet {
315                key_column,
316                key_value,
317                ..
318            } => {
319                assert_eq!(key_column, "Sku", "the catalog spelling names the key");
320                assert_eq!(key_value, SqlValue::String("p1".into()));
321            }
322            other => panic!("expected a point lookup, got {other:?}"),
323        }
324    }
325
326    #[test]
327    fn an_id_column_beside_a_declared_key_stays_a_scan() {
328        let plan = optimize(
329            equality_scan("id", EngineType::DocumentSchemaless),
330            &SkuCatalog,
331        )
332        .unwrap();
333        assert!(matches!(plan, SqlPlan::Scan { .. }), "{plan:?}");
334    }
335
336    #[test]
337    fn a_key_equality_on_an_engine_without_point_lookups_stays_a_scan() {
338        for engine in [EngineType::Timeseries, EngineType::Array] {
339            let plan = optimize(equality_scan("sku", engine), &SkuCatalog).unwrap();
340            assert!(matches!(plan, SqlPlan::Scan { .. }), "{plan:?}");
341        }
342    }
343}
nodedb_sql/optimizer/point_get.rs

nodedb_sql/optimizer/
point_get.rs
