diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index eca32c03d3e0f..fc276da2f9d8c 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -472,6 +472,12 @@ ], "sqlState" : "0A000" }, + "BITMAP_INPUT_TOO_LARGE" : { + "message" : [ + "The input bitmap has bytes, which exceeds the maximum supported size of bytes." + ], + "sqlState" : "22001" + }, "CALL_ON_STREAMING_DATASET_UNSUPPORTED" : { "message" : [ "The method can not be called on streaming Dataset/DataFrame." diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala index 1f943ed7edb08..b41cf6e64c955 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala +++ b/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) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index a5d64ddd314ed..4fe1012633f46 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/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( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/BitmapExpressionsQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/BitmapExpressionsQuerySuite.scala index 458688de3fbde..4ac1f21358b87 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/BitmapExpressionsQuerySuite.scala +++ b/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" + ) + ) + } }