Skip to content

Commit

Permalink
support regexp_replace function
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-bli committed Sep 19, 2023
1 parent 265c417 commit bcdd884
Show file tree
Hide file tree
Showing 4 changed files with 81 additions and 0 deletions.
32 changes: 32 additions & 0 deletions src/main/java/com/snowflake/snowpark_java/Functions.java
Original file line number Diff line number Diff line change
Expand Up @@ -2354,6 +2354,38 @@ public static Column regexp_count(Column strExpr, Column pattern) {
strExpr.toScalaColumn(), pattern.toScalaColumn()));
}

/**
* Returns the subject with the specified pattern (or all occurrences of the pattern) removed. If
* no matches are found, returns the original subject.
*
* @param strExpr The input string
* @param pattern The pattern
* @return The result column
* @since 1.9.0
*/
public static Column regexp_replace(Column strExpr, Column pattern) {
return new Column(
com.snowflake.snowpark.functions.regexp_replace(
strExpr.toScalaColumn(), pattern.toScalaColumn()));
}

/**
* Returns the subject with the specified pattern (or all occurrences of the pattern)
* replaced by a replacement string. If no matches are found, returns the original
* subject.
*
* @param strExpr The input string
* @param pattern The pattern
* @param replacement The replacement string
* @return The result column
* @since 1.9.0
*/
public static Column regexp_replace(Column strExpr, Column pattern, Column replacement) {
return new Column(
com.snowflake.snowpark.functions.regexp_replace(
strExpr.toScalaColumn(), pattern.toScalaColumn(), replacement.toScalaColumn()));
}

/**
* Removes all occurrences of a specified strExpr, and optionally replaces them with replacement.
*
Expand Down
22 changes: 22 additions & 0 deletions src/main/scala/com/snowflake/snowpark/functions.scala
Original file line number Diff line number Diff line change
Expand Up @@ -1796,6 +1796,28 @@ object functions {
def regexp_count(strExpr: Column, pattern: Column): Column =
builtin("regexp_count")(strExpr, pattern)

/**
* Returns the subject with the specified pattern (or all occurrences of the pattern) removed.
* If no matches are found, returns the original subject.
*
* @group str_func
* @since 1.9.0
*/
def regexp_replace(strExpr: Column, pattern: Column): Column =
builtin("regexp_replace")(strExpr, pattern)

/**
* Returns the subject with the specified pattern (or all occurrences of the pattern)
* replaced by a replacement string. If no matches are found,
* returns the original subject.
*
* @group str_func
* @since 1.9.0
*/
def regexp_replace(strExpr: Column, pattern: Column, replacement: Column): Column =
builtin("regexp_replace")(strExpr, pattern, replacement)


/**
* Removes all occurrences of a specified strExpr,
* and optionally replaces them with replacement.
Expand Down
13 changes: 13 additions & 0 deletions src/test/java/com/snowflake/snowpark_test/JavaFunctionSuite.java
Original file line number Diff line number Diff line change
Expand Up @@ -1447,6 +1447,19 @@ public void regexp_count() {
df.select(Functions.regexp_count(df.col("a"), Functions.lit("a"))), expected1, false);
}

@Test
public void regexp_replace() {
DataFrame df = getSession().sql("select * from values('cat'),('dog'),('mouse') as T(a)");
Column pattern = Functions.lit("^ca|^[m|d]o");
Row[] expected = {Row.create("t"), Row.create("g"), Row.create("use")};
checkAnswer(df.select(Functions.regexp_replace(df.col("a"), pattern)), expected, false);

Column replacement = Functions.lit("ch");
Row[] expected1 = {Row.create("cht"), Row.create("chg"), Row.create("chuse")};
checkAnswer(
df.select(Functions.regexp_replace(df.col("a"), pattern, replacement)), expected1, false);
}

@Test
public void replace() {
DataFrame df = getSession().sql("select * from values('apple'),('banana'),('peach') as T(a)");
Expand Down
14 changes: 14 additions & 0 deletions src/test/scala/com/snowflake/snowpark_test/FunctionSuite.scala
Original file line number Diff line number Diff line change
Expand Up @@ -2144,6 +2144,20 @@ trait FunctionSuite extends TestData {
Seq(Row("1'2'3'4'5")))
}

test("regexp_replace") {
val data = Seq("cat", "dog", "mouse").toDF("a")
val pattern = lit("^ca|^[m|d]o")
var expected = Seq(Row("t"), Row("g"), Row("use"))
checkAnswer(data.select(regexp_replace(data("a"), pattern)), expected, sort = false)

val replacement = lit("ch")
expected = Seq(Row("cht"), Row("chg"), Row("chuse"))
checkAnswer(
data.select(regexp_replace(data("a"), pattern, replacement)),
expected,
sort = false)
}

}

class EagerFunctionSuite extends FunctionSuite with EagerSession
Expand Down

0 comments on commit bcdd884

Please sign in to comment.