[SPARK-57561][SQL] Add size validation to bitmap_or_agg and bitmap_and_agg

common/utils/src/main/resources/error/error-conditions.json

@@ -472,6 +472,12 @@
     ],
     "sqlState" : "0A000"
   },
+  "BITMAP_INPUT_TOO_LARGE" : {
+    "message" : [
+      "The input bitmap has <inputNumBytes> bytes, which exceeds the maximum supported size of <maxNumBytes> bytes."
+    ],
+    "sqlState" : "22001"
+  },
   "CALL_ON_STREAMING_DATASET_UNSUPPORTED" : {
     "message" : [
       "The method <methodName> can not be called on streaming Dataset/DataFrame."

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala

@@ -248,7 +248,9 @@ case class BitmapConstructAgg(child: Expression,
 @ExpressionDescription(
   usage = """
     _FUNC_(child) - Returns a bitmap that is the bitwise OR of all of the bitmaps from the child
-    expression. The input should be bitmaps created from bitmap_construct_agg().
+    expression. Input bitmaps shorter than 4096 bytes are merged into the corresponding byte
+    positions of the result; inputs produced by bitmap_construct_agg() are always exactly 4096
+    bytes. Inputs longer than 4096 bytes are not supported and will cause a runtime error.
   """,
   // scalastyle:off line.size.limit
   examples = """
@@ -322,6 +324,10 @@ case class BitmapOrAgg(child: Expression,
   override def update(buffer: InternalRow, input: InternalRow): Unit = {
     val input_bitmap = child.eval(input).asInstanceOf[Array[Byte]]
     if (input_bitmap != null) {
+      if (input_bitmap.length > BitmapExpressionUtils.NUM_BYTES) {
+        throw QueryExecutionErrors.bitmapInputTooLargeError(
+          input_bitmap.length, BitmapExpressionUtils.NUM_BYTES)
+      }
       val bitmap = buffer.getBinary(mutableAggBufferOffset)
       BitmapExpressionUtils.bitmapMerge(bitmap, input_bitmap)
     }
@@ -341,7 +347,11 @@ case class BitmapOrAgg(child: Expression,
 @ExpressionDescription(
   usage = """
     _FUNC_(child) - Returns a bitmap that is the bitwise AND of all of the bitmaps from the child
-    expression. The input should be bitmaps created from bitmap_construct_agg().
+    expression. The aggregation buffer is initialized to all 0xFF bytes (ones). Input bitmaps
+    shorter than 4096 bytes are ANDed into the corresponding byte positions of the result; tail
+    bytes not covered by a short input remain all-ones. Inputs produced by bitmap_construct_agg()
+    are always exactly 4096 bytes. Inputs longer than 4096 bytes are not supported and will cause
+    a runtime error.
   """,
   // scalastyle:off line.size.limit
   examples = """
@@ -415,6 +425,10 @@ case class BitmapAndAgg(
   override def update(buffer: InternalRow, input: InternalRow): Unit = {
     val input_bitmap = child.eval(input).asInstanceOf[Array[Byte]]
     if (input_bitmap != null) {
+      if (input_bitmap.length > BitmapExpressionUtils.NUM_BYTES) {
+        throw QueryExecutionErrors.bitmapInputTooLargeError(
+          input_bitmap.length, BitmapExpressionUtils.NUM_BYTES)
+      }
       val bitmap = buffer.getBinary(mutableAggBufferOffset)
       BitmapExpressionUtils.bitmapAndMerge(bitmap, input_bitmap)
     }

sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala

@@ -257,6 +257,15 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE
       summary = getSummary(context))
   }
 
+  def bitmapInputTooLargeError(inputNumBytes: Int, maxNumBytes: Int): SparkRuntimeException = {
+    new SparkRuntimeException(
+      errorClass = "BITMAP_INPUT_TOO_LARGE",
+      messageParameters = Map(
+        "inputNumBytes" -> inputNumBytes.toString,
+        "maxNumBytes" -> maxNumBytes.toString),
+      cause = null)
+  }
+
   def invalidBitmapPositionError(bitPosition: Long,
                                  bitmapNumBytes: Long): ArrayIndexOutOfBoundsException = {
     new SparkArrayIndexOutOfBoundsException(

sql/core/src/test/scala/org/apache/spark/sql/BitmapExpressionsQuerySuite.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql
 
+import org.apache.spark.SparkRuntimeException
 import org.apache.spark.sql.functions.{bitmap_and_agg, bitmap_bit_position, bitmap_bucket_number, bitmap_construct_agg, bitmap_count, bitmap_or_agg, col, expr, hex, lit, substring, to_binary}
 import org.apache.spark.sql.test.SharedSparkSession
 
@@ -356,4 +357,36 @@ class BitmapExpressionsQuerySuite extends SharedSparkSession {
         "inputType" -> "\"INT\""),
       context = ExpectedContext(fragment = "bitmap_and_agg(a)", start = 0, stop = 16))
   }
+
+  test("bitmap_or_agg rejects input larger than 4096 bytes") {
+    // 4097 bytes exceeds the 4096-byte limit
+    val oversized = Array.fill[Byte](4097)(0xFF.toByte)
+    val df = Seq(oversized).toDF("a")
+    checkError(
+      exception = intercept[SparkRuntimeException] {
+        df.select(bitmap_or_agg(col("a"))).collect()
+      },
+      condition = "BITMAP_INPUT_TOO_LARGE",
+      parameters = Map(
+        "inputNumBytes" -> "4097",
+        "maxNumBytes" -> "4096"
+      )
+    )
+  }
+
+  test("bitmap_and_agg rejects input larger than 4096 bytes") {
+    // 4097 bytes exceeds the 4096-byte limit
+    val oversized = Array.fill[Byte](4097)(0xFF.toByte)
+    val df = Seq(oversized).toDF("a")
+    checkError(
+      exception = intercept[SparkRuntimeException] {
+        df.select(bitmap_and_agg(col("a"))).collect()
+      },
+      condition = "BITMAP_INPUT_TOO_LARGE",
+      parameters = Map(
+        "inputNumBytes" -> "4097",
+        "maxNumBytes" -> "4096"
+      )
+    )
+  }
 }