From 7759f25095fbb27e9aef60c18407d4612b1fcbef Mon Sep 17 00:00:00 2001
From: Soumyava <93540295+somu-imply@users.noreply.github.com>
Date: Thu, 4 Apr 2024 12:49:29 -0700
Subject: [PATCH] Moving bitwise_or to use native calcite operator (#16237)

---
 .../builtin/BitwiseSqlAggregator.java         | 31 +-------
 .../druid/sql/calcite/CalciteQueryTest.java   | 73 +++++++++++++++++++
 2 files changed, 74 insertions(+), 30 deletions(-)

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/aggregation/builtin/BitwiseSqlAggregator.java b/sql/src/main/java/org/apache/druid/sql/calcite/aggregation/builtin/BitwiseSqlAggregator.java
index a5c7fb61cff..489080acd16 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/aggregation/builtin/BitwiseSqlAggregator.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/aggregation/builtin/BitwiseSqlAggregator.java
@@ -22,14 +22,7 @@ package org.apache.druid.sql.calcite.aggregation.builtin;
 import com.google.common.collect.ImmutableSet;
 import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.sql.SqlAggFunction;
-import org.apache.calcite.sql.SqlFunctionCategory;
-import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
-import org.apache.calcite.sql.type.InferTypes;
-import org.apache.calcite.sql.type.OperandTypes;
-import org.apache.calcite.sql.type.ReturnTypes;
-import org.apache.calcite.sql.type.SqlTypeName;
-import org.apache.calcite.util.Optionality;
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.math.expr.ExprMacroTable;
@@ -54,8 +47,6 @@ import java.util.stream.Collectors;
 
 public class BitwiseSqlAggregator implements SqlAggregator
 {
-  private static final SqlAggFunction XOR_FUNCTION = new BitwiseXorSqlAggFunction();
-
   public enum Op
   {
     AND {
@@ -88,8 +79,7 @@ public class BitwiseSqlAggregator implements SqlAggregator
       @Override
       SqlAggFunction getCalciteFunction()
       {
-        // newer versions of calcite have this built-in so someday we can drop this...
-        return XOR_FUNCTION;
+        return SqlStdOperatorTable.BIT_XOR;
       }
 
       @Override
@@ -175,23 +165,4 @@ public class BitwiseSqlAggregator implements SqlAggregator
         )
     );
   }
-
-  private static class BitwiseXorSqlAggFunction extends SqlAggFunction
-  {
-    BitwiseXorSqlAggFunction()
-    {
-      super(
-          "BIT_XOR",
-          null,
-          SqlKind.OTHER_FUNCTION,
-          ReturnTypes.explicit(SqlTypeName.BIGINT),
-          InferTypes.ANY_NULLABLE,
-          OperandTypes.EXACT_NUMERIC,
-          SqlFunctionCategory.NUMERIC,
-          false,
-          false,
-          Optionality.IGNORED
-      );
-    }
-  }
 }
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
index bc21586ac8c..478648088c3 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
@@ -15550,4 +15550,77 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
         ImmutableList.of(NullHandling.sqlCompatible() ? new Object[]{null} : new Object[]{0})
     );
   }
+
+  @Test
+  public void testBitwiseXor()
+  {
+    skipVectorize();
+    cannotVectorize();
+    msqIncompatible();
+    testQuery(
+        "select count(*) from (\n"
+        + "  select __time, cityName, bit_xor(cityName) c2\n"
+        + "  from wikipedia\n"
+        + "  group by __time, cityName\n"
+        + "  having bit_xor(cityName) is null\n"
+        + ")",
+        ImmutableList.of(
+            GroupByQuery.builder()
+                        .setDataSource(
+                            new QueryDataSource(
+                                new GroupByQuery.Builder()
+                                    .setDataSource(CalciteTests.WIKIPEDIA)
+                                    .setInterval(querySegmentSpec(Filtration.eternity()))
+                                    .setGranularity(Granularities.ALL)
+                                    .setDimensions(
+                                        new DefaultDimensionSpec("__time", "d0", ColumnType.LONG),
+                                        new DefaultDimensionSpec("cityName", "d1", ColumnType.STRING)
+                                    )
+                                    .setLimitSpec(
+                                        NoopLimitSpec.instance()
+                                    )
+                                    .setAggregatorSpecs(aggregators(new FilteredAggregatorFactory(
+                                        new ExpressionLambdaAggregatorFactory(
+                                            "a0",
+                                            ImmutableSet.of("cityName"),
+                                            "__acc",
+                                            "0",
+                                            "0",
+                                            NullHandling.sqlCompatible(),
+                                            false,
+                                            false,
+                                            "bitwiseXor(\"__acc\", \"cityName\")",
+                                            "bitwiseXor(\"__acc\", \"a0\")",
+                                            null,
+                                            null,
+                                            ExpressionLambdaAggregatorFactory.DEFAULT_MAX_SIZE_BYTES,
+                                            TestExprMacroTable.INSTANCE
+                                        ),
+                                        notNull("cityName")
+                                    )))
+                                    .setHavingSpec(
+                                        having(
+                                            isNull("a0")
+                                        )
+                                    )
+                                    .setContext(OUTER_LIMIT_CONTEXT)
+                                    .build()
+                            )
+                        )
+                        .setInterval(querySegmentSpec(Filtration.eternity()))
+                        .setGranularity(Granularities.ALL)
+                        .setAggregatorSpecs(
+                            aggregators(
+                                new CountAggregatorFactory("_a0")
+                            )
+                        )
+                        .setContext(QUERY_CONTEXT_DEFAULT)
+                        .build()
+        ),
+        ImmutableList.of(
+            useDefault ?
+            new Object[]{0L} : new Object[]{37091L}
+        )
+    );
+  }
 }