Moving bitwise_or to use native calcite operator (apache#16237)

sergioferragut · Apr 4, 2024 · 7759f25 · 7759f25
1 parent 9729376
commit 7759f25
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 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
@@ -22,14 +22,7 @@
 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 @@
 
 public class BitwiseSqlAggregator implements SqlAggregator
 {
-  private static final SqlAggFunction XOR_FUNCTION = new BitwiseXorSqlAggFunction();
-
   public enum Op
   {
     AND {
@@ -88,8 +79,7 @@ String getDruidFunction()
       @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 Aggregation toDruidAggregation(
         )
     );
   }
-
-  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
@@ -15550,4 +15550,77 @@ public void testIpv4ParseWithNullableType()
         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}
+        )
+    );
+  }
 }