Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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);
}
}

Expand Down Expand Up @@ -1182,6 +1183,55 @@ 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.
*
* <p>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.
*
* <p>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.
*
* <p>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.
Comment thread
nielspardon marked this conversation as resolved.
*
* <p>{@link MaskExpression#getMaintainSingularStruct()} is not read. A relation's record type is
* always a struct, so a mask over its columns has nothing to unwrap, and the only masks that
* could unwrap a nested struct select inside a column, which is refused above.
*
* @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
*/
protected RelNode applyProjection(RelNode relNode, Optional<MaskExpression> projection) {
if (projection.isEmpty()) {
return relNode;
}
List<MaskExpression.StructItem> items = projection.get().getSelect().getStructItems();
RelDataType rowType = relNode.getRowType();
List<RexNode> 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.
Comment thread
nielspardon marked this conversation as resolved.
Expand All @@ -1205,13 +1255,14 @@ protected RelNode applyRelCommon(RelNode relNode, Rel rel, RelNode... inputs) {
* converted into.
*
* <p>The names are applied to the projection this relation's own conversion produced: the one a
* {@link Project} becomes, or the one {@link #applyRemap(RelNode, Optional)} adds for a relation
* with an emit mapping. Anywhere else they are dropped, rather than renaming a node that stands
* for another relation or adding a projection the plan never asked for. A relation converted into
* a bare Calcite operator therefore keeps the names Calcite derives, and so does one Calcite
* builds no operator for at all -- a filter that cannot filter, a sort with no sort fields, an
* identity emit mapping -- where the node handed back is the input's, which is what the given
* inputs are compared against.
* {@link Project} becomes, the one {@link #applyRemap(RelNode, Optional)} adds for a relation
* with an emit mapping, or the one {@link #applyProjection(RelNode, Optional)} adds for a read's
* projection. Anywhere else they are dropped, rather than renaming a node that stands for another
* relation or adding a projection the plan never asked for. A relation converted into a bare
* Calcite operator therefore keeps the names Calcite derives, and so does one Calcite builds no
* operator for at all -- a filter that cannot filter, a sort with no sort fields, an identity
* emit mapping -- where the node handed back is the input's, which is what the given inputs are
* compared against.
*
* <p>They are dropped as well where the columns of that projection are not the columns of the
* relation's record type, type by type. An aggregate over several grouping sets types its
Expand Down
223 changes: 223 additions & 0 deletions isthmus/src/test/java/io/substrait/isthmus/ReadProjectionTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,223 @@
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(0, 1)).build();

RelNode relNode = substraitToCalcite.convert(masked);

assertEquals(List.of("c", "c0"), relNode.getRowType().getFieldNames());
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"));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
Loading