Skip to content
Closed
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,8 @@ private[sql] trait ColumnarBatchScan extends CodegenSupport {

val inMemoryTableScan: InMemoryTableScanExec = null

def vectorTypes: Option[Seq[String]] = None

override lazy val metrics = Map(
"numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
"scanTime" -> SQLMetrics.createTimingMetric(sparkContext, "scan time"))
Expand Down Expand Up @@ -83,13 +85,15 @@ private[sql] trait ColumnarBatchScan extends CodegenSupport {
val batch = ctx.freshName("batch")
ctx.addMutableState(columnarBatchClz, batch, s"$batch = null;")

val columnVectorClz = "org.apache.spark.sql.execution.vectorized.ColumnVector"
val idx = ctx.freshName("batchIdx")
ctx.addMutableState("int", idx, s"$idx = 0;")
val colVars = output.indices.map(i => ctx.freshName("colInstance" + i))
val columnAssigns = colVars.zipWithIndex.map { case (name, i) =>
ctx.addMutableState(columnVectorClz, name, s"$name = null;")
s"$name = $batch.column($i);"
val columnVectorClzs = vectorTypes.getOrElse(
Seq.fill(colVars.size)("org.apache.spark.sql.execution.vectorized.ColumnVector"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: classOf[ColumnVector].getName?

val columnAssigns = colVars.zip(columnVectorClzs).zipWithIndex.map {
case ((name, columnVectorClz), i) =>
ctx.addMutableState(columnVectorClz, name, s"$name = null;")
s"$name = ($columnVectorClz) $batch.column($i);"
}

val nextBatch = ctx.freshName("nextBatch")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,12 @@ case class FileSourceScanExec(
false
}

override def vectorTypes: Option[Seq[String]] =
relation.fileFormat.vectorTypes(
sparkSession = Option(relation.sparkSession),
requiredSchema = requiredSchema,
partitionSchema = relation.partitionSchema)

@transient private lazy val selectedPartitions: Seq[PartitionDirectory] = {
val optimizerMetadataTimeNs = relation.location.metadataOpsTimeNs.getOrElse(0L)
val startTime = System.nanoTime()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,17 @@ trait FileFormat {
false
}

/**
* Returns concrete column vector class names for each column to be used in a columnar batch
* if this format supports returning columnar batch.
*/
def vectorTypes(
sparkSession: Option[SparkSession],
requiredSchema: StructType,
partitionSchema: StructType): Option[Seq[String]] = {
None
}

/**
* Returns whether a file with `path` could be splitted or not.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection
import org.apache.spark.sql.catalyst.parser.LegacyTypeStringParser
import org.apache.spark.sql.execution.datasources._
import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.sources._
import org.apache.spark.sql.types._
Expand Down Expand Up @@ -272,6 +273,14 @@ class ParquetFileFormat
schema.forall(_.dataType.isInstanceOf[AtomicType])
}

override def vectorTypes(
sparkSession: Option[SparkSession],

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we really need the session parameter?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, let's remove this for now.

requiredSchema: StructType,
partitionSchema: StructType): Option[Seq[String]] = {
Option(Seq.fill(requiredSchema.fields.length + partitionSchema.fields.length)(
classOf[OnHeapColumnVector].getName))
}

override def isSplitable(
sparkSession: SparkSession,
options: Map[String, String],
Expand Down