nodedb_sql/planner/array_fn/
table_fn.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! `SELECT * FROM ARRAY_*(...)` table-valued function planning:
4//! slice / project / aggregate / elementwise.
5
6use crate::types::{ArrayAggPlan, ArrayElementwisePlan, ArrayProjectPlan, ArraySlicePlan};
7use sqlparser::ast;
8
9use nodedb_types::Value;
10
11use super::helpers::{
12    collect_args, expect_string_array, expect_string_literal, expect_u32, is_null_literal,
13    require_array_name, value_to_coord_literal,
14};
15use crate::error::{Result, SqlError};
16use crate::parser::object_literal::parse_object_literal_complete;
17use crate::temporal::TemporalScope;
18use crate::types::{SqlCatalog, SqlPlan};
19use crate::types_array::{ArrayBinaryOpAst, ArrayReducerAst, ArraySliceAst, NamedDimRange};
20
21/// Try to intercept a `SELECT * FROM array_xxx(...)` table-valued
22/// function call. Returns `Ok(Some(plan))` on a match, `Ok(None)` if
23/// the FROM is not an array function (caller falls through to normal
24/// catalog resolution).
25///
26/// `temporal` carries any `AS OF SYSTEM TIME` / `AS OF VALID TIME` qualifiers
27/// extracted by the pre-processor. It is propagated verbatim into
28/// `SqlPlan::ArraySlice` and `SqlPlan::ArrayAgg` so the planner-to-physical
29/// conversion can populate the corresponding `ArrayOp` fields. When neither
30/// clause was present, `temporal` is `TemporalScope::default()`, which maps to
31/// the live-state fast path in the Data Plane handler.
32pub fn try_plan_array_table_fn(
33    from: &[ast::TableWithJoins],
34    catalog: &dyn SqlCatalog,
35    temporal: TemporalScope,
36) -> Result<Option<SqlPlan>> {
37    if from.len() != 1 {
38        return Ok(None);
39    }
40    let twj = &from[0];
41    if !twj.joins.is_empty() {
42        return Ok(None);
43    }
44    let (name, args) = match &twj.relation {
45        ast::TableFactor::Table {
46            name,
47            args: Some(args),
48            alias,
49            ..
50        } => {
51            if let Some(alias) = alias {
52                crate::reserved::check_ast_identifier(&alias.name)?;
53            }
54            (name, args)
55        }
56        _ => return Ok(None),
57    };
58    let fn_name = crate::parser::normalize::normalize_object_name_checked(name)?;
59    let arg_exprs = collect_args(&args.args);
60    match fn_name.as_str() {
61        "array_slice" => Ok(Some(plan_slice(&arg_exprs, catalog, temporal)?)),
62        "array_project" => Ok(Some(plan_project(&arg_exprs, catalog)?)),
63        "array_agg" => Ok(Some(plan_agg(&arg_exprs, catalog, temporal)?)),
64        "array_elementwise" => Ok(Some(plan_elementwise(&arg_exprs, catalog)?)),
65        _ => Ok(None),
66    }
67}
68
69fn plan_slice(
70    args: &[ast::Expr],
71    catalog: &dyn SqlCatalog,
72    temporal: TemporalScope,
73) -> Result<SqlPlan> {
74    if args.len() < 2 || args.len() > 4 {
75        return Err(SqlError::Unsupported {
76            detail: format!(
77                "ARRAY_SLICE expects 2..=4 args (name, slice_obj, [attrs], [limit]); got {}",
78                args.len()
79            ),
80        });
81    }
82    let name = require_array_name(args, 0, "ARRAY_SLICE", catalog)?;
83    let view = catalog
84        .lookup_array(&name)
85        .ok_or_else(|| SqlError::Unsupported {
86            detail: format!("ARRAY_SLICE: array '{name}' not found"),
87        })?;
88
89    // Slice-predicate literal: encoded as a quoted string carrying the
90    // brace-form object literal. The PostgreSQL dialect does not accept
91    // bare `{...}` in expression position, so we decode the string
92    // contents here.
93    let slice_str = expect_string_literal(&args[1], "ARRAY_SLICE slice predicate")?;
94    let parsed =
95        parse_object_literal_complete(&slice_str).ok_or_else(|| SqlError::Unsupported {
96            detail: format!("ARRAY_SLICE: slice predicate must be an object literal: {slice_str}"),
97        })?;
98    let map = parsed.map_err(|detail| SqlError::Unsupported {
99        detail: format!("ARRAY_SLICE: slice parse: {detail}"),
100    })?;
101    let mut dim_ranges: Vec<NamedDimRange> = Vec::with_capacity(map.len());
102    for (dim, val) in map {
103        // Verify the dim exists on the array.
104        if !view.dims.iter().any(|d| d.name == dim) {
105            return Err(SqlError::Unsupported {
106                detail: format!("ARRAY_SLICE: array '{name}' has no dim '{dim}'"),
107            });
108        }
109        let arr = match val {
110            Value::Array(a) if a.len() == 2 => a,
111            _ => {
112                return Err(SqlError::Unsupported {
113                    detail: format!(
114                        "ARRAY_SLICE: dim '{dim}' range must be a 2-element array [lo, hi]"
115                    ),
116                });
117            }
118        };
119        let lo = value_to_coord_literal(&arr[0], &dim)?;
120        let hi = value_to_coord_literal(&arr[1], &dim)?;
121        dim_ranges.push(NamedDimRange { dim, lo, hi });
122    }
123
124    let attr_projection = if args.len() >= 3 {
125        match &args[2] {
126            ast::Expr::Value(v) if matches!(v.value, ast::Value::SingleQuotedString(ref s) if s == "*") => {
127                Vec::new()
128            }
129            _ => expect_string_array(&args[2], "ARRAY_SLICE attr projection")?,
130        }
131    } else {
132        Vec::new()
133    };
134    // Validate attr names against the catalog.
135    for attr in &attr_projection {
136        if !view.attrs.iter().any(|a| &a.name == attr) {
137            return Err(SqlError::Unsupported {
138                detail: format!("ARRAY_SLICE: array '{name}' has no attr '{attr}'"),
139            });
140        }
141    }
142
143    let limit = if args.len() >= 4 {
144        expect_u32(&args[3], "ARRAY_SLICE limit")?
145    } else {
146        0
147    };
148
149    Ok(SqlPlan::ArraySlice(ArraySlicePlan {
150        name,
151        slice: ArraySliceAst { dim_ranges },
152        attr_projection,
153        limit,
154        temporal,
155    }))
156}
157
158fn plan_project(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result<SqlPlan> {
159    if args.len() != 2 {
160        return Err(SqlError::Unsupported {
161            detail: format!(
162                "ARRAY_PROJECT expects 2 args (name, [attrs]); got {}",
163                args.len()
164            ),
165        });
166    }
167    let name = require_array_name(args, 0, "ARRAY_PROJECT", catalog)?;
168    let view = catalog
169        .lookup_array(&name)
170        .ok_or_else(|| SqlError::Unsupported {
171            detail: format!("ARRAY_PROJECT: array '{name}' not found"),
172        })?;
173    let attr_projection = expect_string_array(&args[1], "ARRAY_PROJECT attrs")?;
174    if attr_projection.is_empty() {
175        return Err(SqlError::Unsupported {
176            detail: "ARRAY_PROJECT: attr list must not be empty".into(),
177        });
178    }
179    for attr in &attr_projection {
180        if !view.attrs.iter().any(|a| &a.name == attr) {
181            return Err(SqlError::Unsupported {
182                detail: format!("ARRAY_PROJECT: array '{name}' has no attr '{attr}'"),
183            });
184        }
185    }
186    Ok(SqlPlan::ArrayProject(ArrayProjectPlan {
187        name,
188        attr_projection,
189    }))
190}
191
192fn plan_agg(
193    args: &[ast::Expr],
194    catalog: &dyn SqlCatalog,
195    temporal: TemporalScope,
196) -> Result<SqlPlan> {
197    if args.len() < 3 || args.len() > 4 {
198        return Err(SqlError::Unsupported {
199            detail: format!(
200                "ARRAY_AGG expects 3..=4 args (name, attr, reducer, [group_by_dim]); got {}",
201                args.len()
202            ),
203        });
204    }
205    let name = require_array_name(args, 0, "ARRAY_AGG", catalog)?;
206    let view = catalog
207        .lookup_array(&name)
208        .ok_or_else(|| SqlError::Unsupported {
209            detail: format!("ARRAY_AGG: array '{name}' not found"),
210        })?;
211
212    let attr = expect_string_literal(&args[1], "ARRAY_AGG attr")?;
213    if !view.attrs.iter().any(|a| a.name == attr) {
214        return Err(SqlError::Unsupported {
215            detail: format!("ARRAY_AGG: array '{name}' has no attr '{attr}'"),
216        });
217    }
218
219    let reducer_str = expect_string_literal(&args[2], "ARRAY_AGG reducer")?;
220    let reducer = ArrayReducerAst::parse(&reducer_str).ok_or_else(|| SqlError::Unsupported {
221        detail: format!("ARRAY_AGG: unknown reducer '{reducer_str}' (want sum/count/min/max/mean)"),
222    })?;
223
224    let group_by_dim = if args.len() == 4 && !is_null_literal(&args[3]) {
225        let dim = expect_string_literal(&args[3], "ARRAY_AGG group_by_dim")?;
226        if !view.dims.iter().any(|d| d.name == dim) {
227            return Err(SqlError::Unsupported {
228                detail: format!("ARRAY_AGG: array '{name}' has no dim '{dim}'"),
229            });
230        }
231        Some(dim)
232    } else {
233        None
234    };
235
236    Ok(SqlPlan::ArrayAgg(ArrayAggPlan {
237        name,
238        attr,
239        reducer,
240        group_by_dim,
241        temporal,
242    }))
243}
244
245fn plan_elementwise(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result<SqlPlan> {
246    if args.len() != 4 {
247        return Err(SqlError::Unsupported {
248            detail: format!(
249                "ARRAY_ELEMENTWISE expects 4 args (left, right, op, attr); got {}",
250                args.len()
251            ),
252        });
253    }
254    let left = require_array_name(args, 0, "ARRAY_ELEMENTWISE", catalog)?;
255    let right = require_array_name(args, 1, "ARRAY_ELEMENTWISE", catalog)?;
256    let lview = catalog
257        .lookup_array(&left)
258        .ok_or_else(|| SqlError::Unsupported {
259            detail: format!("ARRAY_ELEMENTWISE: array '{left}' not found"),
260        })?;
261    let rview = catalog
262        .lookup_array(&right)
263        .ok_or_else(|| SqlError::Unsupported {
264            detail: format!("ARRAY_ELEMENTWISE: array '{right}' not found"),
265        })?;
266    if lview.dims.len() != rview.dims.len() || lview.attrs.len() != rview.attrs.len() {
267        return Err(SqlError::Unsupported {
268            detail: format!(
269                "ARRAY_ELEMENTWISE: arrays '{left}' and '{right}' must share schema shape"
270            ),
271        });
272    }
273    let op_str = expect_string_literal(&args[2], "ARRAY_ELEMENTWISE op")?;
274    let op = ArrayBinaryOpAst::parse(&op_str).ok_or_else(|| SqlError::Unsupported {
275        detail: format!("ARRAY_ELEMENTWISE: unknown op '{op_str}' (want add/sub/mul/div)"),
276    })?;
277    let attr = expect_string_literal(&args[3], "ARRAY_ELEMENTWISE attr")?;
278    if !lview.attrs.iter().any(|a| a.name == attr) {
279        return Err(SqlError::Unsupported {
280            detail: format!("ARRAY_ELEMENTWISE: array '{left}' has no attr '{attr}'"),
281        });
282    }
283    Ok(SqlPlan::ArrayElementwise(ArrayElementwisePlan {
284        left,
285        right,
286        op,
287        attr,
288    }))
289}
290
291#[cfg(test)]
292mod tests {
293    use crate::catalog::{ArrayCatalogView, SqlCatalogError};
294    use crate::error::Result;
295    use crate::functions::registry::FunctionRegistry;
296    use crate::parser::statement::parse_sql;
297    use crate::types::{ArrayAggPlan, ArrayProjectPlan, ArraySlicePlan};
298    use crate::types::{CollectionInfo, SqlCatalog, SqlPlan};
299    use crate::types_array::{
300        ArrayAttrAst, ArrayAttrType, ArrayDimAst, ArrayDimType, ArrayDomainBound, ArrayReducerAst,
301    };
302
303    struct StubCatalog {
304        view: Option<ArrayCatalogView>,
305        right_view: Option<ArrayCatalogView>,
306    }
307    impl SqlCatalog for StubCatalog {
308        fn get_collection(
309            &self,
310            _: nodedb_types::DatabaseId,
311            _name: &str,
312        ) -> std::result::Result<Option<CollectionInfo>, SqlCatalogError> {
313            Ok(None)
314        }
315        fn lookup_array(&self, name: &str) -> Option<ArrayCatalogView> {
316            if name == "g" || name == "left" {
317                self.view.clone()
318            } else if name == "right" {
319                self.right_view.clone()
320            } else {
321                None
322            }
323        }
324    }
325
326    fn view() -> ArrayCatalogView {
327        ArrayCatalogView {
328            name: "g".into(),
329            dims: vec![
330                ArrayDimAst {
331                    name: "chrom".into(),
332                    dtype: ArrayDimType::Int64,
333                    lo: ArrayDomainBound::Int64(1),
334                    hi: ArrayDomainBound::Int64(23),
335                },
336                ArrayDimAst {
337                    name: "pos".into(),
338                    dtype: ArrayDimType::Int64,
339                    lo: ArrayDomainBound::Int64(0),
340                    hi: ArrayDomainBound::Int64(1_000_000),
341                },
342            ],
343            attrs: vec![
344                ArrayAttrAst {
345                    name: "variant".into(),
346                    dtype: ArrayAttrType::String,
347                    nullable: true,
348                },
349                ArrayAttrAst {
350                    name: "qual".into(),
351                    dtype: ArrayAttrType::Float64,
352                    nullable: true,
353                },
354            ],
355            tile_extents: vec![1, 1_000_000],
356        }
357    }
358
359    fn cat() -> StubCatalog {
360        StubCatalog {
361            view: Some(view()),
362            right_view: Some(view()),
363        }
364    }
365
366    fn plan_one(sql: &str) -> Result<SqlPlan> {
367        let stmts = parse_sql(sql)?;
368        let q = match &stmts[0] {
369            sqlparser::ast::Statement::Query(q) => q,
370            _ => panic!("not a query"),
371        };
372        crate::planner::select::plan_query(
373            q,
374            &cat(),
375            &FunctionRegistry::new(),
376            crate::TemporalScope::default(),
377        )
378    }
379
380    #[test]
381    fn slice_happy() {
382        let p = plan_one(
383            "SELECT * FROM ARRAY_SLICE('g', '{chrom: [1,1], pos: [0, 100]}', ['qual'], 50)",
384        )
385        .unwrap();
386        match p {
387            SqlPlan::ArraySlice(ArraySlicePlan {
388                name,
389                slice,
390                attr_projection,
391                limit,
392                ..
393            }) => {
394                assert_eq!(name, "g");
395                assert_eq!(slice.dim_ranges.len(), 2);
396                assert_eq!(attr_projection, vec!["qual".to_string()]);
397                assert_eq!(limit, 50);
398            }
399            other => panic!("expected ArraySlice, got {other:?}"),
400        }
401    }
402
403    #[test]
404    fn slice_unknown_dim_rejected() {
405        let err = plan_one("SELECT * FROM ARRAY_SLICE('g', '{nope: [1, 2]}')")
406            .err()
407            .unwrap();
408        assert!(format!("{err}").contains("no dim"));
409    }
410
411    #[test]
412    fn project_happy() {
413        let p = plan_one("SELECT * FROM ARRAY_PROJECT('g', ['qual', 'variant'])").unwrap();
414        match p {
415            SqlPlan::ArrayProject(ArrayProjectPlan {
416                name,
417                attr_projection,
418            }) => {
419                assert_eq!(name, "g");
420                assert_eq!(attr_projection, vec!["qual".to_string(), "variant".into()]);
421            }
422            other => panic!("expected ArrayProject, got {other:?}"),
423        }
424    }
425
426    #[test]
427    fn project_empty_rejected() {
428        assert!(plan_one("SELECT * FROM ARRAY_PROJECT('g', ARRAY[])").is_err());
429    }
430
431    #[test]
432    fn agg_scalar() {
433        let p = plan_one("SELECT * FROM ARRAY_AGG('g', 'qual', 'sum')").unwrap();
434        match p {
435            SqlPlan::ArrayAgg(ArrayAggPlan {
436                name,
437                attr,
438                reducer,
439                group_by_dim,
440                ..
441            }) => {
442                assert_eq!(name, "g");
443                assert_eq!(attr, "qual");
444                assert_eq!(reducer, ArrayReducerAst::Sum);
445                assert!(group_by_dim.is_none());
446            }
447            other => panic!("expected ArrayAgg, got {other:?}"),
448        }
449    }
450
451    #[test]
452    fn agg_grouped() {
453        let p = plan_one("SELECT * FROM ARRAY_AGG('g', 'qual', 'mean', 'chrom')").unwrap();
454        match p {
455            SqlPlan::ArrayAgg(ArrayAggPlan {
456                reducer,
457                group_by_dim,
458                ..
459            }) => {
460                assert_eq!(reducer, ArrayReducerAst::Mean);
461                assert_eq!(group_by_dim, Some("chrom".into()));
462            }
463            other => panic!("expected ArrayAgg, got {other:?}"),
464        }
465    }
466
467    #[test]
468    fn agg_unknown_reducer_rejected() {
469        assert!(plan_one("SELECT * FROM ARRAY_AGG('g', 'qual', 'bogus')").is_err());
470    }
471
472    #[test]
473    fn elementwise_happy() {
474        let p =
475            plan_one("SELECT * FROM ARRAY_ELEMENTWISE('left', 'right', 'add', 'qual')").unwrap();
476        assert!(matches!(p, SqlPlan::ArrayElementwise { .. }));
477    }
478
479    #[test]
480    fn elementwise_unknown_op_rejected() {
481        assert!(
482            plan_one("SELECT * FROM ARRAY_ELEMENTWISE('left', 'right', 'wat', 'qual')").is_err()
483        );
484    }
485}
nodedb_sql/planner/array_fn/table_fn.rs

nodedb_sql/planner/array_fn/
table_fn.rs
