query_planning/plan/
gpu_function_options.rs

1//! Shared dense-relation bindings and GPU UDTF option metadata.
2
3use std::collections::BTreeSet;
4
5use arrow::datatypes::DataType;
6use serde_json::{Value, json};
7
8/// Scalar ID types accepted by dense relational operations.
9///
10/// This is deliberately narrower than the set of Arrow integer types. Native
11/// cuVS output/gather paths are only defined for these five stable SQL types.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
13pub enum RelationIdType {
14    /// Signed 32-bit integer identifiers.
15    Int32,
16    /// Signed 64-bit integer identifiers.
17    Int64,
18    /// UTF-8 string identifiers with 32-bit offsets.
19    Utf8,
20    /// UTF-8 string identifiers with 64-bit offsets.
21    LargeUtf8,
22    /// UTF-8 identifiers stored as Arrow string views.
23    Utf8View,
24}
25
26impl RelationIdType {
27    /// Returns the supported binding type, or `None` for other Arrow types.
28    pub fn from_data_type(data_type: &DataType) -> Option<Self> {
29        match data_type {
30            DataType::Int32 => Some(Self::Int32),
31            DataType::Int64 => Some(Self::Int64),
32            DataType::Utf8 => Some(Self::Utf8),
33            DataType::LargeUtf8 => Some(Self::LargeUtf8),
34            DataType::Utf8View => Some(Self::Utf8View),
35            _ => None,
36        }
37    }
38}
39
40/// Stable order used in generated descriptions and validation diagnostics.
41pub const RELATION_ID_TYPES: &[RelationIdType] = &[
42    RelationIdType::Int32,
43    RelationIdType::Int64,
44    RelationIdType::Utf8,
45    RelationIdType::LargeUtf8,
46    RelationIdType::Utf8View,
47];
48
49/// Vector element types for dense relation descriptors.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
51pub enum VectorElementType {
52    /// Single-precision floating-point vector elements.
53    Float32,
54    /// Double-precision floating-point vector elements.
55    Float64,
56    /// Signed 8-bit vector elements.
57    Int8,
58    /// Unsigned 8-bit vector elements.
59    UInt8,
60}
61
62impl VectorElementType {
63    /// Returns the supported binding type, or `None` for other Arrow types.
64    pub fn from_data_type(data_type: &DataType) -> Option<Self> {
65        match data_type {
66            DataType::Float32 => Some(Self::Float32),
67            DataType::Float64 => Some(Self::Float64),
68            DataType::Int8 => Some(Self::Int8),
69            DataType::UInt8 => Some(Self::UInt8),
70            _ => None,
71        }
72    }
73
74    /// Arrow-compatible type label used in option schemas and diagnostics.
75    pub const fn as_str(self) -> &'static str {
76        match self {
77            Self::Float32 => "Float32",
78            Self::Float64 => "Float64",
79            Self::Int8 => "Int8",
80            Self::UInt8 => "UInt8",
81        }
82    }
83}
84
85/// Canonical scalar option metadata used by JSON schema generation and value
86/// validation in the DataFusion adapter.
87#[derive(Debug, Clone, Copy, PartialEq)]
88pub enum GpuOptionKind {
89    /// Object that names an ID column and a supported dense-vector binding.
90    RelationBinding,
91    /// Unsigned integer constrained by inclusive bounds.
92    Integer {
93        /// Inclusive lower bound.
94        minimum: u64,
95        /// Inclusive upper bound.
96        maximum: u64,
97        /// Value inserted when an optional key is absent; `None` leaves it absent.
98        default: Option<u64>,
99    },
100    /// Finite floating-point value greater than or equal to zero.
101    NonnegativeNumber {
102        /// Value inserted when an optional key is absent.
103        default: f64,
104    },
105    /// Finite floating-point value strictly greater than zero.
106    PositiveNumber {
107        /// Value inserted when an optional key is absent.
108        default: f64,
109    },
110    /// Finite floating-point value in the inclusive interval `[0, 1]`.
111    UnitInterval {
112        /// Value inserted when an optional key is absent.
113        default: f64,
114    },
115    /// Finite floating-point value in the interval `(0, 1]`.
116    Fraction {
117        /// Value inserted when an optional key is absent.
118        default: f64,
119    },
120    /// JSON boolean with an optional default.
121    Boolean {
122        /// Value inserted when an optional key is absent; `None` leaves it absent.
123        default: Option<bool>,
124    },
125    /// String selected from the declared values.
126    Enum {
127        /// Accepted strings, matched case-sensitively.
128        values: &'static [&'static str],
129        /// Value inserted when an optional key is absent; `None` leaves it absent.
130        default: Option<&'static str>,
131    },
132    /// Non-empty string of at most `max_chars` Unicode scalar values, the unit
133    /// JSON Schema `maxLength` counts.
134    Text {
135        /// Maximum Unicode scalar values in the string.
136        max_chars: usize,
137    },
138}
139
140impl GpuOptionKind {
141    pub(crate) const fn default_u64(self) -> Option<u64> {
142        match self {
143            Self::Integer { default, .. } => default,
144            _ => None,
145        }
146    }
147
148    pub(crate) const fn default_f64(self) -> Option<f64> {
149        match self {
150            Self::NonnegativeNumber { default }
151            | Self::PositiveNumber { default }
152            | Self::UnitInterval { default }
153            | Self::Fraction { default } => Some(default),
154            _ => None,
155        }
156    }
157
158    pub(crate) const fn default_bool(self) -> Option<bool> {
159        match self {
160            Self::Boolean { default } => default,
161            _ => None,
162        }
163    }
164
165    pub(crate) const fn default_str(self) -> Option<&'static str> {
166        match self {
167            Self::Enum { default, .. } => default,
168            _ => None,
169        }
170    }
171}
172
173/// One declared option of a GPU table function: its JSON key, whether a call
174/// may omit it, and the value contract that validation and schema generation
175/// share.
176#[derive(Debug, Clone, Copy, PartialEq)]
177pub struct GpuOptionDefinition {
178    /// Case-sensitive key accepted in the option object.
179    pub name: &'static str,
180    /// A missing required option is an error; a missing optional one takes
181    /// the kind's default when it has one and is otherwise left absent.
182    pub required: bool,
183    /// Value constraints used by both validation and JSON schema generation.
184    pub kind: GpuOptionKind,
185    /// User-facing semantics projected through the unified GPU descriptor.
186    pub description: &'static str,
187}
188
189/// One documented output field in a GPU relational result.
190#[derive(Debug, Clone, Copy, PartialEq, Eq)]
191pub struct GpuOutputFieldDefinition {
192    /// Output field name or repeated-field template used to construct the result schema.
193    pub name: &'static str,
194    /// SQL type label for the generated reference and descriptor.
195    pub data_type: &'static str,
196    /// Whether the result field permits nulls.
197    pub nullable: bool,
198    /// Literal option that determines how many physical output fields exist.
199    pub repeats: Option<&'static str>,
200    /// Meaning of the output field in the generated descriptor.
201    pub description: &'static str,
202}
203
204/// Rejection of a GPU function's option object. `code` is a stable error
205/// identifier read by the adapter's error mapping; `message` names the
206/// function family and is for people.
207#[derive(Debug, Clone, PartialEq, Eq)]
208pub struct GpuOptionValidationError {
209    code: &'static str,
210    message: String,
211    argument: Option<&'static str>,
212}
213
214impl GpuOptionValidationError {
215    /// Stable error code for adapter classification.
216    pub fn code(&self) -> &'static str {
217        self.code
218    }
219
220    /// The declared option this error concerns, when it is attributable to one.
221    pub fn argument(&self) -> Option<&'static str> {
222        self.argument
223    }
224
225    pub(crate) fn for_argument(mut self, argument: &'static str) -> Self {
226        self.argument = Some(argument);
227        self
228    }
229
230    /// Human-readable explanation of the invalid option.
231    pub fn message(&self) -> &str {
232        &self.message
233    }
234
235    pub(crate) fn new(code: &'static str, message: impl Into<String>) -> Self {
236        Self { code, message: message.into(), argument: None }
237    }
238}
239
240pub(crate) fn dense_relation_binding_schema(elements: &[VectorElementType]) -> Value {
241    let elements_label = elements.iter().map(|element| element.as_str()).collect::<Vec<_>>().join(", ");
242    let list_types = elements
243        .iter()
244        .map(|element| {
245            format!(
246                "FixedSizeList<{}, D>, List<{}>, or LargeList<{}>",
247                element.as_str(),
248                element.as_str(),
249                element.as_str()
250            )
251        })
252        .collect::<Vec<_>>()
253        .join("; ");
254    json!({
255        "type": "object", "additionalProperties": false, "required": ["id", "vector"],
256        "properties": {
257            "id": {"type": "string", "minLength": 1,
258                "description": "Logical row ID column; a null ID fails the query. IDs may repeat; result ordinals disambiguate physical rows."},
259            "vector": {"description": "Choose exactly one dense-vector binding shape.", "oneOf": [
260                {"type": "object", "additionalProperties": false, "required": ["columns"], "properties": {
261                    "columns": {"type": "array", "minItems": 1, "uniqueItems": true,
262                        "description": format!("Ordered, unique {elements_label} feature columns without null values; order defines vector dimensions."),
263                        "items": {"type": "string", "minLength": 1}}}},
264                {"type": "object", "additionalProperties": false, "required": ["column"], "properties": {
265                    "column": {"type": "string", "minLength": 1, "description": format!("One {list_types} column without null rows or null elements.")}}}
266            ]}
267        }
268    })
269}
270
271pub(crate) fn validate_dense_relation_binding_value(
272    family: &str,
273    option: &GpuOptionDefinition,
274    value: &Value,
275    allow_label: bool,
276) -> Result<(), GpuOptionValidationError> {
277    let invalid = || {
278        GpuOptionValidationError::new("missing_or_invalid_option", format!("invalid {family} option '{}'", option.name))
279    };
280    let binding = value.as_object().ok_or_else(invalid)?;
281    if let Some(unknown) =
282        binding.keys().find(|key| !matches!(key.as_str(), "id" | "vector") && !(allow_label && key.as_str() == "label"))
283    {
284        return Err(GpuOptionValidationError::new(
285            "unknown_option",
286            format!("unknown {family} option '{unknown}' in {}", option.name),
287        ));
288    }
289    if binding.get("id").and_then(Value::as_str).filter(|value| !value.is_empty()).is_none() {
290        return Err(GpuOptionValidationError::new(
291            "missing_or_invalid_option",
292            format!("{family} option '{}.id' must be a non-empty string", option.name),
293        ));
294    }
295    if allow_label && binding.get("label").and_then(Value::as_str) != Some("label") {
296        return Err(GpuOptionValidationError::new(
297            "missing_or_invalid_option",
298            format!("{family} option '{}.label' must be 'label'", option.name),
299        ));
300    }
301    let vector = binding.get("vector").and_then(Value::as_object).ok_or_else(|| {
302        GpuOptionValidationError::new(
303            "missing_or_invalid_option",
304            format!("{family} option '{}.vector' must be an object", option.name),
305        )
306    })?;
307    if let Some(unknown) = vector.keys().find(|key| !matches!(key.as_str(), "columns" | "column")) {
308        return Err(GpuOptionValidationError::new(
309            "unknown_option",
310            format!("unknown {family} option '{unknown}' in {}.vector", option.name),
311        ));
312    }
313    match (vector.get("columns"), vector.get("column")) {
314        (Some(columns), None) => {
315            let columns = columns.as_array().filter(|columns| !columns.is_empty()).ok_or_else(|| {
316                GpuOptionValidationError::new(
317                    "invalid_vector_binding",
318                    format!("{family} option '{}.vector.columns' must be a non-empty string array", option.name),
319                )
320            })?;
321            let mut unique = BTreeSet::new();
322            for column in columns {
323                let column = column.as_str().filter(|column| !column.is_empty()).ok_or_else(|| {
324                    GpuOptionValidationError::new(
325                        "invalid_vector_binding",
326                        format!("{family} option '{}.vector.columns' entries must be non-empty strings", option.name),
327                    )
328                })?;
329                if !unique.insert(column) {
330                    return Err(GpuOptionValidationError::new(
331                        "duplicate_vector_column",
332                        format!("{family} option '{}.vector.columns' contains duplicate field '{column}'", option.name),
333                    ));
334                }
335            }
336        },
337        (None, Some(column)) => {
338            if column.as_str().filter(|column| !column.is_empty()).is_none() {
339                return Err(GpuOptionValidationError::new(
340                    "invalid_vector_binding",
341                    format!("{family} option '{}.vector.column' must be a non-empty string", option.name),
342                ));
343            }
344        },
345        _ => {
346            return Err(GpuOptionValidationError::new(
347                "invalid_vector_binding",
348                format!("{family} option '{}.vector' must contain exactly one of columns or column", option.name),
349            ));
350        },
351    }
352    Ok(())
353}
354
355pub(crate) fn option_schema(kind: GpuOptionKind, relation_schema: &Value) -> Value {
356    match kind {
357        GpuOptionKind::RelationBinding => relation_schema.clone(),
358        GpuOptionKind::Integer { minimum, maximum, default } => {
359            let mut value = json!({"type": "integer", "minimum": minimum, "maximum": maximum});
360            if let Some(default) = default {
361                value["default"] = json!(default);
362            }
363            value
364        },
365        GpuOptionKind::NonnegativeNumber { default } => json!({"type": "number", "minimum": 0.0, "default": default}),
366        GpuOptionKind::PositiveNumber { default } => {
367            json!({"type": "number", "exclusiveMinimum": 0, "default": default})
368        },
369        GpuOptionKind::UnitInterval { default } => {
370            json!({"type": "number", "minimum": 0, "maximum": 1, "default": default})
371        },
372        GpuOptionKind::Fraction { default } => {
373            json!({"type": "number", "exclusiveMinimum": 0, "maximum": 1, "default": default})
374        },
375        GpuOptionKind::Boolean { default } => {
376            let mut value = json!({"type": "boolean"});
377            if let Some(default) = default {
378                value["default"] = json!(default);
379            }
380            value
381        },
382        GpuOptionKind::Enum { values, default } => {
383            let mut value = json!({"type": "string", "enum": values});
384            if let Some(default) = default {
385                value["default"] = json!(default);
386            }
387            value
388        },
389        GpuOptionKind::Text { max_chars } => json!({"type": "string", "minLength": 1, "maxLength": max_chars}),
390    }
391}
392
393/// Parses a function's raw JSON options into the canonical object the plan
394/// carries: unknown keys rejected, every optional key with a default filled
395/// in, every present value validated. Filling defaults here makes two calls
396/// that differ only in spelled-out defaults build equal nodes, which the
397/// builder then interns. Each error names the offending option when one is
398/// attributable.
399pub(crate) fn canonicalize_options(
400    options: &[GpuOptionDefinition],
401    raw: &str,
402    family: &str,
403    validate_relation: impl Fn(&GpuOptionDefinition, &Value) -> Result<(), GpuOptionValidationError>,
404) -> Result<serde_json::Map<String, Value>, GpuOptionValidationError> {
405    let parsed: Value = serde_json::from_str(raw).map_err(|_| {
406        GpuOptionValidationError::new("invalid_options_json", format!("{family} options must be valid JSON"))
407    })?;
408    let mut object = parsed.as_object().cloned().ok_or_else(|| {
409        GpuOptionValidationError::new("options_not_object", format!("{family} options must be a JSON object"))
410    })?;
411    if let Some(unknown) = object.keys().find(|key| !options.iter().any(|option| option.name == key.as_str())) {
412        return Err(GpuOptionValidationError::new("unknown_option", format!("unknown {family} option '{unknown}'")));
413    }
414    for option in options {
415        let Some(value) = object.get_mut(option.name) else {
416            if option.required {
417                return Err(GpuOptionValidationError::new(
418                    "missing_or_invalid_option",
419                    format!("{family} option '{}' is required", option.name),
420                )
421                .for_argument(option.name));
422            }
423            let default = match option.kind {
424                GpuOptionKind::Integer { default: Some(value), .. } => json!(value),
425                GpuOptionKind::NonnegativeNumber { default }
426                | GpuOptionKind::PositiveNumber { default }
427                | GpuOptionKind::UnitInterval { default }
428                | GpuOptionKind::Fraction { default } => json!(default),
429                GpuOptionKind::Boolean { default: Some(value) } => json!(value),
430                GpuOptionKind::Enum { default: Some(value), .. } => json!(value),
431                _ => continue,
432            };
433            object.insert(option.name.to_string(), default);
434            continue;
435        };
436        validate_option_value(option, value, family, &validate_relation)
437            .map_err(|error| error.for_argument(option.name))?;
438    }
439    Ok(object)
440}
441
442pub(crate) fn validate_option_value(
443    option: &GpuOptionDefinition,
444    value: &Value,
445    family: &str,
446    validate_relation: &impl Fn(&GpuOptionDefinition, &Value) -> Result<(), GpuOptionValidationError>,
447) -> Result<(), GpuOptionValidationError> {
448    let invalid = || {
449        GpuOptionValidationError::new("missing_or_invalid_option", format!("invalid {family} option '{}'", option.name))
450    };
451    match option.kind {
452        GpuOptionKind::RelationBinding => validate_relation(option, value)?,
453        GpuOptionKind::Integer { minimum, maximum, .. } => {
454            if value.as_u64().filter(|value| *value >= minimum && *value <= maximum).is_none() {
455                return Err(invalid());
456            }
457        },
458        GpuOptionKind::NonnegativeNumber { .. } => {
459            if value.as_f64().filter(|value| value.is_finite() && *value >= 0.0).is_none() {
460                return Err(invalid());
461            }
462        },
463        GpuOptionKind::PositiveNumber { .. } => {
464            if value.as_f64().filter(|value| value.is_finite() && *value > 0.0).is_none() {
465                return Err(invalid());
466            }
467        },
468        GpuOptionKind::UnitInterval { .. } => {
469            if value.as_f64().filter(|value| value.is_finite() && (0.0..=1.0).contains(value)).is_none() {
470                return Err(invalid());
471            }
472        },
473        GpuOptionKind::Fraction { .. } => {
474            if value.as_f64().filter(|value| value.is_finite() && *value > 0.0 && *value <= 1.0).is_none() {
475                return Err(invalid());
476            }
477        },
478        GpuOptionKind::Boolean { .. } => {
479            if !value.is_boolean() {
480                return Err(invalid());
481            }
482        },
483        GpuOptionKind::Enum { values, .. } => {
484            if value.as_str().is_none_or(|value| !values.contains(&value)) {
485                return Err(invalid());
486            }
487        },
488        GpuOptionKind::Text { max_chars } => {
489            if value.as_str().is_none_or(|value| value.is_empty() || value.chars().count() > max_chars) {
490                return Err(invalid());
491            }
492        },
493    }
494    Ok(())
495}
496
497/// Binding from one native relation to a dense vector matrix.
498#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
499pub enum DenseVectorBinding {
500    /// Separate columns of one supported element type packed in the listed order.
501    Columns(Vec<String>),
502    /// One fixed-size or variable-size list column of a supported element type.
503    VectorColumn(String),
504}
505
506impl DenseVectorBinding {
507    pub(crate) fn columns(columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
508        Self::Columns(columns.into_iter().map(Into::into).collect())
509    }
510
511    pub(crate) fn vector_column(column: impl Into<String>) -> Self {
512        Self::VectorColumn(column.into())
513    }
514
515    pub(crate) fn display_label(&self) -> String {
516        match self {
517            Self::Columns(columns) => columns.join("|"),
518            Self::VectorColumn(column) => format!("[{column}]"),
519        }
520    }
521}
query_planning/plan/gpu_function_options.rs

query_planning/plan/
gpu_function_options.rs
