expose large var type to all other api calls

Author: wecharyu

backends-velox/src-iceberg/main/scala/org/apache/gluten/connector/write/IcebergDataWriteFactory.scala

@@ -24,9 +24,8 @@ import org.apache.gluten.runtime.Runtimes
 import org.apache.gluten.utils.ArrowAbiUtil
 
 import org.apache.spark.sql.connector.write.DataWriter
-import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
-import org.apache.spark.sql.utils.SparkArrowUtil
+import org.apache.spark.sql.utils.SparkSchemaUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
 import org.apache.arrow.c.ArrowSchema
@@ -101,7 +100,7 @@ case class IcebergDataWriteFactory(
       partitionSpec: IcebergPartitionSpec,
       field: IcebergNestedField,
       icebergProperties: util.HashMap[String, String]): (Long, IcebergWriteJniWrapper) = {
-    val schema = SparkArrowUtil.toArrowSchema(localSchema, SQLConf.get.sessionLocalTimeZone)
+    val schema = SparkSchemaUtil.toArrowSchema(localSchema)
     val arrowAlloc = ArrowBufferAllocators.contextInstance()
     val cSchema = ArrowSchema.allocateNew(arrowAlloc)
     ArrowAbiUtil.exportSchema(arrowAlloc, schema, cSchema)

backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala

@@ -31,9 +31,8 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.execution.{BroadcastUtils, SparkPlan}
 import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
-import org.apache.spark.sql.utils.SparkArrowUtil
+import org.apache.spark.sql.utils.SparkSchemaUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.task.TaskResources
 import org.apache.spark.unsafe.Platform
@@ -128,8 +127,7 @@ object RowToVeloxColumnarExec {
       return Iterator.empty
     }
 
-    val arrowSchema =
-      SparkArrowUtil.toArrowSchema(schema, SQLConf.get.sessionLocalTimeZone)
+    val arrowSchema = SparkSchemaUtil.toArrowSchema(schema)
     val runtime = Runtimes.contextInstance(BackendsApiManager.getBackendName, "RowToColumnar")
     val jniWrapper = NativeRowToColumnarJniWrapper.create(runtime)
     val arrowAllocator = ArrowBufferAllocators.contextInstance()

backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala

@@ -32,8 +32,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSeq, BindReferences, BoundReference, Expression, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode
 import org.apache.spark.sql.execution.joins.{BuildSideRelation, HashedRelationBroadcastMode}
-import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.utils.SparkArrowUtil
+import org.apache.spark.sql.utils.SparkSchemaUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.task.TaskResources
 import org.apache.spark.util.KnownSizeEstimation
@@ -117,9 +116,8 @@ case class ColumnarBuildSideRelation(
     val serializeHandle: Long = {
       val allocator = ArrowBufferAllocators.contextInstance()
       val cSchema = ArrowSchema.allocateNew(allocator)
-      val arrowSchema = SparkArrowUtil.toArrowSchema(
-        SparkShimLoader.getSparkShims.structFromAttributes(output),
-        SQLConf.get.sessionLocalTimeZone)
+      val arrowSchema = SparkSchemaUtil.toArrowSchema(
+        SparkShimLoader.getSparkShims.structFromAttributes(output))
       ArrowAbiUtil.exportSchema(allocator, arrowSchema, cSchema)
       val handle = jniWrapper
         .init(cSchema.memoryAddress())
@@ -170,9 +168,8 @@ case class ColumnarBuildSideRelation(
         val serializeHandle: Long = {
           val allocator = ArrowBufferAllocators.contextInstance()
           val cSchema = ArrowSchema.allocateNew(allocator)
-          val arrowSchema = SparkArrowUtil.toArrowSchema(
-            SparkShimLoader.getSparkShims.structFromAttributes(output),
-            SQLConf.get.sessionLocalTimeZone)
+          val arrowSchema = SparkSchemaUtil.toArrowSchema(
+            SparkShimLoader.getSparkShims.structFromAttributes(output))
           ArrowAbiUtil.exportSchema(allocator, arrowSchema, cSchema)
           val handle = jniWrapper
             .init(cSchema.memoryAddress())
@@ -362,9 +359,8 @@ case class ColumnarBuildSideRelation(
     val serializeHandle = {
       val allocator = ArrowBufferAllocators.contextInstance()
       val cSchema = ArrowSchema.allocateNew(allocator)
-      val arrowSchema = SparkArrowUtil.toArrowSchema(
-        SparkShimLoader.getSparkShims.structFromAttributes(output),
-        SQLConf.get.sessionLocalTimeZone)
+      val arrowSchema = SparkSchemaUtil.toArrowSchema(
+        SparkShimLoader.getSparkShims.structFromAttributes(output))
       ArrowAbiUtil.exportSchema(allocator, arrowSchema, cSchema)
       val handle = serializerJniWrapper.init(cSchema.memoryAddress())
       cSchema.close()

backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarCachedBatchSerializer.scala

@@ -39,7 +39,7 @@ import org.apache.spark.sql.columnar.SimpleMetricsCachedBatchSerializer
 import org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
-import org.apache.spark.sql.utils.SparkArrowUtil
+import org.apache.spark.sql.utils.SparkSchemaUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.storage.StorageLevel
 import org.apache.spark.unsafe.types.UTF8String
@@ -953,15 +953,14 @@ class ColumnarCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer
       }
       val shouldSelectAttributes = cacheAttributes != selectedAttributes
       val localSchema = toStructType(cacheAttributes)
-      val timezoneId = SQLConf.get.sessionLocalTimeZone
       input.mapPartitions {
         it =>
           val runtime = Runtimes.contextInstance(
             BackendsApiManager.getBackendName,
             "ColumnarCachedBatchSerializer#read")
           val jniWrapper = ColumnarBatchSerializerJniWrapper
             .create(runtime)
-          val schema = SparkArrowUtil.toArrowSchema(localSchema, timezoneId)
+          val schema = SparkSchemaUtil.toArrowSchema(localSchema)
           val arrowAlloc = ArrowBufferAllocators.contextInstance()
           val cSchema = ArrowSchema.allocateNew(arrowAlloc)
           ArrowAbiUtil.exportSchema(arrowAlloc, schema, cSchema)

backends-velox/src/main/scala/org/apache/spark/sql/execution/datasources/velox/VeloxFormatWriterInjects.scala

@@ -29,9 +29,8 @@ import org.apache.gluten.utils.ArrowAbiUtil
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.execution.datasources._
-import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
-import org.apache.spark.sql.utils.SparkArrowUtil
+import org.apache.spark.sql.utils.SparkSchemaUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
 import org.apache.arrow.c.ArrowSchema
@@ -56,8 +55,7 @@ trait VeloxFormatWriterInjects extends GlutenFormatWriterInjectsBase {
       }
     }
 
-    val arrowSchema =
-      SparkArrowUtil.toArrowSchema(dataSchema, SQLConf.get.sessionLocalTimeZone)
+    val arrowSchema = SparkSchemaUtil.toArrowSchema(dataSchema)
     val cSchema = ArrowSchema.allocateNew(ArrowBufferAllocators.contextInstance())
     var dsHandle = -1L
     val runtime = Runtimes.contextInstance(BackendsApiManager.getBackendName, "VeloxWriter")

backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala

@@ -34,8 +34,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSeq, BindR
 import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode
 import org.apache.spark.sql.execution.{BroadcastModeUtils, HashExprSafeBroadcastMode, HashSafeBroadcastMode, IdentitySafeBroadcastMode, SafeBroadcastMode}
 import org.apache.spark.sql.execution.joins.{BuildSideRelation, HashedRelationBroadcastMode}
-import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.utils.SparkArrowUtil
+import org.apache.spark.sql.utils.SparkSchemaUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.task.TaskResources
 import org.apache.spark.util.{KnownSizeEstimation, Utils}
@@ -141,9 +140,8 @@ class UnsafeColumnarBuildSideRelation(
         val serializeHandle: Long = {
           val allocator = ArrowBufferAllocators.contextInstance()
           val cSchema = ArrowSchema.allocateNew(allocator)
-          val arrowSchema = SparkArrowUtil.toArrowSchema(
-            SparkShimLoader.getSparkShims.structFromAttributes(output),
-            SQLConf.get.sessionLocalTimeZone)
+          val arrowSchema = SparkSchemaUtil.toArrowSchema(
+            SparkShimLoader.getSparkShims.structFromAttributes(output))
           ArrowAbiUtil.exportSchema(allocator, arrowSchema, cSchema)
           val handle = jniWrapper
             .init(cSchema.memoryAddress())
@@ -381,9 +379,8 @@ class UnsafeColumnarBuildSideRelation(
     val serializerHandle: Long = {
       val allocator = ArrowBufferAllocators.contextInstance()
       val cSchema = ArrowSchema.allocateNew(allocator)
-      val arrowSchema = SparkArrowUtil.toArrowSchema(
-        SparkShimLoader.getSparkShims.structFromAttributes(output),
-        SQLConf.get.sessionLocalTimeZone)
+      val arrowSchema = SparkSchemaUtil.toArrowSchema(
+        SparkShimLoader.getSparkShims.structFromAttributes(output))
       ArrowAbiUtil.exportSchema(allocator, arrowSchema, cSchema)
       val handle = jniWrapper
         .init(cSchema.memoryAddress())
@@ -430,9 +427,8 @@ class UnsafeColumnarBuildSideRelation(
     val serializerHandle = {
       val allocator = ArrowBufferAllocators.contextInstance()
       val cSchema = ArrowSchema.allocateNew(allocator)
-      val arrowSchema = SparkArrowUtil.toArrowSchema(
-        SparkShimLoader.getSparkShims.structFromAttributes(output),
-        SQLConf.get.sessionLocalTimeZone)
+      val arrowSchema = SparkSchemaUtil.toArrowSchema(
+        SparkShimLoader.getSparkShims.structFromAttributes(output))
       ArrowAbiUtil.exportSchema(allocator, arrowSchema, cSchema)
       val handle = serializerJniWrapper.init(cSchema.memoryAddress())
       cSchema.close()

gluten-arrow/src/main/java/org/apache/gluten/vectorized/ArrowWritableColumnVector.java

@@ -207,11 +207,7 @@ public ArrowWritableColumnVector(int capacity, DataType dataType) {
     super(capacity, dataType);
     vectorCount.getAndIncrement();
     refCnt.getAndIncrement();
-    String timeZoneId = SparkSchemaUtil.getLocalTimezoneID();
-    boolean useLargeVarTypes = SparkSchemaUtil.enableLargeVarTypes();
-    List<Field> fields =
-        Arrays.asList(
-            SparkArrowUtil.toArrowField("col", dataType, true, timeZoneId, useLargeVarTypes));
+    List<Field> fields = Arrays.asList(SparkSchemaUtil.toArrowField("col", dataType, true));
     Schema arrowSchema = new Schema(fields);
     VectorSchemaRoot root =
         VectorSchemaRoot.create(arrowSchema, ArrowBufferAllocators.contextInstance());

gluten-arrow/src/main/scala/org/apache/gluten/utils/ArrowUtil.scala

@@ -18,43 +18,17 @@ package org.apache.gluten.utils
 
 import org.apache.gluten.vectorized.ArrowWritableColumnVector
 
-import org.apache.spark.sql.catalyst.expressions.Attribute
-import org.apache.spark.sql.types._
 import org.apache.spark.sql.utils.{SparkArrowUtil, SparkSchemaUtil}
 import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
 
 import org.apache.arrow.c.{ArrowSchema, CDataDictionaryProvider, Data}
 import org.apache.arrow.memory.BufferAllocator
-import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+import org.apache.arrow.vector.types.pojo.{Field, Schema}
 
 import java.util
 
-import scala.collection.JavaConverters._
-
 object ArrowUtil {
 
-  private val defaultTimeZoneId = SparkSchemaUtil.getLocalTimezoneID
-
-  private def getResultType(dataType: DataType): ArrowType = {
-    getResultType(dataType, defaultTimeZoneId)
-  }
-
-  private def getResultType(dataType: DataType, timeZoneId: String): ArrowType = {
-    dataType match {
-      case other =>
-        SparkArrowUtil.toArrowType(dataType, timeZoneId)
-    }
-  }
-
-  def toArrowSchema(attributes: Seq[Attribute]): Schema = {
-    val fields = attributes.map(
-      attr => {
-        Field
-          .nullable(s"${attr.name}#${attr.exprId.id}", getResultType(attr.dataType))
-      })
-    new Schema(fields.toList.asJava)
-  }
-
   def toArrowSchema(
       cSchema: ArrowSchema,
       allocator: BufferAllocator,
@@ -65,8 +39,7 @@ object ArrowUtil {
     originFields.forEach {
       field =>
         val dt = SparkArrowUtil.fromArrowField(field)
-        fields.add(
-          SparkArrowUtil.toArrowField(field.getName, dt, true, SparkSchemaUtil.getLocalTimezoneID))
+        fields.add(SparkSchemaUtil.toArrowField(field.getName, dt, true))
     }
     new Schema(fields)
   }

gluten-arrow/src/main/scala/org/apache/spark/sql/utils/SparkArrowUtil.scala

@@ -31,7 +31,7 @@ import scala.collection.JavaConverters._
 object SparkArrowUtil {
 
   /** Maps data type from Spark to Arrow. NOTE: timeZoneId required for TimestampTypes */
-  def toArrowType(
+  private[utils] def toArrowType(
       dt: DataType,
       timeZoneId: String,
       largeVarTypes: Boolean = false): ArrowType = dt match {
@@ -100,7 +100,7 @@ object SparkArrowUtil {
   }
 
   /** Maps field from Spark to Arrow. NOTE: timeZoneId required for TimestampType */
-  def toArrowField(
+  private[utils] def toArrowField(
       name: String,
       dt: DataType,
       nullable: Boolean,
@@ -179,7 +179,7 @@ object SparkArrowUtil {
   }
 
   /** Maps schema from Spark to Arrow. NOTE: timeZoneId required for TimestampType in StructType */
-  def toArrowSchema(
+  private[utils] def toArrowSchema(
       schema: StructType,
       timeZoneId: String,
       largeVarTypes: Boolean = false): Schema = {

gluten-arrow/src/main/scala/org/apache/spark/sql/utils/SparkSchemaUtil.scala

@@ -17,14 +17,18 @@
 package org.apache.spark.sql.utils
 
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{DataType, StructType}
 
-import org.apache.arrow.vector.types.pojo.Schema
+import org.apache.arrow.vector.types.pojo.{Field, Schema}
 
 import java.util.{Objects, TimeZone}
 
 object SparkSchemaUtil {
 
+  def toArrowField(name: String, dt: DataType, nullable: Boolean): Field = {
+    SparkArrowUtil.toArrowField(name, dt, nullable, getLocalTimezoneID, enableLargeVarTypes)
+  }
+
   def fromArrowSchema(schema: Schema): StructType = {
     SparkArrowUtil.fromArrowSchema(schema)
   }