diff --git a/isthmus/src/main/java/io/substrait/isthmus/calcite/SubstraitOperatorTable.java b/isthmus/src/main/java/io/substrait/isthmus/calcite/SubstraitOperatorTable.java index c73baed7e..3e71cc837 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/calcite/SubstraitOperatorTable.java +++ b/isthmus/src/main/java/io/substrait/isthmus/calcite/SubstraitOperatorTable.java @@ -2,6 +2,7 @@ import io.substrait.isthmus.AggregateFunctions; import io.substrait.isthmus.expression.CurrentTimezoneFunction; +import io.substrait.isthmus.expression.FunctionMappings; import java.util.EnumSet; import java.util.List; import java.util.Set; @@ -55,7 +56,7 @@ public class SubstraitOperatorTable implements SqlOperatorTable { // feed OVERRIDE_KINDS: they share generic kinds such as OTHER_FUNCTION with many standard // operators, which we must not shadow. private static final SqlOperatorTable SUBSTRAIT_SCALAR_OPERATOR_TABLE = - SqlOperatorTables.of(List.of(CurrentTimezoneFunction.INSTANCE)); + SqlOperatorTables.of(List.of(CurrentTimezoneFunction.INSTANCE, FunctionMappings.RIGHTSHIFT)); // Utilisation of extended library operators available from calcite 1.35+, i.e hyperbolic // functions diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionMappings.java b/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionMappings.java index 85f44a312..fea31418d 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionMappings.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionMappings.java @@ -11,6 +11,7 @@ import org.apache.calcite.sql.fun.SqlLibraryOperators; import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; import org.apache.calcite.sql.type.SqlTypeFamily; import org.apache.calcite.sql.type.SqlTypeName; @@ -50,6 +51,18 @@ public class FunctionMappings { opBinding -> opBinding.getTypeFactory().createSqlType(SqlTypeName.BOOLEAN), OperandTypes.family(SqlTypeFamily.ARRAY, SqlTypeFamily.ANY)); + /** + * The {@code RIGHTSHIFT(value, shift)} function. Calcite provides {@link + * SqlStdOperatorTable#LEFTSHIFT} (and the {@code <<} operator {@link + * SqlStdOperatorTable#BIT_LEFT_SHIFT}) but has no right-shift counterpart, so Isthmus defines one + * to map to the Substrait {@code shift_right} function. + */ + public static final SqlFunction RIGHTSHIFT = + SqlBasicFunction.create( + "RIGHTSHIFT", + ReturnTypes.ARG0_NULLABLE, + OperandTypes.family(SqlTypeFamily.INTEGER, SqlTypeFamily.INTEGER)); + /** Scalar function mappings. */ public static final ImmutableList SCALAR_SIGS = ImmutableList.builder() @@ -128,6 +141,7 @@ public class FunctionMappings { s(SqlLibraryOperators.GREATEST, "greatest"), s(SqlStdOperatorTable.BIT_LEFT_SHIFT, "shift_left"), s(SqlStdOperatorTable.LEFTSHIFT, "shift_left"), + s(RIGHTSHIFT, "shift_right"), s(SqlLibraryOperators.STARTS_WITH, "starts_with"), s(SqlLibraryOperators.ENDS_WITH, "ends_with"), s(SqlLibraryOperators.CONTAINS_SUBSTR, "contains"), diff --git a/isthmus/src/test/java/io/substrait/isthmus/ArithmeticFunctionTest.java b/isthmus/src/test/java/io/substrait/isthmus/ArithmeticFunctionTest.java index 9fbff8a96..e71934b43 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ArithmeticFunctionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ArithmeticFunctionTest.java @@ -220,4 +220,11 @@ void leftshift(String column) throws Exception { String query = String.format("SELECT LEFTSHIFT(%s, 1) FROM numbers", column); assertFullRoundTrip(query, CREATES); } + + @ParameterizedTest + @ValueSource(strings = {"i8", "i16", "i32", "i64"}) + void rightshift(String column) throws Exception { + String query = String.format("SELECT RIGHTSHIFT(%s, 1) FROM numbers", column); + assertFullRoundTrip(query, CREATES); + } }