AggExec: implement columnar accumulator states.

refactor execution context

Author: zhangli20

.github/workflows/tpcds-reusable.yml

@@ -120,23 +120,13 @@ jobs:
       - name: Install Blaze JAR
         run: cp blaze-engine-*${{ inputs.sparkver }}*.jar spark-bin-${{ inputs.sparkver }}/jars/
 
-      - name: Run with BHJ
+      - name: Run
         run: |
           export RUST_LOG=ERROR
           export RUST_BACKTRACE=1
           SPARK_HOME=spark-bin-${{ inputs.sparkver }} dev/run-tpcds-test \
             --data-location dev/tpcds_1g \
             --conf spark.driver.memory=3g \
-            --conf spark.driver.memoryOverhead=2048 \
-            --query-filter ${{ matrix.query }}
-
-      - name: Run without BHJ
-        run: |
-          export RUST_LOG=ERROR
-          export RUST_BACKTRACE=1
-          SPARK_HOME=spark-bin-${{ inputs.sparkver }} dev/run-tpcds-test \
-            --conf spark.sql.autoBroadcastJoinThreshold=-1 \
-            --conf spark.driver.memory=3g \
-            --conf spark.driver.memoryOverhead=2048 \
-            --data-location dev/tpcds_1g \
+            --conf spark.driver.memoryOverhead=2536 \
+            --conf spark.sql.broadcastTimeout=900s \
             --query-filter ${{ matrix.query }}

Cargo.lock

@@ -59,7 +59,7 @@ use datafusion_ext_exprs::{
     string_ends_with::StringEndsWithExpr, string_starts_with::StringStartsWithExpr,
 };
 use datafusion_ext_plans::{
-    agg::{create_agg, AggExecMode, AggExpr, AggFunction, AggMode, GroupingExpr},
+    agg::{agg::create_agg, AggExecMode, AggExpr, AggFunction, AggMode, GroupingExpr},
     agg_exec::AggExec,
     broadcast_join_build_hash_map_exec::BroadcastJoinBuildHashMapExec,
     broadcast_join_exec::BroadcastJoinExec,
@@ -536,7 +536,6 @@ impl TryInto<Arc<dyn ExecutionPlan>> for &protobuf::PhysicalPlanNode {
                     exec_mode,
                     physical_groupings,
                     physical_aggs,
-                    agg.initial_input_buffer_offset as usize,
                     agg.supports_partial_skipping,
                     input,
                 )?))

native-engine/blaze-serde/src/from_proto.rs

@@ -31,24 +31,23 @@ use datafusion::{
     common::Result,
     error::DataFusionError,
     execution::context::TaskContext,
-    physical_plan::{
-        metrics::{BaselineMetrics, ExecutionPlanMetricsSet},
-        ExecutionPlan,
-    },
+    physical_plan::{metrics::ExecutionPlanMetricsSet, ExecutionPlan},
+};
+use datafusion_ext_commons::df_execution_err;
+use datafusion_ext_plans::{
+    common::{execution_context::ExecutionContext, output::TaskOutputter},
+    parquet_sink_exec::ParquetSinkExec,
 };
-use datafusion_ext_commons::{df_execution_err, streams::coalesce_stream::CoalesceInput};
-use datafusion_ext_plans::{common::output::TaskOutputter, parquet_sink_exec::ParquetSinkExec};
 use futures::{FutureExt, StreamExt};
 use jni::objects::{GlobalRef, JObject};
 use tokio::runtime::Runtime;
 
 use crate::{handle_unwinded_scope, metrics::update_spark_metric_node};
 
 pub struct NativeExecutionRuntime {
+    exec_ctx: Arc<ExecutionContext>,
     native_wrapper: GlobalRef,
     plan: Arc<dyn ExecutionPlan>,
-    task_context: Arc<TaskContext>,
-    partition: usize,
     batch_receiver: Receiver<Result<Option<RecordBatch>>>,
     rt: Runtime,
 }
@@ -61,21 +60,23 @@ impl NativeExecutionRuntime {
         context: Arc<TaskContext>,
     ) -> Result<Self> {
         // execute plan to output stream
-        let stream = plan.execute(partition, context.clone())?;
-        let schema = stream.schema();
+        let exec_ctx = ExecutionContext::new(
+            context.clone(),
+            partition,
+            plan.schema(),
+            &ExecutionPlanMetricsSet::new(),
+        );
+        let stream = exec_ctx.execute(&plan)?;
 
         // coalesce
         let mut stream = if plan.as_any().downcast_ref::<ParquetSinkExec>().is_some() {
             stream // cannot coalesce parquet sink output
         } else {
-            context.coalesce_with_default_batch_size(
-                stream,
-                &BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), partition),
-            )?
+            exec_ctx.coalesce_with_default_batch_size(stream)
         };
 
         // init ffi schema
-        let ffi_schema = FFI_ArrowSchema::try_from(schema.as_ref())?;
+        let ffi_schema = FFI_ArrowSchema::try_from(exec_ctx.output_schema().as_ref())?;
         jni_call!(BlazeCallNativeWrapper(native_wrapper.as_obj())
             .importSchema(&ffi_schema as *const FFI_ArrowSchema as i64) -> ()
         )?;
@@ -98,12 +99,11 @@ impl NativeExecutionRuntime {
 
         let (batch_sender, batch_receiver) = std::sync::mpsc::sync_channel(1);
         let nrt = Self {
+            exec_ctx,
             native_wrapper: native_wrapper.clone(),
             plan,
-            partition,
             rt,
             batch_receiver,
-            task_context: context,
         };
 
         // spawn batch producer
@@ -188,7 +188,7 @@ impl NativeExecutionRuntime {
             }
         };
 
-        let partition = self.partition;
+        let partition = self.exec_ctx.partition_id();
         match next_batch() {
             Ok(ret) => return ret,
             Err(err) => {
@@ -203,13 +203,13 @@ impl NativeExecutionRuntime {
     }
 
     pub fn finalize(self) {
-        let partition = self.partition;
+        let partition = self.exec_ctx.partition_id();
 
         log::info!("[partition={partition}] native execution finalizing");
         self.update_metrics().unwrap_or_default();
         drop(self.plan);
 
-        self.task_context.cancel_task(); // cancel all pending streams
+        self.exec_ctx.task_ctx().cancel_task(); // cancel all pending streams
         self.rt.shutdown_background();
         log::info!("[partition={partition}] native execution finalized");
     }

native-engine/blaze/src/rt.rs

@@ -25,7 +25,6 @@ num = "0.4.2"
 once_cell = "1.20.2"
 paste = "1.0.15"
 radsort = "0.1.1"
-slimmer_box = "0.6.5"
 tempfile = "3"
 thrift = "0.17.0"
 tokio = "1.41"

native-engine/datafusion-ext-commons/Cargo.toml

@@ -138,6 +138,22 @@ pub fn read_u8<R: Read>(input: &mut R) -> std::io::Result<u8> {
     Ok(buf[0])
 }
 
+pub fn read_bytes_into_vec<R: Read>(
+    input: &mut R,
+    buf: &mut Vec<u8>,
+    len: usize,
+) -> std::io::Result<()> {
+    buf.reserve(len);
+    unsafe {
+        // safety: space has been reserved
+        input.read_exact(std::slice::from_raw_parts_mut(
+            buf.as_mut_ptr().add(buf.len()),
+            len,
+        ))?;
+        buf.set_len(buf.len() + len);
+    }
+    Ok(())
+}
 pub fn read_bytes_slice<R: Read>(input: &mut R, len: usize) -> std::io::Result<Box<[u8]>> {
     // safety - assume_init() is safe for [u8]
     let mut byte_slice = unsafe { Box::new_uninit_slice(len).assume_init() };