Fix Beam SQL UNNEST of array elements with a nested row through JDBC

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

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<Integer> 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.
+     *
+     * <p>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<id,
+     * author ROW<name, email>>}, 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<Object> getNestedRowBaseValues(Row nestedRow) {
-      return IntStream.range(0, nestedRow.getFieldCount())
-          .mapToObj(
-              (i) -> {
-                List<Object> 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<Object> 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<Object> 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<Row> 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<Object> 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)

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<Row> inputTable =
+        PBegin.in(pipeline).apply(Create.of(input).withRowSchema(inputType));
+
+    PCollection<Row> 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();
+  }
 }

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<ROW< \n"
+                + "              description VARCHAR, \n"
+                + "              nestedRow ROW<id BIGINT, name VARCHAR>>> \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<Row> 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();