Handle UNKNOWN type input in approx_distinct (#11367)

Summary: Pull Request resolved: #11367 Reviewed By: amitkdutta Differential Revision: D65100134
facebookincubator · Oct 29, 2024 · c839f70 · c839f70
1 parent 27d0527
commit c839f70
Show file tree

Hide file tree

Showing 4 changed files with 123 additions and 13 deletions.
diff --git a/velox/exec/SetAccumulator.h b/velox/exec/SetAccumulator.h
@@ -339,6 +339,62 @@ struct ComplexTypeSetAccumulator {
   }
 };
 
+class UnknownTypeSetAccumulator {
+ public:
+  UnknownTypeSetAccumulator(
+      const TypePtr& /*type*/,
+      HashStringAllocator* /*allocator*/) {}
+
+  void addValue(
+      const DecodedVector& decoded,
+      vector_size_t index,
+      HashStringAllocator* /*allocator*/) {
+    VELOX_DCHECK(decoded.isNullAt(index));
+    hasNull_ = true;
+  }
+
+  void addValues(
+      const ArrayVector& arrayVector,
+      vector_size_t index,
+      const DecodedVector& values,
+      HashStringAllocator* allocator) {
+    VELOX_DCHECK(!arrayVector.isNullAt(index));
+    const auto size = arrayVector.sizeAt(index);
+    const auto offset = arrayVector.offsetAt(index);
+    for (auto i = 0; i < size; ++i) {
+      addValue(values, offset + i, allocator);
+    }
+  }
+
+  void addNonNullValue(
+      const DecodedVector& /*decoded*/,
+      vector_size_t /*index*/,
+      HashStringAllocator* /*allocator*/) {}
+
+  void addNonNullValues(
+      const ArrayVector& /*arrayVector*/,
+      vector_size_t /*index*/,
+      const DecodedVector& /*values*/,
+      HashStringAllocator* /*/allocator*/) {}
+
+  size_t size() const {
+    return hasNull_ ? 1 : 0;
+  }
+
+  vector_size_t extractValues(BaseVector& values, vector_size_t offset) {
+    if (!hasNull_) {
+      return 0;
+    }
+    values.setNull(offset, true);
+    return 1;
+  }
+
+  void free(HashStringAllocator& /*allocator*/) {}
+
+ private:
+  bool hasNull_ = false;
+};
+
 template <typename T>
 struct SetAccumulatorTypeTraits {
   using AccumulatorType = SetAccumulator<T>;
@@ -369,6 +425,12 @@ template <>
 struct SetAccumulatorTypeTraits<ComplexType> {
   using AccumulatorType = ComplexTypeSetAccumulator;
 };
+
+template <>
+struct SetAccumulatorTypeTraits<UnknownValue> {
+  using AccumulatorType = UnknownTypeSetAccumulator;
+};
+
 } // namespace detail
 
 // A wrapper around SetAccumulator that overrides hash and equal_to functions to

diff --git a/velox/functions/prestosql/aggregates/ApproxDistinctAggregate.cpp b/velox/functions/prestosql/aggregates/ApproxDistinctAggregate.cpp
@@ -453,19 +453,21 @@ exec::AggregateRegistrationResult registerApproxDistinct(
                              .argumentType("hyperloglog")
                              .build());
   } else {
-    for (const auto& inputType :
-         {"boolean",
-          "tinyint",
-          "smallint",
-          "integer",
-          "bigint",
-          "hugeint",
-          "real",
-          "double",
-          "varchar",
-          "varbinary",
-          "timestamp",
-          "date"}) {
+    for (const auto& inputType : {
+             "boolean",
+             "tinyint",
+             "smallint",
+             "integer",
+             "bigint",
+             "hugeint",
+             "real",
+             "double",
+             "varchar",
+             "varbinary",
+             "timestamp",
+             "date",
+             "unknown",
+         }) {
       signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                                .returnType(returnType)
                                .intermediateType("varbinary")
@@ -505,6 +507,10 @@ exec::AggregateRegistrationResult registerApproxDistinct(
           const TypePtr& resultType,
           const core::QueryConfig& /*config*/)
           -> std::unique_ptr<exec::Aggregate> {
+        if (argTypes[0]->isUnKnown()) {
+          return std::make_unique<ApproxDistinctAggregate<UnknownValue>>(
+              resultType, hllAsFinalResult, hllAsRawInput, defaultError);
+        }
         return VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH(
             createApproxDistinct,
             argTypes[0]->kind(),

diff --git a/velox/functions/prestosql/aggregates/SetAggregates.cpp b/velox/functions/prestosql/aggregates/SetAggregates.cpp
@@ -419,6 +419,9 @@ void registerCountDistinctAggregate(
           case TypeKind::ROW:
             return std::make_unique<CountDistinctAggregate<ComplexType>>(
                 resultType, argTypes[0]);
+          case TypeKind::UNKNOWN:
+            return std::make_unique<CountDistinctAggregate<UnknownValue>>(
+                resultType, argTypes[0]);
           default:
             VELOX_UNREACHABLE(
                 "Unexpected type {}", mapTypeKindToName(typeKind));

diff --git a/velox/functions/prestosql/aggregates/tests/ApproxDistinctTest.cpp b/velox/functions/prestosql/aggregates/tests/ApproxDistinctTest.cpp
@@ -432,5 +432,44 @@ TEST_F(ApproxDistinctTest, toIntermediate) {
       digests, {"c0"}, {"merge(a0)"}, {"c0", "cardinality(a0)"}, {input});
 }
 
+TEST_F(ApproxDistinctTest, unknownType) {
+  constexpr int kSize = 10;
+  auto input = makeRowVector({
+      makeFlatVector<int32_t>(kSize, [](auto i) { return i % 2; }),
+      makeAllNullFlatVector<UnknownValue>(kSize),
+  });
+  testAggregations(
+      {input},
+      {},
+      {"approx_distinct(c1)", "approx_distinct(c1, 0.023)"},
+      {makeRowVector(std::vector<VectorPtr>(2, makeConstant<int64_t>(0, 1)))});
+  testAggregations(
+      {input},
+      {},
+      {"approx_set(c1)", "approx_set(c1, 0.01625)"},
+      {"cardinality(a0)", "cardinality(a1)"},
+      {makeRowVector(
+          std::vector<VectorPtr>(2, makeNullConstant(TypeKind::BIGINT, 1)))});
+  testAggregations(
+      {input},
+      {"c0"},
+      {"approx_distinct(c1)", "approx_distinct(c1, 0.023)"},
+      {makeRowVector({
+          makeFlatVector<int32_t>({0, 1}),
+          makeFlatVector<int64_t>({0, 0}),
+          makeFlatVector<int64_t>({0, 0}),
+      })});
+  testAggregations(
+      {input},
+      {"c0"},
+      {"approx_set(c1)", "approx_set(c1, 0.01625)"},
+      {"c0", "cardinality(a0)", "cardinality(a1)"},
+      {makeRowVector({
+          makeFlatVector<int32_t>({0, 1}),
+          makeNullConstant(TypeKind::BIGINT, 2),
+          makeNullConstant(TypeKind::BIGINT, 2),
+      })});
+}
+
 } // namespace
 } // namespace facebook::velox::aggregate::test