feat(udf)!: switch ScalarFunction.evaluate to ColumnarValue API (closes #62) (#64)

Author: andygrove

core/src/main/java/org/apache/datafusion/ColumnarValue.java

@@ -0,0 +1,81 @@
+/*
+ * 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 org.apache.datafusion;
+
+import java.util.Objects;
+
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+
+/**
+ * The value of a scalar UDF argument or result: either a per-row {@link Array} of length {@code
+ * rowCount}, or a {@link Scalar} (length-1 vector) that the framework broadcasts.
+ *
+ * <p>Mirrors DataFusion's {@code datafusion::logical_expr::ColumnarValue} enum. Use {@link
+ * #array(FieldVector)} and {@link #scalar(FieldVector)} factories rather than constructing the
+ * records directly so the length invariants are enforced consistently.
+ */
+public sealed interface ColumnarValue permits ColumnarValue.Array, ColumnarValue.Scalar {
+
+  /** The underlying Arrow vector. For {@link Scalar} this vector has {@code valueCount == 1}. */
+  FieldVector vector();
+
+  /** Convenience: the vector's declared Arrow type. */
+  default ArrowType dataType() {
+    return vector().getField().getType();
+  }
+
+  /** Wrap an arbitrary-length vector as an {@link Array}. */
+  static ColumnarValue array(FieldVector vector) {
+    return new Array(Objects.requireNonNull(vector, "vector"));
+  }
+
+  /**
+   * Wrap a length-1 vector as a {@link Scalar}.
+   *
+   * @throws IllegalArgumentException if {@code vector.getValueCount() != 1}
+   */
+  static ColumnarValue scalar(FieldVector vector) {
+    Objects.requireNonNull(vector, "vector");
+    if (vector.getValueCount() != 1) {
+      throw new IllegalArgumentException(
+          "Scalar vector must have valueCount == 1, got " + vector.getValueCount());
+    }
+    return new Scalar(vector);
+  }
+
+  /** Per-row Arrow vector of length equal to the batch row count. */
+  record Array(FieldVector vector) implements ColumnarValue {
+    public Array {
+      Objects.requireNonNull(vector, "vector");
+    }
+  }
+
+  /** Length-1 Arrow vector representing a single value broadcast across all rows. */
+  record Scalar(FieldVector vector) implements ColumnarValue {
+    public Scalar {
+      Objects.requireNonNull(vector, "vector");
+      if (vector.getValueCount() != 1) {
+        throw new IllegalArgumentException(
+            "Scalar vector must have valueCount == 1, got " + vector.getValueCount());
+      }
+    }
+  }
+}

core/src/main/java/org/apache/datafusion/ScalarFunction.java

@@ -22,7 +22,6 @@
 import java.util.List;
 
 import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 
 /**
@@ -46,7 +45,9 @@ public interface ScalarFunction {
    */
   List<ArrowType> argTypes();
 
-  /** Declared return type. The returned {@link FieldVector} must have this exact type. */
+  /**
+   * Declared return type. The returned {@link ColumnarValue}'s vector must have this exact type.
+   */
   ArrowType returnType();
 
   /**
@@ -59,14 +60,16 @@ public interface ScalarFunction {
   /**
    * Compute the function result for one input batch.
    *
-   * @param allocator the {@link BufferAllocator} that MUST be used for any new {@link FieldVector}
+   * @param allocator the {@link BufferAllocator} that MUST be used for any new Arrow vector
    *     allocation, including the result. Buffers allocated from other allocators will not survive
    *     the JNI handoff.
-   * @param args one {@link FieldVector} per declared argument, all of the same length. These are
-   *     read-only views; the implementation must NOT close them.
-   * @return a {@link FieldVector} of the declared return type and the same length as the inputs.
-   *     Ownership transfers to the framework on return; the implementation must NOT close the
-   *     returned vector.
+   * @param args the per-arg {@link ColumnarValue}s and the batch row count. Each {@link
+   *     ColumnarValue} is a read-only view; the implementation must NOT close its underlying
+   *     vector.
+   * @return a {@link ColumnarValue} of the declared return type. If {@link ColumnarValue.Array},
+   *     the underlying vector must have length {@code args.rowCount()}; if {@link
+   *     ColumnarValue.Scalar}, length 1. Ownership of the returned vector transfers to the
+   *     framework; the implementation must NOT close it.
    */
-  FieldVector evaluate(BufferAllocator allocator, List<FieldVector> args);
+  ColumnarValue evaluate(BufferAllocator allocator, ScalarFunctionArgs args);
 }

core/src/main/java/org/apache/datafusion/ScalarFunctionArgs.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 org.apache.datafusion;
+
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Bundle of inputs passed to {@link ScalarFunction#evaluate}: the per-arg {@link ColumnarValue}s
+ * (in declared order) and the batch row count DataFusion is driving.
+ *
+ * <p>Mirrors DataFusion's {@code datafusion::logical_expr::ScalarFunctionArgs}. {@code rowCount} is
+ * the only channel by which an array-returning UDF without array-typed inputs (all-scalar args, or
+ * nullary) can size its output. Nullary UDFs that prefer to broadcast a single value should return
+ * {@link ColumnarValue#scalar(org.apache.arrow.vector.FieldVector) ColumnarValue.scalar(...)}
+ * instead, which removes the need to consult {@code rowCount}.
+ */
+public record ScalarFunctionArgs(List<ColumnarValue> args, int rowCount) {
+  public ScalarFunctionArgs {
+    args = List.copyOf(Objects.requireNonNull(args, "args"));
+    if (rowCount < 0) {
+      throw new IllegalArgumentException("rowCount must be >= 0, got " + rowCount);
+    }
+  }
+}

core/src/main/java/org/apache/datafusion/internal/JniBridge.java

@@ -19,6 +19,7 @@
 
 package org.apache.datafusion.internal;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import org.apache.arrow.c.ArrowArray;
@@ -27,7 +28,9 @@
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.datafusion.ColumnarValue;
 import org.apache.datafusion.ScalarFunction;
+import org.apache.datafusion.ScalarFunctionArgs;
 
 /** Internal trampoline invoked from native code on every UDF call. Not part of the public API. */
 public final class JniBridge {
@@ -40,54 +43,100 @@ public final class JniBridge {
 
   private JniBridge() {}
 
+  /** argKind byte signalling a {@link ColumnarValue.Array} arg. */
+  private static final byte KIND_ARRAY = 0;
+
+  /** argKind byte signalling a {@link ColumnarValue.Scalar} arg. */
+  private static final byte KIND_SCALAR = 1;
+
   /**
    * Invoke a scalar UDF for one batch. Called from native code; not for application use.
    *
-   * @param impl the registered {@link ScalarFunction} implementation
-   * @param argsArrayAddr address of a populated {@code FFI_ArrowArray} struct holding the input
-   *     batch as a struct array (one field per UDF argument)
-   * @param argsSchemaAddr address of the matching {@code FFI_ArrowSchema}
-   * @param resultArrayAddr address of an empty {@code FFI_ArrowArray} the bridge writes into
-   * @param resultSchemaAddr address of an empty {@code FFI_ArrowSchema} the bridge writes into
-   * @param expectedRowCount the row count the result vector must have
+   * <p>Args arrive split into two struct arrays: {@code arrayArgs*} of length {@code rowCount}
+   * holding the {@link ColumnarValue.Array} arguments in their relative order, and {@code
+   * scalarArgs*} of length 1 holding the {@link ColumnarValue.Scalar} arguments. {@code argKinds}
+   * records the original positional order so the bridge can interleave them back into a single
+   * {@code List<ColumnarValue>} for the user.
+   *
+   * @return {@link #KIND_ARRAY} if the UDF returned an Array, {@link #KIND_SCALAR} if it returned a
+   *     Scalar. The native caller uses this to reconstruct the right {@code ColumnarValue} variant.
    */
-  public static void invokeScalarUdf(
+  public static byte invokeScalarUdf(
       ScalarFunction impl,
-      long argsArrayAddr,
-      long argsSchemaAddr,
+      long arrayArgsArrayAddr,
+      long arrayArgsSchemaAddr,
+      long scalarArgsArrayAddr,
+      long scalarArgsSchemaAddr,
+      byte[] argKinds,
       long resultArrayAddr,
       long resultSchemaAddr,
-      int expectedRowCount) {
-    ArrowArray argsArr = ArrowArray.wrap(argsArrayAddr);
-    ArrowSchema argsSch = ArrowSchema.wrap(argsSchemaAddr);
+      int rowCount) {
+    ArrowArray arrayArr = ArrowArray.wrap(arrayArgsArrayAddr);
+    ArrowSchema arraySch = ArrowSchema.wrap(arrayArgsSchemaAddr);
+    ArrowArray scalarArr = ArrowArray.wrap(scalarArgsArrayAddr);
+    ArrowSchema scalarSch = ArrowSchema.wrap(scalarArgsSchemaAddr);
     ArrowArray resultArr = ArrowArray.wrap(resultArrayAddr);
     ArrowSchema resultSch = ArrowSchema.wrap(resultSchemaAddr);
 
-    try (VectorSchemaRoot root = Data.importVectorSchemaRoot(ALLOCATOR, argsArr, argsSch, null)) {
-      List<FieldVector> argVectors = root.getFieldVectors();
+    try (VectorSchemaRoot arrayRoot =
+            Data.importVectorSchemaRoot(ALLOCATOR, arrayArr, arraySch, null);
+        VectorSchemaRoot scalarRoot =
+            Data.importVectorSchemaRoot(ALLOCATOR, scalarArr, scalarSch, null)) {
 
-      FieldVector result = impl.evaluate(ALLOCATOR, argVectors);
+      List<FieldVector> arrayFields = arrayRoot.getFieldVectors();
+      List<FieldVector> scalarFields = scalarRoot.getFieldVectors();
+
+      List<ColumnarValue> args = new ArrayList<>(argKinds.length);
+      int arrayIdx = 0;
+      int scalarIdx = 0;
+      for (byte kind : argKinds) {
+        if (kind == KIND_ARRAY) {
+          args.add(ColumnarValue.array(arrayFields.get(arrayIdx++)));
+        } else if (kind == KIND_SCALAR) {
+          args.add(ColumnarValue.scalar(scalarFields.get(scalarIdx++)));
+        } else {
+          throw new IllegalStateException("Unknown argKind byte: " + kind);
+        }
+      }
+
+      ColumnarValue result = impl.evaluate(ALLOCATOR, new ScalarFunctionArgs(args, rowCount));
 
       if (result == null) {
         throw new IllegalStateException("ScalarFunction.evaluate returned null");
       }
-      if (result.getValueCount() != expectedRowCount) {
+
+      FieldVector resultVec = result.vector();
+      byte resultKind;
+      int expectedLen;
+      if (result instanceof ColumnarValue.Array) {
+        resultKind = KIND_ARRAY;
+        expectedLen = rowCount;
+      } else {
+        resultKind = KIND_SCALAR;
+        expectedLen = 1;
+      }
+
+      if (resultVec.getValueCount() != expectedLen) {
         try {
           throw new IllegalStateException(
-              "ScalarFunction.evaluate returned vector with "
-                  + result.getValueCount()
+              "ScalarFunction.evaluate returned "
+                  + (resultKind == KIND_ARRAY ? "Array" : "Scalar")
+                  + " vector with "
+                  + resultVec.getValueCount()
                   + " rows; expected "
-                  + expectedRowCount);
+                  + expectedLen);
         } finally {
-          result.close();
+          resultVec.close();
         }
       }
 
       try {
-        Data.exportVector(ALLOCATOR, result, null, resultArr, resultSch);
+        Data.exportVector(ALLOCATOR, resultVec, null, resultArr, resultSch);
       } finally {
-        result.close();
+        resultVec.close();
       }
+
+      return resultKind;
     }
   }
 }

core/src/test/java/org/apache/datafusion/ColumnarValueTest.java

@@ -0,0 +1,82 @@
+/*
+ * 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 org.apache.datafusion;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.junit.jupiter.api.Test;
+
+class ColumnarValueTest {
+
+  private static final ArrowType INT32 = new ArrowType.Int(32, true);
+
+  @Test
+  void array_factory_returnsArrayVariant() {
+    try (BufferAllocator allocator = new RootAllocator();
+        IntVector v = new IntVector("v", allocator)) {
+      v.allocateNew(3);
+      v.setValueCount(3);
+      ColumnarValue cv = ColumnarValue.array(v);
+      assertSame(v, cv.vector());
+      assertEquals(INT32, cv.dataType());
+    }
+  }
+
+  @Test
+  void scalar_factory_returnsScalarVariant() {
+    try (BufferAllocator allocator = new RootAllocator();
+        IntVector v = new IntVector("v", allocator)) {
+      v.allocateNew(1);
+      v.set(0, 42);
+      v.setValueCount(1);
+      ColumnarValue cv = ColumnarValue.scalar(v);
+      assertSame(v, cv.vector());
+      assertEquals(INT32, cv.dataType());
+    }
+  }
+
+  @Test
+  void scalar_factory_rejectsNonOneLength() {
+    try (BufferAllocator allocator = new RootAllocator();
+        IntVector v = new IntVector("v", allocator)) {
+      v.allocateNew(2);
+      v.setValueCount(2);
+      IllegalArgumentException ex =
+          assertThrows(IllegalArgumentException.class, () -> ColumnarValue.scalar(v));
+      assertEquals("Scalar vector must have valueCount == 1, got 2", ex.getMessage());
+    }
+  }
+
+  @Test
+  void array_factory_rejectsNull() {
+    assertThrows(NullPointerException.class, () -> ColumnarValue.array(null));
+  }
+
+  @Test
+  void scalar_factory_rejectsNull() {
+    assertThrows(NullPointerException.class, () -> ColumnarValue.scalar(null));
+  }
+}