From 17b5c2763c79b0884e9e70172a8040e8dbe8effe Mon Sep 17 00:00:00 2001 From: bvolpato Date: Thu, 3 Sep 2026 19:11:57 +0000 Subject: [PATCH] fix(isthmus): preserve update targets in nested schemas For a table src(s ROW(x INTEGER), x INTEGER, n INTEGER), UPDATE src SET n = 99 converts back to UPDATE src SET x = 99. The transform's top-level column ordinal is incorrectly used to index the flattened depth-first name list [S, X, X, N]. Reconstruct the declared row type before resolving target names so nested fields do not shift top-level update columns. This also preserves multiple-assignment order when nested and top-level names collide. Partially addresses #1175; struct-literal assignment field names are separate from target-column resolution. --- .../isthmus/SubstraitRelNodeConverter.java | 7 ++- .../isthmus/NestedUpdateTargetTest.java | 53 +++++++++++++++++++ 2 files changed, 56 insertions(+), 4 deletions(-) create mode 100644 isthmus/src/test/java/io/substrait/isthmus/NestedUpdateTargetTest.java diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 72d319364..df6f03a17 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -771,8 +771,7 @@ public RelNode visit(NamedUpdate update, Context context) { context.enterScope(AnchoredInput.of(update.getRelAnchor(), relBuilder.peek().getRowType())); RexNode condition = update.getCondition().accept(expressionRexConverter, context); - NamedStruct tableSchema = update.getTableSchema(); - List fieldNames = tableSchema.names(); + List fieldNames = toRowType(update.getTableSchema()).getFieldNames(); List updateColumnList = new ArrayList<>(); List sourceExpressionList = new ArrayList<>(); @@ -1061,10 +1060,10 @@ private RexNode unwrapNullabilityCast(RexNode rexNode) { } /** - * Converts the schema a DDL relation declares into the row type that describes it. + * Converts a declared schema into the row type that describes it. * * @param schema the declared schema, whose names are one per field at every level of the struct - * @return the row type of the object the statement creates + * @return the row type with field names at their corresponding nesting levels */ private RelDataType toRowType(NamedStruct schema) { return typeConverter.toCalcite(typeFactory, schema.struct(), schema.names()); diff --git a/isthmus/src/test/java/io/substrait/isthmus/NestedUpdateTargetTest.java b/isthmus/src/test/java/io/substrait/isthmus/NestedUpdateTargetTest.java new file mode 100644 index 000000000..f92fae583 --- /dev/null +++ b/isthmus/src/test/java/io/substrait/isthmus/NestedUpdateTargetTest.java @@ -0,0 +1,53 @@ +package io.substrait.isthmus; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; + +import io.substrait.isthmus.sql.SubstraitCreateStatementParser; +import io.substrait.plan.Plan; +import io.substrait.plan.PlanProtoConverter; +import io.substrait.plan.ProtoPlanConverter; +import io.substrait.relation.NamedUpdate; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.calcite.prepare.Prepare; +import org.apache.calcite.rel.core.TableModify; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +class NestedUpdateTargetTest { + + static Stream schemasAndTargets() { + return Stream.of( + Arguments.of("x INTEGER, n INTEGER", List.of("N")), + Arguments.of("s ROW(x INTEGER), x INTEGER, n INTEGER", List.of("N")), + Arguments.of("s ROW(x INTEGER), x INTEGER, n INTEGER", List.of("X")), + Arguments.of("s ROW(x INTEGER), x INTEGER, n INTEGER", List.of("N", "X")), + Arguments.of("x INTEGER, s ROW(n INTEGER), n INTEGER", List.of("X", "N")), + Arguments.of("s ROW(a ROW(x INTEGER), n INTEGER), x INTEGER, n INTEGER", List.of("N", "X")), + Arguments.of("s ROW(a INTEGER, b INTEGER), n INTEGER", List.of("N"))); + } + + @ParameterizedTest + @MethodSource("schemasAndTargets") + void preservesTopLevelUpdateTargets(String schema, List targets) throws Exception { + Prepare.CatalogReader catalog = + SubstraitCreateStatementParser.processCreateStatementsToCatalog( + "CREATE TABLE src (" + schema + ")"); + String assignments = + targets.stream().map(name -> name + " = 99").collect(Collectors.joining(", ")); + Plan plan = new SqlToSubstrait().convert("UPDATE src SET " + assignments, catalog); + Plan decoded = new ProtoPlanConverter().from(new PlanProtoConverter().toProto(plan)); + NamedUpdate update = assertInstanceOf(NamedUpdate.class, decoded.getRoots().get(0).getInput()); + + TableModify converted = + assertInstanceOf( + TableModify.class, + new SubstraitToCalcite(ConverterProvider.DEFAULT, catalog).convert(update)); + + assertEquals(targets, converted.getUpdateColumnList()); + assertEquals(targets.size(), converted.getSourceExpressionList().size()); + } +}