Support the null value in bloom_filter_agg Spark aggregate function

Author: weixiuli

velox/functions/sparksql/aggregates/BloomFilterAggAggregate.cpp

@@ -85,14 +85,6 @@ class BloomFilterAggAggregate : public exec::Aggregate {
     }
   }
 
-  static FOLLY_ALWAYS_INLINE void checkBloomFilterNotNull(
-      DecodedVector& decoded,
-      vector_size_t idx) {
-    VELOX_USER_CHECK(
-        !decoded.isNullAt(idx),
-        "First argument of bloom_filter_agg cannot be null");
-  }
-
   void addRawInput(
       char** groups,
       const SelectivityVector& rows,
@@ -102,8 +94,8 @@ class BloomFilterAggAggregate : public exec::Aggregate {
     computeCapacity();
     auto mayHaveNulls = decodedRaw_.mayHaveNulls();
     rows.applyToSelected([&](vector_size_t row) {
-      if (mayHaveNulls) {
-        checkBloomFilterNotNull(decodedRaw_, row);
+      if (mayHaveNulls && decodedRaw_.isNullAt(row)) {
+        return;
       }
       auto group = groups[row];
       auto tracker = trackRowSize(group);
@@ -144,17 +136,18 @@ class BloomFilterAggAggregate : public exec::Aggregate {
     accumulator->init(capacity_);
     if (decodedRaw_.isConstantMapping()) {
       // All values are same, just do for the first.
-      checkBloomFilterNotNull(decodedRaw_, 0);
-      accumulator->insert(decodedRaw_.valueAt<int64_t>(0));
-      return;
-    }
-    auto mayHaveNulls = decodedRaw_.mayHaveNulls();
-    rows.applyToSelected([&](vector_size_t row) {
-      if (mayHaveNulls) {
-        checkBloomFilterNotNull(decodedRaw_, row);
+      if (!decodedRaw_.isNullAt(0)) {
+        accumulator->insert(decodedRaw_.valueAt<int64_t>(0));
       }
-      accumulator->insert(decodedRaw_.valueAt<int64_t>(row));
-    });
+    } else {
+      auto mayHaveNulls = decodedRaw_.mayHaveNulls();
+      rows.applyToSelected([&](vector_size_t row) {
+        if (mayHaveNulls && decodedRaw_.isNullAt(row)) {
+          return;
+        }
+        accumulator->insert(decodedRaw_.valueAt<int64_t>(row));
+      });
+    };
   }
 
   void addSingleGroupIntermediateResults(

velox/functions/sparksql/aggregates/tests/BloomFilterAggAggregateTest.cpp

@@ -71,14 +71,12 @@ TEST_F(BloomFilterAggAggregateTest, emptyInput) {
   testAggregations(vectors, {}, {"bloom_filter_agg(c0, 5, 64)"}, expected);
 }
 
-TEST_F(BloomFilterAggAggregateTest, nullBloomFilter) {
-  auto vectors = {makeRowVector({makeAllNullFlatVector<int64_t>(2)})};
-  auto expectedFake = {makeRowVector(
-      {makeNullableFlatVector<StringView>({std::nullopt}, VARBINARY())})};
-  VELOX_ASSERT_THROW(
-      testAggregations(
-          vectors, {}, {"bloom_filter_agg(c0, 5, 64)"}, expectedFake),
-      "First argument of bloom_filter_agg cannot be null");
+TEST_F(BloomFilterAggAggregateTest, nullInput) {
+  auto vectors = {makeRowVector(
+      {makeFlatVector<int64_t>(100, [](vector_size_t row) { return row % 9; }),
+       makeAllNullFlatVector<int64_t>(1)})};
+  auto expected = {makeRowVector({getSerializedBloomFilter(4)})};
+  testAggregations(vectors, {}, {"bloom_filter_agg(c0, 5, 64)"}, expected);
 }
 
 TEST_F(BloomFilterAggAggregateTest, config) {