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 sqlparser::ast;
7
8use nodedb_types::Value;
9
10use super::helpers::{
11    collect_args, expect_string_array, expect_string_literal, expect_u32, is_null_literal,
12    require_array_name, value_to_coord_literal,
13};
14use crate::error::{Result, SqlError};
15use crate::parser::object_literal::parse_object_literal;
16use crate::temporal::TemporalScope;
17use crate::types::{SqlCatalog, SqlPlan};
18use crate::types_array::{ArrayBinaryOpAst, ArrayReducerAst, ArraySliceAst, NamedDimRange};
19
20/// Try to intercept a `SELECT * FROM array_xxx(...)` table-valued
21/// function call. Returns `Ok(Some(plan))` on a match, `Ok(None)` if
22/// the FROM is not an array function (caller falls through to normal
23/// catalog resolution).
24///
25/// `temporal` carries any `AS OF SYSTEM TIME` / `AS OF VALID TIME` qualifiers
26/// extracted by the pre-processor. It is propagated verbatim into
27/// `SqlPlan::ArraySlice` and `SqlPlan::ArrayAgg` so the planner-to-physical
28/// conversion can populate the corresponding `ArrayOp` fields. When neither
29/// clause was present, `temporal` is `TemporalScope::default()`, which maps to
30/// the live-state fast path in the Data Plane handler.
31pub fn try_plan_array_table_fn(
32    from: &[ast::TableWithJoins],
33    catalog: &dyn SqlCatalog,
34    temporal: TemporalScope,
35) -> Result<Option<SqlPlan>> {
36    if from.len() != 1 {
37        return Ok(None);
38    }
39    let twj = &from[0];
40    if !twj.joins.is_empty() {
41        return Ok(None);
42    }
43    let (name, args) = match &twj.relation {
44        ast::TableFactor::Table {
45            name,
46            args: Some(args),
47            ..
48        } => (name, args),
49        _ => return Ok(None),
50    };
51    let fn_name = crate::parser::normalize::normalize_object_name_checked(name)?;
52    let arg_exprs = collect_args(&args.args);
53    match fn_name.as_str() {
54        "array_slice" => Ok(Some(plan_slice(&arg_exprs, catalog, temporal)?)),
55        "array_project" => Ok(Some(plan_project(&arg_exprs, catalog)?)),
56        "array_agg" => Ok(Some(plan_agg(&arg_exprs, catalog, temporal)?)),
57        "array_elementwise" => Ok(Some(plan_elementwise(&arg_exprs, catalog)?)),
58        _ => Ok(None),
59    }
60}
61
62fn plan_slice(
63    args: &[ast::Expr],
64    catalog: &dyn SqlCatalog,
65    temporal: TemporalScope,
66) -> Result<SqlPlan> {
67    if args.len() < 2 || args.len() > 4 {
68        return Err(SqlError::Unsupported {
69            detail: format!(
70                "ARRAY_SLICE expects 2..=4 args (name, slice_obj, [attrs], [limit]); got {}",
71                args.len()
72            ),
73        });
74    }
75    let name = require_array_name(args, 0, "ARRAY_SLICE", catalog)?;
76    let view = catalog
77        .lookup_array(&name)
78        .ok_or_else(|| SqlError::Unsupported {
79            detail: format!("ARRAY_SLICE: array '{name}' not found"),
80        })?;
81
82    // Slice-predicate literal: encoded as a quoted string carrying the
83    // brace-form object literal. The PostgreSQL dialect does not accept
84    // bare `{...}` in expression position, so we decode the string
85    // contents here.
86    let slice_str = expect_string_literal(&args[1], "ARRAY_SLICE slice predicate")?;
87    let parsed = parse_object_literal(&slice_str).ok_or_else(|| SqlError::Unsupported {
88        detail: format!("ARRAY_SLICE: slice predicate must be an object literal: {slice_str}"),
89    })?;
90    let map = parsed.map_err(|detail| SqlError::Unsupported {
91        detail: format!("ARRAY_SLICE: slice parse: {detail}"),
92    })?;
93    let mut dim_ranges: Vec<NamedDimRange> = Vec::with_capacity(map.len());
94    for (dim, val) in map {
95        // Verify the dim exists on the array.
96        if !view.dims.iter().any(|d| d.name == dim) {
97            return Err(SqlError::Unsupported {
98                detail: format!("ARRAY_SLICE: array '{name}' has no dim '{dim}'"),
99            });
100        }
101        let arr = match val {
102            Value::Array(a) if a.len() == 2 => a,
103            _ => {
104                return Err(SqlError::Unsupported {
105                    detail: format!(
106                        "ARRAY_SLICE: dim '{dim}' range must be a 2-element array [lo, hi]"
107                    ),
108                });
109            }
110        };
111        let lo = value_to_coord_literal(&arr[0], &dim)?;
112        let hi = value_to_coord_literal(&arr[1], &dim)?;
113        dim_ranges.push(NamedDimRange { dim, lo, hi });
114    }
115
116    let attr_projection = if args.len() >= 3 {
117        match &args[2] {
118            ast::Expr::Value(v) if matches!(v.value, ast::Value::SingleQuotedString(ref s) if s == "*") => {
119                Vec::new()
120            }
121            _ => expect_string_array(&args[2], "ARRAY_SLICE attr projection")?,
122        }
123    } else {
124        Vec::new()
125    };
126    // Validate attr names against the catalog.
127    for attr in &attr_projection {
128        if !view.attrs.iter().any(|a| &a.name == attr) {
129            return Err(SqlError::Unsupported {
130                detail: format!("ARRAY_SLICE: array '{name}' has no attr '{attr}'"),
131            });
132        }
133    }
134
135    let limit = if args.len() >= 4 {
136        expect_u32(&args[3], "ARRAY_SLICE limit")?
137    } else {
138        0
139    };
140
141    Ok(SqlPlan::ArraySlice {
142        name,
143        slice: ArraySliceAst { dim_ranges },
144        attr_projection,
145        limit,
146        temporal,
147    })
148}
149
150fn plan_project(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result<SqlPlan> {
151    if args.len() != 2 {
152        return Err(SqlError::Unsupported {
153            detail: format!(
154                "ARRAY_PROJECT expects 2 args (name, [attrs]); got {}",
155                args.len()
156            ),
157        });
158    }
159    let name = require_array_name(args, 0, "ARRAY_PROJECT", catalog)?;
160    let view = catalog
161        .lookup_array(&name)
162        .ok_or_else(|| SqlError::Unsupported {
163            detail: format!("ARRAY_PROJECT: array '{name}' not found"),
164        })?;
165    let attr_projection = expect_string_array(&args[1], "ARRAY_PROJECT attrs")?;
166    if attr_projection.is_empty() {
167        return Err(SqlError::Unsupported {
168            detail: "ARRAY_PROJECT: attr list must not be empty".into(),
169        });
170    }
171    for attr in &attr_projection {
172        if !view.attrs.iter().any(|a| &a.name == attr) {
173            return Err(SqlError::Unsupported {
174                detail: format!("ARRAY_PROJECT: array '{name}' has no attr '{attr}'"),
175            });
176        }
177    }
178    Ok(SqlPlan::ArrayProject {
179        name,
180        attr_projection,
181    })
182}
183
184fn plan_agg(
185    args: &[ast::Expr],
186    catalog: &dyn SqlCatalog,
187    temporal: TemporalScope,
188) -> Result<SqlPlan> {
189    if args.len() < 3 || args.len() > 4 {
190        return Err(SqlError::Unsupported {
191            detail: format!(
192                "ARRAY_AGG expects 3..=4 args (name, attr, reducer, [group_by_dim]); got {}",
193                args.len()
194            ),
195        });
196    }
197    let name = require_array_name(args, 0, "ARRAY_AGG", catalog)?;
198    let view = catalog
199        .lookup_array(&name)
200        .ok_or_else(|| SqlError::Unsupported {
201            detail: format!("ARRAY_AGG: array '{name}' not found"),
202        })?;
203
204    let attr = expect_string_literal(&args[1], "ARRAY_AGG attr")?;
205    if !view.attrs.iter().any(|a| a.name == attr) {
206        return Err(SqlError::Unsupported {
207            detail: format!("ARRAY_AGG: array '{name}' has no attr '{attr}'"),
208        });
209    }
210
211    let reducer_str = expect_string_literal(&args[2], "ARRAY_AGG reducer")?;
212    let reducer = ArrayReducerAst::parse(&reducer_str).ok_or_else(|| SqlError::Unsupported {
213        detail: format!("ARRAY_AGG: unknown reducer '{reducer_str}' (want sum/count/min/max/mean)"),
214    })?;
215
216    let group_by_dim = if args.len() == 4 && !is_null_literal(&args[3]) {
217        let dim = expect_string_literal(&args[3], "ARRAY_AGG group_by_dim")?;
218        if !view.dims.iter().any(|d| d.name == dim) {
219            return Err(SqlError::Unsupported {
220                detail: format!("ARRAY_AGG: array '{name}' has no dim '{dim}'"),
221            });
222        }
223        Some(dim)
224    } else {
225        None
226    };
227
228    Ok(SqlPlan::ArrayAgg {
229        name,
230        attr,
231        reducer,
232        group_by_dim,
233        temporal,
234    })
235}
236
237fn plan_elementwise(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result<SqlPlan> {
238    if args.len() != 4 {
239        return Err(SqlError::Unsupported {
240            detail: format!(
241                "ARRAY_ELEMENTWISE expects 4 args (left, right, op, attr); got {}",
242                args.len()
243            ),
244        });
245    }
246    let left = require_array_name(args, 0, "ARRAY_ELEMENTWISE", catalog)?;
247    let right = require_array_name(args, 1, "ARRAY_ELEMENTWISE", catalog)?;
248    let lview = catalog
249        .lookup_array(&left)
250        .ok_or_else(|| SqlError::Unsupported {
251            detail: format!("ARRAY_ELEMENTWISE: array '{left}' not found"),
252        })?;
253    let rview = catalog
254        .lookup_array(&right)
255        .ok_or_else(|| SqlError::Unsupported {
256            detail: format!("ARRAY_ELEMENTWISE: array '{right}' not found"),
257        })?;
258    if lview.dims.len() != rview.dims.len() || lview.attrs.len() != rview.attrs.len() {
259        return Err(SqlError::Unsupported {
260            detail: format!(
261                "ARRAY_ELEMENTWISE: arrays '{left}' and '{right}' must share schema shape"
262            ),
263        });
264    }
265    let op_str = expect_string_literal(&args[2], "ARRAY_ELEMENTWISE op")?;
266    let op = ArrayBinaryOpAst::parse(&op_str).ok_or_else(|| SqlError::Unsupported {
267        detail: format!("ARRAY_ELEMENTWISE: unknown op '{op_str}' (want add/sub/mul/div)"),
268    })?;
269    let attr = expect_string_literal(&args[3], "ARRAY_ELEMENTWISE attr")?;
270    if !lview.attrs.iter().any(|a| a.name == attr) {
271        return Err(SqlError::Unsupported {
272            detail: format!("ARRAY_ELEMENTWISE: array '{left}' has no attr '{attr}'"),
273        });
274    }
275    Ok(SqlPlan::ArrayElementwise {
276        left,
277        right,
278        op,
279        attr,
280    })
281}
nodedb_sql/planner/array_fn/table_fn.rs

nodedb_sql/planner/array_fn/
table_fn.rs
