[SPARK-52492][SQL] Make InMemoryRelation.convertToColumnarIfPossible customizable

https://issues.apache.org/jira/browse/SPARK-52492

### What changes were proposed in this pull request?

This PR moves `InMemoryRelation.convertToColumnarIfPossible` to as a public API of `CachedBatchSerializer`.

### Why are the changes needed?

TL;DR: So that plugins like Gluten could have the relevant logic customized for their own catch serializers.

Currently, InMemoryRelation.convertToColumnarIfPossible is highly coupled with vanilla Spark's columnar processing logic. It unwraps the input columnar plan by removing the topmost ColumnarToRowExec, the assumes that the outcome RDD after this process can be recognized by the user-customized cache serializer.

But sometimes this assertion is invalid. As in the Apache Gluten project, we may continue distiguishing plans that are all have `supportsColumnar=true` with different columnar batch types. So even the topmost `ColumnarToRowExec` is removed, we still don't know whether the columnar RDD unwrapped can be accepted by Gluten's cache serializer (assuming it only handles one certain type of columnar batch or something).

So in Gluten we had a rule to workaround the logic in InMemoryRelation.convertToColumnarIfPossible: https://github.com/apache/incubator-gluten/blob/c6461b4e0c7d3022a31fa832aeab588b1a3200e6/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/MiscColumnarRules.scala#L192-L217. This is the best way we had thought about to get through the issue but it's still not elegant, especially the rule is even caller-sensitive as it needs to determine whether it's called in the caching planning process or not.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

With the added UTs.

Closes #51189 from zhztheplayer/wip-cache-customize-unwrap.

Lead-authored-by: Hongze Zhang <[email protected]>
Co-authored-by: Kent Yao <[email protected]>
Signed-off-by: Kent Yao <[email protected]>

Author: zhztheplayer

sql/core/src/main/scala/org/apache/spark/sql/columnar/CachedBatchSerializer.scala

@@ -24,6 +24,8 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, BindReferences, EqualNullSafe, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, In, IsNotNull, IsNull, Length, LessThan, LessThanOrEqual, Literal, Or, Predicate, StartsWith}
+import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, SparkPlan, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
 import org.apache.spark.sql.execution.columnar.{ColumnStatisticsSchema, PartitionStatistics}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{AtomicType, BinaryType, StructType}
@@ -58,6 +60,40 @@ trait CachedBatchSerializer extends Serializable {
    */
   def supportsColumnarInput(schema: Seq[Attribute]): Boolean
 
+  /**
+   * Attempt to convert a query plan to its columnar equivalence for columnar caching.
+   * Called on the query plan that is about to be cached once [[supportsColumnarInput]] returns
+   * true on its output schema.
+   *
+   * The default implementation works by stripping the topmost columnar-to-row transition to
+   * expose the columnar-based plan to the serializer.
+   *
+   * @param plan The plan to convert.
+   * @return The output plan. Could either be a columnar plan if the input plan is convertible, or
+   *         the input plan unchanged if no viable conversion can be done.
+   */
+  @DeveloperApi
+  @Since("4.1.0")
+  def convertToColumnarPlanIfPossible(plan: SparkPlan): SparkPlan = plan match {
+    case gen: WholeStageCodegenExec =>
+      gen.child match {
+        case c2r: ColumnarToRowTransition =>
+          c2r.child match {
+            case ia: InputAdapter => ia.child
+            case _ => plan
+          }
+        case _ => plan
+      }
+    case c2r: ColumnarToRowTransition => // This matches when whole stage code gen is disabled.
+      c2r.child
+    case adaptive: AdaptiveSparkPlanExec =>
+      // If AQE is enabled for cached plan and table cache supports columnar in, we should mark
+      // `AdaptiveSparkPlanExec.supportsColumnar` as true to avoid inserting `ColumnarToRow`, so
+      // that `CachedBatchSerializer` can use `convertColumnarBatchToCachedBatch` to cache data.
+      adaptive.copy(supportsColumnar = true)
+    case _ => plan
+  }
+
   /**
    * Convert an `RDD[InternalRow]` into an `RDD[CachedBatch]` in preparation for caching the data.
    * @param input the input `RDD` to be converted.

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala

@@ -381,22 +381,8 @@ object InMemoryRelation {
   /* Visible for testing */
   private[columnar] def clearSerializer(): Unit = synchronized { ser = None }
 
-  def convertToColumnarIfPossible(plan: SparkPlan): SparkPlan = plan match {
-    case gen: WholeStageCodegenExec => gen.child match {
-      case c2r: ColumnarToRowTransition => c2r.child match {
-        case ia: InputAdapter => ia.child
-        case _ => plan
-      }
-      case _ => plan
-    }
-    case c2r: ColumnarToRowTransition => // This matches when whole stage code gen is disabled.
-      c2r.child
-    case adaptive: AdaptiveSparkPlanExec =>
-      // If AQE is enabled for cached plan and table cache supports columnar in, we should mark
-      // `AdaptiveSparkPlanExec.supportsColumnar` as true to avoid inserting `ColumnarToRow`, so
-      // that `CachedBatchSerializer` can use `convertColumnarBatchToCachedBatch` to cache data.
-      adaptive.copy(supportsColumnar = true)
-    case _ => plan
+  def convertToColumnarIfPossible(plan: SparkPlan): SparkPlan = {
+    getSerializer(plan.conf).convertToColumnarPlanIfPossible(plan)
   }
 
   def apply(
@@ -406,7 +392,7 @@ object InMemoryRelation {
     val optimizedPlan = qe.optimizedPlan
     val serializer = getSerializer(optimizedPlan.conf)
     val child = if (serializer.supportsColumnarInput(optimizedPlan.output)) {
-      convertToColumnarIfPossible(qe.executedPlan)
+      serializer.convertToColumnarPlanIfPossible(qe.executedPlan)
     } else {
       qe.executedPlan
     }
@@ -433,8 +419,9 @@ object InMemoryRelation {
 
   def apply(cacheBuilder: CachedRDDBuilder, qe: QueryExecution): InMemoryRelation = {
     val optimizedPlan = qe.optimizedPlan
-    val newBuilder = if (cacheBuilder.serializer.supportsColumnarInput(optimizedPlan.output)) {
-      cacheBuilder.copy(cachedPlan = convertToColumnarIfPossible(qe.executedPlan))
+    val serializer = cacheBuilder.serializer
+    val newBuilder = if (serializer.supportsColumnarInput(optimizedPlan.output)) {
+      cacheBuilder.copy(cachedPlan = serializer.convertToColumnarPlanIfPossible(qe.executedPlan))
     } else {
       cacheBuilder.copy(cachedPlan = qe.executedPlan)
     }

sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala

@@ -25,7 +25,7 @@ import org.apache.spark.sql.{QueryTest, Row}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, UnsafeProjection}
 import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer}
-import org.apache.spark.sql.execution.ColumnarToRowExec
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, WholeStageCodegenExec}
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper}
 import org.apache.spark.sql.execution.columnar.InMemoryRelation.clearSerializer
 import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
@@ -122,6 +122,25 @@ class TestSingleIntColumnarCachedBatchSerializer extends CachedBatchSerializer {
   }
 }
 
+/**
+ * An equivalence of Spark's [[DefaultCachedBatchSerializer]] while the API
+ * [[convertToColumnarPlanIfPossible]] is being tested.
+ */
+class DefaultCachedBatchSerializerNoUnwrap extends DefaultCachedBatchSerializer {
+  override def supportsColumnarInput(schema: Seq[Attribute]): Boolean = {
+    // Return true to let Spark call #convertToColumnarPlanIfPossible to unwrap the input
+    // columnar plan out from the guard of the topmost ColumnarToRowExec.
+    true
+  }
+
+  override def convertToColumnarPlanIfPossible(plan: SparkPlan): SparkPlan = {
+    assert(!plan.supportsColumnar)
+    // Disable the unwrapping code path from default CachedBatchSerializer so
+    // Spark will keep the topmost columnar-to-row plan node.
+    plan
+  }
+}
+
 class CachedBatchSerializerSuite extends QueryTest
   with SharedSparkSession with AdaptiveSparkPlanHelper {
   import testImplicits._
@@ -180,3 +199,41 @@ class CachedBatchSerializerSuite extends QueryTest
     }
   }
 }
+
+
+class CachedBatchSerializerNoUnwrapSuite extends QueryTest
+  with SharedSparkSession with AdaptiveSparkPlanHelper {
+
+  import testImplicits._
+
+  override protected def sparkConf: SparkConf = {
+    super.sparkConf.set(
+      StaticSQLConf.SPARK_CACHE_SERIALIZER.key,
+      classOf[DefaultCachedBatchSerializerNoUnwrap].getName)
+  }
+
+  test("Do not unwrap ColumnarToRowExec") {
+    withTempPath { workDir =>
+      val workDirPath = workDir.getAbsolutePath
+      val input = Seq(100, 200).toDF("count")
+      input.write.parquet(workDirPath)
+      val data = spark.read.parquet(workDirPath)
+      data.cache()
+      val df = data.union(data)
+      assert(df.count() == 4)
+      checkAnswer(df, Row(100) :: Row(200) :: Row(100) :: Row(200) :: Nil)
+
+      val finalPlan = df.queryExecution.executedPlan
+      val cachedPlans = finalPlan.collect {
+        case i: InMemoryTableScanExec => i.relation.cachedPlan
+      }
+      assert(cachedPlans.length == 2)
+      cachedPlans.foreach {
+        cachedPlan =>
+          assert(cachedPlan.isInstanceOf[WholeStageCodegenExec])
+          assert(cachedPlan.asInstanceOf[WholeStageCodegenExec]
+            .child.isInstanceOf[ColumnarToRowExec])
+      }
+    }
+  }
+}