viirya commented on a change in pull request #32049:
URL: https://github.com/apache/spark/pull/32049#discussion_r606829349
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
##########
@@ -687,6 +691,40 @@ object DataSourceStrategy
(nonconvertiblePredicates ++ unhandledPredicates, pushedFilters,
handledFilters)
}
+ protected[sql] def translateAggregate(aggregates: AggregateExpression):
Option[AggregateFunc] = {
+
+ def columnAsString(e: Expression): String = e match {
+ case AttributeReference(name, _, _, _) => name
+ case Cast(child, _, _) => columnAsString(child)
+ case _ => ""
+ }
+
+ aggregates.aggregateFunction match {
+ case min: aggregate.Min =>
+ val colName = columnAsString(min.child)
Review comment:
One suggestion: Can we reuse `PushableColumn` which is used by predicate
pushdown to capture pushed column?
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetUtils.scala
##########
@@ -127,4 +147,328 @@ object ParquetUtils {
file.getName == ParquetFileWriter.PARQUET_COMMON_METADATA_FILE ||
file.getName == ParquetFileWriter.PARQUET_METADATA_FILE
}
+
+ private[sql] def aggResultToSparkInternalRows(
Review comment:
For key methods added here, can you add some descriptive comment?
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetPartitionReaderFactory.scala
##########
@@ -80,43 +88,135 @@ case class ParquetPartitionReaderFactory(
private val datetimeRebaseModeInRead =
parquetOptions.datetimeRebaseModeInRead
private val int96RebaseModeInRead = parquetOptions.int96RebaseModeInRead
+ private def buildAggSchema: StructType = {
+ var aggSchema = new StructType()
+ for (i <- 0 until aggregation.aggregateExpressions.size) {
+ var index = 0
+ aggregation.aggregateExpressions(i) match {
+ case Max(col, _) =>
+ index = dataSchema.fieldNames.toList.indexOf(col)
+ val field = dataSchema.fields(index)
+ aggSchema = aggSchema.add(field.copy("max(" + field.name + ")"))
+ case Min(col, _) =>
+ index = dataSchema.fieldNames.toList.indexOf(col)
+ val field = dataSchema.fields(index)
+ aggSchema = aggSchema.add(field.copy("min(" + field.name + ")"))
+ case Count(col, _, _) =>
+ if (col.equals("1")) {
+ aggSchema = aggSchema.add(new StructField("count(*)", LongType))
+ } else {
+ aggSchema = aggSchema.add(new StructField("count(" + col + ")",
LongType))
+ }
+ case _ =>
+ }
+ }
+ aggSchema
+ }
+
override def supportColumnarReads(partition: InputPartition): Boolean = {
sqlConf.parquetVectorizedReaderEnabled && sqlConf.wholeStageEnabled &&
resultSchema.length <= sqlConf.wholeStageMaxNumFields &&
resultSchema.forall(_.dataType.isInstanceOf[AtomicType])
}
override def buildReader(file: PartitionedFile):
PartitionReader[InternalRow] = {
- val reader = if (enableVectorizedReader) {
- createVectorizedReader(file)
+ val fileReader = if (aggregation.aggregateExpressions.isEmpty) {
+
+ val reader = if (enableVectorizedReader) {
+ createVectorizedReader(file)
Review comment:
So if we use aggregate pushdown for Parquet, we cannot use vectorized
Parquet reader, right? Can you describe it too in the config doc?
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
##########
@@ -687,6 +691,40 @@ object DataSourceStrategy
(nonconvertiblePredicates ++ unhandledPredicates, pushedFilters,
handledFilters)
}
+ protected[sql] def translateAggregate(aggregates: AggregateExpression):
Option[AggregateFunc] = {
+
+ def columnAsString(e: Expression): String = e match {
+ case AttributeReference(name, _, _, _) => name
+ case Cast(child, _, _) => columnAsString(child)
+ case _ => ""
+ }
+
+ aggregates.aggregateFunction match {
+ case min: aggregate.Min =>
+ val colName = columnAsString(min.child)
+ if (colName.nonEmpty) Some(Min(colName, min.dataType)) else None
+ case max: aggregate.Max =>
+ val colName = columnAsString(max.child)
+ if (colName.nonEmpty) Some(Max(colName, max.dataType)) else None
+ case avg: aggregate.Average =>
+ val colName = columnAsString(avg.child)
+ if (colName.nonEmpty) Some(Avg(colName, avg.dataType,
aggregates.isDistinct)) else None
+ case sum: aggregate.Sum =>
+ val colName = columnAsString(sum.child)
+ if (colName.nonEmpty) Some(Sum(colName, sum.dataType,
aggregates.isDistinct)) else None
+ case count: aggregate.Count =>
+ val columnName = count.children.head match {
+ case Literal(_, _) => "1" // SELECT (*) FROM table is translated to
SELECT 1 FROM table
Review comment:
typo? `SELECT (*)`? You mean `SELECT count(*) FROM table` -> `SELECT
count(1) FROM table`?
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetUtils.scala
##########
@@ -127,4 +147,328 @@ object ParquetUtils {
file.getName == ParquetFileWriter.PARQUET_COMMON_METADATA_FILE ||
file.getName == ParquetFileWriter.PARQUET_METADATA_FILE
}
+
+ private[sql] def aggResultToSparkInternalRows(
+ footer: ParquetMetadata,
+ parquetTypes: Seq[PrimitiveType.PrimitiveTypeName],
+ values: Seq[Any],
+ dataSchema: StructType,
+ datetimeRebaseModeInRead: String,
+ int96RebaseModeInRead: String,
+ convertTz: Option[ZoneId]): InternalRow = {
+ val mutableRow = new SpecificInternalRow(dataSchema.fields.map(x =>
x.dataType))
+ val footerFileMetaData = footer.getFileMetaData
+ val datetimeRebaseMode = DataSourceUtils.datetimeRebaseMode(
+ footerFileMetaData.getKeyValueMetaData.get,
+ datetimeRebaseModeInRead)
+ val int96RebaseMode = DataSourceUtils.int96RebaseMode(
+ footerFileMetaData.getKeyValueMetaData.get,
+ int96RebaseModeInRead)
+ parquetTypes.zipWithIndex.map {
+ case (PrimitiveType.PrimitiveTypeName.INT32, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ dataSchema.fields(i).dataType match {
+ case b: ByteType =>
+ mutableRow.setByte(i, values(i).asInstanceOf[Integer].toByte)
+ case s: ShortType =>
+ mutableRow.setShort(i, values(i).asInstanceOf[Integer].toShort)
+ case int: IntegerType =>
+ mutableRow.setInt(i, values(i).asInstanceOf[Integer])
+ case d: DateType =>
+ val dateRebaseFunc = DataSourceUtils.creteDateRebaseFuncInRead(
+ datetimeRebaseMode, "Parquet")
+ mutableRow.update(i,
dateRebaseFunc(values(i).asInstanceOf[Integer]))
+ case d: DecimalType =>
+ val decimal = Decimal(values(i).asInstanceOf[Integer].toLong,
d.precision, d.scale)
+ mutableRow.setDecimal(i, decimal, d.precision)
+ case _ => throw new IllegalArgumentException("Unexpected type for
INT32")
+ }
+ }
+ case (PrimitiveType.PrimitiveTypeName.INT64, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ dataSchema.fields(i).dataType match {
+ case long: LongType =>
+ mutableRow.setLong(i, values(i).asInstanceOf[Long])
+ case d: DecimalType =>
+ val decimal = Decimal(values(i).asInstanceOf[Integer].toLong,
d.precision, d.scale)
Review comment:
values(i).asInstanceOf[Integer]? Or values(i).asInstanceOf[Long]? It is
`PrimitiveTypeName.INT64`.
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetUtils.scala
##########
@@ -127,4 +147,328 @@ object ParquetUtils {
file.getName == ParquetFileWriter.PARQUET_COMMON_METADATA_FILE ||
file.getName == ParquetFileWriter.PARQUET_METADATA_FILE
}
+
+ private[sql] def aggResultToSparkInternalRows(
+ footer: ParquetMetadata,
+ parquetTypes: Seq[PrimitiveType.PrimitiveTypeName],
+ values: Seq[Any],
+ dataSchema: StructType,
+ datetimeRebaseModeInRead: String,
+ int96RebaseModeInRead: String,
+ convertTz: Option[ZoneId]): InternalRow = {
+ val mutableRow = new SpecificInternalRow(dataSchema.fields.map(x =>
x.dataType))
+ val footerFileMetaData = footer.getFileMetaData
+ val datetimeRebaseMode = DataSourceUtils.datetimeRebaseMode(
+ footerFileMetaData.getKeyValueMetaData.get,
+ datetimeRebaseModeInRead)
+ val int96RebaseMode = DataSourceUtils.int96RebaseMode(
+ footerFileMetaData.getKeyValueMetaData.get,
+ int96RebaseModeInRead)
+ parquetTypes.zipWithIndex.map {
+ case (PrimitiveType.PrimitiveTypeName.INT32, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ dataSchema.fields(i).dataType match {
+ case b: ByteType =>
+ mutableRow.setByte(i, values(i).asInstanceOf[Integer].toByte)
+ case s: ShortType =>
+ mutableRow.setShort(i, values(i).asInstanceOf[Integer].toShort)
+ case int: IntegerType =>
+ mutableRow.setInt(i, values(i).asInstanceOf[Integer])
+ case d: DateType =>
+ val dateRebaseFunc = DataSourceUtils.creteDateRebaseFuncInRead(
+ datetimeRebaseMode, "Parquet")
+ mutableRow.update(i,
dateRebaseFunc(values(i).asInstanceOf[Integer]))
+ case d: DecimalType =>
+ val decimal = Decimal(values(i).asInstanceOf[Integer].toLong,
d.precision, d.scale)
+ mutableRow.setDecimal(i, decimal, d.precision)
+ case _ => throw new IllegalArgumentException("Unexpected type for
INT32")
+ }
+ }
+ case (PrimitiveType.PrimitiveTypeName.INT64, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ dataSchema.fields(i).dataType match {
+ case long: LongType =>
+ mutableRow.setLong(i, values(i).asInstanceOf[Long])
+ case d: DecimalType =>
+ val decimal = Decimal(values(i).asInstanceOf[Integer].toLong,
d.precision, d.scale)
+ mutableRow.setDecimal(i, decimal, d.precision)
+ case _ => throw new IllegalArgumentException("Unexpected type for
INT64")
+ }
+ }
+ case (PrimitiveType.PrimitiveTypeName.INT96, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ dataSchema.fields(i).dataType match {
+ case l: LongType =>
+ mutableRow.setLong(i, values(i).asInstanceOf[Long])
+ case d: TimestampType =>
+ val int96RebaseFunc =
DataSourceUtils.creteTimestampRebaseFuncInRead(
+ int96RebaseMode, "Parquet INT96")
+ val julianMicros =
+
ParquetRowConverter.binaryToSQLTimestamp(values(i).asInstanceOf[Binary])
+ val gregorianMicros = int96RebaseFunc(julianMicros)
+ val adjTime =
+ convertTz.map(DateTimeUtils.convertTz(gregorianMicros, _,
ZoneOffset.UTC))
+ .getOrElse(gregorianMicros)
+ mutableRow.setLong(i, adjTime)
+ case _ =>
+ }
+ }
+ case (PrimitiveType.PrimitiveTypeName.FLOAT, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ mutableRow.setFloat(i, values(i).asInstanceOf[Float])
+ }
+ case (PrimitiveType.PrimitiveTypeName.DOUBLE, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ mutableRow.setDouble(i, values(i).asInstanceOf[Double])
+ }
+ case (PrimitiveType.PrimitiveTypeName.BOOLEAN, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ mutableRow.setBoolean(i, values(i).asInstanceOf[Boolean])
+ }
+ case (PrimitiveType.PrimitiveTypeName.BINARY, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ val bytes = values(i).asInstanceOf[Binary].getBytes
+ dataSchema.fields(i).dataType match {
+ case s: StringType =>
+ mutableRow.update(i, UTF8String.fromBytes(bytes))
+ case b: BinaryType =>
+ mutableRow.update(i, bytes)
+ case d: DecimalType =>
+ val decimal =
+ Decimal(new BigDecimal(new BigInteger(bytes), d.scale),
d.precision, d.scale)
+ mutableRow.setDecimal(i, decimal, d.precision)
+ case _ => throw new IllegalArgumentException("Unexpected type for
Binary")
+ }
+ }
+ case (PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY, i) =>
+ if (values(i) == null) {
+ mutableRow.setNullAt(i)
+ } else {
+ val bytes = values(i).asInstanceOf[Binary].getBytes
+ dataSchema.fields(i).dataType match {
+ case d: DecimalType =>
+ val decimal =
+ Decimal(new BigDecimal(new BigInteger(bytes), d.scale),
d.precision, d.scale)
+ mutableRow.setDecimal(i, decimal, d.precision)
+ case _ => throw new IllegalArgumentException("Unexpected type for
FIXED_LEN_BYTE_ARRAY")
+ }
+ }
+ case _ =>
+ throw new IllegalArgumentException("Unexpected parquet type name")
+ }
+ mutableRow
+ }
+
+ private[sql] def aggResultToSparkColumnarBatch(
+ footer: ParquetMetadata,
+ parquetTypes: Seq[PrimitiveType.PrimitiveTypeName],
+ values: Seq[Any],
+ dataSchema: StructType,
+ offHeap: Boolean,
+ datetimeRebaseModeInRead: String,
+ int96RebaseModeInRead: String,
+ convertTz: Option[ZoneId]): ColumnarBatch = {
+ val capacity = 4 * 1024
+ val footerFileMetaData = footer.getFileMetaData
+ val datetimeRebaseMode = DataSourceUtils.datetimeRebaseMode(
+ footerFileMetaData.getKeyValueMetaData.get,
+ datetimeRebaseModeInRead)
+ val int96RebaseMode = DataSourceUtils.int96RebaseMode(
+ footerFileMetaData.getKeyValueMetaData.get,
+ int96RebaseModeInRead)
+ val columnVectors = if (offHeap) {
+ OffHeapColumnVector.allocateColumns(capacity, dataSchema)
+ } else {
+ OnHeapColumnVector.allocateColumns(capacity, dataSchema)
+ }
+
+ parquetTypes.zipWithIndex.map {
+ case (PrimitiveType.PrimitiveTypeName.INT32, i) =>
+ if (values(i) == null) {
+ columnVectors(i).appendNull()
+ } else {
+ dataSchema.fields(i).dataType match {
+ case b: ByteType =>
+
columnVectors(i).appendByte(values(i).asInstanceOf[Integer].toByte)
+ case s: ShortType =>
+
columnVectors(i).appendShort(values(i).asInstanceOf[Integer].toShort)
+ case int: IntegerType =>
+ columnVectors(i).appendInt(values(i).asInstanceOf[Integer])
+ case d: DateType =>
+ val dateRebaseFunc = DataSourceUtils.creteDateRebaseFuncInRead(
+ datetimeRebaseMode, "Parquet")
+
columnVectors(i).appendInt(dateRebaseFunc(values(i).asInstanceOf[Integer]))
+ case _ => throw new IllegalArgumentException("Unexpected type for
INT32")
+ }
+ }
+ case (PrimitiveType.PrimitiveTypeName.INT64, i) =>
+ if (values(i) == null) {
+ columnVectors(i).appendNull()
+ } else {
+ columnVectors(i).appendLong(values(i).asInstanceOf[Long])
+ }
+ case (PrimitiveType.PrimitiveTypeName.INT96, i) =>
+ if (values(i) == null) {
+ columnVectors(i).appendNull()
+ } else {
+ dataSchema.fields(i).dataType match {
+ case l: LongType =>
+ columnVectors(i).appendLong(values(i).asInstanceOf[Long])
+ case d: TimestampType =>
+ val int96RebaseFunc =
DataSourceUtils.creteTimestampRebaseFuncInRead(
+ int96RebaseMode, "Parquet INT96")
+ val julianMicros =
+
ParquetRowConverter.binaryToSQLTimestamp(values(i).asInstanceOf[Binary])
+ val gregorianMicros = int96RebaseFunc(julianMicros)
+ val adjTime =
+ convertTz.map(DateTimeUtils.convertTz(gregorianMicros, _,
ZoneOffset.UTC))
+ .getOrElse(gregorianMicros)
+ columnVectors(i).appendLong(adjTime)
+ case _ => throw new IllegalArgumentException("Unexpected type for
INT96")
+ }
+ }
+ case (PrimitiveType.PrimitiveTypeName.FLOAT, i) =>
+ if (values(i) == null) {
+ columnVectors(i).appendNull()
+ } else {
+ columnVectors(i).appendFloat(values(i).asInstanceOf[Float])
+ }
+ case (PrimitiveType.PrimitiveTypeName.DOUBLE, i) =>
+ if (values(i) == null) {
+ columnVectors(i).appendNull()
+ } else {
+ columnVectors(i).appendDouble(values(i).asInstanceOf[Double])
+ }
+ case (PrimitiveType.PrimitiveTypeName.BINARY, i) =>
+ if (values(i) == null) {
+ columnVectors(i).appendNull()
+ } else {
+ val bytes = values(i).asInstanceOf[Binary].getBytes
+ columnVectors(i).putByteArray(0, bytes, 0, bytes.length)
+ }
+ case (PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY, i) =>
+ if (values(i) == null) {
+ columnVectors(i).appendNull()
+ } else {
+ val bytes = values(i).asInstanceOf[Binary].getBytes
+ columnVectors(i).putByteArray(0, bytes, 0, bytes.length)
+ }
+ case (PrimitiveType.PrimitiveTypeName.BOOLEAN, i) =>
+ if (values(i) == null) {
+ columnVectors(i).appendNull()
+ } else {
+ columnVectors(i).appendBoolean(values(i).asInstanceOf[Boolean])
+ }
+ case _ =>
+ throw new IllegalArgumentException("Unexpected parquet type name")
+ }
+ new ColumnarBatch(columnVectors.asInstanceOf[Array[ColumnVector]], 1)
+ }
+
+ private[sql] def getPushedDownAggResult(
+ footer: ParquetMetadata,
+ dataSchema: StructType,
+ aggregation: Aggregation)
+ : (Array[PrimitiveType.PrimitiveTypeName], Array[Any]) = {
+ val footerFileMetaData = footer.getFileMetaData
+ val fields = footerFileMetaData.getSchema.getFields
+ val blocks = footer.getBlocks()
+ val typesBuilder = ArrayBuilder.make[PrimitiveType.PrimitiveTypeName]
+ val valuesBuilder = ArrayBuilder.make[Any]
+
+ blocks.forEach { block =>
+ val columns = block.getColumns()
+ for (i <- 0 until aggregation.aggregateExpressions.size) {
Review comment:
Correct me if I misunderstand it.
Seems that this method reads each block then get aggregated results for each
aggregate function. The aggregated results are put into an array.
Consider two aggregate functions max(col1) and min(col1), So the array
content looks like [max(col1), min(col2)].
How does this deal with more than one block case? Seems this method puts
aggregated results sequentially like [max(col1) for block1, min(col2) for
block1, max(col1) for block2, min(col2) for block2, ...]?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
