handle y null

wayneguow · Jun 26, 2024 · e5dbb7c · e5dbb7c
1 parent a474b88
commit e5dbb7c
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 9 deletions.
diff --git a/...src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/linearRegression.scala b/...src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/linearRegression.scala
@@ -258,7 +258,8 @@ case class RegrSlope(left: Expression, right: Expression) extends DeclarativeAgg
 
   private val covarPop = new CovPopulation(right, left)
 
-  private val varPop = new VariancePop(right)
+  private val varPop = new VariancePop(If(And(IsNotNull(left), IsNotNull(right)),
+    right, Literal.create(null, right.dataType)))
 
   override def nullable: Boolean = true
 
@@ -311,7 +312,8 @@ case class RegrIntercept(left: Expression, right: Expression) extends Declarativ
 
   private val covarPop = new CovPopulation(right, left)
 
-  private val varPop = new VariancePop(right)
+  private val varPop = new VariancePop(If(And(IsNotNull(left), IsNotNull(right)),
+    right, Literal.create(null, right.dataType)))
 
   override def nullable: Boolean = true
 

diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/linear-regression.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/linear-regression.sql.out
@@ -1,11 +1,11 @@
 -- Automatically generated by SQLQueryTestSuite
 -- !query
 CREATE OR REPLACE TEMPORARY VIEW testRegression AS SELECT * FROM VALUES
-(1, 10, null), (2, 10, 11), (2, 20, 22), (2, 25, null), (2, 30, 35)
+(1, 10, null), (2, 10, 11), (2, 20, 22), (2, 25, null), (2, 30, 35), (2, null, 40)
 AS testRegression(k, y, x)
 -- !query analysis
 CreateViewCommand `testRegression`, SELECT * FROM VALUES
-(1, 10, null), (2, 10, 11), (2, 20, 22), (2, 25, null), (2, 30, 35)
+(1, 10, null), (2, 10, 11), (2, 20, 22), (2, 25, null), (2, 30, 35), (2, null, 40)
 AS testRegression(k, y, x), false, true, LocalTempView, UNSUPPORTED, true
    +- Project [k#x, y#x, x#x]
       +- SubqueryAlias testRegression

diff --git a/sql/core/src/test/resources/sql-tests/inputs/linear-regression.sql b/sql/core/src/test/resources/sql-tests/inputs/linear-regression.sql
@@ -1,6 +1,6 @@
 -- Test data.
 CREATE OR REPLACE TEMPORARY VIEW testRegression AS SELECT * FROM VALUES
-(1, 10, null), (2, 10, 11), (2, 20, 22), (2, 25, null), (2, 30, 35)
+(1, 10, null), (2, 10, 11), (2, 20, 22), (2, 25, null), (2, 30, 35), (2, null, 40)
 AS testRegression(k, y, x);
 
 -- SPARK-37613: Support ANSI Aggregate Function: regr_count

diff --git a/sql/core/src/test/resources/sql-tests/results/linear-regression.sql.out b/sql/core/src/test/resources/sql-tests/results/linear-regression.sql.out
@@ -1,7 +1,7 @@
 -- Automatically generated by SQLQueryTestSuite
 -- !query
 CREATE OR REPLACE TEMPORARY VIEW testRegression AS SELECT * FROM VALUES
-(1, 10, null), (2, 10, 11), (2, 20, 22), (2, 25, null), (2, 30, 35)
+(1, 10, null), (2, 10, 11), (2, 20, 22), (2, 25, null), (2, 30, 35), (2, null, 40)
 AS testRegression(k, y, x)
 -- !query schema
 struct<>
@@ -31,7 +31,7 @@ SELECT k, count(*), regr_count(y, x) FROM testRegression GROUP BY k
 struct<k:int,count(1):bigint,regr_count(y, x):bigint>
 -- !query output
 1	1	0
-2	4	3
+2	5	3
 
 
 -- !query
@@ -40,7 +40,7 @@ SELECT k, count(*) FILTER (WHERE x IS NOT NULL), regr_count(y, x) FROM testRegre
 struct<k:int,count(1) FILTER (WHERE (x IS NOT NULL)):bigint,regr_count(y, x):bigint>
 -- !query output
 1	0	0
-2	3	3
+2	4	3
 
 
 -- !query
@@ -99,7 +99,7 @@ SELECT k, avg(x), avg(y), regr_avgx(y, x), regr_avgy(y, x) FROM testRegression G
 struct<k:int,avg(x):double,avg(y):double,regr_avgx(y, x):double,regr_avgy(y, x):double>
 -- !query output
 1	NULL	10.0	NULL	NULL
-2	22.666666666666668	21.25	22.666666666666668	20.0
+2	27.0	21.25	22.666666666666668	20.0
 
 
 -- !query