add array function (#539)

Author: QiZhang1997

geaflow/geaflow-dsl/geaflow-dsl-plan/src/main/java/com/antgroup/geaflow/dsl/schema/function/BuildInSqlFunctionTable.java

@@ -37,6 +37,10 @@
 import com.antgroup.geaflow.dsl.udf.table.agg.PercentileDouble;
 import com.antgroup.geaflow.dsl.udf.table.agg.PercentileInteger;
 import com.antgroup.geaflow.dsl.udf.table.agg.PercentileLong;
+import com.antgroup.geaflow.dsl.udf.table.array.ArrayAppend;
+import com.antgroup.geaflow.dsl.udf.table.array.ArrayContains;
+import com.antgroup.geaflow.dsl.udf.table.array.ArrayDistinct;
+import com.antgroup.geaflow.dsl.udf.table.array.ArrayUnion;
 import com.antgroup.geaflow.dsl.udf.table.date.AddMonths;
 import com.antgroup.geaflow.dsl.udf.table.date.DateAdd;
 import com.antgroup.geaflow.dsl.udf.table.date.DateDiff;
@@ -145,6 +149,12 @@ public class BuildInSqlFunctionTable extends ListSqlOperatorTable {
             .add(GeaFlowFunction.of(WeekOfYear.class))
             .add(GeaFlowFunction.of(Year.class))
 
+            // udf.table.array
+            .add(GeaFlowFunction.of(ArrayAppend.class))
+            .add(GeaFlowFunction.of(ArrayContains.class))
+            .add(GeaFlowFunction.of(ArrayDistinct.class))
+            .add(GeaFlowFunction.of(ArrayUnion.class))
+
             // udf.table.math
             .add(GeaFlowFunction.of(E.class))
             .add(GeaFlowFunction.of(Log2.class))

geaflow/geaflow-dsl/geaflow-dsl-plan/src/main/java/com/antgroup/geaflow/dsl/udf/table/array/ArrayAppend.java

@@ -0,0 +1,36 @@
+/*
+ * 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 com.antgroup.geaflow.dsl.udf.table.array;
+
+import com.antgroup.geaflow.dsl.common.function.Description;
+import com.antgroup.geaflow.dsl.common.function.UDF;
+
+@Description(name = "array_append", description = "Append other element to input array.")
+public class ArrayAppend extends UDF {
+
+    public Object[] eval(Object[] input, Object otherElement) {
+        Object[] res = new Object[input.length + 1];
+        for (int i = 0; i < input.length; i++) {
+            res[i] = input[i];
+        }
+        res[res.length - 1] = otherElement;
+        return res;
+    }
+}

geaflow/geaflow-dsl/geaflow-dsl-plan/src/main/java/com/antgroup/geaflow/dsl/udf/table/array/ArrayContains.java

@@ -0,0 +1,42 @@
+/*
+ * 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 com.antgroup.geaflow.dsl.udf.table.array;
+
+import com.antgroup.geaflow.dsl.common.function.Description;
+import com.antgroup.geaflow.dsl.common.function.UDF;
+
+import java.util.Objects;
+
+@Description(name = "array_contains", description = "Judge if the input element exists in array.")
+public class ArrayContains extends UDF {
+
+    public Boolean eval(Object[] array, Object e) {
+        if (array == null || e == null) {
+            return null;
+        }
+
+        for (Object a : array) {
+            if (Objects.equals(a, e)) {
+                return true;
+            }
+        }
+        return false;
+    }
+}

geaflow/geaflow-dsl/geaflow-dsl-plan/src/main/java/com/antgroup/geaflow/dsl/udf/table/array/ArrayDistinct.java

@@ -0,0 +1,39 @@
+/*
+ * 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 com.antgroup.geaflow.dsl.udf.table.array;
+
+import com.antgroup.geaflow.dsl.common.function.Description;
+import com.antgroup.geaflow.dsl.common.function.UDF;
+
+import java.util.HashSet;
+import java.util.Set;
+
+@Description(name = "array_distinct", description = "Dedup input array.")
+public class ArrayDistinct extends UDF {
+
+    public Object[] eval(Object[] input) {
+        Set<Object> set = new HashSet<>();
+
+        for (Object o : input) {
+            set.add(o);
+        }
+        return set.toArray(new Object[0]);
+    }
+}

geaflow/geaflow-dsl/geaflow-dsl-plan/src/main/java/com/antgroup/geaflow/dsl/udf/table/array/ArrayUnion.java

@@ -0,0 +1,43 @@
+/*
+ * 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 com.antgroup.geaflow.dsl.udf.table.array;
+
+import com.antgroup.geaflow.dsl.common.function.Description;
+import com.antgroup.geaflow.dsl.common.function.UDF;
+import com.google.common.collect.Sets;
+
+import java.util.Arrays;
+import java.util.Set;
+
+@Description(name = "array_union", description = "Union two input array.")
+public class ArrayUnion extends UDF {
+
+    public Object[] eval(Object[] input1, Object[] input2) {
+        if (input1 == null) {
+            return input2;
+        }
+        if (input2 == null) {
+            return input1;
+        }
+        Set<Object> result = Sets.newHashSet(input1);
+        result.addAll(Arrays.asList(input2));
+        return result.toArray();
+    }
+}

geaflow/geaflow-dsl/geaflow-dsl-plan/src/test/java/com/antgroup/geaflow/dsl/udf/array/ArrayTest.java

@@ -0,0 +1,64 @@
+/*
+ * 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 com.antgroup.geaflow.dsl.udf.array;
+
+import com.antgroup.geaflow.dsl.udf.table.array.ArrayAppend;
+import com.antgroup.geaflow.dsl.udf.table.array.ArrayContains;
+import com.antgroup.geaflow.dsl.udf.table.array.ArrayDistinct;
+import com.antgroup.geaflow.dsl.udf.table.array.ArrayUnion;
+import com.antgroup.geaflow.dsl.udf.table.string.RegExpExtract;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class ArrayTest {
+
+    @Test
+    public void testArrayAppend() throws Exception {
+        ArrayAppend udf = new ArrayAppend();
+        Object[] input = new Object[]{1, 2, 4, -1};
+        Object[] res = udf.eval(input, 6);
+        Assert.assertEquals(res, new Object[]{1, 2, 4, -1, 6});
+    }
+
+    @Test
+    public void testArrayContains() throws Exception {
+        ArrayContains udf = new ArrayContains();
+        Object[] input = new Object[]{1, 2, 4, -1};
+        Assert.assertTrue(udf.eval(input, 2));
+        Assert.assertFalse(udf.eval(input, -3));
+    }
+
+    @Test
+    public void testArrayDistinct() throws Exception {
+        ArrayDistinct udf = new ArrayDistinct();
+        Object[] input = new Object[]{1, 2, 4, -1, 2, 1};
+        Object[] res = udf.eval(input);
+        Assert.assertEquals(res.length, 4);
+    }
+
+    @Test
+    public void testArrayUnion() throws Exception {
+        ArrayUnion udf = new ArrayUnion();
+        Object[] input1 = new Object[]{1, 2, 4, -1};
+        Object[] input2 = new Object[]{1, 3, 5, 2, -4, -6};
+        Object[] res = udf.eval(input1, input2);
+        Assert.assertEquals(res.length, 8);
+    }
+}