From dd8e97da95507ba906e7a355af82eb6d676f6ffa Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 12:17:07 +0200 Subject: [PATCH 1/3] Fix Beam SQL UNNEST of array elements with a nested row through JDBC (#26911) --- CHANGES.md | 1 + .../sql/impl/rel/BeamUnnestRel.java | 82 +++++++++---------- .../sql/BeamSqlDslUnnestRowsTest.java | 41 ++++++++++ .../extensions/sql/impl/JdbcDriverTest.java | 53 ++++++++++++ 4 files changed, 136 insertions(+), 41 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 04ac961fcd93..a671170737cf 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -142,6 +142,7 @@ * (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)). * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)). * (Go) Fixed generic registration (for example `register.Function1x1` or `register.DoFn1x1`) panicking on functions and methods that take or return an anonymous struct. Unnamed struct types are now only registered with the schema registry, since they have no type key for the runtime type registry ([#21789](https://github.com/apache/beam/issues/21789)). +* (Java) Beam SQL: `CROSS JOIN UNNEST` over an array of rows that contain a nested row no longer fails with a `Row expected N fields ... initialized with M fields` error when the query runs through the JDBC driver or the SQL shell ([#26911](https://github.com/apache/beam/issues/26911)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java index dc97ee7cf58c..e73408a84780 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java @@ -20,10 +20,7 @@ import java.util.ArrayList; import java.util.Collection; import java.util.List; -import java.util.function.Function; -import java.util.stream.Collectors; -import java.util.stream.IntStream; -import java.util.stream.Stream; +import java.util.Objects; import org.apache.beam.sdk.extensions.sql.impl.planner.BeamCostModel; import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery; import org.apache.beam.sdk.extensions.sql.impl.planner.NodeStats; @@ -136,32 +133,41 @@ private UnnestFn(Schema outputSchema, List unnestIndices) { } /** - * This is recursive call to get all the values of the nested rows. The recursion is bounded by - * the amount of nesting with in the data. This mirrors the unnest behavior of calcite towards - * schema. * + * Appends the fields of an unnested array element to {@code values}, laid out the way {@code + * outputSchema} expects them. The recursion is bounded by the nesting depth of the element + * schema. + * + *

Whether a nested row field of the element arrives as one ROW column or as one column per + * leaf field depends on which Calcite planner produced the plan. For an element {@code ROW>}, the planner behind {@code SqlTransform} flattens the element into + * {@code id, name, email}, while the planner behind the JDBC driver and the SQL shell keeps + * {@code id, author}. Building the values from the element schema alone matches only one of + * them, and the other then fails in {@link Row.Builder#build()} with a field count mismatch. So + * each nested row is checked against the output column at the same position: a ROW column takes + * the row as is, any other column means the row was flattened. */ - private List getNestedRowBaseValues(Row nestedRow) { - return IntStream.range(0, nestedRow.getFieldCount()) - .mapToObj( - (i) -> { - List values = new ArrayList<>(); - Schema.FieldType fieldType = nestedRow.getSchema().getField(i).getType(); - if (fieldType.getTypeName().equals(Schema.TypeName.ROW)) { - @Nullable Row row = nestedRow.getBaseValue(i, Row.class); - if (row == null) { - return Stream.builder().build(); - } - List rowValues = getNestedRowBaseValues(row); - if (null != rowValues) { - values.addAll(rowValues); - } - } else { - values.add(nestedRow.getBaseValue(i)); - } - return values.stream(); - }) - .flatMap(Function.identity()) - .collect(Collectors.toList()); + private void addElementValues( + Schema elementSchema, @Nullable Row element, List values) { + for (int i = 0; i < elementSchema.getFieldCount(); i++) { + Schema.FieldType fieldType = elementSchema.getField(i).getType(); + Object value = null; + if (Objects.nonNull(element)) { + value = element.getBaseValue(i); + } + + boolean isRow = fieldType.getTypeName().equals(Schema.TypeName.ROW); + if (!isRow || isRowColumn(values.size())) { + values.add(value); + continue; + } + + // a null nested row still fills one null per flattened column + addElementValues(fieldType.getRowSchema(), (Row) value, values); + } + } + + private boolean isRowColumn(int index) { + return outputSchema.getField(index).getType().getTypeName().equals(Schema.TypeName.ROW); } @ProcessElement @@ -178,21 +184,15 @@ public void process(@Element Row row, OutputReceiver out) { if (rawValues == null) { return; } - Schema.TypeName typeName = - schemaWithArrayField - .getField(unnestIndices.get(0)) - .getType() - .getCollectionElementType() - .getTypeName(); + Schema.FieldType elementType = + schemaWithArrayField.getField(unnestIndices.get(0)).getType().getCollectionElementType(); + Schema.TypeName typeName = elementType.getTypeName(); for (Object uncollectedValue : rawValues) { if (typeName.equals(Schema.TypeName.ROW)) { - Row nestedRow = (Row) uncollectedValue; - out.output( - Row.withSchema(outputSchema) - .addValues(row.getBaseValues()) - .addValues(getNestedRowBaseValues(nestedRow)) - .build()); + List values = new ArrayList<>(row.getBaseValues()); + addElementValues(elementType.getRowSchema(), (Row) uncollectedValue, values); + out.output(Row.withSchema(outputSchema).addValues(values).build()); } else { out.output( Row.withSchema(outputSchema) diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUnnestRowsTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUnnestRowsTest.java index d72bc249154a..cf4e465124a9 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUnnestRowsTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUnnestRowsTest.java @@ -139,4 +139,45 @@ public void testUnnestArrayWithNestedRows() { .build()); pipeline.run(); } + + @Test + public void testUnnestArrayWithNullNestedRow() { + Schema nestedType = Schema.builder().addInt32Field("c1").addStringField("c2").build(); + Schema elementType = + Schema.builder() + .addInt32Field("b1") + .addNullableField("b2", Schema.FieldType.row(nestedType)) + .build(); + Schema inputType = + Schema.builder() + .addInt32Field("a1") + .addArrayField("a2", Schema.FieldType.row(elementType)) + .build(); + + Row input = + Row.withSchema(inputType) + .addValues( + 1, + Arrays.asList( + Row.withSchema(elementType) + .addValues(1, Row.withSchema(nestedType).addValues(10, "x").build()) + .build(), + Row.withSchema(elementType).addValues(2, null).build())) + .build(); + PCollection inputTable = + PBegin.in(pipeline).apply(Create.of(input).withRowSchema(inputType)); + + PCollection result = + PCollectionTuple.of(new TupleTag<>("test"), inputTable) + .apply( + SqlTransform.query("select d.b1, d.b2.c1 from test t cross join unnest(t.a2) d")); + + Schema resultSchema = + Schema.builder().addInt32Field("b1").addNullableField("c1", Schema.FieldType.INT32).build(); + PAssert.that(result) + .containsInAnyOrder( + Row.withSchema(resultSchema).addValues(1, 10).build(), + Row.withSchema(resultSchema).addValues(2, null).build()); + pipeline.run(); + } } diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java index b022d94777ba..fa78a1e5c622 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java @@ -351,6 +351,59 @@ public void testSelectsFromExistingComplexTable() throws Exception { assertThat(resultRows, containsInAnyOrder(row(1L, "aaa"), row(2L, "bbb"))); } + @Test + public void testUnnestArrayOfRowsWithNestedRow() throws Exception { + TestTableProvider tableProvider = new TestTableProvider(); + Connection connection = JdbcDriver.connect(tableProvider, PipelineOptionsFactory.create()); + + connection + .createStatement() + .executeUpdate( + "CREATE EXTERNAL TABLE person ( \n" + + "name VARCHAR, \n" + + "items ARRAY>> \n" + + ") \n" + + "TYPE 'test'"); + + Schema tableSchema = + Schema.builder() + .addNullableField("name", Schema.FieldType.STRING) + .addNullableField("items", Schema.FieldType.array(Schema.FieldType.row(COMPLEX_SCHEMA))) + .build(); + tableProvider.addRows( + "person", + row( + tableSchema, + "person1", + Arrays.asList( + row(COMPLEX_SCHEMA, "description1", row(1L, "aaa")), + row(COMPLEX_SCHEMA, "description2", row(2L, "bbb"))))); + + ResultSet selectResult = + connection + .createStatement() + .executeQuery( + "SELECT item.description, item.nestedRow.name " + + "FROM person CROSS JOIN UNNEST(person.items) AS item"); + + Schema resultSchema = + Schema.builder() + .addNullableField("description", Schema.FieldType.STRING) + .addNullableField("name", Schema.FieldType.STRING) + .build(); + List resultRows = + readResultSet(selectResult).stream() + .map(values -> values.stream().collect(toRow(resultSchema))) + .collect(Collectors.toList()); + + assertThat( + resultRows, + containsInAnyOrder( + row(resultSchema, "description1", "aaa"), row(resultSchema, "description2", "bbb"))); + } + @Test public void testInsertIntoCreatedTable() throws Exception { TestTableProvider tableProvider = new TestTableProvider(); From 7a2a63100324864ba3e73ca95b7e54c6f2046220 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 12:35:38 +0200 Subject: [PATCH 2/3] Cover null nested row through JDBC and keep the field count error on layout mismatch --- .../beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java | 7 +++++++ .../beam/sdk/extensions/sql/impl/JdbcDriverTest.java | 7 +++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java index e73408a84780..d49599c4f797 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java @@ -167,6 +167,13 @@ private void addElementValues( } private boolean isRowColumn(int index) { + // past the last output column the element does not fit the planner's layout. returning + // false lets the values run long so Row.Builder#build() reports the field count mismatch, + // instead of an IndexOutOfBoundsException from getField() that names neither schema + if (index >= outputSchema.getFieldCount()) { + return false; + } + return outputSchema.getField(index).getType().getTypeName().equals(Schema.TypeName.ROW); } diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java index fa78a1e5c622..11f41534ee02 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java @@ -379,7 +379,8 @@ public void testUnnestArrayOfRowsWithNestedRow() throws Exception { "person1", Arrays.asList( row(COMPLEX_SCHEMA, "description1", row(1L, "aaa")), - row(COMPLEX_SCHEMA, "description2", row(2L, "bbb"))))); + row(COMPLEX_SCHEMA, "description2", row(2L, "bbb")), + row(COMPLEX_SCHEMA, "description3", null)))); ResultSet selectResult = connection @@ -401,7 +402,9 @@ public void testUnnestArrayOfRowsWithNestedRow() throws Exception { assertThat( resultRows, containsInAnyOrder( - row(resultSchema, "description1", "aaa"), row(resultSchema, "description2", "bbb"))); + row(resultSchema, "description1", "aaa"), + row(resultSchema, "description2", "bbb"), + row(resultSchema, "description3", null))); } @Test From 390bfadd8b39cb21033668ec129021632c04ee9a Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 12:56:44 +0200 Subject: [PATCH 3/3] Use a direct null check in BeamUnnestRel to match the module style --- .../apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java index d49599c4f797..bae9fd60ba85 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java @@ -20,7 +20,6 @@ import java.util.ArrayList; import java.util.Collection; import java.util.List; -import java.util.Objects; import org.apache.beam.sdk.extensions.sql.impl.planner.BeamCostModel; import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery; import org.apache.beam.sdk.extensions.sql.impl.planner.NodeStats; @@ -151,7 +150,7 @@ private void addElementValues( for (int i = 0; i < elementSchema.getFieldCount(); i++) { Schema.FieldType fieldType = elementSchema.getField(i).getType(); Object value = null; - if (Objects.nonNull(element)) { + if (element != null) { value = element.getBaseValue(i); }