[GLUTEN-12447] Add Arrow large var types support

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())
@@ -265,9 +262,8 @@ case class ColumnarBuildSideRelation(
         val serializeHandle: Long = {
           val allocator = ArrowBufferAllocators.globalInstance()
           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 +358,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())
@@ -237,9 +235,8 @@ class UnsafeColumnarBuildSideRelation(
         val serializeHandle: Long = {
           val allocator = ArrowBufferAllocators.globalInstance()
           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 +378,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 +426,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()

backends-velox/src/test/scala/org/apache/gluten/expression/UDFPartialProjectSuite.scala

@@ -158,8 +158,13 @@ class UDFPartialProjectSuite extends WholeStageTransformerSuite {
   }
 
   test("test concat with string") {
-    runQueryAndCompare("SELECT concat_concat(l_comment), hash(l_partkey) from lineitem") {
-      checkGlutenPlan[ColumnarPartialProjectExec]
+    Seq("false", "true").foreach {
+      useLargeVarTypes =>
+        withSQLConf("spark.sql.execution.arrow.useLargeVarTypes" -> useLargeVarTypes) {
+          runQueryAndCompare("SELECT concat_concat(l_comment), hash(l_partkey) from lineitem") {
+            checkGlutenPlan[ColumnarPartialProjectExec]
+          }
+        }
     }
   }
 

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

@@ -20,6 +20,7 @@
 import org.apache.arrow.vector.complex.ListVector;
 import org.apache.arrow.vector.complex.MapVector;
 import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.holders.NullableLargeVarCharHolder;
 import org.apache.arrow.vector.holders.NullableVarCharHolder;
 import org.apache.spark.annotation.DeveloperApi;
 import org.apache.spark.sql.types.DataType;
@@ -166,8 +167,12 @@ void initAccessor(ValueVector vector) {
       accessor = new DecimalAccessor((DecimalVector) vector);
     } else if (vector instanceof VarCharVector) {
       accessor = new StringAccessor((VarCharVector) vector);
+    } else if (vector instanceof LargeVarCharVector) {
+      accessor = new LargeStringAccessor((LargeVarCharVector) vector);
     } else if (vector instanceof VarBinaryVector) {
       accessor = new BinaryAccessor((VarBinaryVector) vector);
+    } else if (vector instanceof LargeVarBinaryVector) {
+      accessor = new LargeBinaryAccessor((LargeVarBinaryVector) vector);
     } else if (vector instanceof DateDayVector) {
       accessor = new DateAccessor((DateDayVector) vector);
     } else if (vector instanceof TimeStampMicroTZVector) {
@@ -414,6 +419,31 @@ final UTF8String getUTF8String(int rowId) {
     }
   }
 
+  static class LargeStringAccessor extends ArrowVectorAccessor {
+
+    private final LargeVarCharVector accessor;
+    private final NullableLargeVarCharHolder stringResult = new NullableLargeVarCharHolder();
+
+    LargeStringAccessor(LargeVarCharVector vector) {
+      super(vector);
+      this.accessor = vector;
+    }
+
+    @Override
+    final UTF8String getUTF8String(int rowId) {
+      accessor.get(rowId, stringResult);
+      if (stringResult.isSet == 0) {
+        return null;
+      } else {
+        return UTF8String.fromAddress(
+            null,
+            stringResult.buffer.memoryAddress() + stringResult.start,
+            // A single string cannot be larger than the max integer size, so the conversion is safe
+            (int) (stringResult.end - stringResult.start));
+      }
+    }
+  }
+
   static class BinaryAccessor extends ArrowVectorAccessor {
 
     private final VarBinaryVector accessor;
@@ -429,6 +459,21 @@ final byte[] getBinary(int rowId) {
     }
   }
 
+  static class LargeBinaryAccessor extends ArrowVectorAccessor {
+
+    private final LargeVarBinaryVector accessor;
+
+    LargeBinaryAccessor(LargeVarBinaryVector vector) {
+      super(vector);
+      this.accessor = vector;
+    }
+
+    @Override
+    final byte[] getBinary(int rowId) {
+      return accessor.getObject(rowId);
+    }
+  }
+
   static class DateAccessor extends ArrowVectorAccessor {
 
     private final DateDayVector accessor;