diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 8fcaf964ca4..2109292025c 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1241,6 +1241,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) .with_exec_id(exec_context_id) .with_sql_text_pool(&exec_context.spark_plan) + .with_jvm_output_sorts(&exec_context.spark_plan) .with_task_context(exec_context.task_context.clone()) .with_class_loader(exec_context.class_loader.clone()) .with_shuffle_partition_pusher( @@ -3657,6 +3658,7 @@ mod native_sort_spill_tests { share: usize, executor_cores: usize, spill_before_output: Option, + jvm_consumer: bool, ) -> OutputTrace { let shape = source.shape; let rows = source.rows; @@ -3708,7 +3710,8 @@ mod native_sort_spill_tests { }, )]) .unwrap(); - let sort: Arc = Arc::new(SortExec::new(ordering, child)); + let sort: Arc = + Arc::new(SortExec::new(ordering, child).with_spill_before_output(jvm_consumer)); let mut stream = sort.execute(0, session.task_ctx()).unwrap(); let mut trace = OutputTrace { reserved: vec![], @@ -3744,8 +3747,14 @@ mod native_sort_spill_tests { let share = SPILL_BEFORE_OUTPUT_SHARE; for rows_per_batch in [8192, 3] { let case = format!("{shape:?} rows={rows} rows_per_batch={rows_per_batch}"); - let held = - sort_and_trace(KibRows::new(shape, rows, rows_per_batch), share, 8, None).await; + let held = sort_and_trace( + KibRows::new(shape, rows, rows_per_batch), + share, + 8, + None, + true, + ) + .await; eprintln!( "{case} off: spill_count={} peak={}MiB reserved at 0%={}MiB 50%={}MiB 90%={}MiB", held.spill_count(), @@ -3762,6 +3771,7 @@ mod native_sort_spill_tests { share, 8, Some(share / 4), + true, ) .await; eprintln!( @@ -3782,11 +3792,31 @@ mod native_sort_spill_tests { share, 8, Some(held.peak_reserved), + true, ) .await; assert_eq!(below.spill_count(), 0, "{case}"); assert_eq!(below.reserved.len(), held.reserved.len(), "{case}"); assert!(below.reserved_at(0.9) >= input_bytes(rows) / 2, "{case}"); + + let native_consumer = sort_and_trace( + KibRows::new(shape, rows, rows_per_batch), + share, + 8, + Some(share / 4), + false, + ) + .await; + assert_eq!(native_consumer.spill_count(), 0, "{case}"); + assert_eq!( + native_consumer.reserved.len(), + held.reserved.len(), + "{case}" + ); + assert!( + native_consumer.reserved_at(0.9) >= input_bytes(rows) / 2, + "{case}" + ); } } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 1febe577337..63a983baef9 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -314,6 +314,29 @@ pub struct PhysicalPlanner { /// Task-owned destination for remote shuffle blocks, registered on the driving Spark task /// thread before native planning. Only explicit RSS destinations may use it. shuffle_partition_pusher: Option>, + jvm_output_sorts: std::collections::HashSet, +} + +pub(crate) fn jvm_output_sorts(root: &Operator) -> std::collections::HashSet { + let mut sorts = std::collections::HashSet::new(); + let mut pending = vec![root]; + while let Some(op) = pending.pop() { + match op.op_struct.as_ref() { + Some(OpStruct::Sort(_)) => { + sorts.insert(op.plan_id); + } + Some( + OpStruct::Projection(_) + | OpStruct::Filter(_) + | OpStruct::Limit(_) + | OpStruct::Sample(_) + | OpStruct::Expand(_) + | OpStruct::Explode(_), + ) => pending.extend(op.children.iter()), + _ => {} + } + } + sorts } impl Default for PhysicalPlanner { @@ -333,12 +356,18 @@ impl PhysicalPlanner { task_context: None, class_loader: None, shuffle_partition_pusher: None, + jvm_output_sorts: std::collections::HashSet::new(), } } /// Load the SQL text pool from the root operator of the plan about to be planned. Must be /// called with the *root* operator: the JVM only populates the pool there, and /// `QueryContext.sql_text_idx` values are indices into it. + pub fn with_jvm_output_sorts(mut self, root: &Operator) -> Self { + self.jvm_output_sorts = jvm_output_sorts(root); + self + } + pub fn with_sql_text_pool(mut self, root: &Operator) -> Self { self.sql_text_pool = root .sql_text_pool @@ -1783,7 +1812,8 @@ impl PhysicalPlanner { LexOrdering::new(exprs?).unwrap(), Arc::clone(&child.native_plan), ) - .with_fetch(fetch), + .with_fetch(fetch) + .with_spill_before_output(self.jvm_output_sorts.contains(&spark_plan.plan_id)), ); if let Some(skip) = sort.skip.filter(|&n| n > 0).map(|n| n as usize) { @@ -5391,7 +5421,7 @@ mod tests { use crate::execution::operators::{ExecutionError, PartitionedRankLimitExec, WindowFnKind}; use crate::execution::planner::{ - convert_spark_types_to_arrow_schema, literal_to_array_ref, + convert_spark_types_to_arrow_schema, jvm_output_sorts, literal_to_array_ref, parse_file_scan_tasks_from_common, }; use crate::execution::shuffle::CometPartitioning; @@ -6755,6 +6785,138 @@ mod tests { } } + fn op(plan_id: u32, op_struct: OpStruct, children: Vec) -> Operator { + Operator { + plan_id, + sql_text_pool: vec![], + children, + op_struct: Some(op_struct), + } + } + + fn sort_op(plan_id: u32, child: Operator) -> Operator { + let order = spark_expression::Expr { + expr_struct: Some(ExprStruct::SortOrder(Box::new( + spark_expression::SortOrder { + child: Some(Box::new(spark_expression::Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(create_proto_datatype()), + })), + query_context: None, + expr_id: None, + })), + direction: spark_expression::SortDirection::Ascending as i32, + null_ordering: spark_expression::NullOrdering::NullsFirst as i32, + }, + ))), + query_context: None, + expr_id: None, + }; + op( + plan_id, + OpStruct::Sort(spark_operator::Sort { + sort_orders: vec![order], + fetch: None, + skip: None, + }), + vec![child], + ) + } + + fn leaf(plan_id: u32) -> Operator { + let mut scan = create_scan(); + scan.plan_id = plan_id; + scan + } + + #[test] + fn jvm_output_sorts_are_sorts_reached_from_the_root_through_streaming_operators() { + use std::collections::HashSet; + let ids = |v: &[u32]| v.iter().copied().collect::>(); + + assert_eq!(jvm_output_sorts(&sort_op(1, leaf(0))), ids(&[1])); + + let projected = op( + 3, + OpStruct::Projection(Default::default()), + vec![op( + 2, + OpStruct::Filter(Default::default()), + vec![sort_op(1, leaf(0))], + )], + ); + assert_eq!(jvm_output_sorts(&projected), ids(&[1])); + + let nested = sort_op( + 2, + op( + 3, + OpStruct::Window(Default::default()), + vec![sort_op(1, leaf(0))], + ), + ); + assert_eq!(jvm_output_sorts(&nested), ids(&[2])); + + for consumer in [ + OpStruct::ShuffleWriter(Default::default()), + OpStruct::Window(Default::default()), + OpStruct::HashAgg(Default::default()), + OpStruct::ParquetWriter(Default::default()), + OpStruct::WindowGroupLimit(Default::default()), + ] { + let plan = op(2, consumer, vec![sort_op(1, leaf(0))]); + assert!(jvm_output_sorts(&plan).is_empty(), "{plan:?}"); + } + + let join = op( + 5, + OpStruct::SortMergeJoin(Default::default()), + vec![sort_op(1, leaf(0)), sort_op(2, leaf(3))], + ); + assert!(jvm_output_sorts(&join).is_empty()); + assert!( + jvm_output_sorts(&op(6, OpStruct::Projection(Default::default()), vec![join])) + .is_empty() + ); + } + + fn sort_execs(plan: &Arc, out: &mut Vec) { + if let Some(sort) = plan.downcast_ref::() { + out.push(sort.spill_before_output()); + } + for child in plan.children() { + sort_execs(child, out); + } + } + + fn planned_sort_flags(root: &Operator) -> Vec { + let planner = PhysicalPlanner::default().with_jvm_output_sorts(root); + let (_, _, plan) = planner.create_plan(root, &mut vec![], 1).unwrap(); + let mut flags = vec![]; + sort_execs(&plan.native_plan, &mut flags); + flags + } + + #[test] + fn planner_lets_only_a_sort_read_by_the_jvm_spill_before_output() { + assert_eq!(planned_sort_flags(&sort_op(1, leaf(0))), vec![true]); + assert_eq!( + planned_sort_flags(&create_filter(sort_op(1, leaf(0)), 3)), + vec![true] + ); + assert_eq!( + planned_sort_flags(&sort_op(2, sort_op(1, leaf(0)))), + vec![true, false] + ); + let sort = sort_op(1, leaf(0)); + let planner = PhysicalPlanner::default(); + let (_, _, plan) = planner.create_plan(&sort, &mut vec![], 1).unwrap(); + let mut flags = vec![]; + sort_execs(&plan.native_plan, &mut flags); + assert_eq!(flags, vec![false]); + } + fn create_scan() -> Operator { Operator { plan_id: 0, diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs index 51514d63b94..ce3173f5f73 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs @@ -1241,6 +1241,7 @@ pub struct SortExec { /// If `fetch` is `Some`, this will also be set and a TopK operator may be used. /// If `fetch` is `None`, this will be `None`. filter: Option>>, + spill_before_output: bool, } impl SortExec { @@ -1260,9 +1261,19 @@ impl SortExec { common_sort_prefix: sort_prefix, cache: Arc::new(cache), filter: None, + spill_before_output: false, } } + pub fn with_spill_before_output(mut self, spill_before_output: bool) -> Self { + self.spill_before_output = spill_before_output; + self + } + + pub fn spill_before_output(&self) -> bool { + self.spill_before_output + } + /// Whether this `SortExec` preserves partitioning of the children pub fn preserve_partitioning(&self) -> bool { self.preserve_partitioning @@ -1336,6 +1347,7 @@ impl SortExec { fetch: self.fetch, cache: Arc::clone(&self.cache), filter: self.filter.clone(), + spill_before_output: self.spill_before_output, } } @@ -1724,10 +1736,14 @@ impl ExecutionPlan for SortExec { execution_options.sort_spill_reservation_bytes; let in_place_bytes = execution_options.sort_in_place_threshold_bytes; - let spill_before_output = context - .session_config() - .get_extension::() - .map_or(0, |threshold| threshold.0); + let spill_before_output = if self.spill_before_output { + context + .session_config() + .get_extension::() + .map_or(0, |threshold| threshold.0) + } else { + 0 + }; let compression = context.session_config().spill_compression(); let runtime = context.runtime_env(); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -1867,7 +1883,8 @@ impl ExecutionPlan for SortExec { Ok(Some(Arc::new( SortExec::new(updated_exprs, make_with_child(projection, self.input())?) .with_fetch(self.fetch()) - .with_preserve_partitioning(self.preserve_partitioning()), + .with_preserve_partitioning(self.preserve_partitioning()) + .with_spill_before_output(self.spill_before_output), ))) } @@ -1947,7 +1964,8 @@ impl ExecutionPlan for SortExec { let new_sort = Arc::new( SortExec::new(self.expr.clone(), new_child) .with_fetch(self.fetch()) - .with_preserve_partitioning(self.preserve_partitioning()), + .with_preserve_partitioning(self.preserve_partitioning()) + .with_spill_before_output(self.spill_before_output), ) as Arc; Ok(FilterPushdownPropagation { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index b79b412cc03..36e2a8c25e8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -136,6 +136,41 @@ class CometExecSuite extends CometTestBase { } } + test("a native sort spills before output only when the JVM reads its output") { + def sortSpills(plan: SparkPlan): Seq[Long] = + collect(plan) { case s: CometSortExec => s.metrics("spill_count").value } + withSQLConf( + CometConf.COMET_EXEC_SORT_SPILL_BEFORE_OUTPUT_THRESHOLD.key -> "1", + CometConf.COMET_EXEC_WINDOW_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withParquetTable((0 until 20000).map(i => (i % 37, s"v${i * 7919 % 20000}", i)), "a") { + withParquetTable((0 until 5000).map(i => (i % 37, s"w$i")), "b") { + val (_, jvmRead) = checkSparkAnswer(sql("SELECT _1, _2, _3 FROM a ORDER BY _2, _3")) + assert(sortSpills(jvmRead).nonEmpty) + assert(sortSpills(jvmRead).forall(_ > 0), jvmRead) + + val (_, filtered) = + checkSparkAnswer(sql("SELECT _2, _3 + 1 FROM a WHERE _1 > 3 ORDER BY _2, _3")) + assert(sortSpills(filtered).nonEmpty) + assert(sortSpills(filtered).forall(_ > 0), filtered) + + val (_, windowRead) = checkSparkAnswer( + sql("SELECT _1, _2, row_number() OVER (PARTITION BY _1 ORDER BY _2, _3) FROM a")) + assert(collect(windowRead) { case w: CometWindowExec => w }.nonEmpty, windowRead) + assert(sortSpills(windowRead).nonEmpty) + assert(sortSpills(windowRead).forall(_ == 0), windowRead) + + val (_, joinRead) = checkSparkAnswer( + sql("SELECT a._1, a._2, b._2 FROM a JOIN b ON a._1 = b._1 AND a._3 < 300")) + assert(collect(joinRead) { case j: CometSortMergeJoinExec => j }.nonEmpty, joinRead) + assert(sortSpills(joinRead).nonEmpty) + assert(sortSpills(joinRead).forall(_ == 0), joinRead) + } + } + } + } + test("sample without replacement") { withParquetTable((0 until 1000).map(i => (i, i + 1)), "tbl") { val df = sql("SELECT * FROM tbl").sample(withReplacement = false, fraction = 0.3, seed = 42)