Check arg type against Py UDF signature at query compile time

engine/table/src/main/java/io/deephaven/engine/table/impl/lang/QueryLanguageParser.java

@@ -153,7 +153,7 @@ public final class QueryLanguageParser extends GenericVisitorAdapter<Class<?>, Q
     private final Map<String, Class<?>> staticImportLookupCache = new HashMap<>();
 
     // We need some class to represent null. We know for certain that this one won't be used...
-    private static final Class<?> NULL_CLASS = QueryLanguageParser.class;
+    public static final Class<?> NULL_CLASS = QueryLanguageParser.class;
 
     /**
      * The result of the QueryLanguageParser for the expression passed given to the constructor.
@@ -1939,7 +1939,7 @@ private static boolean isAssociativitySafeExpression(Expression expr) {
      * @return {@code true} if a conversion from {@code original} to {@code target} is a widening conversion; otherwise,
      *         {@code false}.
      */
-    static boolean isWideningPrimitiveConversion(Class<?> original, Class<?> target) {
+    public static boolean isWideningPrimitiveConversion(Class<?> original, Class<?> target) {
         if (original == null || !original.isPrimitive() || target == null || !target.isPrimitive()
                 || original.equals(void.class) || target.equals(void.class)) {
             throw new IllegalArgumentException("Arguments must be a primitive type (excluding void)!");
@@ -1968,6 +1968,7 @@ static boolean isWideningPrimitiveConversion(Class<?> original, Class<?> target)
         return false;
     }
 
+
     private enum LanguageParserPrimitiveType {
         // Including "Enum" (or really, any differentiating string) in these names is important. They're used
         // in a switch() statement, which apparently does not support qualified names. And we can't use
@@ -2498,13 +2499,37 @@ public Class<?> visit(MethodCallExpr n, VisitArgs printer) {
 
         // Attempt python function call vectorization.
         if (scopeType != null && PyCallableWrapper.class.isAssignableFrom(scopeType)) {
+            verifyPyCallableArguments(n, argTypes);
             tryVectorizePythonCallable(n, scopeType, convertedArgExpressions, argTypes);
         }
 
         return calculateMethodReturnTypeUsingGenerics(scopeType, n.getScope().orElse(null), method, expressionTypes,
                 typeArguments);
     }
 
+    private void verifyPyCallableArguments(@NotNull MethodCallExpr n, @NotNull Class<?>[] argTypes) {
+        final String invokedMethodName = n.getNameAsString();
+
+        if (GET_ATTRIBUTE_METHOD_NAME.equals(invokedMethodName)) {
+            // Currently Python UDF handling is only supported for top module level function(callable) calls.
+            // The getAttribute() calls which is needed to support Python method calls, which is beyond the scope of
+            // current implementation. So we are skipping the argument verification for getAttribute() calls.
+            return;
+        }
+        if (!n.containsData(QueryLanguageParserDataKeys.PY_CALLABLE_DETAILS)) {
+            return;
+        }
+        final PyCallableDetails pyCallableDetails = n.getData(QueryLanguageParserDataKeys.PY_CALLABLE_DETAILS);
+        final String pyMethodName = pyCallableDetails.pythonMethodName;
+        final Object methodVar = queryScopeVariables.get(pyMethodName);
+        if (!(methodVar instanceof PyCallableWrapper)) {
+            return;
+        }
+        final PyCallableWrapper pyCallableWrapper = (PyCallableWrapper) methodVar;
+        pyCallableWrapper.parseSignature();
+        pyCallableWrapper.verifyArguments(argTypes);
+    }
+
     private Optional<CastExpr> makeCastExpressionForPyCallable(Class<?> retType, MethodCallExpr callMethodCall) {
         if (retType.isPrimitive()) {
             return Optional.of(new CastExpr(
@@ -2552,7 +2577,7 @@ private Optional<Class<?>> pyCallableReturnType(@NotNull MethodCallExpr n) {
         }
         final PyCallableWrapper pyCallableWrapper = (PyCallableWrapper) paramValueRaw;
         pyCallableWrapper.parseSignature();
-        return Optional.ofNullable(pyCallableWrapper.getReturnType());
+        return Optional.ofNullable(pyCallableWrapper.getSignature().getReturnType());
     }
 
     @NotNull
@@ -2683,7 +2708,8 @@ private void checkVectorizability(@NotNull final MethodCallExpr n,
         pyCallableWrapper.parseSignature();
         if (!pyCallableWrapper.isVectorizableReturnType()) {
             throw new PythonCallVectorizationFailure(
-                    "Python function return type is not supported: " + pyCallableWrapper.getReturnType());
+                    "Python function return type is not supported: "
+                            + pyCallableWrapper.getSignature().getReturnType());
         }
 
         // Python vectorized functions(numba, DH) return arrays of primitive/Object types. This will break the generated
@@ -2726,11 +2752,10 @@ private void checkVectorizability(@NotNull final MethodCallExpr n,
             }
         }
 
-        List<Class<?>> paramTypes = pyCallableWrapper.getParamTypes();
-        if (paramTypes.size() != expressions.length) {
+        if (pyCallableWrapper.getSignature().getParameters().size() != expressions.length) {
             // note vectorization doesn't handle Python variadic arguments
             throw new PythonCallVectorizationFailure("Python function argument count mismatch: " + n + " "
-                    + paramTypes.size() + " vs. " + expressions.length);
+                    + pyCallableWrapper.getSignature().getParameters().size() + " vs. " + expressions.length);
         }
     }
 
@@ -2739,10 +2764,9 @@ private void prepareVectorizationArgs(
             Expression[] expressions,
             Class<?>[] argTypes,
             PyCallableWrapper pyCallableWrapper) {
-        List<Class<?>> paramTypes = pyCallableWrapper.getParamTypes();
-        if (paramTypes.size() != expressions.length) {
+        if (pyCallableWrapper.getSignature().getParameters().size() != expressions.length) {
             throw new PythonCallVectorizationFailure("Python function argument count mismatch: " + n + " "
-                    + paramTypes.size() + " vs. " + expressions.length);
+                    + pyCallableWrapper.getSignature().getParameters().size() + " vs. " + expressions.length);
         }
 
         pyCallableWrapper.initializeChunkArguments();
@@ -2763,11 +2787,6 @@ private void prepareVectorizationArgs(
             } else {
                 throw new IllegalStateException("Vectorizability check failed: " + n);
             }
-
-            if (!isSafelyCoerceable(argTypes[i], paramTypes.get(i))) {
-                throw new PythonCallVectorizationFailure("Python vectorized function argument type mismatch: " + n + " "
-                        + argTypes[i].getSimpleName() + " -> " + paramTypes.get(i).getSimpleName());
-            }
         }
     }
 

engine/table/src/main/java/io/deephaven/engine/table/impl/select/AbstractConditionFilter.java

@@ -267,7 +267,7 @@ private void checkAndInitializeVectorization(QueryLanguageParser.Result result,
         final PyCallableWrapperJpyImpl pyCallableWrapper = cws[0];
 
         if (pyCallableWrapper.isVectorizable()) {
-            checkReturnType(result, pyCallableWrapper.getReturnType());
+            checkReturnType(result, pyCallableWrapper.getSignature().getReturnType());
 
             for (String variable : result.getVariablesUsed()) {
                 if (variable.equals("i")) {
@@ -284,7 +284,8 @@ private void checkAndInitializeVectorization(QueryLanguageParser.Result result,
             ArgumentsChunked argumentsChunked = pyCallableWrapper.buildArgumentsChunked(usedColumns);
             PyObject vectorized = pyCallableWrapper.vectorizedCallable();
             DeephavenCompatibleFunction dcf = DeephavenCompatibleFunction.create(vectorized,
-                    pyCallableWrapper.getReturnType(), usedColumns.toArray(new String[0]), argumentsChunked, true);
+                    pyCallableWrapper.getSignature().getReturnType(), usedColumns.toArray(new String[0]),
+                    argumentsChunked, true);
             setFilter(new ConditionFilter.ChunkFilter(
                     dcf.toFilterKernel(),
                     dcf.getColumnNames().toArray(new String[0]),

engine/table/src/main/java/io/deephaven/engine/table/impl/select/DhFormulaColumn.java

@@ -238,7 +238,8 @@ private void checkAndInitializeVectorization(Map<String, ColumnDefinition<?>> co
             PyObject vectorized = pyCallableWrapper.vectorizedCallable();
             formulaColumnPython = FormulaColumnPython.create(this.columnName,
                     DeephavenCompatibleFunction.create(vectorized,
-                            pyCallableWrapper.getReturnType(), this.analyzedFormula.sourceDescriptor.sources,
+                            pyCallableWrapper.getSignature().getReturnType(),
+                            this.analyzedFormula.sourceDescriptor.sources,
                             argumentsChunked,
                             true));
             formulaColumnPython.initDef(columnDefinitionMap);

engine/table/src/main/java/io/deephaven/engine/util/PyCallableWrapper.java

@@ -6,6 +6,7 @@
 import org.jpy.PyObject;
 
 import java.util.List;
+import java.util.Set;
 
 /**
  * Created by rbasralian on 8/12/23
@@ -19,8 +20,6 @@ public interface PyCallableWrapper {
 
     Object call(Object... args);
 
-    List<Class<?>> getParamTypes();
-
     boolean isVectorized();
 
     boolean isVectorizable();
@@ -31,7 +30,46 @@ public interface PyCallableWrapper {
 
     void addChunkArgument(ChunkArgument chunkArgument);
 
-    Class<?> getReturnType();
+    Signature getSignature();
+
+    void verifyArguments(Class<?>[] argTypes);
+
+    class Parameter {
+        private final String name;
+        private final Set<Class<?>> possibleTypes;
+
+
+        public Parameter(String name, Set<Class<?>> possibleTypes) {
+            this.name = name;
+            this.possibleTypes = possibleTypes;
+        }
+
+        public Set<Class<?>> getPossibleTypes() {
+            return possibleTypes;
+        }
+
+        public String getName() {
+            return name;
+        }
+    }
+
+    class Signature {
+        private final List<Parameter> parameters;
+        private final Class<?> returnType;
+
+        public Signature(List<Parameter> parameters, Class<?> returnType) {
+            this.parameters = parameters;
+            this.returnType = returnType;
+        }
+
+        public List<Parameter> getParameters() {
+            return parameters;
+        }
+
+        public Class<?> getReturnType() {
+            return returnType;
+        }
+    }
 
     abstract class ChunkArgument {
         private final Class<?> type;
@@ -88,4 +126,5 @@ public Object getValue() {
     }
 
     boolean isVectorizableReturnType();
+
 }