Return list parameter type information through the function registry.

dremio · Nov 17, 2023 · b6593b3 · b6593b3
1 parent b7c5288
commit b6593b3
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 7 deletions.
diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ExpressionRegistry.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ExpressionRegistry.java
@@ -17,9 +17,11 @@
 
 package org.apache.arrow.gandiva.evaluator;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
 
+import org.apache.arrow.flatbuf.Type;
 import org.apache.arrow.gandiva.exceptions.GandivaException;
 import org.apache.arrow.gandiva.ipc.GandivaTypes;
 import org.apache.arrow.gandiva.ipc.GandivaTypes.ExtGandivaType;
@@ -32,7 +34,6 @@
 import org.apache.arrow.vector.types.TimeUnit;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 
-import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import com.google.protobuf.InvalidProtocolBufferException;
 
@@ -117,9 +118,16 @@ private static Set<FunctionSignature> getSupportedFunctionsFromGandiva() throws
         String functionName = protoFunctionSignature.getName();
         ArrowType returnType = getArrowType(protoFunctionSignature.getReturnType());
         ArrowType returnListType = getArrowTypeSimple(protoFunctionSignature.getReturnType().getListType());
-        List<ArrowType> paramTypes = Lists.newArrayList();
+        List<List<ArrowType>> paramTypes = new ArrayList<List<ArrowType>>();
         for (ExtGandivaType type : protoFunctionSignature.getParamTypesList()) {
-          paramTypes.add(getArrowType(type));
+          ArrowType paramType = getArrowType(type);
+          ArrowType paramListType = getArrowTypeSimple(type.getListType());
+          List<ArrowType> paramArrowList = new ArrayList<ArrowType>();
+          paramArrowList.add(paramType);
+          if (paramType.getTypeID().getFlatbufID() == Type.List) {
+            paramArrowList.add(paramListType);
+          }
+          paramTypes.add(paramArrowList);
         }
         FunctionSignature functionSignature = new FunctionSignature(functionName,
                                                                     returnType, returnListType, paramTypes);

diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/FunctionSignature.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/FunctionSignature.java
@@ -17,6 +17,7 @@
 
 package org.apache.arrow.gandiva.evaluator;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -31,7 +32,7 @@ public class FunctionSignature {
   private final String name;
   private final ArrowType returnType;
   private final ArrowType returnListType;
-  private final List<ArrowType> paramTypes;
+  private final List<List<ArrowType>> paramTypes;
 
   public ArrowType getReturnType() {
     return returnType;
@@ -41,7 +42,7 @@ public ArrowType getReturnListType() {
     return returnListType;
   }
 
-  public List<ArrowType> getParamTypes() {
+  public List<List<ArrowType>> getParamTypes() {
     return paramTypes;
   }
 
@@ -56,7 +57,8 @@ public String getName() {
    * @param returnListType optional list type
    * @param paramTypes - data type of input args.
    */
-  public FunctionSignature(String name, ArrowType returnType, ArrowType returnListType, List<ArrowType> paramTypes) {
+  public FunctionSignature(String name, ArrowType returnType, ArrowType returnListType, 
+      List<List<ArrowType>> paramTypes) {
     this.name = name;
     this.returnType = returnType;
     this.returnListType = returnListType;
@@ -73,7 +75,13 @@ public FunctionSignature(String name, ArrowType returnType, List<ArrowType> para
     this.name = name;
     this.returnType = returnType;
     this.returnListType = ArrowType.Null.INSTANCE;
-    this.paramTypes = paramTypes;
+    this.paramTypes = new ArrayList<List<ArrowType>>();
+    for (ArrowType paramType : paramTypes) {
+      List<ArrowType> paramArrowList = new ArrayList<ArrowType>();
+      paramArrowList.add(paramType);
+      this.paramTypes.add(paramArrowList);
+    }
+
   }
 
   /**