add white list for agg opt

apache · Oct 31, 2023 · 6ebb3f2 · 6ebb3f2
1 parent 6503c27
commit 6ebb3f2
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 3 deletions.
diff --git a/...scala/io/glutenproject/execution/GlutenClickHouseTPCHColumnarShuffleParquetAQESuite.scala b/...scala/io/glutenproject/execution/GlutenClickHouseTPCHColumnarShuffleParquetAQESuite.scala
@@ -288,4 +288,43 @@ class GlutenClickHouseTPCHColumnarShuffleParquetAQESuite
       }
     }
   }
+
+  test("collect_set") {
+    val sql =
+      """
+        |select a, b from (
+        |select n_regionkey as a, collect_set(if(n_regionkey=0, n_name, null)) as set from nation group by n_regionkey)
+        |lateral view explode(set) as b
+        |order by a, b
+        |""".stripMargin
+    runQueryAndCompare(sql)(checkOperatorMatch[CHHashAggregateExecTransformer])
+  }
+
+  test("test 'aggregate function collect_list'") {
+    val df = runQueryAndCompare(
+      "select l_orderkey,from_unixtime(l_orderkey, 'yyyy-MM-dd HH:mm:ss') " +
+        "from lineitem order by l_orderkey desc limit 10"
+    )(checkOperatorMatch[ProjectExecTransformer])
+    checkLengthAndPlan(df, 10)
+  }
+
+  test("test max string") {
+    withSQLConf(("spark.gluten.sql.columnar.force.hashagg", "true")) {
+      val sql =
+        """
+          |SELECT
+          |    l_returnflag,
+          |    l_linestatus,
+          |    max(l_comment)
+          |FROM
+          |    lineitem
+          |WHERE
+          |    l_shipdate <= date'1998-09-02' - interval 1 day
+          |GROUP BY
+          |    l_returnflag,
+          |    l_linestatus
+          |""".stripMargin
+      runQueryAndCompare(sql, noFallBack = false) { df => }
+    }
+  }
 }
diff --git a/cpp-ch/local-engine/Builder/SerializedPlanBuilder.cpp b/cpp-ch/local-engine/Builder/SerializedPlanBuilder.cpp
@@ -259,7 +259,7 @@ std::shared_ptr<substrait::Type> SerializedPlanBuilder::buildType(const DB::Data
         res->mutable_i32()->set_nullability(type_nullability);
     else if (which.isInt64())
         res->mutable_i64()->set_nullability(type_nullability);
-    else if (which.isString() || which.isAggregateFunction() || which.isFixedString())
+    else if (which.isStringOrFixedString() || which.isAggregateFunction())
         res->mutable_binary()->set_nullability(type_nullability); /// Spark Binary type is more similiar to CH String type
     else if (which.isFloat32())
         res->mutable_fp32()->set_nullability(type_nullability);

diff --git a/cpp-ch/local-engine/Storages/IO/AggregateSerializationUtils.cpp b/cpp-ch/local-engine/Storages/IO/AggregateSerializationUtils.cpp
@@ -18,6 +18,7 @@
 #include <Common/Arena.h>
 
 #include <Columns/ColumnAggregateFunction.h>
+#include <Columns/ColumnFixedString.h>
 #include <DataTypes/DataTypeAggregateFunction.h>
 #include <DataTypes/DataTypeFixedString.h>
 
@@ -26,20 +27,34 @@ using namespace DB;
 
 namespace local_engine
 {
+
+bool isFixedSizeStateAggregateFunction(const String& name)
+{
+    // TODO max(String) should exclude, but fallback now
+    static const std::set<String> function_set = {"min", "max", "sum", "count", "avg"};
+    return function_set.contains(name);
+}
+
 DB::ColumnWithTypeAndName convertAggregateStateToFixedString(DB::ColumnWithTypeAndName col)
 {
     if (!isAggregateFunction(col.type))
     {
         return col;
     }
     const auto *aggregate_col = checkAndGetColumn<ColumnAggregateFunction>(*col.column);
+    // only support known fixed size aggregate function
+    if (!isFixedSizeStateAggregateFunction(aggregate_col->getAggregateFunction()->getName()))
+    {
+        return col;
+    }
     size_t state_size = aggregate_col->getAggregateFunction()->sizeOfData();
     auto res_type = std::make_shared<DataTypeFixedString>(state_size);
     auto res_col = res_type->createColumn();
-    res_col->reserve(aggregate_col->size());
+    PaddedPODArray<UInt8> & column_chars_t = assert_cast<ColumnFixedString &>(*res_col).getChars();
+    column_chars_t.reserve(aggregate_col->size() * state_size);
     for (const auto & item : aggregate_col->getData())
     {
-        res_col->insertData(item, state_size);
+        column_chars_t.insert_assume_reserved(item, item + state_size);
     }
     return DB::ColumnWithTypeAndName(std::move(res_col), res_type, col.name);
 }