Skip to content

Commit 62e46f0

Browse files
damccormdamccorm
committed
Reconcile schema nullability in BeamCoGBKJoinRel
1 parent 9bcdedd commit 62e46f0

2 files changed

Lines changed: 144 additions & 4 deletions

File tree

‎sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCoGBKJoinRel.java‎

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.apache.beam.sdk.transforms.windowing.Trigger;
3737
import org.apache.beam.sdk.transforms.windowing.Window;
3838
import org.apache.beam.sdk.transforms.windowing.WindowFn;
39+
import org.apache.beam.sdk.util.Preconditions;
3940
import org.apache.beam.sdk.values.PCollection;
4041
import org.apache.beam.sdk.values.PCollectionList;
4142
import org.apache.beam.sdk.values.Row;
@@ -193,11 +194,49 @@ private PCollection<Row> standardJoin(
193194
}
194195

195196
// Flatten the lhs and rhs fields into a single row.
197+
FieldAccessDescriptor flattenFields =
198+
FieldAccessDescriptor.withFieldNames(
199+
org.apache.beam.sdk.schemas.transforms.Join.LHS_TAG + ".*",
200+
org.apache.beam.sdk.schemas.transforms.Join.RHS_TAG + ".*");
201+
202+
// Reconcile the desired output schema (which carries the Calcite-derived field names and the
203+
// correct top-level, outer-join-aware nullability) with the types the data actually carries.
204+
// Calcite's join row-type derivation can report a different nullability than the rows hold --
205+
// in particular it can mark the fields nested inside a struct column as nullable even when the
206+
// joined rows still keep them NOT NULL. Forcing the Calcite schema verbatim then trips Select's
207+
// type-equality guard. The flatten emits the lhs struct's fields followed by the rhs struct's
208+
// fields, so walk the Calcite output positionally against those data fields, keeping Calcite's
209+
// names and top-level nullability but adopting the data's (possibly deeper) field types.
210+
Schema calciteSchema = CalciteUtils.toSchema(getRowType());
211+
Schema joinedSchema = joinedRows.getSchema();
212+
List<Schema.Field> dataFields = new java.util.ArrayList<>();
213+
dataFields.addAll(
214+
Preconditions.checkArgumentNotNull(joinedSchema.getField(0).getType().getRowSchema())
215+
.getFields());
216+
dataFields.addAll(
217+
Preconditions.checkArgumentNotNull(joinedSchema.getField(1).getType().getRowSchema())
218+
.getFields());
219+
220+
if (calciteSchema.getFieldCount() != dataFields.size()) {
221+
throw new IllegalStateException(
222+
String.format(
223+
"Field count mismatch: Calcite schema has %d fields, but data schema has %d fields",
224+
calciteSchema.getFieldCount(), dataFields.size()));
225+
}
226+
227+
Schema.Builder reconciled = Schema.builder();
228+
for (int i = 0; i < calciteSchema.getFieldCount(); i++) {
229+
Schema.Field calciteField = calciteSchema.getField(i);
230+
// Keep the data's type verbatim -- including the nullability of fields nested inside struct
231+
// columns, which Calcite's join row-type derivation can report incorrectly -- but honour
232+
// Calcite's top-level nullability so outer-join columns stay nullable.
233+
reconciled.addField(
234+
calciteField.withType(
235+
dataFields.get(i).getType().withNullable(calciteField.getType().getNullable())));
236+
}
237+
196238
return joinedRows.apply(
197-
Select.<Row>fieldNames(
198-
org.apache.beam.sdk.schemas.transforms.Join.LHS_TAG + ".*",
199-
org.apache.beam.sdk.schemas.transforms.Join.RHS_TAG + ".*")
200-
.withOutputSchema(CalciteUtils.toSchema(getRowType())));
239+
Select.<Row>fieldAccess(flattenFields).withOutputSchema(reconciled.build()));
201240
}
202241

203242
@Override

‎sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslJoinTest.java‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -325,6 +325,107 @@ private PCollection<Row> ordersUnbounded() {
325325
.buildUnbounded();
326326
}
327327

328+
@Test
329+
public void testLeftOuterJoinWithNestedRows() {
330+
Schema nestedSchema =
331+
Schema.builder().addInt32Field("nested_int").addStringField("nested_str").build();
332+
Schema lhsSchema = Schema.builder().addInt32Field("id").addStringField("val").build();
333+
Schema rhsSchema =
334+
Schema.builder().addInt32Field("id").addRowField("nested", nestedSchema).build();
335+
336+
PCollection<Row> lhs =
337+
pipeline.apply(
338+
"lhs",
339+
org.apache.beam.sdk.transforms.Create.of(
340+
Row.withSchema(lhsSchema).addValues(1, "lhs1").build(),
341+
Row.withSchema(lhsSchema).addValues(2, "lhs2").build())
342+
.withRowSchema(lhsSchema));
343+
344+
PCollection<Row> rhs =
345+
pipeline.apply(
346+
"rhs",
347+
org.apache.beam.sdk.transforms.Create.of(
348+
Row.withSchema(rhsSchema)
349+
.addValues(1, Row.withSchema(nestedSchema).addValues(10, "nested1").build())
350+
.build())
351+
.withRowSchema(rhsSchema));
352+
353+
String sql =
354+
"SELECT lhs.id, lhs.val, rhs.nested FROM lhs LEFT OUTER JOIN rhs ON lhs.id = rhs.id";
355+
356+
PCollection<Row> result =
357+
PCollectionTuple.of("lhs", lhs).and("rhs", rhs).apply("sql", SqlTransform.query(sql));
358+
359+
Schema expectedSchema =
360+
Schema.builder()
361+
.addInt32Field("id")
362+
.addStringField("val")
363+
.addNullableField("nested", Schema.FieldType.row(nestedSchema))
364+
.build();
365+
366+
PAssert.that(result)
367+
.containsInAnyOrder(
368+
Row.withSchema(expectedSchema)
369+
.addValues(1, "lhs1", Row.withSchema(nestedSchema).addValues(10, "nested1").build())
370+
.build(),
371+
Row.withSchema(expectedSchema).addValues(2, "lhs2", null).build());
372+
373+
pipeline.run();
374+
}
375+
376+
@Test
377+
public void testLeftOuterJoinWithArrayOfNestedRows() {
378+
Schema nestedSchema =
379+
Schema.builder().addInt32Field("nested_int").addStringField("nested_str").build();
380+
Schema lhsSchema = Schema.builder().addInt32Field("id").addStringField("val").build();
381+
Schema rhsSchema =
382+
Schema.builder()
383+
.addInt32Field("id")
384+
.addArrayField("nested_list", Schema.FieldType.row(nestedSchema))
385+
.build();
386+
387+
PCollection<Row> lhs =
388+
pipeline.apply(
389+
"lhs",
390+
org.apache.beam.sdk.transforms.Create.of(
391+
Row.withSchema(lhsSchema).addValues(1, "lhs1").build(),
392+
Row.withSchema(lhsSchema).addValues(2, "lhs2").build())
393+
.withRowSchema(lhsSchema));
394+
395+
java.util.List<Row> nestedRows =
396+
java.util.Arrays.asList(
397+
Row.withSchema(nestedSchema).addValues(10, "nested1").build(),
398+
Row.withSchema(nestedSchema).addValues(20, "nested2").build());
399+
400+
PCollection<Row> rhs =
401+
pipeline.apply(
402+
"rhs",
403+
org.apache.beam.sdk.transforms.Create.of(
404+
Row.withSchema(rhsSchema).addValues(1, nestedRows).build())
405+
.withRowSchema(rhsSchema));
406+
407+
String sql =
408+
"SELECT lhs.id, lhs.val, rhs.nested_list FROM lhs LEFT OUTER JOIN rhs ON lhs.id = rhs.id";
409+
410+
PCollection<Row> result =
411+
PCollectionTuple.of("lhs", lhs).and("rhs", rhs).apply("sql", SqlTransform.query(sql));
412+
413+
Schema expectedSchema =
414+
Schema.builder()
415+
.addInt32Field("id")
416+
.addStringField("val")
417+
.addNullableField(
418+
"nested_list", Schema.FieldType.array(Schema.FieldType.row(nestedSchema)))
419+
.build();
420+
421+
PAssert.that(result)
422+
.containsInAnyOrder(
423+
Row.withSchema(expectedSchema).addValues(1, "lhs1", nestedRows).build(),
424+
Row.withSchema(expectedSchema).addValues(2, "lhs2", null).build());
425+
426+
pipeline.run();
427+
}
428+
328429
private PCollection<Row> queryFromOrderTables(String sql) {
329430
return tuple(
330431
"ORDER_DETAILS1",

0 commit comments

Comments
 (0)