This is an automated email from the ASF dual-hosted git repository.
viirya 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 2e4ec701 feat: add nullOnDivideByZero for Covariance (#564)
2e4ec701 is described below
commit 2e4ec70166ff4517b04aefaf74ddcf8256dff180
Author: Huaxin Gao <[email protected]>
AuthorDate: Sat Jun 22 10:45:02 2024 -0700
feat: add nullOnDivideByZero for Covariance (#564)
* feat: add nullOnDivideByZero for Covariance
* change test
* fix style
* refactor test to reuse code
* Trigger Build
---
.../datafusion/expressions/correlation.rs | 2 +-
.../execution/datafusion/expressions/covariance.rs | 33 ++-
core/src/execution/datafusion/planner.rs | 43 ++--
core/src/execution/proto/expr.proto | 19 +-
.../org/apache/comet/serde/QueryPlanSerde.scala | 16 +-
.../apache/comet/exec/CometAggregateSuite.scala | 263 +++++++--------------
6 files changed, 149 insertions(+), 227 deletions(-)
diff --git a/core/src/execution/datafusion/expressions/correlation.rs
b/core/src/execution/datafusion/expressions/correlation.rs
index c83341e5..78a3f341 100644
--- a/core/src/execution/datafusion/expressions/correlation.rs
+++ b/core/src/execution/datafusion/expressions/correlation.rs
@@ -150,7 +150,7 @@ impl CorrelationAccumulator {
/// Creates a new `CorrelationAccumulator`
pub fn try_new(null_on_divide_by_zero: bool) -> Result<Self> {
Ok(Self {
- covar: CovarianceAccumulator::try_new(StatsType::Population)?,
+ covar: CovarianceAccumulator::try_new(StatsType::Population,
null_on_divide_by_zero)?,
stddev1: StddevAccumulator::try_new(StatsType::Population,
null_on_divide_by_zero)?,
stddev2: StddevAccumulator::try_new(StatsType::Population,
null_on_divide_by_zero)?,
null_on_divide_by_zero,
diff --git a/core/src/execution/datafusion/expressions/covariance.rs
b/core/src/execution/datafusion/expressions/covariance.rs
index 5d0e550f..151fb61a 100644
--- a/core/src/execution/datafusion/expressions/covariance.rs
+++ b/core/src/execution/datafusion/expressions/covariance.rs
@@ -44,6 +44,7 @@ pub struct Covariance {
expr1: Arc<dyn PhysicalExpr>,
expr2: Arc<dyn PhysicalExpr>,
stats_type: StatsType,
+ null_on_divide_by_zero: bool,
}
impl Covariance {
@@ -54,6 +55,7 @@ impl Covariance {
name: impl Into<String>,
data_type: DataType,
stats_type: StatsType,
+ null_on_divide_by_zero: bool,
) -> Self {
// the result of covariance just support FLOAT64 data type.
assert!(matches!(data_type, DataType::Float64));
@@ -62,6 +64,7 @@ impl Covariance {
expr1,
expr2,
stats_type,
+ null_on_divide_by_zero,
}
}
}
@@ -77,7 +80,10 @@ impl AggregateExpr for Covariance {
}
fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
- Ok(Box::new(CovarianceAccumulator::try_new(self.stats_type)?))
+ Ok(Box::new(CovarianceAccumulator::try_new(
+ self.stats_type,
+ self.null_on_divide_by_zero,
+ )?))
}
fn state_fields(&self) -> Result<Vec<Field>> {
@@ -123,6 +129,7 @@ impl PartialEq<dyn Any> for Covariance {
&& self.expr1.eq(&x.expr1)
&& self.expr2.eq(&x.expr2)
&& self.stats_type == x.stats_type
+ && self.null_on_divide_by_zero == x.null_on_divide_by_zero
})
.unwrap_or(false)
}
@@ -136,17 +143,19 @@ pub struct CovarianceAccumulator {
mean2: f64,
count: f64,
stats_type: StatsType,
+ null_on_divide_by_zero: bool,
}
impl CovarianceAccumulator {
/// Creates a new `CovarianceAccumulator`
- pub fn try_new(s_type: StatsType) -> Result<Self> {
+ pub fn try_new(s_type: StatsType, null_on_divide_by_zero: bool) ->
Result<Self> {
Ok(Self {
algo_const: 0_f64,
mean1: 0_f64,
mean2: 0_f64,
count: 0_f64,
stats_type: s_type,
+ null_on_divide_by_zero,
})
}
@@ -287,22 +296,24 @@ impl Accumulator for CovarianceAccumulator {
}
fn evaluate(&mut self) -> Result<ScalarValue> {
+ if self.count == 0.0 {
+ return Ok(ScalarValue::Float64(None));
+ }
+
let count = match self.stats_type {
StatsType::Population => self.count,
+ StatsType::Sample if self.count > 1.0 => self.count - 1.0,
StatsType::Sample => {
- if self.count > 0.0 {
- self.count - 1.0
+ // self.count == 1.0
+ return if self.null_on_divide_by_zero {
+ Ok(ScalarValue::Float64(None))
} else {
- self.count
- }
+ Ok(ScalarValue::Float64(Some(f64::NAN)))
+ };
}
};
- if count == 0.0 {
- Ok(ScalarValue::Float64(None))
- } else {
- Ok(ScalarValue::Float64(Some(self.algo_const / count)))
- }
+ Ok(ScalarValue::Float64(Some(self.algo_const / count)))
}
fn size(&self) -> usize {
diff --git a/core/src/execution/datafusion/planner.rs
b/core/src/execution/datafusion/planner.rs
index cd9822d6..afdebce3 100644
--- a/core/src/execution/datafusion/planner.rs
+++ b/core/src/execution/datafusion/planner.rs
@@ -1232,29 +1232,32 @@ impl PhysicalPlanner {
let datatype =
to_arrow_datatype(expr.datatype.as_ref().unwrap());
Ok(Arc::new(BitXor::new(child, "bit_xor", datatype)))
}
- AggExprStruct::CovSample(expr) => {
+ AggExprStruct::Covariance(expr) => {
let child1 = self.create_expr(expr.child1.as_ref().unwrap(),
schema.clone())?;
let child2 = self.create_expr(expr.child2.as_ref().unwrap(),
schema.clone())?;
let datatype =
to_arrow_datatype(expr.datatype.as_ref().unwrap());
- Ok(Arc::new(Covariance::new(
- child1,
- child2,
- "covariance",
- datatype,
- StatsType::Sample,
- )))
- }
- AggExprStruct::CovPopulation(expr) => {
- let child1 = self.create_expr(expr.child1.as_ref().unwrap(),
schema.clone())?;
- let child2 = self.create_expr(expr.child2.as_ref().unwrap(),
schema.clone())?;
- let datatype =
to_arrow_datatype(expr.datatype.as_ref().unwrap());
- Ok(Arc::new(Covariance::new(
- child1,
- child2,
- "covariance_pop",
- datatype,
- StatsType::Population,
- )))
+ match expr.stats_type {
+ 0 => Ok(Arc::new(Covariance::new(
+ child1,
+ child2,
+ "covariance",
+ datatype,
+ StatsType::Sample,
+ expr.null_on_divide_by_zero,
+ ))),
+ 1 => Ok(Arc::new(Covariance::new(
+ child1,
+ child2,
+ "covariance_pop",
+ datatype,
+ StatsType::Population,
+ expr.null_on_divide_by_zero,
+ ))),
+ stats_type => Err(ExecutionError::GeneralError(format!(
+ "Unknown StatisticsType {:?} for Variance",
+ stats_type
+ ))),
+ }
}
AggExprStruct::Variance(expr) => {
let child = self.create_expr(expr.child.as_ref().unwrap(),
schema.clone())?;
diff --git a/core/src/execution/proto/expr.proto
b/core/src/execution/proto/expr.proto
index 093b07b3..6b66a307 100644
--- a/core/src/execution/proto/expr.proto
+++ b/core/src/execution/proto/expr.proto
@@ -92,11 +92,10 @@ message AggExpr {
BitAndAgg bitAndAgg = 9;
BitOrAgg bitOrAgg = 10;
BitXorAgg bitXorAgg = 11;
- CovSample covSample = 12;
- CovPopulation covPopulation = 13;
- Variance variance = 14;
- Stddev stddev = 15;
- Correlation correlation = 16;
+ Covariance covariance = 12;
+ Variance variance = 13;
+ Stddev stddev = 14;
+ Correlation correlation = 15;
}
}
@@ -159,18 +158,12 @@ message BitXorAgg {
DataType datatype = 2;
}
-message CovSample {
- Expr child1 = 1;
- Expr child2 = 2;
- bool null_on_divide_by_zero = 3;
- DataType datatype = 4;
-}
-
-message CovPopulation {
+message Covariance {
Expr child1 = 1;
Expr child2 = 2;
bool null_on_divide_by_zero = 3;
DataType datatype = 4;
+ StatisticsType stats_type = 5;
}
message Variance {
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 13abaa0c..1e61ef75 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -428,40 +428,44 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
withInfo(aggExpr, child)
None
}
- case cov @ CovSample(child1, child2, _) =>
+ case cov @ CovSample(child1, child2, nullOnDivideByZero) =>
val child1Expr = exprToProto(child1, inputs, binding)
val child2Expr = exprToProto(child2, inputs, binding)
val dataType = serializeDataType(cov.dataType)
if (child1Expr.isDefined && child2Expr.isDefined &&
dataType.isDefined) {
- val covBuilder = ExprOuterClass.CovSample.newBuilder()
+ val covBuilder = ExprOuterClass.Covariance.newBuilder()
covBuilder.setChild1(child1Expr.get)
covBuilder.setChild2(child2Expr.get)
+ covBuilder.setNullOnDivideByZero(nullOnDivideByZero)
covBuilder.setDatatype(dataType.get)
+ covBuilder.setStatsTypeValue(0)
Some(
ExprOuterClass.AggExpr
.newBuilder()
- .setCovSample(covBuilder)
+ .setCovariance(covBuilder)
.build())
} else {
None
}
- case cov @ CovPopulation(child1, child2, _) =>
+ case cov @ CovPopulation(child1, child2, nullOnDivideByZero) =>
val child1Expr = exprToProto(child1, inputs, binding)
val child2Expr = exprToProto(child2, inputs, binding)
val dataType = serializeDataType(cov.dataType)
if (child1Expr.isDefined && child2Expr.isDefined &&
dataType.isDefined) {
- val covBuilder = ExprOuterClass.CovPopulation.newBuilder()
+ val covBuilder = ExprOuterClass.Covariance.newBuilder()
covBuilder.setChild1(child1Expr.get)
covBuilder.setChild2(child2Expr.get)
+ covBuilder.setNullOnDivideByZero(nullOnDivideByZero)
covBuilder.setDatatype(dataType.get)
+ covBuilder.setStatsTypeValue(1)
Some(
ExprOuterClass.AggExpr
.newBuilder()
- .setCovPopulation(covBuilder)
+ .setCovariance(covBuilder)
.build())
} else {
None
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 93be1111..2acc562b 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -1089,36 +1089,97 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
- test("covar_pop and covar_samp") {
+ def setupAndTestAggregates(
+ table: String,
+ data: Seq[(Any, Any, Any)],
+ dataTypes: (String, String, String),
+ aggregates: String): Unit = {
+ val (type1, type2, type3) = dataTypes
+ withTable(table) {
+ sql(s"create table $table(col1 $type1, col2 $type2, col3 $type3) using
parquet")
+ val values = data
+ .map { case (c1, c2, c3) =>
+ s"($c1, $c2, $c3)"
+ }
+ .mkString(", ")
+ sql(s"insert into $table values $values")
+
+ val expectedNumOfCometAggregates = 2
+
+ checkSparkAnswerWithTolAndNumOfAggregates(
+ s"SELECT $aggregates FROM $table",
+ expectedNumOfCometAggregates)
+
+ checkSparkAnswerWithTolAndNumOfAggregates(
+ s"SELECT $aggregates FROM $table GROUP BY col3",
+ expectedNumOfCometAggregates)
+ }
+ }
+
+ test("covariance & correlation") {
withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
- Seq("native", "jvm").foreach { cometShuffleMode =>
+ Seq("jvm", "native").foreach { cometShuffleMode =>
withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
Seq(true, false).foreach { dictionary =>
withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
- val table = "test"
- withTable(table) {
- sql(
- s"create table $table(col1 int, col2 int, col3 int, col4
float, col5 double," +
- " col6 double, col7 int) using parquet")
- sql(s"insert into $table values(1, 4, null, 1.1, 2.2, null,
1)," +
- " (2, 5, 6, 3.4, 5.6, null, 1), (3, 6, null, 7.9, 2.4, null,
2)")
- val expectedNumOfCometAggregates = 2
- checkSparkAnswerAndNumOfAggregates(
- "SELECT covar_samp(col1, col2), covar_samp(col1, col3),
covar_samp(col4, col5)," +
- " covar_samp(col4, col6) FROM test",
- expectedNumOfCometAggregates)
- checkSparkAnswerAndNumOfAggregates(
- "SELECT covar_pop(col1, col2), covar_pop(col1, col3),
covar_pop(col4, col5)," +
- " covar_pop(col4, col6) FROM test",
- expectedNumOfCometAggregates)
- checkSparkAnswerAndNumOfAggregates(
- "SELECT covar_samp(col1, col2), covar_samp(col1, col3),
covar_samp(col4, col5)," +
- " covar_samp(col4, col6) FROM test GROUP BY col7",
- expectedNumOfCometAggregates)
- checkSparkAnswerAndNumOfAggregates(
- "SELECT covar_pop(col1, col2), covar_pop(col1, col3),
covar_pop(col4, col5)," +
- " covar_pop(col4, col6) FROM test GROUP BY col7",
- expectedNumOfCometAggregates)
+ Seq(true, false).foreach { nullOnDivideByZero =>
+ withSQLConf(
+ "spark.sql.legacy.statisticalAggregate" ->
nullOnDivideByZero.toString) {
+ val table = "test"
+ val aggregates =
+ "covar_samp(col1, col2), covar_pop(col1, col2), corr(col1,
col2)"
+ setupAndTestAggregates(
+ table,
+ Seq((1, 4, 1), (2, 5, 1), (3, 6, 2)),
+ ("double", "double", "double"),
+ aggregates)
+ setupAndTestAggregates(
+ table,
+ Seq((1, 4, 3), (2, -5, 3), (3, 6, 1)),
+ ("double", "double", "double"),
+ aggregates)
+ setupAndTestAggregates(
+ table,
+ Seq((1.1, 4.1, 2.3), (2, 5, 1.5), (3, 6, 2.3)),
+ ("double", "double", "double"),
+ aggregates)
+ setupAndTestAggregates(
+ table,
+ Seq(
+ (1, 4, 1),
+ (2, 5, 2),
+ (3, 6, 3),
+ (1.1, 4.4, 1),
+ (2.2, 5.5, 2),
+ (3.3, 6.6, 3)),
+ ("double", "double", "double"),
+ aggregates)
+ setupAndTestAggregates(
+ table,
+ Seq((1, 4, 1), (2, 5, 2), (3, 6, 3)),
+ ("int", "int", "int"),
+ aggregates)
+ setupAndTestAggregates(
+ table,
+ Seq((1, 4, 2), (null, null, 2), (3, 6, 1), (3, 3, 1)),
+ ("int", "int", "int"),
+ aggregates)
+ setupAndTestAggregates(
+ table,
+ Seq((1, 4, 1), (null, 5, 1), (2, 5, 2), (9, null, 2), (3,
6, 2)),
+ ("int", "int", "int"),
+ aggregates)
+ setupAndTestAggregates(
+ table,
+ Seq((null, null, 1), (1, 2, 1), (null, null, 2)),
+ ("int", "int", "int"),
+ aggregates)
+ setupAndTestAggregates(
+ table,
+ Seq((null, null, 1), (null, null, 1), (null, null, 2)),
+ ("int", "int", "int"),
+ aggregates)
+ }
}
}
}
@@ -1208,156 +1269,6 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
- test("correlation") {
- withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
- Seq("jvm", "native").foreach { cometShuffleMode =>
- withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
- Seq(true, false).foreach { dictionary =>
- withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
- Seq(true, false).foreach { nullOnDivideByZero =>
- withSQLConf(
- "spark.sql.legacy.statisticalAggregate" ->
nullOnDivideByZero.toString) {
- val table = "test"
- withTable(table) {
- sql(
- s"create table $table(col1 double, col2 double, col3
double) using parquet")
- sql(s"insert into $table values(1, 4, 1), (2, 5, 1), (3,
6, 2)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
-
- withTable(table) {
- sql(
- s"create table $table(col1 double, col2 double, col3
double) using parquet")
- sql(s"insert into $table values(1, 4, 3), (2, -5, 3), (3,
6, 1)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
-
- withTable(table) {
- sql(
- s"create table $table(col1 double, col2 double, col3
double) using parquet")
- sql(s"insert into $table values(1.1, 4.1, 2.3), (2, 5,
1.5), (3, 6, 2.3)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
-
- withTable(table) {
- sql(
- s"create table $table(col1 double, col2 double, col3
double) using parquet")
- sql(s"insert into $table values(1, 4, 1), (2, 5, 2), (3,
6, 3), (1.1, 4.4, 1), (2.2, 5.5, 2), (3.3, 6.6, 3)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
-
- withTable(table) {
- sql(s"create table $table(col1 int, col2 int, col3 int)
using parquet")
- sql(s"insert into $table values(1, 4, 1), (2, 5, 2), (3,
6, 3)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
-
- withTable(table) {
- sql(s"create table $table(col1 int, col2 int, col3 int)
using parquet")
- sql(
- s"insert into $table values(1, 4, 2), (null, null, 2),
(3, 6, 1), (3, 3, 1)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
-
- withTable(table) {
- sql(s"create table $table(col1 int, col2 int, col3 int)
using parquet")
- sql(s"insert into $table values(1, 4, 1), (null, 5, 1),
(2, 5, 2), (9, null, 2), (3, 6, 2)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
-
- withTable(table) {
- sql(s"create table $table(col1 int, col2 int, col3 int)
using parquet")
- sql(s"insert into $table values(null, null, 1), (1, 2, 1),
(null, null, 2)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
-
- withTable(table) {
- sql(s"create table $table(col1 int, col2 int, col3 int)
using parquet")
- sql(
- s"insert into $table values(null, null, 1), (null, null,
1), (null, null, 2)")
- val expectedNumOfCometAggregates = 2
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table",
- expectedNumOfCometAggregates)
-
- checkSparkAnswerWithTolAndNumOfAggregates(
- s"SELECT corr(col1, col2) FROM $table GROUP BY col3",
- expectedNumOfCometAggregates)
- }
- }
- }
- }
- }
- }
- }
- }
- }
-
protected def checkSparkAnswerAndNumOfAggregates(query: String,
numAggregates: Int): Unit = {
val df = sql(query)
checkSparkAnswer(df)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
