Fix aggregation masks and validate the size of count arguments (oap-p…

…roject#142) * fix aggregation mask * validate count argument size
zhejiangxiaomai · Apr 20, 2023 · 36797c5 · 36797c5
1 parent 5657cdb
commit 36797c5
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 15 deletions.
diff --git a/velox/substrait/SubstraitToVeloxPlan.cpp b/velox/substrait/SubstraitToVeloxPlan.cpp
@@ -376,21 +376,18 @@ std::shared_ptr<const core::PlanNode> SubstraitVeloxPlanConverter::toVeloxAgg(
   aggregateMasks.reserve(aggRel.measures().size());
   for (const auto& smea : aggRel.measures()) {
     core::FieldAccessTypedExprPtr aggregateMask;
-    ::substrait::Expression substraitAggMask = smea.filter();
     // Get Aggregation Masks.
-    if (smea.has_filter()) {
-      if (substraitAggMask.ByteSizeLong() == 0) {
-        aggregateMask = {};
-      } else {
-        aggregateMask =
-            std::dynamic_pointer_cast<const core::FieldAccessTypedExpr>(
-                exprConverter_->toVeloxExpr(substraitAggMask, inputType));
-        VELOX_CHECK(
-            aggregateMask != nullptr,
-            " the agg filter expression in Aggregate Operator only support field");
-      }
-      aggregateMasks.push_back(aggregateMask);
+    if (!smea.has_filter()) {
+      aggregateMask = {};
+    } else {
+      aggregateMask =
+          std::dynamic_pointer_cast<const core::FieldAccessTypedExpr>(
+              exprConverter_->toVeloxExpr(smea.filter(), inputType));
+      VELOX_CHECK(
+          aggregateMask != nullptr,
+          " the agg filter expression in Aggregate Operator only support field");
     }
+    aggregateMasks.push_back(aggregateMask);
     const auto& aggFunction = smea.measure();
     std::string funcName = subParser_->findVeloxFunction(
         functionMap_, aggFunction.function_reference());

diff --git a/velox/substrait/SubstraitToVeloxPlanValidator.cpp b/velox/substrait/SubstraitToVeloxPlanValidator.cpp
@@ -718,9 +718,15 @@ bool SubstraitToVeloxPlanValidator::validate(
       }
 
       const auto& aggFunction = smea.measure();
-      funcSpecs.emplace_back(
-          planConverter_->findFuncSpec(aggFunction.function_reference()));
+      const auto& functionSpec = planConverter_->findFuncSpec(aggFunction.function_reference());
+      funcSpecs.emplace_back(functionSpec);
       toVeloxType(subParser_->parseType(aggFunction.output_type())->type);
+      // Validate the size of arguments.
+      if (subParser_->getSubFunctionName(functionSpec) == "count" &&
+          aggFunction.arguments().size() > 1) {
+        // Count accepts only one argument.
+        return false;
+      }
       for (const auto& arg : aggFunction.arguments()) {
         auto typeCase = arg.value().rex_type_case();
         switch (typeCase) {