diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 72d319364..4e939f8bc 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -4,6 +4,7 @@ import io.substrait.expression.Expression; import io.substrait.expression.Expression.SortDirection; import io.substrait.expression.FunctionArg; +import io.substrait.expression.MaskExpression; import io.substrait.extension.FunctionBindingResolver; import io.substrait.extension.ResolvedAggregateBinding; import io.substrait.extension.ResolvedArgument; @@ -204,7 +205,7 @@ public RelNode visit(Filter filter, Context context) throws RuntimeException { @Override public RelNode visit(NamedScan namedScan, Context context) throws RuntimeException { RelNode node = relBuilder.scan(namedScan.getNames()).build(); - return applyRelCommon(node, namedScan); + return applyRelCommon(applyProjection(node, namedScan.getProjection()), namedScan); } @Override @@ -866,11 +867,6 @@ public RelNode visit(NamedDdl namedDdl, Context context) { @Override public RelNode visit(VirtualTableScan virtualTableScan, Context context) { - if (virtualTableScan.getProjection().isPresent()) { - throw new UnsupportedOperationException( - "Projection on a VirtualTableScan is not supported: its columns would have to be " - + "masked before an emit mapping selects from them"); - } // A schema's names are one per field at every level of the struct, in depth-first order, so // they have to be handed to the conversion rather than paired with the row type afterwards: // with a nested struct anywhere in the schema the two lists do not even have the same length. @@ -921,7 +917,9 @@ public RelNode visit(VirtualTableScan virtualTableScan, Context context) { tuplesBuilder.add(tupleBuilder.build()); } return applyRelCommon( - LogicalValues.create(relBuilder.getCluster(), rowType, tuplesBuilder.build()), + applyProjection( + LogicalValues.create(relBuilder.getCluster(), rowType, tuplesBuilder.build()), + virtualTableScan.getProjection()), virtualTableScan); } else { // A row that does not fit a LogicalValues tuple keeps its expressions, in a relation of our @@ -930,7 +928,10 @@ public RelNode visit(VirtualTableScan virtualTableScan, Context context) { // consumer whose planner only knows Calcite's own relations can expand it with // VirtualTableExpansionRule. return applyRelCommon( - VirtualTable.create(relBuilder.getCluster(), rowType, convertedRows), virtualTableScan); + applyProjection( + VirtualTable.create(relBuilder.getCluster(), rowType, convertedRows), + virtualTableScan.getProjection()), + virtualTableScan); } } @@ -1182,6 +1183,51 @@ public RelNode visitFallback(Rel rel, Context context) throws RuntimeException { rel, rel.getClass().getCanonicalName(), this.getClass().getCanonicalName())); } + /** + * Applies a read relation's projection to the node its scan was converted into. + * + *

A projection masks the columns of the initial schema, and the read produces the ones it + * leaves: {@link io.substrait.relation.AbstractReadRel#deriveRecordType()} applies it to that + * schema, so an emit mapping's indices, and every field reference a parent relation makes, count + * the masked columns. Building the scan from the schema alone leaves the node one column list and + * the relation another. + * + *

The columns come out in the order the mask lists them. That is the order {@link + * io.substrait.expression.MaskExpressionTypeProjector} builds the record type in, and the node + * has to carry the columns the relation says it carries. Spec v0.102.0 describes a mask as + * removing columns and asks whether reordering should be supported at all, so this order follows + * the record type the model derives rather than a rule the specification settles. + * + *

Only a mask that selects whole columns is converted. A mask can also select inside a column + * -- some of a struct's fields, some of a list's elements, some of a map's entries -- and + * applying that would mean rebuilding the column's value from the parts the mask keeps, which + * this conversion does not do. Such a mask is reported rather than applied to the columns it + * selects whole, which would drop the rest of what it says. + * + * @param relNode the node the read was converted into + * @param projection the projection the read carries, if any + * @return the node, with the masked columns projected out of it + */ + private RelNode applyProjection(RelNode relNode, Optional projection) { + if (projection.isEmpty()) { + return relNode; + } + List items = projection.get().getSelect().getStructItems(); + RelDataType rowType = relNode.getRowType(); + List rexList = new ArrayList<>(items.size()); + for (MaskExpression.StructItem item : items) { + if (item.getChild().isPresent()) { + throw new UnsupportedOperationException( + "A read projection that selects inside a column is not supported: only a mask that " + + "selects whole columns is applied, and pruning a struct, a list or a map would " + + "have to rebuild the column's value"); + } + rexList.add( + new RexInputRef(item.getField(), rowType.getFieldList().get(item.getField()).getType())); + } + return relBuilder.push(relNode).project(rexList).build(); + } + /** * Applies the parts of a relation's {@code RelCommon} that Calcite can hold: its emit mapping * first, and then the alternative output field names of its hint. diff --git a/isthmus/src/test/java/io/substrait/isthmus/ReadProjectionTest.java b/isthmus/src/test/java/io/substrait/isthmus/ReadProjectionTest.java new file mode 100644 index 000000000..1d4b283cf --- /dev/null +++ b/isthmus/src/test/java/io/substrait/isthmus/ReadProjectionTest.java @@ -0,0 +1,222 @@ +package io.substrait.isthmus; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.substrait.expression.Expression; +import io.substrait.expression.ImmutableMaskExpression; +import io.substrait.expression.MaskExpression; +import io.substrait.hint.Hint; +import io.substrait.relation.Filter; +import io.substrait.relation.NamedScan; +import io.substrait.relation.Rel; +import io.substrait.relation.VirtualTableScan; +import io.substrait.type.NamedStruct; +import java.util.List; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.TableScan; +import org.junit.jupiter.api.Test; + +/** + * The projection a read relation carries masks the columns of its initial schema, and the read + * produces the ones the mask leaves. + */ +class ReadProjectionTest extends PlanTestBase { + + private final NamedScan scan = + (NamedScan) + sb.namedScan(List.of("t"), List.of("a", "b", "c"), List.of(R.I64, N.STRING, R.FP64)); + + /** A mask selecting whole columns, by the index each has in the initial schema. */ + private static MaskExpression columns(int... fields) { + ImmutableMaskExpression.StructSelect.Builder select = MaskExpression.StructSelect.builder(); + for (int field : fields) { + select.addStructItems(MaskExpression.StructItem.of(field)); + } + return MaskExpression.builder().select(select.build()).build(); + } + + @Test + void aProjectionMasksTheColumnsANamedScanReads() { + NamedScan masked = NamedScan.builder().from(scan).projection(columns(0, 2)).build(); + + RelNode relNode = substraitToCalcite.convert(masked); + + assertEquals(List.of("a", "c"), relNode.getRowType().getFieldNames()); + assertEquals( + masked.getRecordType(), + SubstraitRelVisitor.convert(relNode, converterProvider).getRecordType()); + } + + /** + * A mask that lists its columns out of schema order produces them in the order it lists. That is + * the order the model derives the record type in, and the node has to carry the columns the + * relation says it carries. Spec v0.102.0 describes a mask as removing columns and asks whether + * reordering should be supported at all, so this pins what the model already derives rather than + * a rule the specification settles. + */ + @Test + void theColumnsComeOutInTheOrderTheMaskListsThem() { + NamedScan reordered = NamedScan.builder().from(scan).projection(columns(2, 0)).build(); + + RelNode relNode = substraitToCalcite.convert(reordered); + + assertEquals(List.of("c", "a"), relNode.getRowType().getFieldNames()); + assertEquals( + reordered.getRecordType(), + SubstraitRelVisitor.convert(relNode, converterProvider).getRecordType()); + } + + /** + * An emit mapping selects from the columns the mask leaves, not from the schema it masked: on + * this scan index 1 is the second column the mask keeps, which is the schema's third. + */ + @Test + void anEmitMappingSelectsFromTheMaskedColumns() { + NamedScan masked = + NamedScan.builder().from(scan).projection(columns(0, 2)).remap(sb.remap(1)).build(); + + RelNode relNode = substraitToCalcite.convert(masked); + + assertEquals(List.of("c"), relNode.getRowType().getFieldNames()); + assertEquals( + masked.getRecordType(), + SubstraitRelVisitor.convert(relNode, converterProvider).getRecordType()); + } + + /** + * So does a field reference a parent relation makes against the read: index 1 is the fp64 column + * the mask leaves there, and not the string the schema has at that index. + */ + @Test + void aParentRelationReferencesTheMaskedColumns() { + NamedScan masked = NamedScan.builder().from(scan).projection(columns(0, 2)).build(); + Filter filter = sb.filter(input -> sb.equal(sb.fieldReference(input, 1), sb.fp64(5)), masked); + + RelNode relNode = substraitToCalcite.convert(filter); + + Rel converted = SubstraitRelVisitor.convert(relNode, converterProvider); + assertEquals(filter.getRecordType(), converted.getRecordType()); + assertEquals(filter.getCondition(), assertInstanceOf(Filter.class, converted).getCondition()); + } + + /** A mask that selects every column in order leaves the scan the columns it already reads. */ + @Test + void aProjectionSelectingEveryColumnLeavesTheScanAlone() { + NamedScan masked = NamedScan.builder().from(scan).projection(columns(0, 1, 2)).build(); + + assertInstanceOf(TableScan.class, substraitToCalcite.convert(masked)); + } + + /** + * The masked columns are the ones the relation produces, so they are the ones an output name from + * a hint names -- and the projection the mask adds is where those names can go. + */ + @Test + void anOutputNameFromAHintNamesAMaskedColumn() { + NamedScan masked = + NamedScan.builder() + .from(scan) + .projection(columns(0, 2)) + .hint(Hint.builder().addOutputNames("x", "y").build()) + .build(); + + assertEquals( + List.of("x", "y"), substraitToCalcite.convert(masked).getRowType().getFieldNames()); + } + + /** + * A mask selects by index: a schema's names are not uniquified, so a name can stand for more than + * one column. + */ + @Test + void aProjectionSelectsByIndexWhereTwoColumnsShareAName() { + NamedScan sharedNames = + (NamedScan) + sb.namedScan(List.of("t"), List.of("c", "c", "d"), List.of(R.I64, R.STRING, R.FP64)); + NamedScan masked = NamedScan.builder().from(sharedNames).projection(columns(1, 2)).build(); + + RelNode relNode = substraitToCalcite.convert(masked); + + assertEquals( + masked.getRecordType(), + SubstraitRelVisitor.convert(relNode, converterProvider).getRecordType()); + } + + @Test + void aProjectionMasksTheColumnsAVirtualTableReads() { + NamedStruct schema = + NamedStruct.of(List.of("col1", "col2", "col3"), R.struct(R.I32, R.STRING, R.BOOLEAN)); + VirtualTableScan table = + VirtualTableScan.builder() + .initialSchema(schema) + .addRows( + Expression.NestedStruct.builder() + .addFields(sb.i32(2), sb.str("a"), sb.bool(true)) + .build()) + .projection(columns(1, 2)) + .build(); + + RelNode relNode = substraitToCalcite.convert(table); + + assertEquals(List.of("col2", "col3"), relNode.getRowType().getFieldNames()); + assertEquals( + table.getRecordType(), + SubstraitRelVisitor.convert(relNode, converterProvider).getRecordType()); + } + + /** A row that no {@code LogicalValues} tuple holds is computed, and masked the same way. */ + @Test + void aProjectionMasksTheColumnsAComputedVirtualTableReads() { + NamedStruct schema = + NamedStruct.of(List.of("col1", "col2", "col3"), R.struct(R.I32, R.FP64, R.STRING)); + VirtualTableScan table = + VirtualTableScan.builder() + .initialSchema(schema) + .addRows( + Expression.NestedStruct.builder() + .addFields(sb.i32(2), sb.add(sb.fp64(4.4), sb.fp64(4.5)), sb.str("a")) + .build()) + .projection(columns(1, 2)) + .build(); + + RelNode relNode = substraitToCalcite.convert(table); + + assertEquals(List.of("col2", "col3"), relNode.getRowType().getFieldNames()); + assertEquals( + table.getRecordType(), + SubstraitRelVisitor.convert(relNode, converterProvider).getRecordType()); + } + + /** + * A mask can select inside a column as well, keeping some of a struct's fields or some of a + * list's elements. Calcite reads a column as it stands, so such a mask is reported rather than + * applied to the column it selects from. + */ + @Test + void aProjectionThatSelectsInsideAColumnIsRefused() { + NamedScan structScan = + (NamedScan) + sb.namedScan(List.of("t"), List.of("s", "x", "y"), List.of(R.struct(R.I64, R.STRING))); + MaskExpression insideAColumn = + MaskExpression.builder() + .select( + MaskExpression.StructSelect.builder() + .addStructItems( + MaskExpression.StructItem.of( + 0, + MaskExpression.StructSelect.builder() + .addStructItems(MaskExpression.StructItem.of(1)) + .build())) + .build()) + .build(); + NamedScan masked = NamedScan.builder().from(structScan).projection(insideAColumn).build(); + + assertTrue( + assertThrows(UnsupportedOperationException.class, () -> substraitToCalcite.convert(masked)) + .getMessage() + .contains("selects inside a column")); + } +} diff --git a/isthmus/src/test/java/io/substrait/isthmus/VirtualTableScanTest.java b/isthmus/src/test/java/io/substrait/isthmus/VirtualTableScanTest.java index 5373da93e..31448f89c 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/VirtualTableScanTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/VirtualTableScanTest.java @@ -9,7 +9,6 @@ import com.google.common.collect.ImmutableList; import io.substrait.expression.Expression; import io.substrait.expression.ExpressionCreator; -import io.substrait.expression.MaskExpression; import io.substrait.hint.Hint; import io.substrait.relation.Rel; import io.substrait.relation.VirtualTableScan; @@ -489,34 +488,6 @@ void outputNamesWithoutAMappingAreLeftAlone() { assertEquals(List.of("col1", "col2"), relNode.getRowType().getFieldNames()); } - /** - * A projection masks a read relation's columns before anything else selects from them -- {@link - * io.substrait.relation.AbstractReadRel#deriveRecordType()} applies it to the initial schema -- - * so an emit mapping's indices count the columns it leaves. Isthmus builds the row type from the - * unmasked schema and reads the projection nowhere, so a scan carrying one is refused rather than - * converted against the wrong columns. - */ - @Test - void aProjectionOnAVirtualTableIsRefused() { - NamedStruct schema = NamedStruct.of(List.of("col1", "col2"), R.struct(R.I32, R.STRING)); - VirtualTableScan table = - VirtualTableScan.builder() - .from(virtualTable(schema, List.of(sb.i32(2), sb.str("a")))) - .projection( - MaskExpression.builder() - .select( - MaskExpression.StructSelect.builder() - .addStructItems(MaskExpression.StructItem.of(1)) - .build()) - .build()) - .build(); - - assertTrue( - assertThrows(UnsupportedOperationException.class, () -> substraitToCalcite.convert(table)) - .getMessage() - .contains("Projection on a VirtualTableScan is not supported")); - } - /** * A virtual table's row type carries the names its schema gives it, which nothing uniquifies, so * the mapping has to select its columns by index: resolving a field by name would give the