subfield bloom filter

Author: zhztheplayer

cpp/velox/operators/functions/SparkExprToSubfieldFilterParser.cc

@@ -16,11 +16,104 @@
  */
 #include "operators/functions/SparkExprToSubfieldFilterParser.h"
 
+#include "utils/Exception.h"
+#include "velox/common/base/BloomFilter.h"
+#include "velox/functions/sparksql/XxHash64.h"
+#include "velox/expression/Expr.h"
+#include "velox/vector/ComplexVector.h"
+
 namespace gluten {
 
 using namespace facebook::velox;
 
 namespace {
+
+// Evaluates an expression as a constant. Returns nullptr if the expression is
+// not constant or evaluation fails. Errors are intentionally swallowed because
+// a non-evaluable expression simply means the filter cannot be pushed down.
+VectorPtr toConstant(const core::TypedExprPtr& expr, core::ExpressionEvaluator* evaluator) {
+  auto exprSet = evaluator->compile(expr);
+  if (!exprSet->exprs()[0]->isConstantExpr()) {
+    return nullptr;
+  }
+  RowVector input(evaluator->pool(), ROW({}, {}), nullptr, 1, std::vector<VectorPtr>{});
+  SelectivityVector rows(1);
+  VectorPtr result;
+  try {
+    evaluator->evaluate(exprSet.get(), rows, input, result);
+  } catch (const VeloxUserError&) {
+    return nullptr;
+  }
+  return result;
+}
+
+/// Subfield filter backed by Velox's BloomFilter from bloom_filter_agg / might_contain.
+/// Values are hashed with Spark-compatible XXH64 using the seed extracted from
+/// xxhash64_with_seed, then re-hashed with folly hasher for bloom filter bucket
+/// selection, matching bloom_filter_agg's insertion path.
+class SparkMightContain final : public common::BigintValuesUsingBloomFilter {
+ public:
+  SparkMightContain(VectorPtr constantVector, bool nullAllowed, int64_t seed)
+      : common::BigintValuesUsingBloomFilter(0, nullAllowed),
+        constantVector_(std::move(constantVector)),
+        seed_(seed) {
+    auto sv = constantVector_->as<SimpleVector<StringView>>()->valueAt(0);
+    view_ = std::make_unique<BloomFilterView>(sv.data());
+  }
+
+  bool testInt64(int64_t value) const override {
+    return view_->mayContain(folly::hasher<int64_t>()(
+        facebook::velox::functions::sparksql::XxHash64::hashInt64(value, seed_)));
+  }
+
+  bool testDouble(double value) const override {
+    return view_->mayContain(folly::hasher<int64_t>()(
+        facebook::velox::functions::sparksql::XxHash64::hashDouble(value, seed_)));
+  }
+
+  bool testFloat(float value) const override {
+    return view_->mayContain(folly::hasher<int64_t>()(
+        facebook::velox::functions::sparksql::XxHash64::hashFloat(value, seed_)));
+  }
+
+  bool testBytes(const char* data, int32_t len) const override {
+    return view_->mayContain(folly::hasher<int64_t>()(
+        facebook::velox::functions::sparksql::XxHash64::hashBytes(StringView(data, len), seed_)));
+  }
+
+  bool testTimestamp(const Timestamp& value) const override {
+    return view_->mayContain(folly::hasher<int64_t>()(
+        facebook::velox::functions::sparksql::XxHash64::hashTimestamp(value, seed_)));
+  }
+
+  bool testInt128(const int128_t& value) const override {
+    return view_->mayContain(folly::hasher<int64_t>()(
+        facebook::velox::functions::sparksql::XxHash64::hashLongDecimal(value, seed_)));
+  }
+
+  bool testInt64Range(int64_t /*min*/, int64_t /*max*/, bool /*hasNull*/) const override {
+    return true;
+  }
+
+  std::unique_ptr<Filter> clone(std::optional<bool> nullAllowed) const override {
+    return std::make_unique<SparkMightContain>(
+        constantVector_, nullAllowed.value_or(nullAllowed_), seed_);
+  }
+
+  bool testingEquals(const Filter& other) const override {
+    return dynamic_cast<const SparkMightContain*>(&other) != nullptr;
+  }
+
+  folly::dynamic serialize() const override {
+    VELOX_UNSUPPORTED("Serialization is not supported for SparkMightContain");
+  }
+
+ private:
+  VectorPtr constantVector_;
+  std::unique_ptr<BloomFilterView> view_;
+  int64_t seed_;
+};
+
 std::optional<std::pair<facebook::velox::common::Subfield, std::unique_ptr<facebook::velox::common::Filter>>> combine(
     facebook::velox::common::Subfield& subfield,
     std::unique_ptr<facebook::velox::common::Filter>& filter) {
@@ -30,6 +123,7 @@ std::optional<std::pair<facebook::velox::common::Subfield, std::unique_ptr<faceb
 
   return std::nullopt;
 }
+
 } // namespace
 
 std::optional<std::pair<facebook::velox::common::Subfield, std::unique_ptr<facebook::velox::common::Filter>>>
@@ -93,6 +187,34 @@ SparkExprToSubfieldFilterParser::leafCallToSubfieldFilter(
       }
       return std::make_pair(std::move(subfield), facebook::velox::exec::isNotNull());
     }
+  } else if (call.name() == "might_contain" && !negated) {
+    // Matches: might_contain(bloomFilter, xxhash64_with_seed(seed, field)).
+    GLUTEN_CHECK(call.inputs().size() == 2,
+        "might_contain expects 2 arguments: bloomFilter and xxhash64_with_seed(seed, field)");
+    const auto* hashCall =
+        dynamic_cast<const core::CallTypedExpr*>(call.inputs()[1].get());
+    if (hashCall && hashCall->name() == "xxhash64_with_seed") {
+      GLUTEN_CHECK(hashCall->inputs().size() == 2, "xxhash64_with_seed expects 2 arguments");
+      auto seedValue = toConstant(hashCall->inputs()[0], evaluator);
+      if (!seedValue || seedValue->isNullAt(0)) {
+        LOG(WARNING) << "might_contain: seed value is null or not constant, "
+                     << "cannot push down to subfield filter";
+        return std::nullopt;
+      }
+      auto seed = seedValue->as<SimpleVector<int64_t>>()->valueAt(0);
+      if (!toSubfield(hashCall->inputs()[1].get(), subfield)) {
+        LOG(WARNING) << "might_contain: second argument to xxhash64_with_seed "
+                     << "is not a subfield, cannot push down to subfield filter";
+        return std::nullopt;
+      }
+      auto bloomFilterValue = toConstant(call.inputs()[0], evaluator);
+      if (bloomFilterValue && !bloomFilterValue->isNullAt(0)) {
+        std::unique_ptr<common::Filter> filter =
+            std::make_unique<SparkMightContain>(bloomFilterValue, false /*nullAllowed*/, seed);
+        return combine(subfield, filter);
+      }
+    }
+    LOG(WARNING) << "might_contain could not be converted to a subfield filter";
   }
   return std::nullopt;
 }

ep/build-velox/src/get-velox.sh

@@ -166,6 +166,17 @@ function apply_compilation_fixes {
   git add ${VELOX_HOME}/CMake/resolve_dependency_modules/arrow/modify_arrow.patch # to avoid the file from being deleted by git clean -dffx :/
 }
 
+function apply_bloomfilter_final_patch {
+  echo "Applying bloom filter final-removal patch..."
+  pushd $VELOX_HOME
+  git apply --check ${CURRENT_DIR}/remove-bloomfilter-final.patch && \
+    git apply ${CURRENT_DIR}/remove-bloomfilter-final.patch || {
+    echo "Failed to apply bloom filter final-removal patch"
+    exit 1
+  }
+  popd
+}
+
 function setup_linux {
   local LINUX_DISTRIBUTION=$(. /etc/os-release && echo ${ID})
   local LINUX_VERSION_ID=$(. /etc/os-release && echo ${VERSION_ID})
@@ -248,4 +259,6 @@ apply_provided_velox_patch
 
 apply_compilation_fixes
 
+apply_bloomfilter_final_patch
+
 echo "Finished getting Velox code"

ep/build-velox/src/remove-bloomfilter-final.patch

@@ -0,0 +1,31 @@
+diff --git a/velox/type/Filter.h b/velox/type/Filter.h
+index 82ecd8eb2..cc8e6b31b 100644
+--- a/velox/type/Filter.h
++++ b/velox/type/Filter.h
+@@ -1290,7 +1290,7 @@ class BigintValuesUsingBitmask final : public Filter {
+   const int64_t max_;
+ };
+ 
+-class BigintValuesUsingBloomFilter final : public Filter {
++class BigintValuesUsingBloomFilter : public Filter {
+  public:
+   static int64_t numBlocks(int64_t capacity) {
+     return SplitBlockBloomFilter::numBlocks(capacity, 0.01);
+@@ -1301,7 +1301,7 @@ class BigintValuesUsingBloomFilter final : public Filter {
+         blocks_(numBlocks(capacity)),
+         filter_(blocks_) {}
+ 
+-  bool testInt64(int64_t value) const final {
++  bool testInt64(int64_t value) const override {
+     return filter_.mayContain(hash(value));
+   }
+ 
+@@ -1318,7 +1318,7 @@ class BigintValuesUsingBloomFilter final : public Filter {
+   }
+ 
+   bool testInt64Range(int64_t /*min*/, int64_t /*max*/, bool /*hasNull*/)
+-      const final {
++      const override {
+     return true;
+   }
+ 

gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenBloomFilterAggregateQuerySuite.scala

@@ -18,7 +18,7 @@ package org.apache.spark.sql
 
 import org.apache.gluten.backendsapi.BackendsApiManager
 import org.apache.gluten.config.GlutenConfig
-import org.apache.gluten.execution.HashAggregateExecBaseTransformer
+import org.apache.gluten.execution.{FileSourceScanExecTransformer, HashAggregateExecBaseTransformer}
 
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
@@ -112,6 +112,44 @@ class GlutenBloomFilterAggregateQuerySuite
     }
   }
 
+  testGluten("Test might_contain constant bloom filter pushdown to scan") {
+    val table = "bloom_filter_test"
+    withTable(table) {
+      (1L to 200L).toDF("col").write.format("parquet").saveAsTable(table)
+
+      // Build a bloom filter over xxhash64(col) so it matches the probe side.
+      val hexBf = spark.sql(
+        s"SELECT hex(bloom_filter_agg(xxhash64(col), " +
+          s"cast(200 as long), cast($veloxBloomFilterMaxNumBits as long))) " +
+          s"FROM $table")
+        .collect()(0).getString(0)
+
+      val sqlString =
+        s"SELECT * FROM $table WHERE might_contain(X'$hexBf', xxhash64(col))"
+
+      val df = spark.sql(sqlString)
+      val result = df.collect()
+      // All 200 rows should match since the bloom filter covers them.
+      assert(result.length == 200)
+
+      // Verify the filter was pushed to the scan.
+      val scans = collect(df.queryExecution.executedPlan) {
+        case s: FileSourceScanExecTransformer => s
+      }
+      assert(scans.nonEmpty)
+      scans.foreach { scan =>
+        val metrics = scan.metrics
+        assert(metrics.contains("rawInputRows"))
+        assert(metrics.contains("numOutputRows"))
+        assert(metrics.contains("remainingFilterTime"))
+        assert(metrics("rawInputRows").value > 0)
+        assert(metrics("numOutputRows").value > 0)
+        assert(metrics("numOutputRows").value < metrics("rawInputRows").value)
+        assert(metrics("remainingFilterTime").value == 0)
+      }
+    }
+  }
+
   testGluten("Test bloom_filter_agg agg fallback") {
     val table = "bloom_filter_test"
     val numEstimatedItems = 5000000L