From 4f41159b311f4385bf5355f789d70986e9fd4b0d Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Thu, 16 May 2024 09:15:50 -0700 Subject: [PATCH 1/2] Add SQL DIV function. This function has been documented for some time, but lacked a binding, so it wasn't usable. --- .../builtin/DivOperatorConversion.java | 45 +++++++++++++++++++ .../calcite/planner/DruidOperatorTable.java | 2 + .../druid/sql/calcite/CalciteQueryTest.java | 29 ++++++++++++ 3 files changed, 76 insertions(+) create mode 100644 sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/DivOperatorConversion.java diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/DivOperatorConversion.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/DivOperatorConversion.java new file mode 100644 index 000000000000..fd73689779f6 --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/DivOperatorConversion.java @@ -0,0 +1,45 @@ +/* + * 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.druid.sql.calcite.expression.builtin; + +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.InferTypes; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.druid.sql.calcite.expression.DirectOperatorConversion; +import org.apache.druid.sql.calcite.expression.OperatorConversions; + +public class DivOperatorConversion extends DirectOperatorConversion +{ + private static final SqlOperator SQL_OPERATOR = + OperatorConversions + .operatorBuilder("DIV") + .operandTypeChecker(OperandTypes.DIVISION_OPERATOR) + .operandTypeInference(InferTypes.FIRST_KNOWN) + .returnTypeCascadeNullable(SqlTypeName.BIGINT) + .functionCategory(SqlFunctionCategory.NUMERIC) + .build(); + + public DivOperatorConversion() + { + super(SQL_OPERATOR, "div"); + } +} diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java index 723d5f45a39d..27efe16270ed 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java @@ -81,6 +81,7 @@ import org.apache.druid.sql.calcite.expression.builtin.ContainsOperatorConversion; import org.apache.druid.sql.calcite.expression.builtin.DateTruncOperatorConversion; import org.apache.druid.sql.calcite.expression.builtin.DecodeBase64UTFOperatorConversion; +import org.apache.druid.sql.calcite.expression.builtin.DivOperatorConversion; import org.apache.druid.sql.calcite.expression.builtin.ExtractOperatorConversion; import org.apache.druid.sql.calcite.expression.builtin.FloorOperatorConversion; import org.apache.druid.sql.calcite.expression.builtin.GreatestOperatorConversion; @@ -369,6 +370,7 @@ public class DruidOperatorTable implements SqlOperatorTable .add(new AliasedOperatorConversion(CHARACTER_LENGTH_CONVERSION, "STRLEN")) .add(new DirectOperatorConversion(SqlStdOperatorTable.CONCAT, "concat")) .add(new DirectOperatorConversion(SqlStdOperatorTable.EXP, "exp")) + .add(new DivOperatorConversion()) .add(new DirectOperatorConversion(SqlStdOperatorTable.DIVIDE_INTEGER, "div")) .add(new DirectOperatorConversion(SqlStdOperatorTable.LN, "log")) .add(new DirectOperatorConversion(SqlStdOperatorTable.LOWER, "lower")) 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 7a3df49eedfa..cd1be7ad9d99 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 @@ -602,6 +602,35 @@ public void testSafeDivide() ); } + @Test + public void testDiv() + { + cannotVectorize(); + final Map context = new HashMap<>(QUERY_CONTEXT_DEFAULT); + + testQuery( + "select m1, div(m1, 2) from foo", + context, + ImmutableList.of( + newScanQueryBuilder() + .dataSource(CalciteTests.DATASOURCE1) + .intervals(querySegmentSpec(Filtration.eternity())) + .virtualColumns(expressionVirtualColumn("v0", "div(\"m1\",2)", ColumnType.LONG)) + .columns(ImmutableList.of("m1", "v0")) + .context(QUERY_CONTEXT_DEFAULT) + .build() + ), + ImmutableList.of( + new Object[]{1.0f, 0L}, + new Object[]{2.0f, 1L}, + new Object[]{3.0f, 1L}, + new Object[]{4.0f, 2L}, + new Object[]{5.0f, 2L}, + new Object[]{6.0f, 3L} + ) + ); + } + @Test public void testGroupByLimitWrappingOrderByAgg() { From 1ed4b550e1a7b2c434befb5e29c0177f365ebbde Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Thu, 16 May 2024 20:24:02 -0700 Subject: [PATCH 2/2] Add a case with two expression inputs. --- .../druid/sql/calcite/CalciteQueryTest.java | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) 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 cd1be7ad9d99..5d5aab111358 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 @@ -609,24 +609,27 @@ public void testDiv() final Map context = new HashMap<>(QUERY_CONTEXT_DEFAULT); testQuery( - "select m1, div(m1, 2) from foo", + "select cnt, m1, div(m1, 2), div(cnt+2, cnt+1) from foo", context, ImmutableList.of( newScanQueryBuilder() .dataSource(CalciteTests.DATASOURCE1) .intervals(querySegmentSpec(Filtration.eternity())) - .virtualColumns(expressionVirtualColumn("v0", "div(\"m1\",2)", ColumnType.LONG)) - .columns(ImmutableList.of("m1", "v0")) + .virtualColumns( + expressionVirtualColumn("v0", "div(\"m1\",2)", ColumnType.LONG), + expressionVirtualColumn("v1", "div((\"cnt\" + 2),(\"cnt\" + 1))", ColumnType.LONG) + ) + .columns(ImmutableList.of("cnt", "m1", "v0", "v1")) .context(QUERY_CONTEXT_DEFAULT) .build() ), ImmutableList.of( - new Object[]{1.0f, 0L}, - new Object[]{2.0f, 1L}, - new Object[]{3.0f, 1L}, - new Object[]{4.0f, 2L}, - new Object[]{5.0f, 2L}, - new Object[]{6.0f, 3L} + new Object[]{1L, 1.0f, 0L, 1L}, + new Object[]{1L, 2.0f, 1L, 1L}, + new Object[]{1L, 3.0f, 1L, 1L}, + new Object[]{1L, 4.0f, 2L, 1L}, + new Object[]{1L, 5.0f, 2L, 1L}, + new Object[]{1L, 6.0f, 3L, 1L} ) ); }