Support streaming SQL IN subqueries

dozer-sql/src/builder/mod.rs

@@ -2,11 +2,16 @@ use crate::aggregation::factory::AggregationProcessorFactory;
 use crate::builder::PipelineError::InvalidQuery;
 use crate::errors::PipelineError;
 use crate::selection::factory::SelectionProcessorFactory;
+use crate::selection::in_subquery::{
+    InSubqueryProcessorFactory, LEFT_IN_SUBQUERY_PORT, RIGHT_IN_SUBQUERY_PORT,
+};
 use dozer_core::app::AppPipeline;
 use dozer_core::node::PortHandle;
 use dozer_core::DEFAULT_PORT_HANDLE;
 use dozer_sql_expression::builder::{ExpressionBuilder, NameOrAlias};
-use dozer_sql_expression::sqlparser::ast::{SetOperator, SetQuantifier, TableFactor};
+use dozer_sql_expression::sqlparser::ast::{
+    Expr as SqlExpr, SetOperator, SetQuantifier, TableFactor,
+};
 use dozer_types::models::udf_config::UdfConfig;
 
 use dozer_sql_expression::sqlparser::{
@@ -292,38 +297,27 @@ fn select_to_pipeline(
 
     pipeline.add_processor(Box::new(aggregation), gen_agg_name.clone());
 
-    // Where clause
-    if let Some(selection) = select.selection {
-        let selection = SelectionProcessorFactory::new(
-            gen_selection_name.clone(),
+    let (aggregation_input_node, aggregation_input_port) = if let Some(selection) = select.selection
+    {
+        insert_selection_to_pipeline(
             selection,
-            query_ctx.udfs.clone(),
-            query_ctx.runtime.clone(),
-        );
-
-        pipeline.add_processor(Box::new(selection), gen_selection_name.clone());
-
-        pipeline.connect_nodes(
             gen_product_name,
             product_output_port,
-            gen_selection_name.clone(),
-            DEFAULT_PORT_HANDLE,
-        );
-
-        pipeline.connect_nodes(
             gen_selection_name,
-            DEFAULT_PORT_HANDLE,
-            gen_agg_name.clone(),
-            DEFAULT_PORT_HANDLE,
-        );
+            pipeline,
+            query_ctx,
+            pipeline_idx,
+        )?
     } else {
-        pipeline.connect_nodes(
-            gen_product_name,
-            product_output_port,
-            gen_agg_name.clone(),
-            DEFAULT_PORT_HANDLE,
-        );
-    }
+        (gen_product_name, product_output_port)
+    };
+
+    pipeline.connect_nodes(
+        aggregation_input_node,
+        aggregation_input_port,
+        gen_agg_name.clone(),
+        DEFAULT_PORT_HANDLE,
+    );
 
     query_ctx.pipeline_map.insert(
         (pipeline_idx, table_info.name.0.to_string()),
@@ -360,6 +354,87 @@ fn select_to_pipeline(
     Ok(gen_agg_name)
 }
 
+#[allow(clippy::too_many_arguments)]
+fn insert_selection_to_pipeline(
+    selection: SqlExpr,
+    input_node: String,
+    input_port: PortHandle,
+    selection_node: String,
+    pipeline: &mut AppPipeline,
+    query_ctx: &mut QueryContext,
+    pipeline_idx: usize,
+) -> Result<(String, PortHandle), PipelineError> {
+    match selection {
+        SqlExpr::InSubquery {
+            expr,
+            subquery,
+            negated,
+        } => {
+            let subquery_name = format!("in_subquery_{}", query_ctx.get_next_processor_id());
+            query_to_pipeline(
+                TableInfo {
+                    name: NameOrAlias(subquery_name.clone(), None),
+                    override_name: None,
+                },
+                *subquery,
+                pipeline,
+                query_ctx,
+                pipeline_idx,
+                false,
+            )?;
+
+            let subquery_output = query_ctx
+                .pipeline_map
+                .get(&(pipeline_idx, subquery_name.clone()))
+                .ok_or_else(|| {
+                    PipelineError::InvalidQuery(format!(
+                        "Invalid IN subquery output: {subquery_name}"
+                    ))
+                })?
+                .clone();
+
+            let selection = InSubqueryProcessorFactory::new(
+                selection_node.clone(),
+                *expr,
+                negated,
+                query_ctx.udfs.clone(),
+                query_ctx.runtime.clone(),
+            );
+            pipeline.add_processor(Box::new(selection), selection_node.clone());
+            pipeline.connect_nodes(
+                input_node,
+                input_port,
+                selection_node.clone(),
+                LEFT_IN_SUBQUERY_PORT,
+            );
+            pipeline.connect_nodes(
+                subquery_output.node,
+                subquery_output.port,
+                selection_node.clone(),
+                RIGHT_IN_SUBQUERY_PORT,
+            );
+            Ok((selection_node, DEFAULT_PORT_HANDLE))
+        }
+        selection => {
+            let selection = SelectionProcessorFactory::new(
+                selection_node.clone(),
+                selection,
+                query_ctx.udfs.clone(),
+                query_ctx.runtime.clone(),
+            );
+
+            pipeline.add_processor(Box::new(selection), selection_node.clone());
+            pipeline.connect_nodes(
+                input_node,
+                input_port,
+                selection_node.clone(),
+                DEFAULT_PORT_HANDLE,
+            );
+            Ok((selection_node, DEFAULT_PORT_HANDLE))
+        }
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 fn set_to_pipeline(
     table_info: TableInfo,