[SPARK-27879][SQL] Add support for bit_and and bit_or aggregates

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

@@ -525,6 +525,8 @@ object FunctionRegistry {
     expression[BitwiseOr]("|"),
     expression[BitwiseXor]("^"),
     expression[BitwiseCount]("bit_count"),
+    expression[BitAndAgg]("bit_and"),
+    expression[BitOrAgg]("bit_or"),
 
     // json
     expression[StructsToJson]("to_json"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/bitwiseAggregates.scala

@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, ExpectsInputTypes, Expression, ExpressionDescription, If, IsNull, Literal}
+import org.apache.spark.sql.types.{AbstractDataType, DataType, IntegralType}
+
+@ExpressionDescription(
+  usage = "_FUNC_(expr) - Returns the bitwise AND of all non-null input values, or null if none.",
+  examples = """
+    Examples:
+      > SELECT _FUNC_(col) FROM VALUES (3), (5) AS tab(col);
+       1
+  """,
+  since = "3.0.0")
+case class BitAndAgg(child: Expression) extends DeclarativeAggregate with ExpectsInputTypes {
+
+  override def nodeName: String = "bit_and"
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  override def dataType: DataType = child.dataType
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(IntegralType)
+
+  private lazy val bitAnd = AttributeReference("bit_and", child.dataType)()
+
+  override lazy val aggBufferAttributes: Seq[AttributeReference] = bitAnd :: Nil
+
+  override lazy val initialValues: Seq[Literal] = Literal.create(null, dataType) :: Nil
+
+  override lazy val updateExpressions: Seq[Expression] =
+    If(IsNull(bitAnd),
+      child,
+      If(IsNull(child), bitAnd, BitwiseAnd(bitAnd, child))) :: Nil
+
+  override lazy val mergeExpressions: Seq[Expression] =
+    If(IsNull(bitAnd.left),
+      bitAnd.right,
+      If(IsNull(bitAnd.right), bitAnd.left, BitwiseAnd(bitAnd.left, bitAnd.right))) :: Nil
+
+  override lazy val evaluateExpression: AttributeReference = bitAnd
+}
+
+@ExpressionDescription(
+  usage = "_FUNC_(expr) - Returns the bitwise OR of all non-null input values, or null if none.",
+  examples = """
+    Examples:
+      > SELECT _FUNC_(col) FROM VALUES (3), (5) AS tab(col);
+       7
+  """,
+  since = "3.0.0")
+case class BitOrAgg(child: Expression) extends DeclarativeAggregate with ExpectsInputTypes {
+
+  override def nodeName: String = "bit_or"
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  override def dataType: DataType = child.dataType
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(IntegralType)
+
+  private lazy val bitOr = AttributeReference("bit_or", child.dataType)()
+
+  override lazy val aggBufferAttributes: Seq[AttributeReference] = bitOr :: Nil
+
+  override lazy val initialValues: Seq[Literal] = Literal.create(null, dataType) :: Nil
+
+  override lazy val updateExpressions: Seq[Expression] =
+    If(IsNull(bitOr),
+      child,
+      If(IsNull(child), bitOr, BitwiseOr(bitOr, child))) :: Nil
+
+  override lazy val mergeExpressions: Seq[Expression] =
+    If(IsNull(bitOr.left),
+      bitOr.right,
+      If(IsNull(bitOr.right), bitOr.left, BitwiseOr(bitOr.left, bitOr.right))) :: Nil
+
+  override lazy val evaluateExpression: AttributeReference = bitOr
+}

sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part2.sql

@@ -41,42 +41,37 @@ create temporary view int4_tbl as select * from values
 --
 -- test for bitwise integer aggregates
 --
--- CREATE TEMPORARY TABLE bitwise_test(
---   i2 INT2,
---   i4 INT4,
---   i8 INT8,
---   i INTEGER,
---   x INT2,
---   y BIT(4)
--- );
+CREATE OR REPLACE TEMPORARY VIEW bitwise_test AS SELECT * FROM VALUES
+  (1, 1, 1, 1L),
+  (3, 3, 3, null),
+  (7, 7, 7, 3L) AS bitwise_test(b1, b2, b3, b4);
 
 -- empty case
--- SELECT
---   BIT_AND(i2) AS "?",
---   BIT_OR(i4)  AS "?"
--- FROM bitwise_test;
-
--- COPY bitwise_test FROM STDIN NULL 'null';
--- 1	1	1	1	1	B0101
--- 3	3	3	null	2	B0100
--- 7	7	7	3	4	B1100
--- \.
-
--- SELECT
---   BIT_AND(i2) AS "1",
---   BIT_AND(i4) AS "1",
---   BIT_AND(i8) AS "1",
---   BIT_AND(i)  AS "?",
---   BIT_AND(x)  AS "0",
---   BIT_AND(y)  AS "0100",
---
---   BIT_OR(i2)  AS "7",
---   BIT_OR(i4)  AS "7",
---   BIT_OR(i8)  AS "7",
---   BIT_OR(i)   AS "?",
---   BIT_OR(x)   AS "7",
---   BIT_OR(y)   AS "1101"
--- FROM bitwise_test;
+SELECT BIT_AND(b1) AS n1, BIT_OR(b2)  AS n2 FROM bitwise_test where 1 = 0;
+
+-- null case
+SELECT BIT_AND(b4) AS n1, BIT_OR(b4)  AS n2 FROM bitwise_test where b4 is null;
+
+SELECT
+ BIT_AND(cast(b1 as tinyint)) AS a1,
+ BIT_AND(cast(b2 as smallint)) AS b1,
+ BIT_AND(b3) AS c1,
+ BIT_AND(b4) AS d1,
+ BIT_OR(cast(b1 as tinyint))  AS e7,
+ BIT_OR(cast(b2 as smallint))  AS f7,
+ BIT_OR(b3)  AS g7,
+ BIT_OR(b4)  AS h3
+FROM bitwise_test;
+
+-- group by
+SELECT b1 , bit_and(b2), bit_or(b4) FROM bitwise_test GROUP BY b1;
+
+--having
+SELECT b1, bit_and(b2) FROM bitwise_test GROUP BY b1 HAVING bit_and(b2) < 7;
+
+-- window
+SELECT b1, b2, bit_and(b2) OVER (PARTITION BY b1 ORDER BY b2) FROM bitwise_test;
+SELECT b1, b2, bit_or(b2) OVER (PARTITION BY b1 ORDER BY b2) FROM bitwise_test;
 
 --
 -- test boolean aggregates

sql/core/src/test/resources/sql-tests/inputs/udf/postgreSQL/udf-aggregates_part2.sql

@@ -43,42 +43,28 @@ create temporary view int4_tbl as select * from values
 --
 -- test for bitwise integer aggregates
 --
--- CREATE TEMPORARY TABLE bitwise_test(
---   i2 INT2,
---   i4 INT4,
---   i8 INT8,
---   i INTEGER,
---   x INT2,
---   y BIT(4)
--- );
+CREATE OR REPLACE TEMPORARY VIEW bitwise_test AS SELECT * FROM VALUES
+  (1, 1, 1, 1L),
+  (3, 3, 3, null),
+  (7, 7, 7, 3L) AS bitwise_test(b1, b2, b3, b4);
 
 -- empty case
--- SELECT
---   BIT_AND(i2) AS "?",
---   BIT_OR(i4)  AS "?"
--- FROM bitwise_test;
-
--- COPY bitwise_test FROM STDIN NULL 'null';
--- 1	1	1	1	1	B0101
--- 3	3	3	null	2	B0100
--- 7	7	7	3	4	B1100
--- \.
-
--- SELECT
---   BIT_AND(i2) AS "1",
---   BIT_AND(i4) AS "1",
---   BIT_AND(i8) AS "1",
---   BIT_AND(i)  AS "?",
---   BIT_AND(x)  AS "0",
---   BIT_AND(y)  AS "0100",
---
---   BIT_OR(i2)  AS "7",
---   BIT_OR(i4)  AS "7",
---   BIT_OR(i8)  AS "7",
---   BIT_OR(i)   AS "?",
---   BIT_OR(x)   AS "7",
---   BIT_OR(y)   AS "1101"
--- FROM bitwise_test;
+SELECT BIT_AND(b1) AS n1, BIT_OR(b2)  AS n2 FROM bitwise_test where 1 = 0;
+
+-- null case
+SELECT BIT_AND(b4) AS n1, BIT_OR(b4)  AS n2 FROM bitwise_test where b4 is null;
+
+
+SELECT
+ BIT_AND(cast(b1 as tinyint)) AS a1,
+ BIT_AND(cast(b2 as smallint)) AS b1,
+ BIT_AND(b3) AS c1,
+ BIT_AND(b4) AS d1,
+ BIT_OR(cast(b1 as tinyint))  AS e7,
+ BIT_OR(cast(b2 as smallint))  AS f7,
+ BIT_OR(b3)  AS g7,
+ BIT_OR(b4)  AS h3
+FROM bitwise_test;
 
 --
 -- test boolean aggregates