jiangxt2 commented on code in PR #58117:
URL: https://github.com/apache/spark/pull/58117#discussion_r4488200429
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala:
##########
@@ -375,6 +383,92 @@ case class BitmapXor(left: Expression, right: Expression)
extends BitmapBinaryEx
copy(left = newLeft, right = newRight)
}
+@ExpressionDescription(
+ usage = """
+ _FUNC_(bitmap, bit_position) - Returns true if the bit at the given
bucket-local position is
+ set in the bitmap, false otherwise.
+ """,
+ arguments = """
+ Arguments:
+ * bitmap - The bitmap to test.
+ An expression that evaluates to a binary. A NULL bitmap produces a
NULL result.
+ * bit_position - The bucket-local bit position to test.
+ An expression that evaluates to a numeric value and is cast to a long.
A NULL position
+ produces a NULL result. A negative position, a position at or above
32768, or a position
Review Comment:
`058768d` updates the SQL, Scala, and PySpark docs to explain the cast to
`BIGINT`, the bounds check that follows it, and ANSI versus legacy cast
behavior. The two new regression cases cover `-0.5` and `2^64` through SQL
only. Both read bit 0 in legacy mode.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala:
##########
@@ -375,6 +383,92 @@ case class BitmapXor(left: Expression, right: Expression)
extends BitmapBinaryEx
copy(left = newLeft, right = newRight)
}
+@ExpressionDescription(
+ usage = """
+ _FUNC_(bitmap, bit_position) - Returns true if the bit at the given
bucket-local position is
+ set in the bitmap, false otherwise.
+ """,
+ arguments = """
+ Arguments:
+ * bitmap - The bitmap to test.
+ An expression that evaluates to a binary. A NULL bitmap produces a
NULL result.
+ * bit_position - The bucket-local bit position to test.
+ An expression that evaluates to a numeric value and is cast to a long.
A NULL position
+ produces a NULL result. A negative position, a position at or above
32768, or a position
+ beyond the actual bitmap length returns false. To query an original
value, use
+ bitmap_bit_position(value) and match bitmap_bucket_number(value) to
the bitmap bucket.
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(X '01', 0L);
+ true
+ > SELECT _FUNC_(X '01', 1L);
+ false
+ > SELECT _FUNC_(X '10', 4L);
+ true
+ """,
+ since = "4.4.0",
+ group = "misc_funcs")
+case class BitmapContains(left: Expression, right: Expression)
+ extends BinaryExpression with RuntimeReplaceable {
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ if (left.dataType != BinaryType && left.dataType != NullType) {
+ DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_INPUT_TYPE",
+ messageParameters = Map(
+ "paramIndex" -> ordinalNumber(0),
+ "requiredType" -> toSQLType(BinaryType),
+ "inputSql" -> toSQLExpr(left),
+ "inputType" -> toSQLType(left.dataType)
+ )
+ )
+ } else if (!right.dataType.isInstanceOf[NumericType] && right.dataType !=
NullType) {
+ DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_INPUT_TYPE",
+ messageParameters = Map(
+ "paramIndex" -> ordinalNumber(1),
+ "requiredType" -> toSQLType(NumericType),
+ "inputSql" -> toSQLExpr(right),
+ "inputType" -> toSQLType(right.dataType)
+ )
+ )
+ } else {
+ TypeCheckSuccess
+ }
+ }
+
+ override def dataType: DataType = BooleanType
+
+ override def prettyName: String = "bitmap_contains"
+
+ override lazy val replacement: Expression = {
+ // StaticInvoke needs an exact BinaryType child, including for a bare NULL
literal.
+ val bitmap = if (left.dataType == NullType) Cast(left, BinaryType) else
left
+ def invoke(bitmap: Expression): Expression = StaticInvoke(
+ classOf[BitmapExpressionUtils],
+ BooleanType,
+ "bitmapContains",
+ Seq(bitmap, Cast(right, LongType)),
+ Seq(BinaryType, LongType),
+ returnNullable = false)
+
+ if (bitmap.nullable) {
Review Comment:
Fixed in `058768d`. Nullable bitmaps that are foldable after
`prepareForEval` now use a guarded `If`. Added regression tests for `SELECT …
VALUES`, `INSERT … VALUES`, and column `DEFAULT`.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtilsSuite.scala:
##########
@@ -93,6 +93,45 @@ class BitmapExpressionUtilsSuite extends SparkFunSuite {
assert(BitmapExpressionUtils.bitmapCount(bitmap) == 15L)
}
+ test("bitmap_contains with set and unset bits") {
+ val bitmap = createBitmap()
+ setBitmapBits(bitmap, 0, 0x81)
+ setBitmapBits(bitmap, 1, 0x01)
+ val original = bitmap.clone()
+
+ assert(BitmapExpressionUtils.bitmapContains(bitmap, 0))
+ assert(!BitmapExpressionUtils.bitmapContains(bitmap, 1))
+ assert(BitmapExpressionUtils.bitmapContains(bitmap, 7))
+ assert(BitmapExpressionUtils.bitmapContains(bitmap, 8))
+ assert(!BitmapExpressionUtils.bitmapContains(bitmap, 9))
+ assert(bitmap.sameElements(original))
+ }
+
+ test("bitmap_contains with empty and short bitmaps") {
+ val emptyBitmap = Array.empty[Byte]
+ assert(!BitmapExpressionUtils.bitmapContains(emptyBitmap, 0))
+
+ val shortBitmap = Array[Byte](0x01, 0x80.toByte)
+ assert(BitmapExpressionUtils.bitmapContains(shortBitmap, 0))
+ assert(BitmapExpressionUtils.bitmapContains(shortBitmap, 15))
+ assert(!BitmapExpressionUtils.bitmapContains(shortBitmap, 16))
+ assert(!BitmapExpressionUtils.bitmapContains(shortBitmap,
BitmapExpressionUtils.NUM_BITS - 1))
+ }
+
+ test("bitmap_contains at fixed-range and long boundaries") {
+ val bitmap = createBitmap()
+ bitmap(bitmap.length - 1) = 0x80.toByte
+ assert(BitmapExpressionUtils.bitmapContains(bitmap,
BitmapExpressionUtils.NUM_BITS - 1))
+ assert(!BitmapExpressionUtils.bitmapContains(bitmap,
BitmapExpressionUtils.NUM_BITS))
Review Comment:
Added a 4097-byte bitmap with bits 32767 and 32768 both set. The test
asserts `true` at 32767 and `false` at 32768.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
