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..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,10 +20,6 @@ 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 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 +132,48 @@ 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 (element != null) { + 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) { + // 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); } @ProcessElement @@ -178,21 +190,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..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 @@ -351,6 +351,62 @@ 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")), + row(COMPLEX_SCHEMA, "description3", null)))); + + 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"), + row(resultSchema, "description3", null))); + } + @Test public void testInsertIntoCreatedTable() throws Exception { TestTableProvider tableProvider = new TestTableProvider();