[CALCITE-6812] Support base64 and unbase64 functions for Hive

apache · Feb 1, 2025 · 4b7c684 · 4b7c684
1 parent f10d0ce
commit 4b7c684
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 1 deletion.
diff --git a/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java b/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java
@@ -158,6 +158,7 @@
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.ASINH;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.ATAND;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.ATANH;
+import static org.apache.calcite.sql.fun.SqlLibraryOperators.BASE64;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.BITAND_AGG;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.BITOR_AGG;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.BIT_COUNT_BIG_QUERY;
@@ -335,6 +336,7 @@
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.UNIX_MICROS;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.UNIX_MILLIS;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.UNIX_SECONDS;
+import static org.apache.calcite.sql.fun.SqlLibraryOperators.UN_BASE64;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.URL_DECODE;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.URL_ENCODE;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.XML_TRANSFORM;
@@ -685,6 +687,8 @@ void populate1() {
       defineMethod(INITCAP, BuiltInMethod.INITCAP.method, NullPolicy.STRICT);
       defineMethod(TO_BASE64, BuiltInMethod.TO_BASE64.method, NullPolicy.STRICT);
       defineMethod(FROM_BASE64, BuiltInMethod.FROM_BASE64.method, NullPolicy.STRICT);
+      defineMethod(BASE64, BuiltInMethod.BASE64.method, NullPolicy.STRICT);
+      defineMethod(UN_BASE64, BuiltInMethod.UN_BASE64.method, NullPolicy.STRICT);
       defineMethod(TO_BASE32, BuiltInMethod.TO_BASE32.method, NullPolicy.STRICT);
       defineMethod(FROM_BASE32, BuiltInMethod.FROM_BASE32.method, NullPolicy.STRICT);
       defineMethod(HEX, BuiltInMethod.HEX.method, NullPolicy.STRICT);

diff --git a/core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java b/core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java
@@ -330,6 +330,28 @@ private static String toBase64_(byte[] bytes) {
     }
   }
 
+  /** SQL BASE64(string) function for string. */
+  public static String base64(String string) {
+    return toBase64_(string.getBytes(UTF_8));
+  }
+
+  /** SQL BASE64(string) function for bytestring. */
+  public static String base64(ByteString byteString) {
+    byte[] bytes = byteString.getBytes();
+    return toBase64_(bytes);
+  }
+
+  /** SQL UNBASE64(string) function. */
+  public static @Nullable String unBase64(String base64) {
+    try {
+      base64 = FROM_BASE64_REGEXP.matcher(base64).replaceAll("");
+      //CHECKSTYLE: IGNORE 1
+      return new String(Base64.getDecoder().decode(base64), UTF_8);
+    } catch (IllegalArgumentException e) {
+      return null;
+    }
+  }
+
   /** SQL TO_BASE32(string) function. */
   public static String toBase32(String string) {
     return toBase32_(string.getBytes(UTF_8));

diff --git a/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java b/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
@@ -1855,6 +1855,19 @@ private static RelDataType deriveTypeMapFromEntries(SqlOperatorBinding opBinding
           OperandTypes.STRING.or(OperandTypes.BINARY),
           SqlFunctionCategory.STRING);
 
+  @LibraryOperator(libraries = {HIVE})
+  public static final SqlFunction BASE64 =
+      SqlBasicFunction.create("BASE64",
+          ReturnTypes.VARCHAR_NULLABLE,
+          OperandTypes.STRING.or(OperandTypes.BINARY),
+          SqlFunctionCategory.STRING);
+
+  @LibraryOperator(libraries = {HIVE})
+  public static final SqlFunction UN_BASE64 =
+      SqlBasicFunction.create("UNBASE64",
+          ReturnTypes.VARCHAR_FORCE_NULLABLE,
+          OperandTypes.STRING, SqlFunctionCategory.STRING);
+
   @LibraryOperator(libraries = {BIG_QUERY})
   public static final SqlFunction FROM_BASE32 =
       SqlBasicFunction.create("FROM_BASE32",

diff --git a/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java b/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java
@@ -433,6 +433,8 @@ public enum BuiltInMethod {
   RIGHT(SqlFunctions.class, "right", String.class, int.class),
   TO_BASE64(SqlFunctions.class, "toBase64", String.class),
   FROM_BASE64(SqlFunctions.class, "fromBase64", String.class),
+  BASE64(SqlFunctions.class, "base64", String.class),
+  UN_BASE64(SqlFunctions.class, "unBase64", String.class),
   TO_BASE32(SqlFunctions.class, "toBase32", String.class),
   FROM_BASE32(SqlFunctions.class, "fromBase32", String.class),
   HEX(SqlFunctions.class, "hex", String.class),

diff --git a/site/_docs/reference.md b/site/_docs/reference.md
@@ -2922,7 +2922,9 @@ In the following:
 | b | FROM_BASE32(string)                            | Returns the decoded result of a base-32 *string* as a string
 | m | TO_BASE64(string)                              | Converts the *string* to base-64 encoded form and returns a encoded string
 | b m | FROM_BASE64(string)                          | Returns the decoded result of a base-64 *string* as a string. If the input argument is an invalid base-64 *string* the function returns `NULL`
-| h s | HEX(string)                                   | Converts *string* into a hexadecimal varchar
+| h | BASE64(string)                                 | Converts the *string* to base-64 encoded form and returns a encoded string
+| h | UNBASE64(string)                               | Returns the decoded result of a base-64 *string* as a string. If the input argument is an invalid base-64 *string* the function returns `NULL`
+| h s | HEX(string)                                  | Converts *string* into a hexadecimal varchar
 | b | TO_HEX(binary)                                 | Converts *binary* into a hexadecimal varchar
 | b | FROM_HEX(varchar)                              | Converts a hexadecimal-encoded *varchar* into bytes
 | b o p r s h | LTRIM(string)                        | Returns *string* with all blanks removed from the start

diff --git a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
@@ -5173,6 +5173,26 @@ void testBitGetFunc(SqlOperatorFixture f, String functionName) {
     f.checkString("to_base64(x'61')", "YQ==", "VARCHAR NOT NULL");
   }
 
+  @Test void testBase64() {
+    final SqlOperatorFixture f = fixture().withLibrary(SqlLibrary.HIVE);
+    f.setFor(SqlLibraryOperators.BASE64);
+    f.checkString("base64(cast('a' as binary))", "YQ==", "VARCHAR NOT NULL");
+    f.checkString("base64('')", "", "VARCHAR NOT NULL");
+    f.checkString("base64('This is a test String.')",
+        "VGhpcyBpcyBhIHRlc3QgU3RyaW5nLg==",
+        "VARCHAR NOT NULL");
+  }
+
+  @Test void testUnBase64() {
+    final SqlOperatorFixture f = fixture().withLibrary(SqlLibrary.HIVE);
+    f.setFor(SqlLibraryOperators.UN_BASE64);
+    f.checkString("unbase64('VGhpcyBpcyBhIHRlc3QgU3RyaW5nLg==')",
+              "This is a test String.",
+              "VARCHAR");
+    f.checkString("unbase64('YQ==')", "a", "VARCHAR");
+    f.checkString("unbase64('')", "", "VARCHAR");
+  }
+
   @Test void testToChar() {
     final SqlOperatorFixture f0 = fixture().setFor(SqlLibraryOperators.TO_CHAR);
     final Consumer<SqlOperatorFixture> consumer = f -> {