executor sink

datafusion-integration/src/main.rs

@@ -1,14 +1,15 @@
 use datafusion::error::Result;
 use datafusion::physical_plan::ExecutionPlan;
-use datafusion::prelude::*;
+use datafusion::prelude::{CsvReadOptions, SessionContext};
 use std::sync::Arc;
 mod util;
 use vayu::{sinks, VayuExecutionEngine};
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    // test_scan_filter_project().await?;
+    test_scan_filter_project().await?;
     test_hash_join().await?;
+    test_store_record_batch().await?;
     Ok(())
 }
 
@@ -72,21 +73,52 @@ async fn test_scan_filter_project() -> Result<()> {
     )
     .await?;
     let mut executor = VayuExecutionEngine::new();
-    let uuid = 1;
-    let sql1 = "SELECT c1,c4,c13  FROM aggregate_test_100 WHERE c3 < 0 AND c1='a'";
+    let sql = "SELECT c1,c3 as neg,c4 as pos,c13  FROM aggregate_test_100 WHERE (c3 < 0 AND c1='a') OR ( c4 > 0 AND c1='b' ) ";
+
+    let plan = util::get_execution_plan_from_sql(&ctx, sql).await?;
+
+    run_pipeline(&mut executor, plan, sinks::SchedulerSinkType::PrintOutput).await;
+
+    Ok(())
+}
 
-    let plan1 = util::get_execution_plan_from_sql(&ctx, sql1).await?;
+async fn test_store_record_batch() -> Result<()> {
+    // create local execution context
+    let ctx: SessionContext = SessionContext::new();
+    // register csv file with the execution context
 
+    ctx.register_csv(
+        "aggregate_test_100",
+        &format!("./testing/data/csv/aggregate_test_100.csv"),
+        CsvReadOptions::new(),
+    )
+    .await?;
+    let mut executor = VayuExecutionEngine::new();
+    let sql = "SELECT c1,c3 as neg,c4 as pos,c13  FROM aggregate_test_100 WHERE c3 < 0 AND c1='a' ";
+
+    let plan = util::get_execution_plan_from_sql(&ctx, sql).await?;
+    let uuid = 42;
     run_pipeline(
         &mut executor,
-        plan1,
+        plan,
         sinks::SchedulerSinkType::StoreRecordBatch(uuid),
     )
     .await;
+    let sql =
+        "SELECT c1,c3 as neg,c4 as pos,c13  FROM aggregate_test_100 WHERE  c4 > 0 AND c1='b' ";
+
+    let plan = util::get_execution_plan_from_sql(&ctx, sql).await?;
+    let uuid = 42;
+    run_pipeline(
+        &mut executor,
+        plan,
+        sinks::SchedulerSinkType::StoreRecordBatch(uuid),
+    )
+    .await;
+
+    // get executor to print the value
+    executor.sink(uuid);
 
-    let sql2 = "SELECT c1,c4,c13  FROM aggregate_test_100 WHERE c4 > 0 AND c1='b'";
-    let plan2 = util::get_execution_plan_from_sql(&ctx, sql2).await?;
-    run_pipeline(&mut executor, plan2, sinks::SchedulerSinkType::PrintOutput).await;
     Ok(())
 }
 

vayu/src/lib.rs

@@ -9,10 +9,11 @@ pub mod sinks;
 
 pub mod store;
 use crate::sinks::SchedulerSinkType;
+use crate::store::Blob::{HashMapBlob, RecordBatchBlob};
 use crate::store::Store;
+use core::panic;
 use datafusion::physical_plan::ExecutionPlan;
 use std::sync::Arc;
-
 pub struct VayuExecutionEngine {
     pub store: Store,
 }
@@ -26,7 +27,6 @@ impl VayuExecutionEngine {
 
     pub fn execute(&mut self, scheduler_pipeline: SchedulerPipeline) {
         let plan = scheduler_pipeline.plan;
-        let sink_type = scheduler_pipeline.sink;
         // convert execution plan to a pipeline
 
         let pipeline = Pipeline::new(plan, &mut self.store, 1);
@@ -37,7 +37,8 @@ impl VayuExecutionEngine {
         // do the sinking - very simple API
         // no need to create a seperate class and introduce indirection unless it moves out of hands
         // to call one function we would need 30+ lines otherwise
-        match sink_type {
+        let sink: SchedulerSinkType = scheduler_pipeline.sink;
+        match sink {
             SchedulerSinkType::PrintOutput => {
                 pretty::print_batches(&result).unwrap();
             }
@@ -51,6 +52,20 @@ impl VayuExecutionEngine {
             }
         };
     }
+    pub fn sink(&mut self, uuid: i32) {
+        let blob = self.store.remove(uuid);
+        match blob {
+            Some(blob) => match blob {
+                RecordBatchBlob(result) => {
+                    pretty::print_batches(&result).unwrap();
+                }
+                HashMapBlob(results) => {
+                    pretty::print_batches(&[results.batch().clone()]).unwrap();
+                }
+            },
+            None => panic!("no blob for {uuid} found"),
+        }
+    }
 }
 pub struct SchedulerPipeline {
     pub plan: Arc<dyn ExecutionPlan>,

vayu/src/pipeline.rs

@@ -3,6 +3,7 @@ use crate::operators::join::HashProbeOperator;
 use crate::operators::projection::ProjectionOperator;
 use crate::store::Store;
 use arrow::array::BooleanBufferBuilder;
+use core::panic;
 use datafusion::arrow::array::RecordBatch;
 use datafusion::datasource::physical_plan::CsvExec;
 use datafusion::error::Result;
@@ -74,28 +75,27 @@ fn make_pipeline(pipeline: &mut Pipeline, plan: Arc<dyn ExecutionPlan>, store: &
         make_pipeline(pipeline, exec.right().clone(), store);
         println!("adding hashprobe");
         let mut hashjoinstream = exec.get_hash_join_stream(0, context).unwrap();
+        // using uuid but this value would be present in HashProbeExec itself
         let build_map = store.remove(pipeline.state.uuid).unwrap();
         let left_data = Arc::new(build_map.get_map());
         let visited_left_side = BooleanBufferBuilder::new(0);
         hashjoinstream.build_side = BuildSide::Ready(BuildSideReadyState {
             left_data,
             visited_left_side,
         });
-        // println!("{:?}", left_data);
         let tt = Box::new(HashProbeOperator::new(hashjoinstream));
         pipeline.operators.push(tt);
         return;
     }
     if let Some(exec) = p.downcast_ref::<RepartitionExec>() {
         make_pipeline(pipeline, exec.input().clone(), store);
         return;
-        // pipeline.operators.push(FilterOperator::new())
     }
     if let Some(exec) = p.downcast_ref::<CoalesceBatchesExec>() {
         make_pipeline(pipeline, exec.input().clone(), store);
         return;
-        // pipeline.operators.push(FilterOperator::new())
     }
+    panic!("should never reach the end");
 }
 pub trait PhysicalOperator {
     fn name(&self) -> String;

vayu/src/pipeline_executor.rs

@@ -40,7 +40,7 @@ impl PipelineExecutor {
         operators: &mut Vec<Box<dyn IntermediateOperator>>,
         mut data: RecordBatch,
     ) -> RecordBatch {
-        for mut x in operators {
+        for x in operators {
             println!(
                 "running operator {} size {}x{}",
                 x.name(),