diff --git a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java index 5f43ace2b..cb793fd10 100644 --- a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java +++ b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java @@ -1013,15 +1013,20 @@ protected Join newJoin(JoinRel rel) { Type.Struct unionedStruct = Type.Struct.builder().from(leftStruct).from(rightStruct).build(); ProtoExpressionConverter converter = new ProtoExpressionConverter(lookup, extensions, unionedStruct, this); + Join.JoinType joinType = Join.JoinType.fromProto(rel.getType()); ImmutableJoin.Builder builder = Join.builder() .left(left) .right(right) .condition(converter.from(rel.getExpression())) - .joinType(Join.JoinType.fromProto(rel.getType())) - .postJoinFilter( - Optional.ofNullable( - rel.hasPostJoinFilter() ? converter.from(rel.getPostJoinFilter()) : null)); + .joinType(joinType); + + if (rel.hasPostJoinFilter()) { + ProtoExpressionConverter outputConverter = + new ProtoExpressionConverter( + lookup, extensions, Join.deriveRecordType(joinType, left, right), this); + builder.postJoinFilter(outputConverter.from(rel.getPostJoinFilter())); + } if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); diff --git a/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java index c8c2ddad3..f18202093 100644 --- a/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java +++ b/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java @@ -10,6 +10,8 @@ import java.util.Arrays; import java.util.List; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; class JoinRoundtripTest extends TestBase { @@ -25,6 +27,39 @@ class JoinRoundtripTest extends TestBase { Arrays.asList("d", "e", "f"), Arrays.asList(R.FP64, R.STRING, R.I64)); + @ParameterizedTest + @EnumSource( + value = Join.JoinType.class, + names = { + "INNER", + "LEFT", + "RIGHT", + "OUTER", + "LEFT_SEMI", + "LEFT_ANTI", + "RIGHT_SEMI", + "RIGHT_ANTI" + }) + void postJoinFilterUsesOutputSchemaBeforeEmit(Join.JoinType joinType) { + Join join = + sb.join( + input -> sb.equal(sb.fieldReference(input, 0), sb.fieldReference(input, 5)), + joinType, + leftTable, + rightTable); + Join filtered = + Join.builder() + .from(join) + .postJoinFilter( + sb.and( + sb.isNull(sb.fieldReference(join, 0)), + sb.isNull(sb.fieldReference(join, join.getRecordType().fields().size() - 1)))) + .remap(sb.remap(0)) + .build(); + + verifyRoundTrip(filtered); + } + @Test void hashJoin() { List leftKeys = Arrays.asList(0, 1); diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 72d319364..d73d0d0d4 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -204,9 +204,22 @@ 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(); + node = applyFilter(node, namedScan.getFilter(), context); return applyRelCommon(node, namedScan); } + private RelNode applyFilter(RelNode input, Optional filter, Context context) { + if (filter.isEmpty()) { + return input; + } + // Embedded predicates use the operator's direct row, before its emit mapping. This is an + // internal input, not another anchored Substrait relation; enclosing scopes still own any + // outer references used by the predicate. + context.enterScope(AnchoredInput.of(Optional.empty(), input.getRowType())); + RexNode condition = filter.get().accept(expressionRexConverter, context); + return relBuilder.push(input).filter(context.exitScope(), condition).build(); + } + @Override public RelNode visit(LocalFiles localFiles, Context context) throws RuntimeException { return visitFallback(localFiles, context); @@ -256,6 +269,7 @@ public RelNode visit(Join join, Context context) throws RuntimeException { JoinRelType joinType = asJoinRelType(join); RelNode node = relBuilder.push(left).push(right).join(joinType, condition, context.exitScope()).build(); + node = applyFilter(node, join.getPostJoinFilter(), context); return applyRelCommon(node, join, left, right); } @@ -921,7 +935,10 @@ public RelNode visit(VirtualTableScan virtualTableScan, Context context) { tuplesBuilder.add(tupleBuilder.build()); } return applyRelCommon( - LogicalValues.create(relBuilder.getCluster(), rowType, tuplesBuilder.build()), + applyFilter( + LogicalValues.create(relBuilder.getCluster(), rowType, tuplesBuilder.build()), + virtualTableScan.getFilter(), + context), virtualTableScan); } else { // A row that does not fit a LogicalValues tuple keeps its expressions, in a relation of our @@ -930,7 +947,11 @@ 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); + applyFilter( + VirtualTable.create(relBuilder.getCluster(), rowType, convertedRows), + virtualTableScan.getFilter(), + context), + virtualTableScan); } } diff --git a/isthmus/src/test/java/io/substrait/isthmus/EmbeddedPredicateTest.java b/isthmus/src/test/java/io/substrait/isthmus/EmbeddedPredicateTest.java new file mode 100644 index 000000000..294f89bd4 --- /dev/null +++ b/isthmus/src/test/java/io/substrait/isthmus/EmbeddedPredicateTest.java @@ -0,0 +1,292 @@ +package io.substrait.isthmus; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.substrait.expression.Expression; +import io.substrait.expression.ExpressionCreator; +import io.substrait.expression.FieldReference; +import io.substrait.extension.ExtensionCollector; +import io.substrait.isthmus.calcite.rel.rules.VirtualTableExpansionRule; +import io.substrait.relation.Join; +import io.substrait.relation.Join.JoinType; +import io.substrait.relation.NamedScan; +import io.substrait.relation.ProtoRelConverter; +import io.substrait.relation.Rel; +import io.substrait.relation.RelProtoConverter; +import io.substrait.relation.VirtualTableScan; +import io.substrait.type.NamedStruct; +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.calcite.DataContext; +import org.apache.calcite.adapter.java.JavaTypeFactory; +import org.apache.calcite.interpreter.Interpreter; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.linq4j.QueryProvider; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Filter; +import org.apache.calcite.rel.core.Project; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rel.core.Values; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexSubQuery; +import org.apache.calcite.schema.SchemaPlus; +import org.apache.calcite.tools.Frameworks; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.ValueSource; + +class EmbeddedPredicateTest extends PlanTestBase { + + @Test + void namedScanFiltersBeforeEmit() { + NamedScan scan = + sb.namedScan(List.of("example"), List.of("id", "keep"), List.of(R.I32, N.BOOLEAN)); + NamedScan filtered = + NamedScan.builder() + .from(scan) + .filter(sb.fieldReference(scan, 1)) + .remap(sb.remap(0)) + .build(); + + Project project = assertInstanceOf(Project.class, substraitToCalcite.convert(filtered)); + Filter filter = assertInstanceOf(Filter.class, project.getInput()); + assertInstanceOf(TableScan.class, filter.getInput()); + assertEquals(1, assertInstanceOf(RexInputRef.class, filter.getCondition()).getIndex()); + assertRowMatch(project.getRowType(), R.I32); + } + + @Test + void namedScanFalseAndNullFiltersProduceNoRows() { + NamedScan scan = sb.namedScan(List.of("example"), List.of("id"), List.of(R.I32)); + for (Expression condition : List.of(sb.bool(false), ExpressionCreator.typedNull(N.BOOLEAN))) { + NamedScan filtered = NamedScan.builder().from(scan).filter(condition).build(); + + Values values = assertInstanceOf(Values.class, substraitToCalcite.convert(filtered)); + assertTrue(values.getTuples().isEmpty()); + assertRowMatch(values.getRowType(), R.I32); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void virtualTableFiltersLiteralAndComputedRowsBeforeEmit(boolean computed) { + VirtualTableScan scan = + VirtualTableScan.builder() + .initialSchema(NamedStruct.of(List.of("id", "keep"), R.struct(R.I32, N.BOOLEAN))) + .addRows( + ExpressionCreator.nestedStruct( + false, + List.of( + computed ? sb.add(sb.i32(1), sb.i32(1)) : sb.i32(2), + ExpressionCreator.bool(true, true))), + ExpressionCreator.nestedStruct( + false, List.of(sb.i32(3), ExpressionCreator.bool(true, false))), + ExpressionCreator.nestedStruct( + false, List.of(sb.i32(4), ExpressionCreator.typedNull(N.BOOLEAN)))) + .build(); + VirtualTableScan filtered = + VirtualTableScan.builder() + .from(scan) + .filter(sb.fieldReference(scan, 1)) + .remap(sb.remap(0)) + .build(); + + assertEquals(List.of(List.of(2)), rows(substraitToCalcite.convert(filtered))); + } + + @Test + void bestEffortReadFilterMayBeIgnored() { + VirtualTableScan scan = + VirtualTableScan.builder().from(integers(1, 2)).bestEffortFilter(sb.bool(false)).build(); + + assertEquals(List.of(List.of(1), List.of(2)), rows(substraitToCalcite.convert(scan))); + } + + @ParameterizedTest + @EnumSource( + value = JoinType.class, + names = {"LEFT", "RIGHT", "OUTER"}) + void postJoinFilterSeesNullExtendedRows(JoinType joinType) { + Join join = equalityJoin(joinType); + int nullField = joinType == JoinType.RIGHT ? 0 : 1; + Join filtered = + Join.builder() + .from(join) + .postJoinFilter(sb.isNull(sb.fieldReference(join, nullField))) + .build(); + + RelNode converted = substraitToCalcite.convert(filtered); + Filter filter = assertInstanceOf(Filter.class, converted); + assertInstanceOf(org.apache.calcite.rel.core.Join.class, filter.getInput()); + List expected = + joinType == JoinType.RIGHT ? Arrays.asList(null, 3) : Arrays.asList(1, null); + assertEquals(List.of(expected), rows(converted)); + } + + @Test + void postJoinFilterRunsBeforeEmit() { + Join join = equalityJoin(JoinType.LEFT); + Join filtered = + Join.builder() + .from(join) + .postJoinFilter(sb.isNull(sb.fieldReference(join, 1))) + .remap(sb.remap(0)) + .build(); + + assertEquals(List.of(List.of(1)), rows(substraitToCalcite.convert(filtered))); + } + + @ParameterizedTest + @EnumSource( + value = JoinType.class, + names = {"LEFT", "RIGHT", "OUTER"}) + void protoPostJoinFilterPreservesUnmatchedRows(JoinType joinType) { + Join join = equalityJoin(joinType); + Join filtered = + Join.builder() + .from(join) + .postJoinFilter( + sb.or(sb.isNull(sb.fieldReference(join, 0)), sb.isNull(sb.fieldReference(join, 1)))) + .remap(sb.remap(1, 0)) + .build(); + ExtensionCollector collector = new ExtensionCollector(); + io.substrait.proto.Rel proto = new RelProtoConverter(collector).toProto(filtered); + Rel decoded = new ProtoRelConverter(collector, extensions).from(proto); + + List> expected = + joinType == JoinType.LEFT + ? List.of(Arrays.asList(null, 1)) + : joinType == JoinType.RIGHT + ? List.of(Arrays.asList(3, null)) + : List.of(Arrays.asList(null, 1), Arrays.asList(3, null)); + assertEquals(expected, rows(substraitToCalcite.convert(decoded))); + } + + @ParameterizedTest + @EnumSource( + value = JoinType.class, + names = {"INNER", "LEFT", "LEFT_SEMI", "LEFT_ANTI"}) + void falseAndNullPostJoinFiltersProduceNoRows(JoinType joinType) { + Join join = equalityJoin(joinType); + for (Expression condition : List.of(sb.bool(false), ExpressionCreator.typedNull(N.BOOLEAN))) { + Join filtered = Join.builder().from(join).postJoinFilter(condition).build(); + + assertEquals(List.of(), rows(substraitToCalcite.convert(filtered))); + } + } + + @ParameterizedTest + @EnumSource( + value = JoinType.class, + names = {"LEFT_SEMI", "LEFT_ANTI"}) + void postJoinFilterUsesSemiAndAntiOutput(JoinType joinType) { + Join join = equalityJoin(joinType); + Join filtered = + Join.builder() + .from(join) + .postJoinFilter(sb.equal(sb.fieldReference(join, 0), sb.i32(1))) + .build(); + + assertEquals( + joinType == JoinType.LEFT_ANTI ? List.of(List.of(1)) : List.of(), + rows(substraitToCalcite.convert(filtered))); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void readFilterRetainsEnclosingCorrelation(boolean named) { + Rel outer = integers(1, 2).withRelAnchor(7); + FieldReference outerRef = + FieldReference.newRootStructOuterReferenceByRelReference(0, outer.getRecordType(), 7); + VirtualTableScan virtual = integers(1, 2); + Expression condition = sb.equal(sb.fieldReference(virtual, 0), outerRef); + Rel inner = + named + ? NamedScan.builder() + .initialSchema(virtual.getInitialSchema()) + .addNames("example") + .filter(condition) + .build() + : VirtualTableScan.builder().from(virtual).filter(condition).build(); + Rel plan = sb.filter(input -> sb.exists(inner), outer); + + Filter converted = assertInstanceOf(Filter.class, substraitToCalcite.convert(plan)); + assertFalse(converted.getVariablesSet().isEmpty()); + RexSubQuery exists = assertInstanceOf(RexSubQuery.class, converted.getCondition()); + Filter innerFilter = assertInstanceOf(Filter.class, exists.rel); + assertTrue(innerFilter.getCondition().toString().contains("$cor")); + } + + @Test + void postJoinFilterRetainsEnclosingCorrelation() { + Rel outer = integers(1, 2).withRelAnchor(7); + FieldReference outerRef = + FieldReference.newRootStructOuterReferenceByRelReference(0, outer.getRecordType(), 7); + Join join = equalityJoin(JoinType.INNER); + Join inner = + Join.builder() + .from(join) + .postJoinFilter(sb.equal(sb.fieldReference(join, 0), outerRef)) + .build(); + + Rel plan = sb.filter(input -> sb.exists(inner), outer); + Filter converted = assertInstanceOf(Filter.class, substraitToCalcite.convert(plan)); + assertFalse(converted.getVariablesSet().isEmpty()); + RexSubQuery exists = assertInstanceOf(RexSubQuery.class, converted.getCondition()); + Filter innerFilter = assertInstanceOf(Filter.class, exists.rel); + assertInstanceOf(org.apache.calcite.rel.core.Join.class, innerFilter.getInput()); + assertTrue(innerFilter.getCondition().toString().contains("$cor")); + } + + private Join equalityJoin(JoinType joinType) { + return sb.join( + input -> sb.equal(sb.fieldReference(input, 0), sb.fieldReference(input, 1)), + joinType, + integers(1, 2), + integers(2, 3)); + } + + private VirtualTableScan integers(int... values) { + return VirtualTableScan.builder() + .initialSchema(NamedStruct.of(List.of("id"), R.struct(R.I32))) + .rows( + Arrays.stream(values) + .mapToObj(value -> ExpressionCreator.nestedStruct(false, List.of(sb.i32(value)))) + .collect(Collectors.toList())) + .build(); + } + + private List> rows(RelNode rel) { + DataContext dataContext = + new DataContext() { + @Override + public SchemaPlus getRootSchema() { + return Frameworks.createRootSchema(true); + } + + @Override + public JavaTypeFactory getTypeFactory() { + return new JavaTypeFactoryImpl(); + } + + @Override + public QueryProvider getQueryProvider() { + return null; + } + + @Override + public Object get(String name) { + return null; + } + }; + RelNode executable = plan(rel, VirtualTableExpansionRule.instance()); + try (Interpreter interpreter = new Interpreter(dataContext, executable)) { + return interpreter.toList().stream().map(Arrays::asList).collect(Collectors.toList()); + } + } +} diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala index c3fa30dba..ae808c13b 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala @@ -227,7 +227,7 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes throw new UnsupportedOperationException(s"Unsupported join type $other") } val plan = Join(left, right, joinType, condition, hint = JoinHint.NONE) - remap(plan, join.getRemap) + remap(applyFilter(plan, join.getPostJoinFilter, context), join.getRemap) } } @@ -415,7 +415,7 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes case _ => LocalRelation(ToSparkType.toAttributeSeq(virtualTableScan.getInitialSchema), rows) } - remap(plan, virtualTableScan.getRemap) + remap(applyFilter(plan, virtualTableScan.getFilter, context), virtualTableScan.getRemap) } override def visit( @@ -425,7 +425,7 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes case m: MultiInstanceRelation => m.newInstance() case other => other } - remap(plan, namedScan.getRemap) + remap(applyFilter(plan, namedScan.getFilter, context), namedScan.getRemap) } override def visit(localFiles: LocalFiles, context: EmptyVisitationContext): LogicalPlan = { @@ -458,7 +458,20 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes catalogTable = None, isStreaming = false ) - remap(plan, localFiles.getRemap) + remap(applyFilter(plan, localFiles.getFilter, context), localFiles.getRemap) + } + + private def applyFilter( + plan: LogicalPlan, + predicate: Optional[SExpression], + context: EmptyVisitationContext): LogicalPlan = { + if (predicate.isPresent) { + withChild(plan) { + Filter(predicate.get.accept(expressionConverter, context), plan) + } + } else { + plan + } } def convertFileFormat(fileFormat: FileFormat): (SparkFileFormat, Map[String, String]) = { diff --git a/spark/src/test/scala/io/substrait/spark/MandatoryPredicatesSuite.scala b/spark/src/test/scala/io/substrait/spark/MandatoryPredicatesSuite.scala new file mode 100644 index 000000000..45e402b55 --- /dev/null +++ b/spark/src/test/scala/io/substrait/spark/MandatoryPredicatesSuite.scala @@ -0,0 +1,188 @@ +package io.substrait.spark + +import io.substrait.spark.logical.{ToLogicalPlan, ToSubstraitRel} + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.Row +import org.apache.spark.sql.classic.DatasetUtil +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.StructType + +import io.substrait.`type`.TypeCreator +import io.substrait.dsl.SubstraitBuilder +import io.substrait.expression.{Expression, ExpressionCreator} +import io.substrait.relation.{AbstractReadRel, Join, LocalFiles => LocalFilesRel, NamedScan, Rel, VirtualTableScan} +import io.substrait.util.EmptyVisitationContext + +class MandatoryPredicatesSuite extends SparkFunSuite with SharedSparkSession { + + private val builder = new SubstraitBuilder + + override def beforeAll(): Unit = { + super.beforeAll() + sparkContext.setLogLevel("WARN") + } + + private def assertRows(rel: Rel, expected: Row*): Unit = { + val plan = rel.accept(new ToLogicalPlan(spark), EmptyVisitationContext.INSTANCE) + assert(plan.resolved) + val actual = DatasetUtil.fromLogicalPlan(spark, plan).collect().toSeq + assertResult(expected.sortBy(_.toString))(actual.sortBy(_.toString)) + } + + private def withScan(kind: String)(body: AbstractReadRel => Unit): Unit = { + val data = spark.sql("SELECT * FROM VALUES (1, true), (2, false), (3, NULL) AS t(id, keep)") + kind match { + case "named table" => + withTempView("predicate_scan") { + data.createOrReplaceTempView("predicate_scan") + body( + NamedScan + .builder() + .addNames("predicate_scan") + .initialSchema(ToSubstraitType.toNamedStruct(data.schema)) + .build()) + } + case "virtual table" => + body( + new ToSubstraitRel() + .visit(data.queryExecution.optimizedPlan) + .asInstanceOf[VirtualTableScan]) + case "local files" => + withTempPath { + path => + data.write.parquet(path.getAbsolutePath) + val read = spark.read.parquet(path.getAbsolutePath) + body( + new ToSubstraitRel() + .visit(read.queryExecution.optimizedPlan) + .asInstanceOf[LocalFilesRel]) + } + case other => throw new IllegalArgumentException(s"Unknown scan kind: $other") + } + } + + private def filteredScan(scan: AbstractReadRel, predicate: Expression): Rel = { + val remap = Rel.Remap.offset(0, 1) + scan match { + case named: NamedScan => + NamedScan.builder().from(named).filter(predicate).remap(remap).build() + case virtual: VirtualTableScan => + VirtualTableScan.builder().from(virtual).filter(predicate).remap(remap).build() + case files: LocalFilesRel => + LocalFilesRel.builder().from(files).filter(predicate).remap(remap).build() + case other => throw new IllegalArgumentException(s"Unknown scan: $other") + } + } + + Seq("named table", "virtual table", "local files").foreach { + kind => + test(s"mandatory read predicates on $kind run before emit") { + withScan(kind) { + scan => + assertRows(scan, Row(1, true), Row(2, false), Row(3, null)) + assertRows(filteredScan(scan, builder.bool(false))) + assertRows( + filteredScan(scan, ExpressionCreator.typedNull(TypeCreator.NULLABLE.BOOLEAN))) + // The predicate column is omitted from the output, and NULL must be rejected. + assertRows(filteredScan(scan, builder.fieldReference(scan, 1)), Row(1)) + } + } + } + + test("mandatory predicate on a zero-column virtual table") { + val scan = VirtualTableScan + .builder() + .initialSchema(ToSubstraitType.toNamedStruct(new StructType())) + .addRows(ExpressionCreator.nestedStruct(false)) + .build() + assertRows(scan, Row()) + assertRows(VirtualTableScan.builder().from(scan).filter(builder.bool(false)).build()) + } + + private def join(joinType: Join.JoinType): Join = { + val left = new ToSubstraitRel().visit( + spark.sql("SELECT * FROM VALUES (1), (2) AS l(id)").queryExecution.optimizedPlan) + val right = new ToSubstraitRel().visit( + spark.sql("SELECT * FROM VALUES (1), (3) AS r(id)").queryExecution.optimizedPlan) + builder.join( + (input: SubstraitBuilder.JoinInput) => + builder.equal(builder.fieldReference(input, 0), builder.fieldReference(input, 1)), + joinType, + left, + right) + } + + Seq(Join.JoinType.INNER, Join.JoinType.LEFT, Join.JoinType.RIGHT, Join.JoinType.OUTER).foreach { + joinType => + test(s"mandatory false post-join predicate on $joinType") { + val input = join(joinType) + assertRows(Join.builder().from(input).postJoinFilter(builder.bool(false)).build()) + } + } + + test("left outer post-join predicate sees null-extended output before emit") { + val input = join(Join.JoinType.LEFT) + assertRows(input, Row(1, 1), Row(2, null)) + val filtered = Join + .builder() + .from(input) + .postJoinFilter(builder.isNull(builder.fieldReference(input, 1))) + .remap(Rel.Remap.offset(0, 1)) + .build() + assertRows(filtered, Row(2)) + } + + test("right outer post-join predicate sees null-extended output before emit") { + val input = join(Join.JoinType.RIGHT) + assertRows(input, Row(1, 1), Row(null, 3)) + val filtered = Join + .builder() + .from(input) + .postJoinFilter(builder.isNull(builder.fieldReference(input, 0))) + .remap(Rel.Remap.offset(1, 1)) + .build() + assertRows(filtered, Row(3)) + } + + test("full outer post-join predicate filters both sides of the join") { + val input = join(Join.JoinType.OUTER) + val filtered = Join + .builder() + .from(input) + .postJoinFilter( + builder.or( + builder.isNull(builder.fieldReference(input, 0)), + builder.isNull(builder.fieldReference(input, 1)))) + .build() + assertRows(filtered, Row(2, null), Row(null, 3)) + } + + test("post-join predicate rejects null-extended rows") { + val input = join(Join.JoinType.LEFT) + val filtered = Join + .builder() + .from(input) + .postJoinFilter(builder.equal(builder.fieldReference(input, 1), builder.i32(1))) + .remap(Rel.Remap.offset(0, 1)) + .build() + assertRows(filtered, Row(1)) + } + + Seq(Join.JoinType.LEFT_SEMI, Join.JoinType.LEFT_ANTI).foreach { + joinType => + test(s"post-join predicate uses $joinType output") { + val input = join(joinType) + val filtered = Join + .builder() + .from(input) + .postJoinFilter(builder.equal(builder.fieldReference(input, 0), builder.i32(2))) + .build() + if (joinType == Join.JoinType.LEFT_ANTI) { + assertRows(filtered, Row(2)) + } else { + assertRows(filtered) + } + } + } +}