[CORE] Add schema validation for broadcast exchange (#4608)

apache · Feb 2, 2024 · b0d50c0 · b0d50c0
1 parent 503afb1
commit b0d50c0
Show file tree

Hide file tree

Showing 7 changed files with 21 additions and 14 deletions.
diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/ValidatorApiImpl.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/ValidatorApiImpl.scala
@@ -81,7 +81,7 @@ class ValidatorApiImpl extends ValidatorApi {
       case array: ArrayType =>
         doSchemaValidate(array.elementType)
       case _ =>
-        Some(s"do not support data type: $schema")
+        Some(s"Schema / data type not supported: $schema")
     }
   }
 

diff --git a/backends-velox/src/main/scala/io/glutenproject/execution/RowToVeloxColumnarExec.scala b/backends-velox/src/main/scala/io/glutenproject/execution/RowToVeloxColumnarExec.scala
@@ -16,7 +16,7 @@
  */
 package io.glutenproject.execution
 
-import io.glutenproject.backendsapi.velox.ValidatorApiImpl
+import io.glutenproject.backendsapi.BackendsApiManager
 import io.glutenproject.columnarbatch.ColumnarBatches
 import io.glutenproject.exec.Runtimes
 import io.glutenproject.memory.arrowalloc.ArrowBufferAllocators
@@ -45,7 +45,7 @@ import scala.collection.mutable.ListBuffer
 case class RowToVeloxColumnarExec(child: SparkPlan) extends RowToColumnarExecBase(child = child) {
 
   override def doExecuteColumnarInternal(): RDD[ColumnarBatch] = {
-    new ValidatorApiImpl().doSchemaValidate(schema).foreach {
+    BackendsApiManager.getValidatorApiInstance.doSchemaValidate(schema).foreach {
       reason =>
         throw new UnsupportedOperationException(
           s"Input schema contains unsupported type when convert row to columnar for $schema " +

diff --git a/...s-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarCachedBatchSerializer.scala b/...s-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarCachedBatchSerializer.scala
@@ -18,7 +18,6 @@ package org.apache.spark.sql.execution
 
 import io.glutenproject.GlutenConfig
 import io.glutenproject.backendsapi.BackendsApiManager
-import io.glutenproject.backendsapi.velox.ValidatorApiImpl
 import io.glutenproject.columnarbatch.ColumnarBatches
 import io.glutenproject.exec.Runtimes
 import io.glutenproject.execution.{RowToVeloxColumnarExec, VeloxColumnarToRowExec}
@@ -91,7 +90,7 @@ class ColumnarCachedBatchSerializer extends CachedBatchSerializer with SQLConfHe
   }
 
   private def validateSchema(schema: StructType): Boolean = {
-    val reason = new ValidatorApiImpl().doSchemaValidate(schema)
+    val reason = BackendsApiManager.getValidatorApiInstance.doSchemaValidate(schema)
     if (reason.isDefined) {
       logInfo(s"Columnar cache does not support schema $schema, due to ${reason.get}")
       false

diff --git a/...en-core/src/main/scala/org/apache/spark/sql/execution/ColumnarBroadcastExchangeExec.scala b/...en-core/src/main/scala/org/apache/spark/sql/execution/ColumnarBroadcastExchangeExec.scala
@@ -131,12 +131,23 @@ case class ColumnarBroadcastExchangeExec(mode: BroadcastMode, child: SparkPlan)
     ColumnarBroadcastExchangeExec(mode.canonicalized, child.canonicalized)
   }
 
-  override protected def doValidateInternal(): ValidationResult = mode match {
-    case _: HashedRelationBroadcastMode =>
-      ValidationResult.ok
-    case _ =>
-      // TODO IdentityBroadcastMode not supported. Need to support BroadcastNestedLoopJoin first.
-      ValidationResult.notOk("Only support HashedRelationBroadcastMode for now.")
+  override protected def doValidateInternal(): ValidationResult = {
+    mode match {
+      case _: HashedRelationBroadcastMode =>
+      case _ =>
+        // TODO IdentityBroadcastMode not supported. Need to support BroadcastNestedLoopJoin first.
+        return ValidationResult.notOk("Only support HashedRelationBroadcastMode for now.")
+    }
+    BackendsApiManager.getValidatorApiInstance
+      .doSchemaValidate(schema)
+      .map {
+        reason =>
+          {
+            ValidationResult.notOk(
+              s"Unsupported schema in broadcast exchange: $schema, reason: $reason")
+          }
+      }
+      .getOrElse(ValidationResult.ok)
   }
 
   override def doPrepare(): Unit = {

diff --git a/gluten-ut/spark32/src/test/scala/io/glutenproject/utils/velox/VeloxTestSettings.scala b/gluten-ut/spark32/src/test/scala/io/glutenproject/utils/velox/VeloxTestSettings.scala
@@ -293,7 +293,6 @@ class VeloxTestSettings extends BackendTestSettings {
   enableSuite[GlutenDataFrameImplicitsSuite]
   enableSuite[GlutenGeneratorFunctionSuite]
   enableSuite[GlutenDataFrameTimeWindowingSuite]
-    .exclude("time window joins") // FIXME hongze
   enableSuite[GlutenDataFrameSessionWindowingSuite]
   enableSuite[GlutenBroadcastExchangeSuite]
   enableSuite[GlutenDataFramePivotSuite]

diff --git a/gluten-ut/spark33/src/test/scala/io/glutenproject/utils/velox/VeloxTestSettings.scala b/gluten-ut/spark33/src/test/scala/io/glutenproject/utils/velox/VeloxTestSettings.scala
@@ -998,7 +998,6 @@ class VeloxTestSettings extends BackendTestSettings {
       "SPARK-9083: sort with non-deterministic expressions"
     )
   enableSuite[GlutenDataFrameTimeWindowingSuite]
-    .exclude("time window joins") // FIXME hongze
   enableSuite[GlutenDataFrameTungstenSuite]
   enableSuite[GlutenDataFrameWindowFramesSuite]
     // Local window fixes are not added.

diff --git a/gluten-ut/spark34/src/test/scala/io/glutenproject/utils/velox/VeloxTestSettings.scala b/gluten-ut/spark34/src/test/scala/io/glutenproject/utils/velox/VeloxTestSettings.scala
@@ -1001,7 +1001,6 @@ class VeloxTestSettings extends BackendTestSettings {
     // test for sort node not present but gluten uses shuffle hash join
     .exclude("SPARK-41048: Improve output partitioning and ordering with AQE cache")
   enableSuite[GlutenDataFrameTimeWindowingSuite]
-    .exclude("time window joins") // FIXME hongze
   enableSuite[GlutenDataFrameTungstenSuite]
   enableSuite[GlutenDataFrameWindowFramesSuite]
     // Local window fixes are not added.