From ad20520208c6e09ff4758ab80745c0efc81b0219 Mon Sep 17 00:00:00 2001 From: Pranav Bhole Date: Tue, 5 Sep 2023 17:14:11 -0700 Subject: [PATCH] Adding new function decode_base64_utf8 and expr macro --- .../druid/math/expr/BuiltInExprMacros.java | 69 +++++++++++++++++++ .../druid/math/expr/ExprMacroTable.java | 3 +- .../apache/druid/math/expr/FunctionTest.java | 9 +++ .../DecodeBase64UTFOperatorConversion.java | 52 ++++++++++++++ .../calcite/planner/DruidOperatorTable.java | 2 + .../druid/sql/calcite/CalciteQueryTest.java | 28 ++++++++ 6 files changed, 162 insertions(+), 1 deletion(-) create mode 100644 sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/DecodeBase64UTFOperatorConversion.java diff --git a/processing/src/main/java/org/apache/druid/math/expr/BuiltInExprMacros.java b/processing/src/main/java/org/apache/druid/math/expr/BuiltInExprMacros.java index e47ef742c309a..39b56f7966fba 100644 --- a/processing/src/main/java/org/apache/druid/math/expr/BuiltInExprMacros.java +++ b/processing/src/main/java/org/apache/druid/math/expr/BuiltInExprMacros.java @@ -144,4 +144,73 @@ public Object getLiteralValue() } } } + + public static class StringDecodeBase64UTFExprMacro implements ExprMacroTable.ExprMacro + { + + public static final String NAME = "decode_base64_utf8"; + + @Override + public Expr apply(List args) + { + return new StringDecodeBase64UTFExpression(args); + } + + @Override + public String name() + { + return NAME; + } + + final class StringDecodeBase64UTFExpression extends ExprMacroTable.BaseScalarMacroFunctionExpr + { + public StringDecodeBase64UTFExpression(List args) + { + super(NAME, args); + validationHelperCheckArgumentCount(args, 1); + } + + @Override + public ExprEval eval(ObjectBinding bindings) + { + ExprEval toDecode = args.get(0).eval(bindings); + if (toDecode.value() == null) { + return ExprEval.of(null); + } + return new StringExpr(StringUtils.fromUtf8(StringUtils.decodeBase64String(toDecode.asString()))).eval(bindings); + } + + @Override + public Expr visit(Shuttle shuttle) + { + return shuttle.visit(apply(shuttle.visitAll(args))); + } + + @Nullable + @Override + public ExpressionType getOutputType(InputBindingInspector inspector) + { + return ExpressionType.STRING; + } + + @Override + public boolean isLiteral() + { + return args.get(0).isLiteral(); + } + + @Override + public boolean isNullLiteral() + { + return args.get(0).isNullLiteral(); + } + + @Nullable + @Override + public Object getLiteralValue() + { + return eval(InputBindings.nilBindings()).value(); + } + } + } } diff --git a/processing/src/main/java/org/apache/druid/math/expr/ExprMacroTable.java b/processing/src/main/java/org/apache/druid/math/expr/ExprMacroTable.java index 6f11ff3e69fd0..d3cc6461c51b2 100644 --- a/processing/src/main/java/org/apache/druid/math/expr/ExprMacroTable.java +++ b/processing/src/main/java/org/apache/druid/math/expr/ExprMacroTable.java @@ -44,7 +44,8 @@ public class ExprMacroTable { private static final List BUILT_IN = ImmutableList.of( - new BuiltInExprMacros.ComplexDecodeBase64ExprMacro() + new BuiltInExprMacros.ComplexDecodeBase64ExprMacro(), + new BuiltInExprMacros.StringDecodeBase64UTFExprMacro() ); private static final ExprMacroTable NIL = new ExprMacroTable(Collections.emptyList()); diff --git a/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java b/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java index e3a7fa909c6cc..c16b12372b3b0 100644 --- a/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java +++ b/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java @@ -925,6 +925,15 @@ public void testRepeat() assertExpr("repeat(nonexistent, 10)", null); } + @Test + public void testDecodeBase64UTF() + { + assertExpr("decode_base64_utf8('aGVsbG8=')", "hello"); + assertExpr("decode_base64_utf8('V2hlbiBhbiBvbmlvbiBpcyBjdXQsIGNlcnRhaW4gKGxhY2hyeW1hdG9yKSBjb21wb3VuZHMgYXJlIHJlbGVhc2VkIGNhdXNpbmcgdGhlIG5lcnZlcyBhcm91bmQgdGhlIGV5ZXMgKGxhY3JpbWFsIGdsYW5kcykgdG8gYmVjb21lIGlycml0YXRlZC4=')", "When an onion is cut, certain (lachrymator) compounds are released causing the nerves around the eyes (lacrimal glands) to become irritated."); + assertExpr("decode_base64_utf8('eyJ0ZXN0IjogMX0=')", "{\"test\": 1}"); + assertExpr("decode_base64_utf8('')", NullHandling.sqlCompatible() ? "" : null); + } + @Test public void testComplexDecode() { diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/DecodeBase64UTFOperatorConversion.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/DecodeBase64UTFOperatorConversion.java new file mode 100644 index 0000000000000..f0a0bf40a8cfe --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/DecodeBase64UTFOperatorConversion.java @@ -0,0 +1,52 @@ +/* + * 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.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.SqlTypeFamily; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.math.expr.BuiltInExprMacros; +import org.apache.druid.sql.calcite.expression.DirectOperatorConversion; +import org.apache.druid.sql.calcite.expression.OperatorConversions; + +public class DecodeBase64UTFOperatorConversion extends DirectOperatorConversion +{ + + private static final SqlFunction SQL_FUNCTION = OperatorConversions + .operatorBuilder(StringUtils.toUpperCase(BuiltInExprMacros.StringDecodeBase64UTFExprMacro.NAME)) + .operandTypes(SqlTypeFamily.CHARACTER) + .returnTypeNullable(SqlTypeName.VARCHAR) + .functionCategory(SqlFunctionCategory.STRING) + .build(); + + public DecodeBase64UTFOperatorConversion() + { + super(SQL_FUNCTION, SQL_FUNCTION.getName()); + } + + @Override + public SqlOperator calciteOperator() + { + return SQL_FUNCTION; + } +} 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 e8ab0cc71d845..f94a88c1927a6 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 @@ -79,6 +79,7 @@ import org.apache.druid.sql.calcite.expression.builtin.ConcatOperatorConversion; 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.ExtractOperatorConversion; import org.apache.druid.sql.calcite.expression.builtin.FloorOperatorConversion; import org.apache.druid.sql.calcite.expression.builtin.GreatestOperatorConversion; @@ -232,6 +233,7 @@ public class DruidOperatorTable implements SqlOperatorTable .add(new CastOperatorConversion()) .add(new ReinterpretOperatorConversion()) .add(new ComplexDecodeBase64OperatorConversion()) + .add(new DecodeBase64UTFOperatorConversion()) .build(); private static final List ARRAY_OPERATOR_CONVERSIONS = 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 582408003719a..d2a59a4e33ac9 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 @@ -2215,6 +2215,34 @@ public void testFilterOnFloat() ); } + @Test + public void testDECODE_BASE64_UTF8() + { + testQuery( + "SELECT DECODE_BASE64_UTF8('aGVsbG8=') FROM druid.foo limit 1", + ImmutableList.of( + Druids.newScanQueryBuilder() + .dataSource(CalciteTests.DATASOURCE1) + .intervals(querySegmentSpec(Filtration.eternity())) + .virtualColumns( + expressionVirtualColumn( + "v0", + "'hello'", + ColumnType.STRING + ) + ) + .limit(1) + .columns(ImmutableList.of("v0")) + .resultFormat(ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .legacy(false) + .context(QUERY_CONTEXT_DEFAULT) + .build() + ), + ImmutableList.of( + new Object[]{"hello"} + ) + ); + } @Test public void testFilterOnDouble() {