address comments

apache · Mar 13, 2024 · 2315dd0 · 2315dd0
1 parent 425a0af
commit 2315dd0
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -188,7 +188,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde {
 
   private def bitwiseAggTypeSupported(dt: DataType): Boolean = {
     dt match {
-      case _: IntegerType => true
+      case _: IntegerType | LongType | ShortType | ByteType => true
       case _ => false
     }
   }

diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -955,25 +955,41 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
           val table = "test"
           withTable(table) {
-            sql(s"create table $table(col1 int, col2 int, col3 int) using parquet")
+            sql(s"create table $table(col1 long, col2 int, col3 short, col4 byte) using parquet")
             sql(
-              s"insert into $table values(4, 1, 1), (4, 1, 1), (3, 3, 1)," +
-                " (2, 4, 2), (1, 3, 2), (null, 1, 1)")
+              s"insert into $table values(4, 1, 1, 3), (4, 1, 1, 3), (3, 3, 1, 4)," +
+                " (2, 4, 2, 5), (1, 3, 2, 6), (null, 1, 1, 7)")
             val expectedNumOfCometAggregates = 2
             checkSparkAnswerAndNumOfAggregates(
-              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1) FROM test",
+              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1)," +
+                " BIT_AND(col2), BIT_OR(col2), BIT_XOR(col2)," +
+                " BIT_AND(col3), BIT_OR(col3), BIT_XOR(col3)," +
+                " BIT_AND(col4), BIT_OR(col4), BIT_XOR(col4) FROM test",
               expectedNumOfCometAggregates)
 
+            // Make sure the combination of BITWISE aggregates and other aggregates work OK
             checkSparkAnswerAndNumOfAggregates(
-              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1), MIN(col1), COUNT(col1) FROM test",
+              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1)," +
+                " BIT_AND(col2), BIT_OR(col2), BIT_XOR(col2)," +
+                " BIT_AND(col3), BIT_OR(col3), BIT_XOR(col3)," +
+                " BIT_AND(col4), BIT_OR(col4), BIT_XOR(col4), MIN(col1), COUNT(col1) FROM test",
               expectedNumOfCometAggregates)
 
             checkSparkAnswerAndNumOfAggregates(
-              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1), col3 FROM test GROUP BY col3",
+              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1)," +
+                " BIT_AND(col2), BIT_OR(col2), BIT_XOR(col2)," +
+                " BIT_AND(col3), BIT_OR(col3), BIT_XOR(col3)," +
+                " BIT_AND(col4), BIT_OR(col4), BIT_XOR(col4), col3 FROM test GROUP BY col3",
               expectedNumOfCometAggregates)
 
+            // Make sure the combination of BITWISE aggregates and other aggregates work OK
+            // with group by
             checkSparkAnswerAndNumOfAggregates(
-              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1), MIN(col1), COUNT(col1), col3 FROM test GROUP BY col3",
+              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1)," +
+                " BIT_AND(col2), BIT_OR(col2), BIT_XOR(col2)," +
+                " BIT_AND(col3), BIT_OR(col3), BIT_XOR(col3)," +
+                " BIT_AND(col4), BIT_OR(col4), BIT_XOR(col4)," +
+                " MIN(col1), COUNT(col1), col3 FROM test GROUP BY col3",
               expectedNumOfCometAggregates)
           }
         }