sql: add bit support for bit_and aggregate function

docs/generated/sql/aggregates.md

@@ -43,6 +43,8 @@
 </span></td></tr>
 <tr><td><a name="bit_and"></a><code>bit_and(arg1: <a href="int.html">int</a>) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Calculates the bitwise AND of all non-null input values, or null if none.</p>
 </span></td></tr>
+<tr><td><a name="bit_and"></a><code>bit_and(arg1: varbit) &rarr; varbit</code></td><td><span class="funcdesc"><p>Calculates the bitwise AND of all non-null input values, or null if none.</p>
+</span></td></tr>
 <tr><td><a name="bit_or"></a><code>bit_or(arg1: <a href="int.html">int</a>) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Calculates the bitwise OR of all non-null input values, or null if none.</p>
 </span></td></tr>
 <tr><td><a name="bool_and"></a><code>bool_and(arg1: <a href="bool.html">bool</a>) &rarr; <a href="bool.html">bool</a></code></td><td><span class="funcdesc"><p>Calculates the boolean value of <code>AND</code>ing all selected values.</p>

pkg/sql/logictest/testdata/logic_test/aggregate

@@ -2241,3 +2241,107 @@ CREATE TABLE t45453(c INT)
 query I
 SELECT count(*) FROM t45453 GROUP BY 0 + 0
 ----
+
+# Tests for the bit_and aggregate function.
+
+subtest bit_and
+
+statement ok
+DROP TABLE IF EXISTS vals
+
+statement ok
+CREATE TABLE vals (
+  v VARBIT,
+  b BIT(8)
+)
+
+# Testing that bit_and returns NULL if there are no rows.
+
+query T
+SELECT bit_and(v) FROM vals
+----
+NULL
+
+# Testing that bit_and does not trigger aggregation on a constant with a source
+# that has no rows.
+
+query T
+SELECT bit_and('1000'::varbit) FROM vals
+----
+NULL
+
+# Testing that bit_and triggers aggregation and computation on a constant
+# with no source.
+
+query TTT
+SELECT bit_and('1'::varbit), bit_and('1000'::bit(4)), bit_and('1010'::varbit)
+----
+1 1000 1010
+
+# Testing that bit_and returns null given a null.
+
+query T
+SELECT bit_and(NULL::varbit)
+----
+NULL
+
+# Testing a successful bit_and over a sequence of non-nulls.
+
+statement ok
+INSERT INTO vals VALUES
+('11111111'::varbit, '11111111'::bit(8)),
+('01111111'::varbit, '01111111'::bit(8)),
+('10111111'::varbit, '10111111'::bit(8)),
+('11011111'::varbit, '11011111'::bit(8)),
+('11101111'::varbit, '11101111'::bit(8))
+
+query TT
+SELECT bit_and(v), bit_and(b) FROM vals
+----
+00001111 00001111
+
+# Testing bit_and over a sequence with nulls and non-nulls.
+
+statement ok
+INSERT INTO vals VALUES
+(NULL::varbit, NULL::bit),
+(NULL::varbit, NULL::bit)
+
+query TT
+SELECT bit_and(v), bit_and(b) FROM vals
+----
+00001111 00001111
+
+# Testing bit_and over a sequence with all nulls.
+
+statement ok
+DELETE FROM vals
+
+statement ok
+INSERT INTO vals VALUES
+(NULL::varbit),
+(NULL::varbit),
+(NULL::varbit),
+(NULL::varbit)
+
+query T
+SELECT bit_and(v) FROM vals
+----
+NULL
+
+# Testing that bit_and returns an error when given an uncasted null.
+
+statement error ambiguous call: bit_and\(unknown\), candidates are
+SELECT bit_and(NULL)
+
+# Testing that an error is returned when bit_and is called on bit arrays of
+# different sizes.
+
+statement error cannot AND bit strings of different sizes
+SELECT bit_and(x::varbit) FROM (VALUES ('1'), ('11')) t(x)
+
+statement error cannot AND bit strings of different sizes
+SELECT bit_and(x) FROM (VALUES ('100'::bit(3)), ('101010111'::varbit)) t(x)
+
+statement error cannot AND bit strings of different sizes
+SELECT bit_and(x) FROM (VALUES (''::varbit), ('1'::varbit)) t(x)

pkg/sql/sem/builtins/aggregate_builtins.go

@@ -24,6 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/arith"
+	"github.com/cockroachdb/cockroach/pkg/util/bitarray"
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/json"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
@@ -131,7 +132,9 @@ var aggregates = map[string]builtinDefinition{
 	),
 
 	"bit_and": makeBuiltin(aggProps(),
-		makeAggOverload([]*types.T{types.Int}, types.Int, newBitAndAggregate,
+		makeAggOverload([]*types.T{types.Int}, types.Int, newIntBitAndAggregate,
+			"Calculates the bitwise AND of all non-null input values, or null if none."),
+		makeAggOverload([]*types.T{types.VarBit}, types.VarBit, newBitBitAndAggregate,
 			"Calculates the bitwise AND of all non-null input values, or null if none."),
 	),
 
@@ -443,7 +446,8 @@ var _ tree.AggregateFunc = &boolOrAggregate{}
 var _ tree.AggregateFunc = &bytesXorAggregate{}
 var _ tree.AggregateFunc = &intXorAggregate{}
 var _ tree.AggregateFunc = &jsonAggregate{}
-var _ tree.AggregateFunc = &bitAndAggregate{}
+var _ tree.AggregateFunc = &intBitAndAggregate{}
+var _ tree.AggregateFunc = &bitBitAndAggregate{}
 var _ tree.AggregateFunc = &bitOrAggregate{}
 
 const sizeOfArrayAggregate = int64(unsafe.Sizeof(arrayAggregate{}))
@@ -474,7 +478,8 @@ const sizeOfBoolOrAggregate = int64(unsafe.Sizeof(boolOrAggregate{}))
 const sizeOfBytesXorAggregate = int64(unsafe.Sizeof(bytesXorAggregate{}))
 const sizeOfIntXorAggregate = int64(unsafe.Sizeof(intXorAggregate{}))
 const sizeOfJSONAggregate = int64(unsafe.Sizeof(jsonAggregate{}))
-const sizeOfBitAndAggregate = int64(unsafe.Sizeof(bitAndAggregate{}))
+const sizeOfIntBitAndAggregate = int64(unsafe.Sizeof(intBitAndAggregate{}))
+const sizeOfBitBitAndAggregate = int64(unsafe.Sizeof(bitBitAndAggregate{}))
 const sizeOfBitOrAggregate = int64(unsafe.Sizeof(bitOrAggregate{}))
 
 // singleDatumAggregateBase is a utility struct that helps aggregate builtins
@@ -865,17 +870,17 @@ func (a *concatAggregate) Size() int64 {
 	return sizeOfConcatAggregate
 }
 
-type bitAndAggregate struct {
+type intBitAndAggregate struct {
 	sawNonNull bool
 	result     int64
 }
 
-func newBitAndAggregate(_ []*types.T, _ *tree.EvalContext, _ tree.Datums) tree.AggregateFunc {
-	return &bitAndAggregate{}
+func newIntBitAndAggregate(_ []*types.T, _ *tree.EvalContext, _ tree.Datums) tree.AggregateFunc {
+	return &intBitAndAggregate{}
 }
 
 // Add inserts one value into the running bitwise AND.
-func (a *bitAndAggregate) Add(_ context.Context, datum tree.Datum, _ ...tree.Datum) error {
+func (a *intBitAndAggregate) Add(_ context.Context, datum tree.Datum, _ ...tree.Datum) error {
 	if datum == tree.DNull {
 		return nil
 	}
@@ -893,25 +898,80 @@ func (a *bitAndAggregate) Add(_ context.Context, datum tree.Datum, _ ...tree.Dat
 }
 
 // Result returns the bitwise AND.
-func (a *bitAndAggregate) Result() (tree.Datum, error) {
+func (a *intBitAndAggregate) Result() (tree.Datum, error) {
 	if !a.sawNonNull {
 		return tree.DNull, nil
 	}
 	return tree.NewDInt(tree.DInt(a.result)), nil
 }
 
 // Reset implements tree.AggregateFunc interface.
-func (a *bitAndAggregate) Reset(context.Context) {
+func (a *intBitAndAggregate) Reset(context.Context) {
 	a.sawNonNull = false
 	a.result = 0
 }
 
 // Close is part of the tree.AggregateFunc interface.
-func (a *bitAndAggregate) Close(context.Context) {}
+func (a *intBitAndAggregate) Close(context.Context) {}
+
+// Size is part of the tree.AggregateFunc interface.
+func (a *intBitAndAggregate) Size() int64 {
+	return sizeOfIntBitAndAggregate
+}
+
+type bitBitAndAggregate struct {
+	sawNonNull bool
+	result     bitarray.BitArray
+}
+
+func newBitBitAndAggregate(_ []*types.T, _ *tree.EvalContext, _ tree.Datums) tree.AggregateFunc {
+	return &bitBitAndAggregate{}
+}
+
+// Add inserts one value into the running bitwise AND.
+func (a *bitBitAndAggregate) Add(_ context.Context, datum tree.Datum, _ ...tree.Datum) error {
+	if datum == tree.DNull {
+		return nil
+	}
+	bits := &tree.MustBeDBitArray(datum).BitArray
+	if !a.sawNonNull {
+		// This is the first non-null datum, so we simply store
+		// the provided value for the aggregation.
+		a.result = *bits
+		a.sawNonNull = true
+		return nil
+	}
+	// If the length of the current bit array is different from that of the
+	// stored value, we return an error.
+	if a.result.BitLen() != bits.BitLen() {
+		return tree.NewCannotMixBitArraySizesError("AND")
+	}
+	// This is not the first non-null datum, so we actually AND it with the
+	// aggregate so far.
+	a.result = bitarray.And(a.result, *bits)
+	return nil
+}
+
+// Result returns the bitwise AND.
+func (a *bitBitAndAggregate) Result() (tree.Datum, error) {
+	if !a.sawNonNull {
+		return tree.DNull, nil
+	}
+	return &tree.DBitArray{BitArray: a.result}, nil
+}
+
+// Reset implements tree.AggregateFunc interface.
+func (a *bitBitAndAggregate) Reset(context.Context) {
+	a.sawNonNull = false
+	a.result = bitarray.BitArray{}
+}
+
+// Close is part of the tree.AggregateFunc interface.
+func (a *bitBitAndAggregate) Close(context.Context) {}
 
 // Size is part of the tree.AggregateFunc interface.
-func (a *bitAndAggregate) Size() int64 {
-	return sizeOfBitAndAggregate
+func (a *bitBitAndAggregate) Size() int64 {
+	return sizeOfBitBitAndAggregate
 }
 
 type bitOrAggregate struct {

pkg/sql/sem/builtins/aggregate_builtins_test.go

@@ -86,13 +86,28 @@ func TestAvgIntervalResultDeepCopy(t *testing.T) {
 func TestBitAndIntResultDeepCopy(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	t.Run("all null", func(t *testing.T) {
-		testAggregateResultDeepCopy(t, newBitAndAggregate, makeNullTestDatum(10))
+		testAggregateResultDeepCopy(t, newIntBitAndAggregate, makeNullTestDatum(10))
 	})
 	t.Run("with null", func(t *testing.T) {
-		testAggregateResultDeepCopy(t, newBitAndAggregate, makeTestWithNullDatum(10, makeIntTestDatum))
+		testAggregateResultDeepCopy(t, newIntBitAndAggregate, makeTestWithNullDatum(10, makeIntTestDatum))
 	})
 	t.Run("without null", func(t *testing.T) {
-		testAggregateResultDeepCopy(t, newBitAndAggregate, makeIntTestDatum(10))
+		testAggregateResultDeepCopy(t, newIntBitAndAggregate, makeIntTestDatum(10))
+	})
+}
+
+func TestBitAndBitResultDeepCopy(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	t.Run("all null", func(t *testing.T) {
+		testAggregateResultDeepCopy(t, newBitBitAndAggregate, makeNullTestDatum(10))
+	})
+	t.Run("with null", func(t *testing.T) {
+		testAggregateResultDeepCopy(t, newBitBitAndAggregate, makeTestWithNullDatum(10, makeBitTestDatum))
+	})
+	t.Run("without null", func(t *testing.T) {
+		for i := 0; i < 1000; i++ {
+			testAggregateResultDeepCopy(t, newBitBitAndAggregate, makeBitTestDatum(10))
+		}
 	})
 }
 
@@ -239,6 +254,19 @@ func makeIntTestDatum(count int) []tree.Datum {
 	return vals
 }
 
+func makeBitTestDatum(count int) []tree.Datum {
+	rng, _ := randutil.NewPseudoRand()
+
+	// Compute randWidth outside the loop so that all bit arrays are the same
+	// length. Generate widths in the range [0, 64].
+	vals := make([]tree.Datum, count)
+	randWidth := uint(rng.Intn(65))
+	for i := range vals {
+		vals[i], _ = tree.NewDBitArrayFromInt(rng.Int63(), randWidth)
+	}
+	return vals
+}
+
 // makeTestWithNullDatum will call the maker function
 // to generate an array of datums, and then a null datum
 // will be placed randomly in the array of datums and

pkg/sql/sem/tree/eval.go

@@ -63,6 +63,13 @@ var (
 	big10E10 = big.NewInt(1e10)
 )
 
+// NewCannotMixBitArraySizesError creates an error for the case when a bitwise
+// aggregate function is called on bit arrays with different sizes.
+func NewCannotMixBitArraySizesError(op string) error {
+	return pgerror.Newf(pgcode.StringDataLengthMismatch,
+		"cannot %s bit strings of different sizes", op)
+}
+
 // UnaryOp is a unary operator.
 type UnaryOp struct {
 	Typ        *types.T
@@ -382,11 +389,6 @@ func getJSONPath(j DJSON, ary DArray) (Datum, error) {
 	return &DJSON{result}, nil
 }
 
-func newCannotMixBitArraySizesError(op string) error {
-	return pgerror.Newf(pgcode.StringDataLengthMismatch,
-		"cannot %s bit strings of different sizes", op)
-}
-
 // BinOps contains the binary operations indexed by operation type.
 var BinOps = map[BinaryOperator]binOpOverload{
 	Bitand: {
@@ -406,7 +408,7 @@ var BinOps = map[BinaryOperator]binOpOverload{
 				lhs := MustBeDBitArray(left)
 				rhs := MustBeDBitArray(right)
 				if lhs.BitLen() != rhs.BitLen() {
-					return nil, newCannotMixBitArraySizesError("AND")
+					return nil, NewCannotMixBitArraySizesError("AND")
 				}
 				return &DBitArray{
 					BitArray: bitarray.And(lhs.BitArray, rhs.BitArray),
@@ -443,7 +445,7 @@ var BinOps = map[BinaryOperator]binOpOverload{
 				lhs := MustBeDBitArray(left)
 				rhs := MustBeDBitArray(right)
 				if lhs.BitLen() != rhs.BitLen() {
-					return nil, newCannotMixBitArraySizesError("OR")
+					return nil, NewCannotMixBitArraySizesError("OR")
 				}
 				return &DBitArray{
 					BitArray: bitarray.Or(lhs.BitArray, rhs.BitArray),
@@ -480,7 +482,7 @@ var BinOps = map[BinaryOperator]binOpOverload{
 				lhs := MustBeDBitArray(left)
 				rhs := MustBeDBitArray(right)
 				if lhs.BitLen() != rhs.BitLen() {
-					return nil, newCannotMixBitArraySizesError("XOR")
+					return nil, NewCannotMixBitArraySizesError("XOR")
 				}
 				return &DBitArray{
 					BitArray: bitarray.Xor(lhs.BitArray, rhs.BitArray),