From 883c2692d23554f5f45aa847b46c5ced25601ab5 Mon Sep 17 00:00:00 2001 From: Pranav Date: Wed, 20 Sep 2023 17:06:34 -0700 Subject: [PATCH] Adding new function decode_base64_utf8 and expr macro (#14943) * Adding new function decode_base64_utf8 and expr macro * using BaseScalarUnivariateMacroFunctionExpr * Print stack trace in case of debug in ChainedExecutionQueryRunner * fix static check --- .../druid/math/expr/BuiltInExprMacros.java | 70 +++++++++++++++++++ .../druid/math/expr/ExprMacroTable.java | 3 +- .../query/ChainedExecutionQueryRunner.java | 6 +- .../apache/druid/math/expr/FunctionTest.java | 9 +++ .../DecodeBase64UTFOperatorConversion.java | 52 ++++++++++++++ .../calcite/planner/DruidOperatorTable.java | 2 + .../druid/sql/calcite/CalciteQueryTest.java | 28 ++++++++ 7 files changed, 168 insertions(+), 2 deletions(-) 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 e47ef742c309..c1298cd79297 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 @@ -23,6 +23,7 @@ import org.apache.druid.segment.column.TypeStrategy; import javax.annotation.Nullable; +import java.util.Collections; import java.util.List; import java.util.stream.Collectors; @@ -144,4 +145,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) + { + validationHelperCheckArgumentCount(args, 1); + return new StringDecodeBase64UTFExpression(args.get(0)); + } + + @Override + public String name() + { + return NAME; + } + + final class StringDecodeBase64UTFExpression extends ExprMacroTable.BaseScalarUnivariateMacroFunctionExpr + { + public StringDecodeBase64UTFExpression(Expr arg) + { + super(NAME, arg); + } + + @Override + public ExprEval eval(ObjectBinding bindings) + { + ExprEval toDecode = arg.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(Collections.singletonList(arg)))); + } + + @Nullable + @Override + public ExpressionType getOutputType(InputBindingInspector inspector) + { + return ExpressionType.STRING; + } + + @Override + public boolean isLiteral() + { + return arg.isLiteral(); + } + + @Override + public boolean isNullLiteral() + { + return arg.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 6f11ff3e69fd..d3cc6461c51b 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/main/java/org/apache/druid/query/ChainedExecutionQueryRunner.java b/processing/src/main/java/org/apache/druid/query/ChainedExecutionQueryRunner.java index 5ff044f65641..c8485814b595 100644 --- a/processing/src/main/java/org/apache/druid/query/ChainedExecutionQueryRunner.java +++ b/processing/src/main/java/org/apache/druid/query/ChainedExecutionQueryRunner.java @@ -123,7 +123,11 @@ public Iterable call() throw e; } catch (Exception e) { - log.noStackTrace().error(e, "Exception with one of the sequences!"); + if (query.context().isDebug()) { + log.error(e, "Exception with one of the sequences!"); + } else { + log.noStackTrace().error(e, "Exception with one of the sequences!"); + } Throwables.propagateIfPossible(e); throw new RuntimeException(e); } 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 e3a7fa909c6c..c16b12372b3b 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 000000000000..f0a0bf40a8cf --- /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 e392ad8a47bf..af6c2fc30d72 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 @@ -78,6 +78,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; @@ -231,6 +232,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 6516358b1b59..85c690ebba7c 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 @@ -2216,6 +2216,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() {