fix

Signed-off-by: Yuan Zhou <yuan.zhou@intel.com>
oap-project · Aug 31, 2022 · 2964705 · 2964705
1 parent fb775a1
commit 2964705
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 11 deletions.
diff --git a/native-sql-engine/cpp/src/codegen/arrow_compute/ext/actions_impl.cc b/native-sql-engine/cpp/src/codegen/arrow_compute/ext/actions_impl.cc
@@ -473,14 +473,30 @@ class CountDistinctAction : public ActionBase {
       cache_.resize(max_group_id + 1, 0);
       length_ = cache_.size();
     }
-
+    in_list_ = in_list;
+    row_id = 0;
+    assert(in_list.size() > 1);
+    int gid = in_list.size() - 1;
+    typed_key_in = std::dynamic_pointer_cast<arrow::BooleanArray>(in_list[gid]);
     // prepare evaluate lambda
     *on_valid = [this](int dest_group_id) {
-      cache_[dest_group_id] += 1;
+      if (typed_key_in->GetView(row_id) != 0) {
+        bool foundNull = false;
+        for (int colId = 0; colId < in_list_.size() - 1; colId++) {
+          if (in_list_[colId]->IsNull(row_id)) {
+            foundNull = true;
+            break;
+          }
+        }
+        if (!foundNull) {
+          cache_[dest_group_id] += 1;
+        }
+      }
+      row_id++;
       return arrow::Status::OK();
     };
 
-    *on_null = [this]() { return arrow::Status::OK(); };
+    *on_null = [this]() { row_id++; return arrow::Status::OK(); };
     return arrow::Status::OK();
   }
 
@@ -587,13 +603,16 @@ class CountDistinctAction : public ActionBase {
   using ResBuilderType = typename arrow::TypeTraits<DataType>::BuilderType;
   // input
   arrow::compute::ExecContext* ctx_;
+  int row_id;
   // for debug only
   int32_t localGid_ = -1;
   // result
   using CType = typename arrow::TypeTraits<DataType>::CType;
   std::vector<CType> cache_;
   std::unique_ptr<ResBuilderType> builder_;
   uint64_t length_ = 0;
+  std::shared_ptr<arrow::BooleanArray> typed_key_in;
+  ArrayList in_list_;
 };
 
 //////////////// CountLiteralAction ///////////////

diff --git a/native-sql-engine/cpp/src/codegen/arrow_compute/ext/hash_aggregate_kernel.cc b/native-sql-engine/cpp/src/codegen/arrow_compute/ext/hash_aggregate_kernel.cc
@@ -876,10 +876,6 @@ class HashAggregateKernel::Impl {
         for (auto idx : action_prepare_index_list_[i]) {
           cols.push_back(in[idx]);
         }
-        if (action->getName() == "CountDistinctAction") {
-            RETURN_NOT_OK(action->EvaluateCountDistinct(cols));
-            continue;
-        }
         std::function<arrow::Status(int)> func;
         std::function<arrow::Status()> null_func;
         action->Submit(cols, max_group_id_, &func, &null_func);
@@ -1039,10 +1035,6 @@ class HashAggregateKernel::Impl {
         for (auto idx : action_prepare_index_list_[i]) {
           cols.push_back(in[idx]);
         }
-        if (action->getName() == "CountDistinctAction") {
-            RETURN_NOT_OK(action->EvaluateCountDistinct(cols));
-            continue;
-        }
         std::function<arrow::Status(int)> func;
         std::function<arrow::Status()> null_func;
         action->Submit(cols, max_group_id_, &func, &null_func);