[NSE-931] Reuse partition vectors for arrow scan (#935)

* reuse partition vectors

* remove extra setValueCount and fix read null value

Author: jackylee-ch

arrow-data-source/standard/src/main/scala/com/intel/oap/spark/sql/execution/datasources/arrow/ArrowFileFormat.scala

@@ -40,13 +40,12 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.execution.datasources.{FileFormat, OutputWriterFactory, PartitionedFile}
 import org.apache.spark.sql.execution.datasources.OutputWriter
+import org.apache.spark.sql.execution.datasources.v2.arrow.{SparkMemoryUtils, SparkVectorUtils}
 import org.apache.spark.sql.execution.datasources.v2.arrow.SparkMemoryUtils.UnsafeItr
-import org.apache.spark.sql.execution.datasources.v2.arrow.SparkVectorUtils
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.sources.{DataSourceRegister, Filter}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
-import org.apache.spark.sql.vectorized.ColumnarBatch;
 
 class ArrowFileFormat extends FileFormat with DataSourceRegister with Serializable {
 
@@ -175,11 +174,17 @@ class ArrowFileFormat extends FileFormat with DataSourceRegister with Serializab
         factory.close()
       }))
 
+      val partitionVectors =
+        ArrowUtils.loadPartitionColumns(batchSize, partitionSchema, file.partitionValues)
+
+      SparkMemoryUtils.addLeakSafeTaskCompletionListener[Unit]((_: TaskContext) => {
+        partitionVectors.foreach(_.close())
+      })
+
       val itr = itrList
         .toIterator
         .flatMap(itr => itr.asScala)
-        .map(batch => ArrowUtils.loadBatch(batch, file.partitionValues, partitionSchema,
-          requiredSchema))
+        .map(batch => ArrowUtils.loadBatch(batch, requiredSchema, partitionVectors))
       new UnsafeItr(itr).asInstanceOf[Iterator[InternalRow]]
     }
   }

arrow-data-source/standard/src/main/scala/com/intel/oap/spark/sql/execution/datasources/v2/arrow/ArrowPartitionReaderFactory.scala

@@ -26,11 +26,13 @@ import com.intel.oap.spark.sql.execution.datasources.v2.arrow.ArrowSQLConf._
 import org.apache.arrow.dataset.scanner.ScanOptions
 import org.apache.arrow.vector.types.pojo.Schema
 
+import org.apache.spark.TaskContext
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader}
 import org.apache.spark.sql.execution.datasources.PartitionedFile
 import org.apache.spark.sql.execution.datasources.v2.FilePartitionReaderFactory
+import org.apache.spark.sql.execution.datasources.v2.arrow.SparkMemoryUtils
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.sources.Filter
 import org.apache.spark.sql.types.StructType
@@ -99,11 +101,17 @@ case class ArrowPartitionReaderFactory(
     val vsrItrList = taskList
       .map(task => task.execute())
 
+    val partitionVectors = ArrowUtils.loadPartitionColumns(
+      batchSize, readPartitionSchema, partitionedFile.partitionValues)
+
+    SparkMemoryUtils.addLeakSafeTaskCompletionListener[Unit]((_: TaskContext) => {
+      partitionVectors.foreach(_.close())
+    })
+
     val batchItr = vsrItrList
       .toIterator
       .flatMap(itr => itr.asScala)
-      .map(batch => ArrowUtils.loadBatch(batch, partitionedFile.partitionValues,
-        readPartitionSchema, readDataSchema))
+      .map(batch => ArrowUtils.loadBatch(batch, readDataSchema, partitionVectors))
 
     new PartitionReader[ColumnarBatch] {
       val holder = new ColumnarBatchRetainer()

arrow-data-source/standard/src/main/scala/com/intel/oap/spark/sql/execution/datasources/v2/arrow/ArrowUtils.scala

@@ -88,6 +88,43 @@ object ArrowUtils {
     SparkSchemaUtils.toArrowSchema(t, SparkSchemaUtils.getLocalTimezoneID())
   }
 
+  def loadPartitionColumns(
+      rowCount: Int,
+      partitionSchema: StructType,
+      partitionValues: InternalRow): Array[ArrowWritableColumnVector] = {
+    val partitionColumns = ArrowWritableColumnVector.allocateColumns(rowCount, partitionSchema)
+    (0 until partitionColumns.length).foreach(i => {
+      ArrowColumnVectorUtils.populate(partitionColumns(i), partitionValues, i)
+      partitionColumns(i).setValueCount(rowCount)
+      partitionColumns(i).setIsConstant()
+    })
+    partitionColumns
+  }
+
+  def loadBatch(
+      input: ArrowRecordBatch,
+      dataSchema: StructType,
+      partitionVectors: Array[ArrowWritableColumnVector]): ColumnarBatch = {
+    val rowCount: Int = input.getLength
+
+    val vectors = try {
+      ArrowWritableColumnVector.loadColumns(rowCount, toArrowSchema(dataSchema), input)
+    } finally {
+      input.close()
+    }
+
+    val batch = new ColumnarBatch(
+      vectors.map(_.asInstanceOf[ColumnVector]) ++
+        partitionVectors
+          .map { vector =>
+            // The vector should call retain() whenever reuse it.
+            vector.retain()
+            vector.asInstanceOf[ColumnVector]
+          },
+      rowCount)
+    batch
+  }
+
   def toArrowField(t: StructField): Field = {
     SparkSchemaUtils.toArrowField(
       t.name, t.dataType, t.nullable, SparkSchemaUtils.getLocalTimezoneID())