[SPARK-53877] Introduce BITMAP_AND_AGG function

Currently, Spark has two bitmap aggregation functions `bitmap_construct_agg` and `bitmap_or_agg` for constructing a bitmap out of set of integers and performing union on two sets represented by bitmaps, respectively. However, efficient intersect operation (bitwise AND) is missing.

## What changes were proposed in this pull request?
- **Implemented `bitmap_and_agg` expression**: New aggregation function that performs bitwise AND operations on binary column inputs.

### Design Decisions
- **Result on empty input is identity element for the operation**: Empty input groups return all-ones bitmaps (AND identity).
- **Missing bytes handling**: For AND operations, missing bytes in input are treated as zeros to maintain intersection semantics.

## How was this patch tested?
Added new test cases to cover `bitmap_and_agg` functionality:
- **`BitmapExpressionsQuerySuite`**: Added test cases for basic AND operations, edge cases, empty group handling, and integration with other bitmap functions.

### Does this PR introduce _any_ user-facing change?
No.

### Was this patch authored or co-authored using generative AI tooling?
No.

Closes apache#52586 from uros7251brick/add-bitmap-and-agg.

Authored-by: Uros Stojkovic <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>

Author: uros7251brick

python/pyspark/sql/connect/functions/builtin.py

@@ -4619,6 +4619,13 @@ def bitmap_or_agg(col: "ColumnOrName") -> Column:
 bitmap_or_agg.__doc__ = pysparkfuncs.bitmap_or_agg.__doc__
 
 
+def bitmap_and_agg(col: "ColumnOrName") -> Column:
+    return _invoke_function_over_columns("bitmap_and_agg", col)
+
+
+bitmap_and_agg.__doc__ = pysparkfuncs.bitmap_and_agg.__doc__
+
+
 # Call Functions
 
 

python/pyspark/sql/functions/__init__.py

@@ -349,6 +349,7 @@
     "bit_and",
     "bit_or",
     "bit_xor",
+    "bitmap_and_agg",
     "bitmap_construct_agg",
     "bitmap_or_agg",
     "bool_and",

python/pyspark/sql/functions/builtin.py

@@ -27441,6 +27441,7 @@ def bitmap_or_agg(col: "ColumnOrName") -> Column:
     :meth:`pyspark.sql.functions.bitmap_bucket_number`
     :meth:`pyspark.sql.functions.bitmap_construct_agg`
     :meth:`pyspark.sql.functions.bitmap_count`
+    :meth:`pyspark.sql.functions.bitmap_and_agg`
 
     Parameters
     ----------
@@ -27461,6 +27462,41 @@ def bitmap_or_agg(col: "ColumnOrName") -> Column:
     return _invoke_function_over_columns("bitmap_or_agg", col)
 
 
+@_try_remote_functions
+def bitmap_and_agg(col: "ColumnOrName") -> Column:
+    """
+    Returns a bitmap that is the bitwise AND of all of the bitmaps from the input column.
+    The input column should be bitmaps created from bitmap_construct_agg().
+
+    .. versionadded:: 4.1.0
+
+    See Also
+    --------
+    :meth:`pyspark.sql.functions.bitmap_bit_position`
+    :meth:`pyspark.sql.functions.bitmap_bucket_number`
+    :meth:`pyspark.sql.functions.bitmap_construct_agg`
+    :meth:`pyspark.sql.functions.bitmap_count`
+    :meth:`pyspark.sql.functions.bitmap_or_agg`
+
+    Parameters
+    ----------
+    col : :class:`~pyspark.sql.Column` or column name
+        The input column should be bitmaps created from bitmap_construct_agg().
+
+    Examples
+    --------
+    >>> from pyspark.sql import functions as sf
+    >>> df = spark.createDataFrame([("F0",),("70",),("30",)], ["a"])
+    >>> df.select(sf.bitmap_and_agg(sf.to_binary(df.a, sf.lit("hex")))).show()
+    +---------------------------------+
+    |bitmap_and_agg(to_binary(a, hex))|
+    +---------------------------------+
+    |             [30 00 00 00 00 0...|
+    +---------------------------------+
+    """
+    return _invoke_function_over_columns("bitmap_and_agg", col)
+
+
 # ---------------------------- User Defined Function ----------------------------------
 
 

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

@@ -4236,6 +4236,15 @@ object functions {
    */
   def bitmap_or_agg(col: Column): Column = Column.fn("bitmap_or_agg", col)
 
+  /**
+   * Returns a bitmap that is the bitwise AND of all of the bitmaps from the input column. The
+   * input column should be bitmaps created from bitmap_construct_agg().
+   *
+   * @group agg_funcs
+   * @since 4.1.0
+   */
+  def bitmap_and_agg(col: Column): Column = Column.fn("bitmap_and_agg", col)
+
   //////////////////////////////////////////////////////////////////////////////////////////////
   // String functions
   //////////////////////////////////////////////////////////////////////////////////////////////

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtils.java

@@ -56,4 +56,17 @@ public static void bitmapMerge(byte[] bitmap1, byte[] bitmap2) {
       bitmap1[i] = (byte) ((bitmap1[i] & 0x0FF) | (bitmap2[i] & 0x0FF));
     }
   }
+
+  /** Performs bitwise AND on both bitmaps and writes the result into bitmap1. */
+  public static void bitmapAndMerge(byte[] bitmap1, byte[] bitmap2) {
+    int minLen = java.lang.Math.min(bitmap1.length, bitmap2.length);
+    for (int i = 0; i < minLen; ++i) {
+      bitmap1[i] = (byte) ((bitmap1[i] & 0x0FF) & (bitmap2[i] & 0x0FF));
+    }
+    // For AND operation, any bytes beyond the input length should be set to 0
+    // since they represent bits that don't exist in the input bitmap
+    for (int i = bitmap2.length; i < bitmap1.length; ++i) {
+      bitmap1[i] = 0;
+    }
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -853,6 +853,7 @@ object FunctionRegistry {
     expression[BitmapConstructAgg]("bitmap_construct_agg"),
     expression[BitmapCount]("bitmap_count"),
     expression[BitmapOrAgg]("bitmap_or_agg"),
+    expression[BitmapAndAgg]("bitmap_and_agg"),
 
     // json
     expression[StructsToJson]("to_json"),

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

@@ -322,3 +322,96 @@ case class BitmapOrAgg(child: Expression,
     buffer.getBinary(mutableAggBufferOffset)
   }
 }
+
+@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().
+  """,
+  // scalastyle:off line.size.limit
+  examples = """
+    Examples:
+      > SELECT substring(hex(_FUNC_(col)), 0, 6) FROM VALUES (X 'F0'), (X '70'), (X '30') AS tab(col);
+       300000
+      > SELECT substring(hex(_FUNC_(col)), 0, 6) FROM VALUES (X 'FF'), (X 'FF'), (X 'FF') AS tab(col);
+       FF0000
+  """,
+  // scalastyle:on line.size.limit
+  since = "4.1.0",
+  group = "agg_funcs")
+case class BitmapAndAgg(
+    child: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0)
+    extends ImperativeAggregate
+    with UnaryLike[Expression] {
+
+  def this(child: Expression) = {
+    this(child = child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+  }
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    if (child.dataType != BinaryType) {
+      DataTypeMismatch(
+        errorSubClass = "UNEXPECTED_INPUT_TYPE",
+        messageParameters = Map(
+          "paramIndex" -> ordinalNumber(0),
+          "requiredType" -> toSQLType(BinaryType),
+          "inputSql" -> toSQLExpr(child),
+          "inputType" -> toSQLType(child.dataType)))
+    } else {
+      TypeCheckSuccess
+    }
+  }
+
+  override def dataType: DataType = BinaryType
+
+  override def prettyName: String = "bitmap_and_agg"
+
+  override protected def withNewChildInternal(newChild: Expression): BitmapAndAgg =
+    copy(child = newChild)
+
+  override def withNewMutableAggBufferOffset(
+      newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+  override def nullable: Boolean = false
+
+  override def aggBufferSchema: StructType = DataTypeUtils.fromAttributes(aggBufferAttributes)
+
+  // The aggregation buffer is a fixed size binary.
+  private val bitmapAttr = AttributeReference("bitmap", BinaryType, false)()
+
+  override def aggBufferAttributes: Seq[AttributeReference] = bitmapAttr :: Nil
+
+  override def defaultResult: Option[Literal] =
+    Option(Literal(Array.fill[Byte](BitmapExpressionUtils.NUM_BYTES)(-1)))
+
+  override val inputAggBufferAttributes: Seq[AttributeReference] =
+    aggBufferAttributes.map(_.newInstance())
+
+  override def initialize(buffer: InternalRow): Unit = {
+    buffer.update(mutableAggBufferOffset, Array.fill[Byte](BitmapExpressionUtils.NUM_BYTES)(-1))
+  }
+
+  override def update(buffer: InternalRow, input: InternalRow): Unit = {
+    val input_bitmap = child.eval(input).asInstanceOf[Array[Byte]]
+    if (input_bitmap != null) {
+      val bitmap = buffer.getBinary(mutableAggBufferOffset)
+      BitmapExpressionUtils.bitmapAndMerge(bitmap, input_bitmap)
+    }
+  }
+
+  override def merge(buffer1: InternalRow, buffer2: InternalRow): Unit = {
+    val bitmap1 = buffer1.getBinary(mutableAggBufferOffset)
+    val bitmap2 = buffer2.getBinary(inputAggBufferOffset)
+    BitmapExpressionUtils.bitmapAndMerge(bitmap1, bitmap2)
+  }
+
+  override def eval(buffer: InternalRow): Any = {
+    buffer.getBinary(mutableAggBufferOffset)
+  }
+}

sql/core/src/test/resources/sql-functions/sql-expression-schema.md

@@ -48,6 +48,7 @@
 | org.apache.spark.sql.catalyst.expressions.Between | between | SELECT 0.5 between 0.1 AND 1.0 | struct<between(0.5, 0.1, 1.0):boolean> |
 | org.apache.spark.sql.catalyst.expressions.Bin | bin | SELECT bin(13) | struct<bin(13):string> |
 | org.apache.spark.sql.catalyst.expressions.BitLength | bit_length | SELECT bit_length('Spark SQL') | struct<bit_length(Spark SQL):int> |
+| org.apache.spark.sql.catalyst.expressions.BitmapAndAgg | bitmap_and_agg | SELECT substring(hex(bitmap_and_agg(col)), 0, 6) FROM VALUES (X 'F0'), (X '70'), (X '30') AS tab(col) | struct<substring(hex(bitmap_and_agg(col)), 0, 6):string> |
 | org.apache.spark.sql.catalyst.expressions.BitmapBitPosition | bitmap_bit_position | SELECT bitmap_bit_position(1) | struct<bitmap_bit_position(1):bigint> |
 | org.apache.spark.sql.catalyst.expressions.BitmapBucketNumber | bitmap_bucket_number | SELECT bitmap_bucket_number(123) | struct<bitmap_bucket_number(123):bigint> |
 | org.apache.spark.sql.catalyst.expressions.BitmapConstructAgg | bitmap_construct_agg | SELECT substring(hex(bitmap_construct_agg(bitmap_bit_position(col))), 0, 6) FROM VALUES (1), (2), (3) AS tab(col) | struct<substring(hex(bitmap_construct_agg(bitmap_bit_position(col))), 0, 6):string> |

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

@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql
 
-import org.apache.spark.sql.functions.{bitmap_bit_position, bitmap_bucket_number, bitmap_construct_agg, bitmap_count, bitmap_or_agg, col, hex, lit, substring, to_binary}
+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
 
 class BitmapExpressionsQuerySuite extends QueryTest with SharedSparkSession {
@@ -208,6 +208,95 @@ class BitmapExpressionsQuerySuite extends QueryTest with SharedSparkSession {
     )
   }
 
+  test("bitmap_and_agg") {
+    // Test basic AND functionality: F0 & 70 & 30 = 30
+    val df = Seq("F0", "70", "30").toDF("a")
+    checkAnswer(
+      df.selectExpr("substring(hex(bitmap_and_agg(to_binary(a, 'hex'))), 0, 6)"),
+      Seq(Row("300000")))
+    checkAnswer(
+      df.select(substring(hex(bitmap_and_agg(to_binary(col("a"), lit("hex")))), 0, 6)),
+      Seq(Row("300000")))
+
+    // Test with all 1s - should return FF
+    val df2 = Seq("FF", "FF", "FF").toDF("a")
+    checkAnswer(
+      df2.selectExpr("substring(hex(bitmap_and_agg(to_binary(a, 'hex'))), 0, 6)"),
+      Seq(Row("FF0000")))
+
+    // Test with mixed values - A0 & F0 & 80 = 80
+    val df3 = Seq("A0", "F0", "80").toDF("a")
+    checkAnswer(
+      df3.selectExpr("substring(hex(bitmap_and_agg(to_binary(a, 'hex'))), 0, 6)"),
+      Seq(Row("800000")))
+
+    // Test with one zero - anything & 00 = 00
+    val df4 = Seq("FF", "00", "FF").toDF("a")
+    checkAnswer(
+      df4.selectExpr("substring(hex(bitmap_and_agg(to_binary(a, 'hex'))), 0, 6)"),
+      Seq(Row("000000")))
+
+    // Test with binary values of different lengths - "FF" & "FFFF" & "FFFFFF" = "FF0000"
+    val df5 = Seq("FF", "FFFF", "FFFFFF").toDF("a")
+    checkAnswer(
+      df5.selectExpr("substring(hex(bitmap_and_agg(to_binary(a, 'hex'))), 0, 6)"),
+      Seq(Row("FF0000")))
+
+    // Test empty result (no rows) - should return all 1s as AND identity
+    val emptyDf = Seq.empty[String].toDF("a")
+    checkAnswer(
+      emptyDf.selectExpr("substring(hex(bitmap_and_agg(to_binary(a, 'hex'))), 0, 6)"),
+      Seq(Row("FFFFFF")))
+
+    val emptyDf2 = Seq.empty[(String, Int)].toDF("a", "b")
+    checkAnswer(
+      emptyDf2
+        .selectExpr("bitmap_and_agg(to_binary(a, 'hex')) " +
+          "filter (where b = 1) as and_agg")
+        .select(substring(hex(col("and_agg")), 0, 6)),
+      Seq(Row("FFFFFF")))
+
+    // Test empty result (no rows) - should return empty DataFrame
+    val emptyDf3 = Seq.empty[(String, Int, Int)].toDF("a", "b", "c")
+    checkAnswer(
+      emptyDf3
+        .groupBy("c")
+        .agg(expr("bitmap_and_agg(to_binary(a, 'hex')) " +
+          "filter (where b = 1) as and_agg").alias("and_agg"))
+        .select(substring(hex(col("and_agg")), 0, 6)),
+      Seq())
+  }
+
+  test("bitmap_and_agg with complex bitmaps from bitmap_construct_agg") {
+    val table = "bitmap_and_test_table"
+    withTable(table) {
+      // Create test data with overlapping bit positions
+      spark.sql(s"""
+           | CREATE TABLE $table (group_id INT, bit_pos LONG)
+           | """.stripMargin)
+      spark.sql(s"""
+           | INSERT INTO $table VALUES
+           | (1, 1), (1, 3), (1, 5),  -- Group 1: bits 1,3,5 set
+           | (2, 1), (2, 2), (2, 3),  -- Group 2: bits 1,2,3 set
+           | (3, 3), (3, 4), (3, 5)   -- Group 3: bits 3,4,5 set
+           | """.stripMargin)
+      // Each group should have their respective bit counts, but when we AND them together
+      // we should get the intersection
+      val intersectionResult = spark.sql(s"""
+           | SELECT bitmap_count(
+           |   bitmap_and_agg(group_bitmap)
+           | ) as intersection_count
+           | FROM (
+           |   SELECT bitmap_construct_agg(bitmap_bit_position(bit_pos)) as group_bitmap
+           |   FROM $table
+           |   GROUP BY group_id
+           | )
+           | """.stripMargin)
+      // The intersection should be 1 (only bit 3 is common)
+      checkAnswer(intersectionResult, Seq(Row(1)))
+    }
+  }
+
   test("bitmap_count called with non-binary type") {
     val df = Seq(12).toDF("a")
     checkError(
@@ -251,4 +340,20 @@ class BitmapExpressionsQuerySuite extends QueryTest with SharedSparkSession {
       )
     )
   }
+
+  test("bitmap_and_agg called with non-binary type") {
+    val df = Seq(12).toDF("a")
+    checkError(
+      exception = intercept[AnalysisException] {
+        df.selectExpr("bitmap_and_agg(a)")
+      },
+      condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE",
+      parameters = Map(
+        "sqlExpr" -> "\"bitmap_and_agg(a)\"",
+        "paramIndex" -> "first",
+        "requiredType" -> "\"BINARY\"",
+        "inputSql" -> "\"a\"",
+        "inputType" -> "\"INT\""),
+      context = ExpectedContext(fragment = "bitmap_and_agg(a)", start = 0, stop = 16))
+  }
 }