fix: cap bloom_filter_agg numItems/numBits and skip null inputs

The Spark 4.0 BloomFilterAggregateQuerySuite CI job aborted the executor
with a multi-exabyte native allocation, and the Spark 3.4 CometExecRuleSuite
job failed analysis. Three bloom-filter issues surfaced once this branch
let bloom_filter_agg execute natively:

- Spark's BloomFilterAggregate caps numItems/numBits at maxNumItems/
  maxNumBits, but CometBloomFilterAggregate forwarded the raw literals.
  Comet's native aggregate stores them as i32, so an oversized Long
  (e.g. the Long.MaxValue cases in BloomFilterAggregateQuerySuite) wrapped
  to a negative size and triggered a 2^61-byte allocation. Apply the same
  cap in the serde so the native side receives Spark-equivalent values.
- update_batch hit `unreachable!()` on a null input value. Spark's
  BloomFilterAggregate.update ignores nulls; skip them, and return an
  error rather than panicking on a genuinely unexpected type.
- The new CometExecRuleSuite BloomFilter cases used an int column, which
  Spark 3.4's bloom_filter_agg rejects (it only accepts a long-typed
  first argument); cast to bigint.

Adds a CometExec3_4PlusSuite regression test covering oversized
numItems/numBits with a null-containing input.

Author: andygrove

native/spark-expr/src/bloom_filter/bloom_filter_agg.rs

@@ -25,7 +25,7 @@ use crate::bloom_filter::spark_bloom_filter::{SparkBloomFilter, SparkBloomFilter
 use arrow::array::ArrayRef;
 use arrow::array::BinaryArray;
 use datafusion::common::{downcast_value, ScalarValue};
-use datafusion::error::Result;
+use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion::logical_expr::{AggregateUDFImpl, Signature};
 use datafusion::physical_expr::expressions::Literal;
@@ -141,8 +141,16 @@ impl Accumulator for SparkBloomFilter {
                 ScalarValue::Utf8(Some(value)) => {
                     self.put_binary(value.as_bytes());
                 }
-                _ => {
-                    unreachable!()
+                // Spark's BloomFilterAggregate.update ignores null inputs.
+                ScalarValue::Int8(None)
+                | ScalarValue::Int16(None)
+                | ScalarValue::Int32(None)
+                | ScalarValue::Int64(None)
+                | ScalarValue::Utf8(None) => {}
+                other => {
+                    return Err(DataFusionError::Internal(format!(
+                        "bloom_filter_agg received an unsupported input type: {other:?}"
+                    )));
                 }
             }
             Ok(())

spark/src/main/scala/org/apache/comet/serde/aggregates.scala

@@ -21,7 +21,7 @@ package org.apache.comet.serde
 
 import scala.jdk.CollectionConverters._
 
-import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Literal}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectSet, Corr, Count, Covariance, CovPopulation, CovSample, First, Last, Max, Min, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{ByteType, DataTypes, DecimalType, IntegerType, LongType, ShortType, StringType}
@@ -649,8 +649,20 @@ object CometBloomFilterAggregate extends CometAggregateExpressionSerde[BloomFilt
     // We ignore mutableAggBufferOffset and inputAggBufferOffset because they are
     // implementation details for Spark's ObjectHashAggregate.
     val childExpr = exprToProto(bloomFilter.child, inputs, binding)
-    val numItemsExpr = exprToProto(bloomFilter.estimatedNumItemsExpression, inputs, binding)
-    val numBitsExpr = exprToProto(bloomFilter.numBitsExpression, inputs, binding)
+    // Spark's BloomFilterAggregate caps numItems / numBits at the configured maxima
+    // (its `estimatedNumItems` / `numBits` lazy vals). Comet's native aggregate stores
+    // these as i32, so an uncapped Long literal (e.g. the Long.MaxValue cases in
+    // BloomFilterAggregateQuerySuite) would wrap to a bogus negative size and abort the
+    // executor with a multi-exabyte allocation. Apply the same cap here so the native
+    // side always receives a sane, Spark-equivalent value.
+    val numItems = math.min(
+      bloomFilter.estimatedNumItemsExpression.eval().asInstanceOf[Number].longValue,
+      conf.getConf(SQLConf.RUNTIME_BLOOM_FILTER_MAX_NUM_ITEMS))
+    val numBits = math.min(
+      bloomFilter.numBitsExpression.eval().asInstanceOf[Number].longValue,
+      conf.getConf(SQLConf.RUNTIME_BLOOM_FILTER_MAX_NUM_BITS))
+    val numItemsExpr = exprToProto(Literal(numItems, LongType), inputs, binding)
+    val numBitsExpr = exprToProto(Literal(numBits, LongType), inputs, binding)
     val dataType = serializeDataType(bloomFilter.dataType)
 
     if (childExpr.isDefined &&

spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala

@@ -29,6 +29,7 @@ import org.scalatest.Tag
 import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.catalyst.FunctionIdentifier
 import org.apache.spark.sql.catalyst.expressions.{BloomFilterMightContain, Expression, ExpressionInfo}
+import org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate
 import org.apache.spark.sql.functions.{col, lit}
 import org.apache.spark.util.sketch.BloomFilter
 
@@ -42,6 +43,7 @@ class CometExec3_4PlusSuite extends CometTestBase {
   import testImplicits._
 
   val func_might_contain = new FunctionIdentifier("might_contain")
+  val func_bloom_filter_agg = new FunctionIdentifier("bloom_filter_agg")
 
   override def beforeAll(): Unit = {
     super.beforeAll()
@@ -51,12 +53,23 @@ class CometExec3_4PlusSuite extends CometTestBase {
         func_might_contain,
         new ExpressionInfo(classOf[BloomFilterMightContain].getName, "might_contain"),
         (children: Seq[Expression]) => BloomFilterMightContain(children.head, children(1)))
+      // Register 'bloom_filter_agg' to builtin.
+      spark.sessionState.functionRegistry.registerFunction(
+        func_bloom_filter_agg,
+        new ExpressionInfo(classOf[BloomFilterAggregate].getName, "bloom_filter_agg"),
+        (children: Seq[Expression]) =>
+          children.size match {
+            case 1 => new BloomFilterAggregate(children.head)
+            case 2 => new BloomFilterAggregate(children.head, children(1))
+            case 3 => new BloomFilterAggregate(children.head, children(1), children(2))
+          })
     }
   }
 
   override def afterAll(): Unit = {
     if (!isSpark42Plus) {
       spark.sessionState.functionRegistry.dropFunction(func_might_contain)
+      spark.sessionState.functionRegistry.dropFunction(func_bloom_filter_agg)
     }
     super.afterAll()
   }
@@ -185,6 +198,24 @@ class CometExec3_4PlusSuite extends CometTestBase {
     }
   }
 
+  test("bloom_filter_agg caps oversized numItems / numBits like Spark") {
+    assume(!isSpark42Plus, "https://github.com/apache/datafusion-comet/issues/4142")
+    val table = "test"
+    withTable(table) {
+      sql(s"create table $table(col1 long) using parquet")
+      sql(s"insert into $table values (1), (2), (3), (201), (null)")
+      // numItems / numBits exceed the Int range. Spark's BloomFilterAggregate caps
+      // them at maxNumItems / maxNumBits; Comet must apply the same cap, otherwise the
+      // oversized values truncate to a negative i32 and abort the executor with a
+      // multi-exabyte allocation.
+      checkSparkAnswerAndOperator(s"""
+           |SELECT bloom_filter_agg(col1,
+           |  cast(9223372036854775807 as long),
+           |  cast(9223372036854775807 as long)) FROM $table
+           |""".stripMargin)
+    }
+  }
+
   private def bloomFilterFromRandomInput(
       expectedItems: Long,
       expectedBits: Long): (Seq[Long], Array[Byte]) = {

spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala

@@ -247,7 +247,10 @@ class CometExecRuleSuite extends CometTestBase {
       withTempView("test_data") {
         createTestDataFrame.createOrReplaceTempView("test_data")
 
-        val sparkPlan = createSparkPlan(spark, "SELECT bloom_filter_agg(id) FROM test_data")
+        // Cast to bigint: Spark 3.4's bloom_filter_agg only accepts a long-typed first
+        // argument; later versions widened it to any integral type.
+        val sparkPlan =
+          createSparkPlan(spark, "SELECT bloom_filter_agg(CAST(id AS BIGINT)) FROM test_data")
 
         val originalObjectAggCount = countOperators(sparkPlan, classOf[ObjectHashAggregateExec])
         assert(originalObjectAggCount == 2)
@@ -283,7 +286,10 @@ class CometExecRuleSuite extends CometTestBase {
       withTempView("test_data") {
         createTestDataFrame.createOrReplaceTempView("test_data")
 
-        val sparkPlan = createSparkPlan(spark, "SELECT bloom_filter_agg(id) FROM test_data")
+        // Cast to bigint: Spark 3.4's bloom_filter_agg only accepts a long-typed first
+        // argument; later versions widened it to any integral type.
+        val sparkPlan =
+          createSparkPlan(spark, "SELECT bloom_filter_agg(CAST(id AS BIGINT)) FROM test_data")
 
         val originalObjectAggCount = countOperators(sparkPlan, classOf[ObjectHashAggregateExec])
         assert(originalObjectAggCount == 2)