andygrove commented on code in PR #59:
URL: https://github.com/apache/datafusion-java/pull/59#discussion_r3284025988
##########
core/src/test/java/org/apache/datafusion/ScalarUdfTest.java:
##########
@@ -612,4 +651,189 @@ void volatilityBytesRoundTrip_forAllThreeKinds() throws
Exception {
}
}
}
+
+ // ---------------------------------------------------------------------
+ // Nested-type UDF tests. These pinpoint the regression #58 fixed.
+ // ---------------------------------------------------------------------
+
+ /**
+ * UDF taking a {@code List<Int32>} argument and returning its length as
Int32. Exercises that
+ * nested Arrow types -- whose element / member types live on the parent
{@link Field}'s {@code
+ * children} list, not inside {@link ArrowType} -- can be declared and
registered.
+ */
+ static final class ListLength extends AbstractScalarFunction {
+ ListLength() {
+ super(
+ "java_list_length",
+ List.of(
+ new Field(
+ "vals",
+ FieldType.nullable(new ArrowType.List()),
+ List.of(Field.nullable("item", INT32)))),
+ Field.nullable("len", INT32),
+ Volatility.IMMUTABLE);
+ }
+
+ @Override
+ public ColumnarValue evaluate(BufferAllocator allocator,
ScalarFunctionArgs args) {
+ ListVector in = (ListVector) args.args().get(0).vector();
+ IntVector out = new IntVector("len_out", allocator);
+ int n = in.getValueCount();
+ out.allocateNew(n);
+ for (int i = 0; i < n; i++) {
+ if (in.isNull(i)) {
+ out.setNull(i);
+ } else {
+ out.set(i, in.getElementEndIndex(i) - in.getElementStartIndex(i));
+ }
+ }
+ out.setValueCount(n);
+ return ColumnarValue.array(out);
+ }
+ }
+
+ @Test
+ void udfWithListArg_canBeRegistered() {
+ // Smoke test: registration alone must succeed for nested-type UDFs.
Without #58's fix the
+ // schema-IPC writer rejects List with no children.
+ try (SessionContext ctx = new SessionContext()) {
+ ctx.registerUdf(new ScalarUdf(new ListLength()));
+ }
+ }
+
+ @Test
+ void udfWithListArg_canBeInvokedFromSql() throws Exception {
+ // End-to-end: the registered UDF is callable from SQL with literal list
arguments and the
+ // body sees the right element type. SELECT java_list_length([10, 20, 30])
-> 3.
+ try (SessionContext ctx = new SessionContext();
+ BufferAllocator allocator = new RootAllocator()) {
+ ctx.registerUdf(new ScalarUdf(new ListLength()));
+
+ try (DataFrame df =
+ ctx.sql(
+ "SELECT java_list_length(make_array(CAST(10 AS INT), CAST(20
AS INT),"
+ + " CAST(30 AS INT))) AS n");
+ ArrowReader r = df.collect(allocator)) {
+ assertEquals(true, r.loadNextBatch());
+ IntVector n = (IntVector) r.getVectorSchemaRoot().getVector("n");
+ assertEquals(1, n.getValueCount());
+ assertEquals(3, n.get(0));
+ }
+ }
+ }
+
+ /**
+ * UDF taking a {@code Struct<a: Int32, b: Int32>} and returning the sum
{@code a + b} as Int64.
+ * Confirms the fix is structural rather than List-specific.
+ */
+ static final class SumStructFields extends AbstractScalarFunction {
+ SumStructFields() {
+ super(
+ "java_sum_struct",
+ List.of(
+ new Field(
+ "ab",
+ FieldType.nullable(new ArrowType.Struct()),
+ List.of(Field.nullable("a", INT32), Field.nullable("b",
INT32)))),
+ Field.nullable("s", INT64),
+ Volatility.IMMUTABLE);
+ }
+
+ @Override
+ public ColumnarValue evaluate(BufferAllocator allocator,
ScalarFunctionArgs args) {
+ StructVector in = (StructVector) args.args().get(0).vector();
+ IntVector a = (IntVector) in.getChild("a");
+ IntVector b = (IntVector) in.getChild("b");
+ org.apache.arrow.vector.BigIntVector out =
+ new org.apache.arrow.vector.BigIntVector("sum_out", allocator);
+ int n = in.getValueCount();
+ out.allocateNew(n);
+ for (int i = 0; i < n; i++) {
+ if (in.isNull(i) || a.isNull(i) || b.isNull(i)) {
+ out.setNull(i);
+ } else {
+ out.set(i, (long) a.get(i) + (long) b.get(i));
+ }
+ }
+ out.setValueCount(n);
+ return ColumnarValue.array(out);
+ }
+ }
+
+ @Test
+ void udfWithStructArg_canBeRegistered() {
+ // Struct child fields ride through the same Field children list as List
elements; if the fix
+ // works for List it should work for Struct. This pins that.
+ try (SessionContext ctx = new SessionContext()) {
+ ctx.registerUdf(new ScalarUdf(new SumStructFields()));
+ }
+ }
+
+ @Test
+ void udfWithStructArg_canBeInvokedFromSql() throws Exception {
+ try (SessionContext ctx = new SessionContext();
+ BufferAllocator allocator = new RootAllocator()) {
+ ctx.registerUdf(new ScalarUdf(new SumStructFields()));
+
+ try (DataFrame df =
+ ctx.sql(
+ "SELECT java_sum_struct(named_struct('a', CAST(3 AS INT),
'b', CAST(4 AS INT)))"
+ + " AS s");
+ ArrowReader r = df.collect(allocator)) {
+ assertEquals(true, r.loadNextBatch());
+ org.apache.arrow.vector.BigIntVector s =
+ (org.apache.arrow.vector.BigIntVector)
r.getVectorSchemaRoot().getVector("s");
+ assertEquals(1, s.getValueCount());
+ assertEquals(7L, s.get(0));
+ }
+ }
+ }
+
+ /**
+ * UDF declaring a non-nullable return Field. DataFusion's default {@code
return_field_from_args}
+ * wraps the return type in a fresh always-nullable Field, so without the
{@code
+ * JavaScalarUdf::return_field_from_args} override the planner sees this
UDF's output as nullable
+ * even though the Java caller said otherwise.
+ */
+ static final class NonNullableConstOne extends AbstractScalarFunction {
+ NonNullableConstOne() {
+ super(
+ "java_const_one_nn",
+ List.of(),
+ new Field("v", new FieldType(false, INT32, null), null),
Review Comment:
```suggestion
new Field("v", FieldType.notNullable(INT32)),
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
