This is an automated email from the ASF dual-hosted git repository.
huaxingao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new c9af3f41 feat: Support sum in window function (#802)
c9af3f41 is described below
commit c9af3f4127ba9f8da8edbb5d6e196fd1359d4f9e
Author: Huaxin Gao <[email protected]>
AuthorDate: Sat Aug 17 10:10:13 2024 -0700
feat: Support sum in window function (#802)
* feat: Support sum in window function
* remove extra line
* remove unused import
* Trigger Build
* Trigger Build
* address comments
---
native/core/src/execution/datafusion/planner.rs | 67 ++++++++--------
.../org/apache/comet/serde/QueryPlanSerde.scala | 29 ++++++-
.../apache/comet/exec/CometAggregateSuite.scala | 88 ----------------------
.../org/apache/comet/exec/CometExecSuite.scala | 60 ++++++++++++++-
.../scala/org/apache/spark/sql/CometTestBase.scala | 87 +++++++++++++++++++++
5 files changed, 210 insertions(+), 121 deletions(-)
diff --git a/native/core/src/execution/datafusion/planner.rs
b/native/core/src/execution/datafusion/planner.rs
index 24ea4142..fe6ef9f7 100644
--- a/native/core/src/execution/datafusion/planner.rs
+++ b/native/core/src/execution/datafusion/planner.rs
@@ -1538,12 +1538,17 @@ impl PhysicalPlanner {
partition_by: &[Arc<dyn PhysicalExpr>],
sort_exprs: &[PhysicalSortExpr],
) -> Result<Arc<dyn WindowExpr>, ExecutionError> {
- let (mut window_func_name, mut window_func_args) = (String::new(),
Vec::new());
+ let window_func_name: String;
+ let window_args: Vec<Arc<dyn PhysicalExpr>>;
if let Some(func) = &spark_expr.built_in_window_function {
match &func.expr_struct {
Some(ExprStruct::ScalarFunc(f)) => {
- window_func_name.clone_from(&f.func);
- window_func_args.clone_from(&f.args);
+ window_func_name = f.func.clone();
+ window_args = f
+ .args
+ .iter()
+ .map(|expr| self.create_expr(expr,
input_schema.clone()))
+ .collect::<Result<Vec<_>, ExecutionError>>()?;
}
other => {
return Err(ExecutionError::GeneralError(format!(
@@ -1552,9 +1557,9 @@ impl PhysicalPlanner {
}
};
} else if let Some(agg_func) = &spark_expr.agg_func {
- let result = Self::process_agg_func(agg_func)?;
+ let result = self.process_agg_func(agg_func,
input_schema.clone())?;
window_func_name = result.0;
- window_func_args = result.1;
+ window_args = result.1;
} else {
return Err(ExecutionError::GeneralError(
"Both func and agg_func are not set".to_string(),
@@ -1570,11 +1575,6 @@ impl PhysicalPlanner {
}
};
- let window_args = window_func_args
- .iter()
- .map(|expr| self.create_expr(expr, input_schema.clone()))
- .collect::<Result<Vec<_>, ExecutionError>>()?;
-
let spark_window_frame = match spark_expr
.spec
.as_ref()
@@ -1644,32 +1644,39 @@ impl PhysicalPlanner {
.map_err(|e| ExecutionError::DataFusionError(e.to_string()))
}
- fn process_agg_func(agg_func: &AggExpr) -> Result<(String, Vec<Expr>),
ExecutionError> {
- fn optional_expr_to_vec(expr_option: &Option<Expr>) -> Vec<Expr> {
- expr_option
- .as_ref()
- .cloned()
- .map_or_else(Vec::new, |e| vec![e])
- }
-
- fn int_to_stats_type(value: i32) -> Option<StatsType> {
- match value {
- 0 => Some(StatsType::Sample),
- 1 => Some(StatsType::Population),
- _ => None,
- }
- }
-
+ fn process_agg_func(
+ &self,
+ agg_func: &AggExpr,
+ schema: SchemaRef,
+ ) -> Result<(String, Vec<Arc<dyn PhysicalExpr>>), ExecutionError> {
match &agg_func.expr_struct {
Some(AggExprStruct::Count(expr)) => {
- let args = &expr.children;
- Ok(("count".to_string(), args.to_vec()))
+ let children = expr
+ .children
+ .iter()
+ .map(|child| self.create_expr(child, schema.clone()))
+ .collect::<Result<Vec<_>, _>>()?;
+ Ok(("count".to_string(), children))
}
Some(AggExprStruct::Min(expr)) => {
- Ok(("min".to_string(), optional_expr_to_vec(&expr.child)))
+ let child = self.create_expr(expr.child.as_ref().unwrap(),
schema.clone())?;
+ Ok(("min".to_string(), vec![child]))
}
Some(AggExprStruct::Max(expr)) => {
- Ok(("max".to_string(), optional_expr_to_vec(&expr.child)))
+ let child = self.create_expr(expr.child.as_ref().unwrap(),
schema.clone())?;
+ Ok(("max".to_string(), vec![child]))
+ }
+ Some(AggExprStruct::Sum(expr)) => {
+ let child = self.create_expr(expr.child.as_ref().unwrap(),
schema.clone())?;
+ let arrow_type =
to_arrow_datatype(expr.datatype.as_ref().unwrap());
+ let datatype = child.data_type(&schema)?;
+
+ let child = if datatype != arrow_type {
+ Arc::new(CastExpr::new(child, arrow_type.clone(), None))
+ } else {
+ child
+ };
+ Ok(("sum".to_string(), vec![child]))
}
other => Err(ExecutionError::GeneralError(format!(
"{other:?} not supported for window function"
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index ca17f6f0..caeda516 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -211,10 +211,35 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
expr match {
case agg: AggregateExpression =>
agg.aggregateFunction match {
- case _: Min | _: Max | _: Count =>
+ case _: Count =>
Some(agg)
+ case min: Min =>
+ if (minMaxDataTypeSupported(min.dataType)) {
+ Some(agg)
+ } else {
+ withInfo(windowExpr, s"datatype ${min.dataType} is not
supported", expr)
+ None
+ }
+ case max: Max =>
+ if (minMaxDataTypeSupported(max.dataType)) {
+ Some(agg)
+ } else {
+ withInfo(windowExpr, s"datatype ${max.dataType} is not
supported", expr)
+ None
+ }
+ case s: Sum =>
+ if (sumDataTypeSupported(s.dataType) &&
!s.dataType.isInstanceOf[DecimalType]) {
+ Some(agg)
+ } else {
+ withInfo(windowExpr, s"datatype ${s.dataType} is not
supported", expr)
+ None
+ }
case _ =>
- withInfo(windowExpr, "Unsupported aggregate", expr)
+ withInfo(
+ windowExpr,
+ s"aggregate ${agg.aggregateFunction}" +
+ " is not supported for window function",
+ expr)
None
}
case _ =>
diff --git
a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
index 2acc562b..5467bd44 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -22,8 +22,6 @@ package org.apache.comet.exec
import scala.util.Random
import org.apache.hadoop.fs.Path
-import org.apache.parquet.example.data.simple.SimpleGroup
-import org.apache.parquet.schema.MessageTypeParser
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
import org.apache.spark.sql.catalyst.optimizer.EliminateSorts
import org.apache.spark.sql.comet.CometHashAggregateExec
@@ -1295,90 +1293,4 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
sparkPlan.collect { case s: CometHashAggregateExec => s }.size
}
- def makeParquetFile(
- path: Path,
- total: Int,
- numGroups: Int,
- dictionaryEnabled: Boolean): Unit = {
- val schemaStr =
- """
- |message root {
- | optional INT32 _1(INT_8);
- | optional INT32 _2(INT_16);
- | optional INT32 _3;
- | optional INT64 _4;
- | optional FLOAT _5;
- | optional DOUBLE _6;
- | optional INT32 _7(DECIMAL(5, 2));
- | optional INT64 _8(DECIMAL(18, 10));
- | optional FIXED_LEN_BYTE_ARRAY(16) _9(DECIMAL(38, 37));
- | optional INT64 _10(TIMESTAMP(MILLIS,true));
- | optional INT64 _11(TIMESTAMP(MICROS,true));
- | optional INT32 _12(DATE);
- | optional INT32 _g1(INT_8);
- | optional INT32 _g2(INT_16);
- | optional INT32 _g3;
- | optional INT64 _g4;
- | optional FLOAT _g5;
- | optional DOUBLE _g6;
- | optional INT32 _g7(DECIMAL(5, 2));
- | optional INT64 _g8(DECIMAL(18, 10));
- | optional FIXED_LEN_BYTE_ARRAY(16) _g9(DECIMAL(38, 37));
- | optional INT64 _g10(TIMESTAMP(MILLIS,true));
- | optional INT64 _g11(TIMESTAMP(MICROS,true));
- | optional INT32 _g12(DATE);
- | optional BINARY _g13(UTF8);
- | optional BINARY _g14;
- |}
- """.stripMargin
-
- val schema = MessageTypeParser.parseMessageType(schemaStr)
- val writer = createParquetWriter(schema, path, dictionaryEnabled = true)
-
- val rand = scala.util.Random
- val expected = (0 until total).map { i =>
- // use a single value for the first page, to make sure dictionary
encoding kicks in
- if (rand.nextBoolean()) None
- else {
- if (dictionaryEnabled) Some(i % 10) else Some(i)
- }
- }
-
- expected.foreach { opt =>
- val record = new SimpleGroup(schema)
- opt match {
- case Some(i) =>
- record.add(0, i.toByte)
- record.add(1, i.toShort)
- record.add(2, i)
- record.add(3, i.toLong)
- record.add(4, rand.nextFloat())
- record.add(5, rand.nextDouble())
- record.add(6, i)
- record.add(7, i.toLong)
- record.add(8, (i % 10).toString * 16)
- record.add(9, i.toLong)
- record.add(10, i.toLong)
- record.add(11, i)
- record.add(12, i.toByte % numGroups)
- record.add(13, i.toShort % numGroups)
- record.add(14, i % numGroups)
- record.add(15, i.toLong % numGroups)
- record.add(16, rand.nextFloat())
- record.add(17, rand.nextDouble())
- record.add(18, i)
- record.add(19, i.toLong)
- record.add(20, (i % 10).toString * 16)
- record.add(21, i.toLong)
- record.add(22, i.toLong)
- record.add(23, i)
- record.add(24, (i % 10).toString * 24)
- record.add(25, (i % 10).toString * 36)
- case _ =>
- }
- writer.write(record)
- }
-
- writer.close()
- }
}
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 15864a33..ed9f2c34 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -1579,6 +1579,58 @@ class CometExecSuite extends CometTestBase {
}
}
+ test("aggregate window function for all types") {
+ val numValues = 2048
+
+ Seq(1, 100, numValues).foreach { numGroups =>
+ Seq(true, false).foreach { dictionaryEnabled =>
+ withTempPath { dir =>
+ val path = new Path(dir.toURI.toString, "test.parquet")
+ makeParquetFile(path, numValues, numGroups, dictionaryEnabled)
+ withParquetTable(path.toUri.toString, "tbl") {
+ Seq(128, numValues + 100).foreach { batchSize =>
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key ->
batchSize.toString) {
+ (1 to 11).foreach { col =>
+ val aggregateFunctions =
+ List(s"COUNT(_$col)", s"MAX(_$col)", s"MIN(_$col)",
s"SUM(_$col)")
+ aggregateFunctions.foreach { function =>
+ val df1 = sql(s"SELECT $function OVER() FROM tbl")
+ checkSparkAnswerWithTol(df1, 1e-6)
+
+ val df2 = sql(s"SELECT $function OVER(order by _2) FROM
tbl")
+ checkSparkAnswerWithTol(df2, 1e-6)
+
+ val df3 = sql(s"SELECT $function OVER(order by _2 desc)
FROM tbl")
+ checkSparkAnswerWithTol(df3, 1e-6)
+
+ val df4 = sql(s"SELECT $function OVER(partition by _2
order by _2) FROM tbl")
+ checkSparkAnswerWithTol(df4, 1e-6)
+ }
+ }
+
+ // SUM doesn't work for Date type.
org.apache.spark.sql.AnalysisException will be thrown.
+ val aggregateFunctionsWithoutSum = List("COUNT(_12)",
"MAX(_12)", "MIN(_12)")
+ aggregateFunctionsWithoutSum.foreach { function =>
+ val df1 = sql(s"SELECT $function OVER() FROM tbl")
+ checkSparkAnswerWithTol(df1, 1e-6)
+
+ val df2 = sql(s"SELECT $function OVER(order by _2) FROM tbl")
+ checkSparkAnswerWithTol(df2, 1e-6)
+
+ val df3 = sql(s"SELECT $function OVER(order by _2 desc) FROM
tbl")
+ checkSparkAnswerWithTol(df3, 1e-6)
+
+ val df4 = sql(s"SELECT $function OVER(partition by _2 order
by _2) FROM tbl")
+ checkSparkAnswerWithTol(df4, 1e-6)
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
test("Windows support") {
Seq("true", "false").foreach(aqeEnabled =>
withSQLConf(
@@ -1586,7 +1638,13 @@ class CometExecSuite extends CometTestBase {
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled) {
withParquetTable((0 until 10).map(i => (i, 10 - i)), "t1") { // TODO:
test nulls
val aggregateFunctions =
- List("COUNT(_1)", "COUNT(*)", "MAX(_1)", "MIN(_1)") // TODO: Test
all the aggregates
+ List(
+ "COUNT(_1)",
+ "COUNT(*)",
+ "MAX(_1)",
+ "MIN(_1)",
+ "SUM(_1)"
+ ) // TODO: Test all the aggregates
aggregateFunctions.foreach { function =>
val queries = Seq(
diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
index a88f9f45..c445aae6 100644
--- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
@@ -816,4 +816,91 @@ abstract class CometTestBase
vertical: Boolean = false): String = {
df.showString(_numRows, truncate, vertical)
}
+
+ def makeParquetFile(
+ path: Path,
+ total: Int,
+ numGroups: Int,
+ dictionaryEnabled: Boolean): Unit = {
+ val schemaStr =
+ """
+ |message root {
+ | optional INT32 _1(INT_8);
+ | optional INT32 _2(INT_16);
+ | optional INT32 _3;
+ | optional INT64 _4;
+ | optional FLOAT _5;
+ | optional DOUBLE _6;
+ | optional INT32 _7(DECIMAL(5, 2));
+ | optional INT64 _8(DECIMAL(18, 10));
+ | optional FIXED_LEN_BYTE_ARRAY(16) _9(DECIMAL(38, 37));
+ | optional INT64 _10(TIMESTAMP(MILLIS,true));
+ | optional INT64 _11(TIMESTAMP(MICROS,true));
+ | optional INT32 _12(DATE);
+ | optional INT32 _g1(INT_8);
+ | optional INT32 _g2(INT_16);
+ | optional INT32 _g3;
+ | optional INT64 _g4;
+ | optional FLOAT _g5;
+ | optional DOUBLE _g6;
+ | optional INT32 _g7(DECIMAL(5, 2));
+ | optional INT64 _g8(DECIMAL(18, 10));
+ | optional FIXED_LEN_BYTE_ARRAY(16) _g9(DECIMAL(38, 37));
+ | optional INT64 _g10(TIMESTAMP(MILLIS,true));
+ | optional INT64 _g11(TIMESTAMP(MICROS,true));
+ | optional INT32 _g12(DATE);
+ | optional BINARY _g13(UTF8);
+ | optional BINARY _g14;
+ |}
+ """.stripMargin
+
+ val schema = MessageTypeParser.parseMessageType(schemaStr)
+ val writer = createParquetWriter(schema, path, dictionaryEnabled = true)
+
+ val rand = scala.util.Random
+ val expected = (0 until total).map { i =>
+ // use a single value for the first page, to make sure dictionary
encoding kicks in
+ if (rand.nextBoolean()) None
+ else {
+ if (dictionaryEnabled) Some(i % 10) else Some(i)
+ }
+ }
+
+ expected.foreach { opt =>
+ val record = new SimpleGroup(schema)
+ opt match {
+ case Some(i) =>
+ record.add(0, i.toByte)
+ record.add(1, i.toShort)
+ record.add(2, i)
+ record.add(3, i.toLong)
+ record.add(4, rand.nextFloat())
+ record.add(5, rand.nextDouble())
+ record.add(6, i)
+ record.add(7, i.toLong)
+ record.add(8, (i % 10).toString * 16)
+ record.add(9, i.toLong)
+ record.add(10, i.toLong)
+ record.add(11, i)
+ record.add(12, i.toByte % numGroups)
+ record.add(13, i.toShort % numGroups)
+ record.add(14, i % numGroups)
+ record.add(15, i.toLong % numGroups)
+ record.add(16, rand.nextFloat())
+ record.add(17, rand.nextDouble())
+ record.add(18, i)
+ record.add(19, i.toLong)
+ record.add(20, (i % 10).toString * 16)
+ record.add(21, i.toLong)
+ record.add(22, i.toLong)
+ record.add(23, i)
+ record.add(24, (i % 10).toString * 24)
+ record.add(25, (i % 10).toString * 36)
+ case _ =>
+ }
+ writer.write(record)
+ }
+
+ writer.close()
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
