[CALCITE-6663] Support SPLIT_PART function for postgresql

zhuqi-lucas · Nov 1, 2024 · e087a3b · e087a3b
1 parent 2e99d51
commit e087a3b
Show file tree

Hide file tree

Showing 7 changed files with 85 additions and 0 deletions.
diff --git a/babel/src/test/resources/sql/postgresql.iq b/babel/src/test/resources/sql/postgresql.iq
@@ -39,6 +39,19 @@ EXPR$0
 false
 !ok
 
+#Test string function split_part
+select split_part('abc~@~def~@~ghi', '~@~', 2);
+-- Expected output: def
+EXPR$0
+def
+!ok
+
+select split_part('abc,def,ghi,jkl', ',', -2);
+-- Expected output: ghi
+EXPR$0
+ghi
+!ok
+
 # Test coercion string to array inside INSERT
 create table sal_emp (name varchar, pay_by_quarter int array, schedule varchar array array);
 (0 rows modified)

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
@@ -294,6 +294,7 @@
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.SOUNDEX_SPARK;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.SPACE;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.SPLIT;
+import static org.apache.calcite.sql.fun.SqlLibraryOperators.SPLIT_PART;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.STARTS_WITH;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.STRCMP;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.STR_TO_MAP;
@@ -746,6 +747,7 @@ void populate1() {
           BuiltInMethod.ARRAY_REVERSE.method);
       defineMethod(LEVENSHTEIN, BuiltInMethod.LEVENSHTEIN.method, NullPolicy.STRICT);
       defineMethod(SPLIT, BuiltInMethod.SPLIT.method, NullPolicy.STRICT);
+      defineMethod(SPLIT_PART, BuiltInMethod.SPLIT_PART.method, NullPolicy.STRICT);
       defineReflective(PARSE_URL, BuiltInMethod.PARSE_URL2.method,
           BuiltInMethod.PARSE_URL3.method);
       defineReflective(REGEXP, BuiltInMethod.RLIKE.method);

diff --git a/core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java b/core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java
@@ -914,6 +914,29 @@ public static List<String> split(String s, String delimiter) {
     }
   }
 
+
+  /** SQL {@code SPLIT_PART(string, string, int)} function. */
+  public static String splitPart(String s, String delimiter, int n) {
+    if (s == null || delimiter == null || delimiter.isEmpty() || n == 0) {
+      return null;
+    }
+
+    String[] parts = s.split(Pattern.quote(delimiter), -1);
+    int partCount = parts.length;
+
+    if (n < 0) {
+      n = partCount + n + 1;
+    }
+
+    if (n <= 0 || n > partCount) {
+      return null;
+    }
+
+    return parts[n - 1];
+  }
+
+
+
   /** SQL {@code SPLIT(string)} function. */
   public static List<String> split(String s) {
     return split(s, ",");

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
@@ -389,6 +389,14 @@ static RelDataType deriveTypeSplit(SqlOperatorBinding operatorBinding,
     return operatorBinding.getTypeFactory().createSqlType(typeName);
   }
 
+  /** The "SPLIT_PART(string, delimiter, n)" function. */
+  @LibraryOperator(libraries = {POSTGRESQL})
+  public static final SqlFunction SPLIT_PART =
+      SqlBasicFunction.create("SPLIT_PART",
+          ReturnTypes.ARG0_NULLABLE_VARYING,
+          OperandTypes.STRING_STRING_INTEGER,
+          SqlFunctionCategory.STRING);
+
   /** The "STRPOS(string, substring)" function. */
   @LibraryOperator(libraries = {BIG_QUERY, POSTGRESQL})
   public static final SqlFunction STRPOS = new SqlPositionFunction("STRPOS");

diff --git a/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java b/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java
@@ -414,6 +414,8 @@ public enum BuiltInMethod {
   REPEAT(SqlFunctions.class, "repeat", String.class, int.class),
   SPACE(SqlFunctions.class, "space", int.class),
   SPLIT(SqlFunctions.class, "split", String.class),
+  SPLIT_PART(SqlFunctions.class, "splitPart", String.class, String.class,
+      int.class),
   SOUNDEX(SqlFunctions.class, "soundex", String.class),
   SOUNDEX_SPARK(SqlFunctions.class, "soundexSpark", String.class),
   STRCMP(SqlFunctions.class, "strcmp", String.class, String.class),

diff --git a/core/src/test/java/org/apache/calcite/test/SqlFunctionsTest.java b/core/src/test/java/org/apache/calcite/test/SqlFunctionsTest.java
@@ -1055,6 +1055,22 @@ private void checkCeil(int x, int y, int result) {
         is(list(s, racad, r, empty)));
   }
 
+  @Test void testSplitPart() {
+    assertThat(SqlFunctions.splitPart("abc~@~def~@~ghi", "~@~", 2), is("def"));
+    assertThat(SqlFunctions.splitPart("abc,def,ghi,jkl", ",", -2), is("ghi"));
+
+    assertThat(SqlFunctions.splitPart("abc,,ghi", ",", 2), is(""));
+    assertThat(SqlFunctions.splitPart("", ",", 1), is(""));
+    assertThat(SqlFunctions.splitPart("abc", "", 1), is(nullValue()));
+
+    assertThat(SqlFunctions.splitPart(null, ",", 1), is(nullValue()));
+    assertThat(SqlFunctions.splitPart("abc,def", null, 1), is(nullValue()));
+    assertThat(SqlFunctions.splitPart("abc,def", ",", 0), is(nullValue()));
+
+    assertThat(SqlFunctions.splitPart("abc,def", ",", 3), is(nullValue()));
+    assertThat(SqlFunctions.splitPart("abc,def", ",", -3), is(nullValue()));
+  }
+
   @Test void testByteString() {
     final byte[] bytes = {(byte) 0xAB, (byte) 0xFF};
     final ByteString byteString = new ByteString(bytes);

diff --git a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
@@ -10901,6 +10901,27 @@ void checkEndsWith(SqlOperatorFixture f0, FunctionAlias functionAlias) {
             + "requires extra delimiter argument", false);
   }
 
+  /** Tests the {@code SPLIT_PART} operator. */
+  @Test
+  void testSplitPartFunction() {
+    final SqlOperatorFixture f0 = fixture().setFor(SqlLibraryOperators.SPLIT_PART);
+    f0.checkFails("^split_part('hello', 1)^",
+        "No match found for function signature SPLIT_PART\\(<CHARACTER>, <NUMERIC>\\)",
+        false);
+    final SqlOperatorFixture f = f0.withLibrary(SqlLibrary.POSTGRESQL);
+
+    f.checkScalar("SPLIT_PART('abc~@~def~@~ghi', '~@~', 2)", "def", "VARCHAR(15) NOT NULL");
+    f.checkScalar("SPLIT_PART('abc,def,ghi,jkl', ',', 3)", "ghi", "VARCHAR(15) NOT NULL");
+
+    f.checkScalar("SPLIT_PART('abc~@~def~@~ghi', '~@~', -1)", "ghi", "VARCHAR(15) NOT NULL");
+    f.checkScalar("SPLIT_PART('abc,def,ghi,jkl', ',', -2)", "ghi", "VARCHAR(15) NOT NULL");
+
+    f.checkScalar("SPLIT_PART('h,e,l,l,o', ',', 7)", "", "VARCHAR(9) NOT NULL");
+    f.checkScalar("SPLIT_PART('h,e,l,l,o', ',', -7)", "", "VARCHAR(9) NOT NULL");
+
+    f.checkScalar("SPLIT_PART('abc,,ghi', ',', 2)", "", "VARCHAR(8) NOT NULL");
+    f.checkScalar("SPLIT_PART('', ',', 1)", "", "VARCHAR(0) NOT NULL");
+  }
 
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-5811">[CALCITE-5811]