Support memory tables in SQL.

crates/arroyo-df/src/lib.rs

@@ -604,7 +604,7 @@ pub async fn parse_and_get_arrow_program(
     }
 
     let mut used_connections = HashSet::new();
-    let mut plan_to_graph_visitor = PlanToGraphVisitor::new(&schema_provider);
+    let mut extensions = vec![];
 
     for insert in inserts {
         let (plan, sink_name) = match insert {
@@ -630,17 +630,29 @@ pub async fn parse_and_get_arrow_program(
         let sink = match sink_name {
             Some(sink_name) => {
                 let table = schema_provider
-                    .get_table(&sink_name)
+                    .get_table_mut(&sink_name)
                     .ok_or_else(|| anyhow!("Connection {} not found", sink_name))?;
-                let Table::ConnectorTable(_connector_table) = table else {
-                    bail!("expected connector table");
-                };
-                SinkExtension::new(
-                    OwnedTableReference::bare(sink_name),
-                    table.clone(),
-                    plan_rewrite.schema().clone(),
-                    Arc::new(plan_rewrite),
-                )
+                match table {
+                    Table::ConnectorTable(_) => SinkExtension::new(
+                        OwnedTableReference::bare(sink_name),
+                        table.clone(),
+                        plan_rewrite.schema().clone(),
+                        Arc::new(plan_rewrite),
+                    ),
+                    Table::MemoryTable { logical_plan, .. } => {
+                        if logical_plan.is_some() {
+                            bail!("Can only insert into a memory table once");
+                        }
+                        logical_plan.replace(plan_rewrite);
+                        continue;
+                    }
+                    Table::TableFromQuery { .. } => {
+                        bail!("Shouldn't be inserting more data into a table made with CREATE TABLE AS");
+                    }
+                    Table::PreviewSink { .. } => {
+                        bail!("queries shouldn't be able insert into preview sink.")
+                    }
+                }
             }
             None => SinkExtension::new(
                 OwnedTableReference::parse_str("preview"),
@@ -651,9 +663,13 @@ pub async fn parse_and_get_arrow_program(
                 Arc::new(plan_rewrite),
             ),
         };
-        plan_to_graph_visitor.add_plan(LogicalPlan::Extension(Extension {
+        extensions.push(LogicalPlan::Extension(Extension {
             node: Arc::new(sink),
-        }))?;
+        }));
+    }
+    let mut plan_to_graph_visitor = PlanToGraphVisitor::new(&schema_provider);
+    for extension in extensions {
+        plan_to_graph_visitor.add_plan(extension)?;
     }
     let graph = plan_to_graph_visitor.into_graph();
     let program = LogicalProgram {

crates/arroyo-df/src/rewriters.rs

@@ -2,6 +2,7 @@ use crate::extension::remote_table::RemoteTableExtension;
 use crate::extension::sink::SinkExtension;
 use crate::extension::table_source::TableSourceExtension;
 use crate::extension::watermark_node::WatermarkNode;
+use crate::schemas::add_timestamp_field;
 use crate::tables::ConnectorTable;
 use crate::tables::FieldSpec;
 use crate::tables::Table;
@@ -10,6 +11,7 @@ use crate::ArroyoSchemaProvider;
 use arrow_schema::DataType;
 use arroyo_rpc::TIMESTAMP_FIELD;
 
+use datafusion_common::plan_err;
 use datafusion_common::tree_node::{
     Transformed, TreeNode, TreeNodeRewriter, TreeNodeVisitor, VisitRecursion,
 };
@@ -222,7 +224,7 @@ impl<'a> TreeNodeRewriter for SourceRewriter<'a> {
     type N = LogicalPlan;
 
     fn mutate(&mut self, node: Self::N) -> DFResult<Self::N> {
-        let LogicalPlan::TableScan(table_scan) = node else {
+        let LogicalPlan::TableScan(mut table_scan) = node else {
             return Ok(node);
         };
 
@@ -234,9 +236,25 @@ impl<'a> TreeNodeRewriter for SourceRewriter<'a> {
 
         match table {
             Table::ConnectorTable(table) => self.mutate_connector_table(&table_scan, table),
-            Table::MemoryTable { .. } => Err(DataFusionError::Plan(
-                "Memory tables are not supported yet".to_string(),
-            )),
+            Table::MemoryTable {
+                name,
+                fields: _,
+                logical_plan,
+            } => {
+                let Some(logical_plan) = logical_plan else {
+                    return plan_err!(
+                        "Can't query from memory table {} without first inserting into it.",
+                        name
+                    );
+                };
+                // this can only be done here, otherwise the query planner will be upset about the timestamp column.
+                table_scan.projected_schema = add_timestamp_field(
+                    table_scan.projected_schema.clone(),
+                    Some(table_scan.table_name.clone()),
+                )?;
+
+                self.mutate_table_from_query(&table_scan, logical_plan)
+            }
             Table::TableFromQuery {
                 name: _,
                 logical_plan,

crates/arroyo-df/src/tables.rs

@@ -367,6 +367,7 @@ pub enum Table {
     MemoryTable {
         name: String,
         fields: Vec<FieldRef>,
+        logical_plan: Option<LogicalPlan>,
     },
     TableFromQuery {
         name: String,
@@ -513,6 +514,7 @@ impl Table {
                             .into_iter()
                             .map(|f| Arc::new(f.field().clone()))
                             .collect(),
+                        logical_plan: None,
                     }))
                 }
                 Some(connector) => {
@@ -588,7 +590,7 @@ impl Table {
 
     pub fn set_inferred_fields(&mut self, fields: Vec<DFField>) -> Result<()> {
         let Table::ConnectorTable(t) = self else {
-            bail!("can only infer schema for connector tables");
+            return Ok(());
         };
 
         if !t.fields.is_empty() {

crates/arroyo-df/src/test/queries/error_memory_table_double_write.sql

@@ -0,0 +1,31 @@
+--fail=Can only insert into a memory table once
+CREATE TABLE cars (
+      timestamp TIMESTAMP,
+      driver_id BIGINT,
+      event_type TEXT,
+      location TEXT
+    ) WITH (
+      connector = 'single_file',
+      path = '$input_dir/cars.json',
+      format = 'json',
+      type = 'source'
+    );
+
+    CREATE TABLE memory (
+      event_type TEXT,
+      location TEXT,
+      driver_id BIGINT,
+    );
+
+    CREATE TABLE cars_output (
+      driver_id BIGINT,
+      event_type TEXT
+    ) WITH (
+      connector = 'single_file',
+      path = '$output_path',
+      format = 'json',
+      type = 'sink'
+    );
+    INSERT INTO memory SELECT event_type, location, driver_id as other_driver_id FROM cars;
+    INSERT INTO memory SELECT location, event_type, driver_id FROM cars; 
+    INSERT INTO cars_output SELECT driver_id as other_driver_id, event_type FROM memory;

crates/arroyo-df/src/test/queries/error_read_from_unwritten_memory_table.sql

@@ -0,0 +1,8 @@
+--fail=Error during planning: Can't query from memory table memory without first inserting into it.
+    CREATE TABLE memory (
+      event_type TEXT,
+      location TEXT,
+      driver_id BIGINT,
+    );
+
+    SELECT * FROM memory;