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
apache · Sep 21, 2023 · 883c269 · 883c269
1 parent 22abc10
commit 883c269
Show file tree

Hide file tree

Showing 7 changed files with 168 additions and 2 deletions.
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
@@ -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<Expr> 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
@@ -44,7 +44,8 @@
 public class ExprMacroTable
 {
   private static final List<ExprMacro> 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
@@ -123,7 +123,11 @@ public Iterable<T> 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
@@ -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/...va/org/apache/druid/sql/calcite/expression/builtin/DecodeBase64UTFOperatorConversion.java b/...va/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
@@ -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<SqlOperatorConversion> 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
@@ -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()
   {