Skip to main content

laminar_sql/planner/
lookup_table.rs

1//! Lookup table planning with declared or resolved reader fields.
2
3use super::{
4    object_name_to_string, Arc, Field, LookupTableInfo, ObjectName, PlanningError, Schema,
5    SchemaRef, StreamingPlan, StreamingPlanner,
6};
7use crate::parser::lookup_table::validate_properties;
8use crate::parser::CreateLookupTableStatement;
9
10impl StreamingPlanner {
11    /// Plans a CREATE LOOKUP TABLE statement.
12    pub(super) fn plan_create_lookup_table(
13        &mut self,
14        lt: &CreateLookupTableStatement,
15    ) -> Result<StreamingPlan, PlanningError> {
16        let name = object_name_to_string(&lt.name);
17
18        if !lt.or_replace && !lt.if_not_exists && self.lookup_tables.contains_key(&name) {
19            return Err(PlanningError::InvalidQuery(format!(
20                "Lookup table '{}' already exists",
21                name
22            )));
23        }
24
25        let columns: Vec<(String, String)> = lt
26            .columns
27            .iter()
28            .map(|c| (c.name.value.clone(), c.data_type.to_string()))
29            .collect();
30
31        let properties = validate_properties(&lt.with_options).map_err(|e| {
32            PlanningError::InvalidQuery(format!("Invalid lookup table properties: {e}"))
33        })?;
34
35        // Compute Arrow schema from column definitions
36        let arrow_fields: Vec<Field> = lt
37            .columns
38            .iter()
39            .map(|c| {
40                let dt = crate::translator::streaming_ddl::sql_type_to_arrow(&c.data_type)
41                    .map_err(|e| PlanningError::InvalidQuery(e.to_string()))?;
42                let nullable = !c
43                    .options
44                    .iter()
45                    .any(|opt| matches!(opt.option, sqlparser::ast::ColumnOption::NotNull));
46                Ok(Field::new(&c.name.value, dt, nullable))
47            })
48            .collect::<Result<_, PlanningError>>()?;
49        let arrow_schema = Arc::new(Schema::new(arrow_fields));
50
51        let info = LookupTableInfo {
52            name: name.clone(),
53            columns,
54            primary_key: lt.primary_key.clone(),
55            properties,
56            arrow_schema,
57            raw_options: lt.with_options.clone(),
58        };
59
60        self.lookup_tables.insert(name, info.clone());
61
62        Ok(StreamingPlan::RegisterLookupTable(info))
63    }
64
65    /// Bind a lookup table to already resolved Arrow fields without a lossy SQL round trip.
66    ///
67    /// # Errors
68    /// Returns the existing property/planner errors, or an empty resolved schema error.
69    pub fn plan_lookup_table_with_schema(
70        &mut self,
71        create: &CreateLookupTableStatement,
72        schema: SchemaRef,
73    ) -> Result<StreamingPlan, PlanningError> {
74        if schema.fields().is_empty() {
75            return Err(PlanningError::InvalidQuery(
76                "resolved lookup reader is empty".into(),
77            ));
78        }
79        let StreamingPlan::RegisterLookupTable(mut info) = self.plan_create_lookup_table(create)?
80        else {
81            return Err(PlanningError::InvalidQuery(
82                "lookup planner returned a different statement".into(),
83            ));
84        };
85        info.columns = schema
86            .fields()
87            .iter()
88            .map(|field| (field.name().clone(), field.data_type().to_string()))
89            .collect();
90        info.arrow_schema = schema;
91        self.lookup_tables.insert(info.name.clone(), info.clone());
92        Ok(StreamingPlan::RegisterLookupTable(info))
93    }
94
95    /// Plans a DROP LOOKUP TABLE statement.
96    pub(super) fn plan_drop_lookup_table(
97        &mut self,
98        name: &ObjectName,
99        if_exists: bool,
100    ) -> Result<StreamingPlan, PlanningError> {
101        let name_str = object_name_to_string(name);
102
103        if !if_exists && !self.lookup_tables.contains_key(&name_str) {
104            return Err(PlanningError::InvalidQuery(format!(
105                "Lookup table '{}' does not exist",
106                name_str
107            )));
108        }
109
110        self.lookup_tables.remove(&name_str);
111
112        Ok(StreamingPlan::DropLookupTable { name: name_str })
113    }
114}