chore: show a bug in Expand conversion

substrait-io · Oct 27, 2024 · b258427 · b258427
1 parent e3139c6
commit b258427
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/spark/src/main/scala/io/substrait/spark/SparkExtension.scala b/spark/src/main/scala/io/substrait/spark/SparkExtension.scala
@@ -34,6 +34,8 @@ object SparkExtension {
   private val EXTENSION_COLLECTION: SimpleExtension.ExtensionCollection =
     SimpleExtension.loadDefaults()
 
+  val COLLECTION: SimpleExtension.ExtensionCollection = EXTENSION_COLLECTION.merge(SparkImpls)
+
   lazy val SparkScalarFunctions: Seq[SimpleExtension.ScalarFunctionVariant] = {
     val ret = new collection.mutable.ArrayBuffer[SimpleExtension.ScalarFunctionVariant]()
     ret.appendAll(EXTENSION_COLLECTION.scalarFunctions().asScala)

diff --git a/spark/src/test/scala/io/substrait/spark/SubstraitPlanTestBase.scala b/spark/src/test/scala/io/substrait/spark/SubstraitPlanTestBase.scala
@@ -17,16 +17,14 @@
 package io.substrait.spark
 
 import io.substrait.spark.logical.{ToLogicalPlan, ToSubstraitRel}
-
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.util.resourceToString
 import org.apache.spark.sql.test.SharedSparkSession
-
 import io.substrait.debug.TreePrinter
 import io.substrait.extension.ExtensionCollector
 import io.substrait.plan.{Plan, PlanProtoConverter, ProtoPlanConverter}
 import io.substrait.proto
-import io.substrait.relation.RelProtoConverter
+import io.substrait.relation.{ProtoRelConverter, RelProtoConverter}
 import org.scalactic.Equality
 import org.scalactic.source.Position
 import org.scalatest.Succeeded
@@ -93,6 +91,10 @@ trait SubstraitPlanTestBase { self: SharedSparkSession =>
     require(logicalPlan2.resolved);
     val pojoRel2 = new ToSubstraitRel().visit(logicalPlan2)
 
+    val extensionCollector = new ExtensionCollector;
+    val proto = new RelProtoConverter(extensionCollector).toProto(pojoRel)
+    new ProtoRelConverter(extensionCollector, SparkExtension.COLLECTION).from(proto)
+
     pojoRel2.shouldEqualPlainly(pojoRel)
     logicalPlan2
   }