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..c9ae3b210 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala @@ -333,17 +333,18 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes } else { projectExprs.map(toNamedExpression) } + val inheritedExpressions = if (createProject) output.map(_.toAttribute) else output + val allExpressions = inheritedExpressions ++ projectList + val remapped = if (project.getRemap.isPresent) { + project.getRemap.get().indices().asScala.map(allExpressions(_)).toSeq + } else { + allExpressions + } if (createProject) { - val allExpressions = output.map(_.toAttribute) ++ projectList - val remapped = if (project.getRemap.isPresent) { - project.getRemap.get().indices().asScala.map(allExpressions(_)).toSeq - } else { - allExpressions - } Project(remapped, child) } else { val aggregate: Aggregate = child.asInstanceOf[Aggregate] - aggregate.copy(aggregateExpressions = projectList) + aggregate.copy(aggregateExpressions = remapped) } } } diff --git a/spark/src/test/scala/io/substrait/spark/ProjectAggregateSuite.scala b/spark/src/test/scala/io/substrait/spark/ProjectAggregateSuite.scala new file mode 100644 index 000000000..f19cf28d7 --- /dev/null +++ b/spark/src/test/scala/io/substrait/spark/ProjectAggregateSuite.scala @@ -0,0 +1,108 @@ +package io.substrait.spark + +import io.substrait.spark.logical.ToLogicalPlan + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate => SparkAggregate} +import org.apache.spark.sql.classic.DatasetUtil +import org.apache.spark.sql.test.SharedSparkSession + +import io.substrait.`type`.{NamedStruct, TypeCreator} +import io.substrait.dsl.SubstraitBuilder +import io.substrait.expression.ExpressionCreator +import io.substrait.extension.DefaultExtensionCatalog +import io.substrait.relation.{Aggregate, Project, Rel, VirtualTableScan} + +import java.util.Arrays + +class ProjectAggregateSuite + extends SparkFunSuite + with SharedSparkSession + with SubstraitPlanTestBase { + + private val builder = new SubstraitBuilder(DefaultExtensionCatalog.DEFAULT_COLLECTION) + + private def projectOverAggregate(): Project = { + val required = TypeCreator.REQUIRED + val input = VirtualTableScan + .builder() + .initialSchema(NamedStruct + .of(Arrays.asList("group", "value"), required.struct(required.I32, required.I64))) + .addRows(ExpressionCreator + .nestedStruct(false, ExpressionCreator.i32(false, 1), ExpressionCreator.i64(false, 10))) + .addRows(ExpressionCreator + .nestedStruct(false, ExpressionCreator.i32(false, 1), ExpressionCreator.i64(false, 20))) + .addRows(ExpressionCreator + .nestedStruct(false, ExpressionCreator.i32(false, 2), ExpressionCreator.i64(false, 7))) + .build() + val aggregate = Aggregate + .builder() + .input(input) + .addGroupings( + Aggregate.Grouping + .builder() + .addExpressions(builder.fieldReference(input, 0)) + .build()) + .addMeasures(builder.sum(input, 1)) + .build() + + Project + .builder() + .input(aggregate) + .addExpressions(ExpressionCreator.i32(false, 99)) + .addExpressions(builder.fieldReference(aggregate, 0)) + .addExpressions( + builder.add(builder.fieldReference(aggregate, 1), ExpressionCreator.i64(false, 1))) + .build() + } + + private def assertRows(project: Project, expected: Seq[Row]): Unit = { + val converted = new ToLogicalPlan(spark).convert(project) + assert(converted.isInstanceOf[SparkAggregate]) + assertResult(project.getRecordType.fields().size())(converted.output.size) + val actual = DatasetUtil.fromLogicalPlan(spark, converted).collect().toSeq + assertResult(expected.sortBy(_.toString))(actual.sortBy(_.toString)) + } + + test("project retains inherited grouping and measure outputs") { + assertRows(projectOverAggregate(), Seq(Row(1, 30L, 99, 1, 31L), Row(2, 7L, 99, 2, 8L))) + } + + test("project emit can select only inherited aggregate outputs") { + val project = Project + .builder() + .from(projectOverAggregate()) + .remap(Rel.Remap.of(Arrays.asList(0, 1))) + .build() + assertRows(project, Seq(Row(1, 30L), Row(2, 7L))) + } + + test("project emit can reorder and duplicate inherited and appended outputs") { + val project = Project + .builder() + .from(projectOverAggregate()) + .remap(Rel.Remap.of(Arrays.asList(4, 1, 0, 1, 2))) + .build() + assertRows(project, Seq(Row(31L, 30L, 1, 30L, 99), Row(8L, 7L, 2, 7L, 99))) + } + + test("project emit can remove all aggregate outputs") { + val project = Project + .builder() + .from(projectOverAggregate()) + .remap(Rel.Remap.of(Arrays.asList())) + .build() + assertRows(project, Seq(Row(), Row())) + } + + test("Spark projected aggregate preserves roundtrip shape and rows") { + val query = + "select group_id + 1 as group_key, sum(value) + 1 as total " + + "from (values (1, 10), (1, 20), (2, 7)) as input(group_id, value) " + + "group by group_id + 1" + val converted = assertSqlSubstraitRelRoundTrip(query) + val actual = DatasetUtil.fromLogicalPlan(spark, converted).collect().toSeq + assertResult(Seq(Row(2, 31L), Row(3, 8L)))(actual.sortBy(_.getInt(0))) + } +}