From 0253af6cd2684f0f7719507fb45702b56cbb9a58 Mon Sep 17 00:00:00 2001
From: wenchengyao <wenchengyao.wcy@antgroup.com>
Date: Tue, 29 Oct 2024 14:46:38 +0800
Subject: [PATCH] feat(reasoner): add udf split_part (#378)

---
 .../reasoner/udf/builtin/udf/SplitPart.java   | 32 +++++++++++++++++++
 .../openspg/reasoner/udf/test/UdfTest.java    | 17 ++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 reasoner/udf/src/main/java/com/antgroup/openspg/reasoner/udf/builtin/udf/SplitPart.java

diff --git a/reasoner/udf/src/main/java/com/antgroup/openspg/reasoner/udf/builtin/udf/SplitPart.java b/reasoner/udf/src/main/java/com/antgroup/openspg/reasoner/udf/builtin/udf/SplitPart.java
new file mode 100644
index 00000000..dc1f5dc2
--- /dev/null
+++ b/reasoner/udf/src/main/java/com/antgroup/openspg/reasoner/udf/builtin/udf/SplitPart.java
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2023 OpenSPG Authors
+ *
+ * Licensed 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.
+ */
+
+package com.antgroup.openspg.reasoner.udf.builtin.udf;
+
+import com.antgroup.openspg.reasoner.udf.model.UdfDefine;
+import org.apache.commons.lang3.StringUtils;
+
+public class SplitPart {
+  @UdfDefine(name = "split_part", compatibleName = "SplitPart")
+  public String split(String str, String separator, Integer partNum) throws Exception {
+    String[] strings = StringUtils.split(str, separator);
+    int length = strings.length;
+    if (partNum < 0) {
+      partNum = length - 1;
+    }
+    if (partNum >= length) {
+      return "";
+    }
+    return strings[partNum];
+  }
+}
diff --git a/reasoner/udf/src/test/java/com/antgroup/openspg/reasoner/udf/test/UdfTest.java b/reasoner/udf/src/test/java/com/antgroup/openspg/reasoner/udf/test/UdfTest.java
index 82d0cb24..201076b1 100644
--- a/reasoner/udf/src/test/java/com/antgroup/openspg/reasoner/udf/test/UdfTest.java
+++ b/reasoner/udf/src/test/java/com/antgroup/openspg/reasoner/udf/test/UdfTest.java
@@ -169,6 +169,23 @@ public class UdfTest {
     Assert.assertEquals(rst2, "岁");
   }
 
+  @Test
+  public void testSplitPart() {
+    UdfMng mng = UdfMngFactory.getUdfMng();
+    IUdfMeta udfMeta =
+        mng.getUdfMeta(
+            "split_part",
+            Lists.newArrayList(KTString$.MODULE$, KTString$.MODULE$, KTInteger$.MODULE$));
+    Object rst1 = udfMeta.invoke("Hello,World!", ",", 0);
+    Assert.assertEquals("Hello", rst1);
+    Object rst2 = udfMeta.invoke("Hello,Ni,Hao", ",", -1);
+    Assert.assertEquals("Hao", rst2);
+    Object rst3 = udfMeta.invoke("A省B市C村XXX", "村", 0);
+    Assert.assertEquals("A省B市C", rst3);
+    Object rst4 = udfMeta.invoke("A省B市C村XXX", "村", 5);
+    Assert.assertEquals("", rst4);
+  }
+
   @Test
   public void testCast() {
     UdfMng mng = UdfMngFactory.getUdfMng();