Skip to content

Commit bcb152a

Browse files
zhuxiangyizhuxiangyi
committed
[AURON #2358] Support native last / last(ignoreNulls) aggregate
Implement native last / last(ignoreNulls) aggregates, mirroring the existing first implementation with "later value wins" semantics: - native: add AggLast / AggLastIgnoresNull (agg/last.rs, agg/last_ignores_null.rs); wire through the AggFunction enum, create_agg, the protobuf contract (LAST / LAST_IGNORES_NULL), the protobuf->AggFunction conversion, and the window-agg mapping. - spark-extension: add the Last expression conversion in NativeConverters; declare the Last native aggregate buffer schema in NativeAggBase.computeNativeAggBufferDataTypes ([dataType] for ignoreNulls, [dataType, Boolean] otherwise) so the partial -> shuffle -> final buffer schema matches the native side. Tests: - Rust unit test agg_exec::test::test_agg_last (partial -> final, nulls). - Scala e2e AuronDataFrameAggregateSuite "native last / last(ignoreNulls) aggregate" (spark34 + spark35), covering the partial -> shuffle -> final native path and asserting NativeAggBase offload.
1 parent d5af99b commit bcb152a

12 files changed

Lines changed: 766 additions & 3 deletions

File tree

‎auron-spark-tests/spark34/src/test/scala/org/apache/spark/sql/AuronDataFrameAggregateSuite.scala‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import scala.util.Random
2121
import org.apache.spark.sql.execution.WholeStageCodegenExec
2222
import org.apache.spark.sql.execution.aggregate.HashAggregateExec
2323
import org.apache.spark.sql.execution.auron.plan.NativeAggBase
24-
import org.apache.spark.sql.functions.{collect_list, monotonically_increasing_id, rand, randn, spark_partition_id, sum}
24+
import org.apache.spark.sql.functions.{collect_list, last, monotonically_increasing_id, rand, randn, spark_partition_id, sum}
2525
import org.apache.spark.sql.internal.SQLConf
2626

2727
class AuronDataFrameAggregateSuite extends DataFrameAggregateSuite with SparkQueryTestsBase {
@@ -75,4 +75,31 @@ class AuronDataFrameAggregateSuite extends DataFrameAggregateSuite with SparkQue
7575
rand(Random.nextLong()),
7676
randn(Random.nextLong())).foreach(assertNoExceptions)
7777
}
78+
79+
testAuron("native last / last(ignoreNulls) aggregate") {
80+
// The grouped aggregate is reliably offloaded to NativeAggBase, and the data
81+
// is deterministic by construction (no intra-group ordering dependence):
82+
// k=1 -> all values 10 => last=10, last(ignoreNulls)=10
83+
// k=2 -> all values null => last=null, last(ignoreNulls)=null
84+
// k=3 -> single row 30 => last=30, last(ignoreNulls)=30
85+
val df = Seq[(Int, Option[Int])](
86+
(1, Some(10)),
87+
(1, Some(10)),
88+
(2, None),
89+
(2, None),
90+
(3, Some(30)))
91+
.toDF("k", "v")
92+
93+
val aggDF = df
94+
.groupBy("k")
95+
.agg(last($"v").as("last_v"), last($"v", ignoreNulls = true).as("last_v_ign"))
96+
97+
checkAnswer(aggDF, Seq(Row(1, 10, 10), Row(2, null, null), Row(3, 30, 30)))
98+
99+
// the aggregate must be offloaded to the native engine
100+
assert(getExecutedPlan(aggDF).exists {
101+
case _: NativeAggBase => true
102+
case _ => false
103+
})
104+
}
78105
}

‎auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameAggregateSuite.scala‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import scala.util.Random
2121
import org.apache.spark.sql.execution.WholeStageCodegenExec
2222
import org.apache.spark.sql.execution.aggregate.HashAggregateExec
2323
import org.apache.spark.sql.execution.auron.plan.NativeAggBase
24-
import org.apache.spark.sql.functions.{collect_list, monotonically_increasing_id, rand, randn, spark_partition_id, sum}
24+
import org.apache.spark.sql.functions.{collect_list, last, monotonically_increasing_id, rand, randn, spark_partition_id, sum}
2525
import org.apache.spark.sql.internal.SQLConf
2626

2727
class AuronDataFrameAggregateSuite extends DataFrameAggregateSuite with SparkQueryTestsBase {
@@ -75,4 +75,31 @@ class AuronDataFrameAggregateSuite extends DataFrameAggregateSuite with SparkQue
7575
rand(Random.nextLong()),
7676
randn(Random.nextLong())).foreach(assertNoExceptions)
7777
}
78+
79+
testAuron("native last / last(ignoreNulls) aggregate") {
80+
// The grouped aggregate is reliably offloaded to NativeAggBase, and the data
81+
// is deterministic by construction (no intra-group ordering dependence):
82+
// k=1 -> all values 10 => last=10, last(ignoreNulls)=10
83+
// k=2 -> all values null => last=null, last(ignoreNulls)=null
84+
// k=3 -> single row 30 => last=30, last(ignoreNulls)=30
85+
val df = Seq[(Int, Option[Int])](
86+
(1, Some(10)),
87+
(1, Some(10)),
88+
(2, None),
89+
(2, None),
90+
(3, Some(30)))
91+
.toDF("k", "v")
92+
93+
val aggDF = df
94+
.groupBy("k")
95+
.agg(last($"v").as("last_v"), last($"v", ignoreNulls = true).as("last_v_ign"))
96+
97+
checkAnswer(aggDF, Seq(Row(1, 10, 10), Row(2, null, null), Row(3, 30, 30)))
98+
99+
// the aggregate must be offloaded to the native engine
100+
assert(getExecutedPlan(aggDF).exists {
101+
case _: NativeAggBase => true
102+
case _ => false
103+
})
104+
}
78105
}

‎native-engine/auron-planner/proto/auron.proto‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,8 @@ enum AggFunction {
148148
FIRST = 7;
149149
FIRST_IGNORES_NULL = 8;
150150
BLOOM_FILTER = 9;
151+
LAST = 10;
152+
LAST_IGNORES_NULL = 11;
151153
BRICKHOUSE_COLLECT = 1000;
152154
BRICKHOUSE_COMBINE_UNIQUE = 1001;
153155
UDAF = 1002;

‎native-engine/auron-planner/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,8 @@ impl From<protobuf::AggFunction> for AggFunction {
135135
protobuf::AggFunction::CollectSet => AggFunction::CollectSet,
136136
protobuf::AggFunction::First => AggFunction::First,
137137
protobuf::AggFunction::FirstIgnoresNull => AggFunction::FirstIgnoresNull,
138+
protobuf::AggFunction::Last => AggFunction::Last,
139+
protobuf::AggFunction::LastIgnoresNull => AggFunction::LastIgnoresNull,
138140
protobuf::AggFunction::BloomFilter => AggFunction::BloomFilter,
139141
protobuf::AggFunction::BrickhouseCollect => AggFunction::BrickhouseCollect,
140142
protobuf::AggFunction::BrickhouseCombineUnique => AggFunction::BrickhouseCombineUnique,

‎native-engine/auron-planner/src/planner.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -680,6 +680,12 @@ impl PhysicalPlanner {
680680
protobuf::AggFunction::FirstIgnoresNull => {
681681
WindowFunction::Agg(AggFunction::FirstIgnoresNull)
682682
}
683+
protobuf::AggFunction::Last => {
684+
WindowFunction::Agg(AggFunction::Last)
685+
}
686+
protobuf::AggFunction::LastIgnoresNull => {
687+
WindowFunction::Agg(AggFunction::LastIgnoresNull)
688+
}
683689
protobuf::AggFunction::BloomFilter => {
684690
WindowFunction::Agg(AggFunction::BloomFilter)
685691
}

‎native-engine/datafusion-ext-plans/src/agg/agg.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@ use crate::agg::{
3333
count::AggCount,
3434
first::AggFirst,
3535
first_ignores_null::AggFirstIgnoresNull,
36+
last::AggLast,
37+
last_ignores_null::AggLastIgnoresNull,
3638
maxmin::{AggMax, AggMin},
3739
spark_udaf_wrapper::SparkUDAFWrapper,
3840
sum::AggSum,
@@ -212,6 +214,14 @@ pub fn create_agg(
212214
let dt = children[0].data_type(input_schema)?;
213215
Arc::new(AggFirstIgnoresNull::try_new(children[0].clone(), dt)?)
214216
}
217+
AggFunction::Last => {
218+
let dt = children[0].data_type(input_schema)?;
219+
Arc::new(AggLast::try_new(children[0].clone(), dt)?)
220+
}
221+
AggFunction::LastIgnoresNull => {
222+
let dt = children[0].data_type(input_schema)?;
223+
Arc::new(AggLastIgnoresNull::try_new(children[0].clone(), dt)?)
224+
}
215225
AggFunction::BloomFilter => {
216226
let dt = children[0].data_type(input_schema)?;
217227
let empty_batch = RecordBatch::new_empty(Arc::new(Schema::empty()));

0 commit comments

Comments
 (0)