Skip to main content

nodedb_sql/resolver/
array_tvf.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Array table-valued function resolution.
4//!
5//! An `ARRAY_*(...)` factor in a FROM clause names an array, not a
6//! collection. This module synthesizes the relation its dims and attrs
7//! expose so column references and equi-join keys against it resolve.
8
9use crate::error::{Result, SqlError};
10use crate::parser::normalize::normalize_object_name_checked;
11use crate::resolver::columns::ResolvedTable;
12use crate::types::{
13    ArrayCatalogView, CollectionInfo, ColumnInfo, EngineType, SqlCatalog, SqlDataType,
14};
15use crate::types_array::{ArrayAttrType, ArrayDimType};
16
17/// If `factor` is `ARRAY_*(name, ...)`, look up the array via the
18/// catalog and build a `ResolvedTable` whose columns mirror the array's
19/// dims + attrs. Returns `Ok(None)` for any non-array-TVF factor.
20pub(super) fn resolve_array_tvf(
21    catalog: &dyn SqlCatalog,
22    factor: &sqlparser::ast::TableFactor,
23) -> Result<Option<ResolvedTable>> {
24    let (fn_name, args, alias) = match factor {
25        sqlparser::ast::TableFactor::Table {
26            name,
27            args: Some(args),
28            alias,
29            ..
30        } => (
31            normalize_object_name_checked(name)?,
32            args,
33            alias
34                .as_ref()
35                .map(|alias| crate::reserved::check_ast_identifier(&alias.name))
36                .transpose()?,
37        ),
38        _ => return Ok(None),
39    };
40    if !matches!(
41        fn_name.as_str(),
42        "array_slice" | "array_project" | "array_agg" | "array_elementwise"
43    ) {
44        return Ok(None);
45    }
46
47    // First positional arg is the array name as a string literal.
48    let first = args.args.first().ok_or_else(|| SqlError::Unsupported {
49        detail: format!("{fn_name}: missing array-name argument"),
50    })?;
51    let array_name = extract_string_literal_arg(first).ok_or_else(|| SqlError::Unsupported {
52        detail: format!("{fn_name}: array-name argument must be a string literal"),
53    })?;
54    let view = catalog
55        .lookup_array(&array_name)
56        .ok_or_else(|| SqlError::UnknownTable {
57            name: array_name.clone(),
58        })?;
59
60    let info = CollectionInfo {
61        name: view.name.clone(),
62        engine: EngineType::Array,
63        columns: array_columns(&view),
64        primary_key: None,
65        has_auto_tier: false,
66        indexes: Vec::new(),
67        bitemporal: false,
68        primary: nodedb_types::PrimaryEngine::Document,
69        vector_primary: None,
70        partition_strategy: nodedb_types::PartitionStrategy::CollectionHomed,
71        open_schema: false,
72    };
73    Ok(Some(ResolvedTable {
74        name: view.name,
75        alias,
76        info,
77    }))
78}
79
80fn array_columns(view: &ArrayCatalogView) -> Vec<ColumnInfo> {
81    let mut cols = Vec::with_capacity(view.dims.len() + view.attrs.len());
82    for d in &view.dims {
83        cols.push(ColumnInfo {
84            name: d.name.clone(),
85            data_type: dim_type_to_sql(d.dtype),
86            nullable: false,
87            is_primary_key: false,
88            default: None,
89            raw_type: None,
90            int_width: None,
91            float_width: None,
92        });
93    }
94    for a in &view.attrs {
95        cols.push(ColumnInfo {
96            name: a.name.clone(),
97            data_type: attr_type_to_sql(a.dtype),
98            nullable: a.nullable,
99            is_primary_key: false,
100            default: None,
101            raw_type: None,
102            int_width: None,
103            float_width: None,
104        });
105    }
106    cols
107}
108
109fn dim_type_to_sql(t: ArrayDimType) -> SqlDataType {
110    match t {
111        ArrayDimType::Int64 => SqlDataType::Int64,
112        ArrayDimType::Float64 => SqlDataType::Float64,
113        ArrayDimType::TimestampMs => SqlDataType::Timestamp,
114        ArrayDimType::String => SqlDataType::String,
115    }
116}
117
118fn attr_type_to_sql(t: ArrayAttrType) -> SqlDataType {
119    match t {
120        ArrayAttrType::Int64 => SqlDataType::Int64,
121        ArrayAttrType::Float64 => SqlDataType::Float64,
122        ArrayAttrType::String => SqlDataType::String,
123        ArrayAttrType::Bytes => SqlDataType::Bytes,
124    }
125}
126
127fn extract_string_literal_arg(arg: &sqlparser::ast::FunctionArg) -> Option<String> {
128    use sqlparser::ast::{Expr, FunctionArg, FunctionArgExpr, Value};
129    let expr = match arg {
130        FunctionArg::Unnamed(FunctionArgExpr::Expr(e)) => e,
131        FunctionArg::Named {
132            arg: FunctionArgExpr::Expr(e),
133            ..
134        } => e,
135        _ => return None,
136    };
137    match expr {
138        Expr::Value(v) => match &v.value {
139            Value::SingleQuotedString(s) => Some(s.clone()),
140            _ => None,
141        },
142        _ => None,
143    }
144}