[HUDI-5807] Read partition values from file and create infra to support reading only a subset of columns #11770
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -30,14 +30,18 @@ import org.apache.hudi.common.table.read.HoodieFileGroupReader | |
| import org.apache.hudi.internal.schema.InternalSchema | ||
| import org.apache.hudi.storage.StorageConfiguration | ||
| import org.apache.hudi.storage.hadoop.{HadoopStorageConfiguration, HoodieHadoopStorage} | ||
|
|
||
| import org.apache.spark.sql.HoodieCatalystExpressionUtils.generateUnsafeProjection | ||
| import org.apache.spark.sql.SparkSession | ||
| import org.apache.spark.sql.catalyst.InternalRow | ||
| import org.apache.spark.sql.catalyst.expressions.JoinedRow | ||
| import org.apache.spark.sql.execution.datasources.PartitionedFile | ||
| import org.apache.spark.sql.execution.vectorized.{OffHeapColumnVector, OnHeapColumnVector} | ||
| import org.apache.spark.sql.internal.SQLConf | ||
| import org.apache.spark.sql.internal.SQLConf.PARQUET_VECTORIZED_READER_ENABLED | ||
| import org.apache.spark.sql.sources.Filter | ||
| import org.apache.spark.sql.types.StructType | ||
| import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnarBatchUtils} | ||
| import org.apache.spark.util.SerializableConfiguration | ||
|
|
||
| import java.io.Closeable | ||
|
|
@@ -87,6 +91,32 @@ class HoodieFileGroupReaderBasedParquetFileFormat(tableState: HoodieTableState, | |
| supportBatchResult | ||
| } | ||
|
|
||
| //for partition columns that we read from the file, we don't want them to be constant column vectors so we | ||
| //modify the vector types in this scenario | ||
| override def vectorTypes(requiredSchema: StructType, | ||
| partitionSchema: StructType, | ||
yihua marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| sqlConf: SQLConf): Option[Seq[String]] = { | ||
| val originalVectorTypes = super.vectorTypes(requiredSchema, partitionSchema, sqlConf) | ||
| if (mandatoryFields.isEmpty) { | ||
| originalVectorTypes | ||
| } else { | ||
| val regularVectorType = if (!sqlConf.offHeapColumnVectorEnabled) { | ||
| classOf[OnHeapColumnVector].getName | ||
| } else { | ||
| classOf[OffHeapColumnVector].getName | ||
| } | ||
| originalVectorTypes.map { | ||
| o: Seq[String] => o.zipWithIndex.map(a => { | ||
| if (a._2 >= requiredSchema.length && mandatoryFields.contains(partitionSchema.fields(a._2 - requiredSchema.length).name)) { | ||
| regularVectorType | ||
| } else { | ||
| a._1 | ||
| } | ||
| }) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private lazy val internalSchemaOpt: org.apache.hudi.common.util.Option[InternalSchema] = if (tableSchema.internalSchema.isEmpty) { | ||
| org.apache.hudi.common.util.Option.empty() | ||
| } else { | ||
|
|
@@ -104,16 +134,21 @@ class HoodieFileGroupReaderBasedParquetFileFormat(tableState: HoodieTableState, | |
| filters: Seq[Filter], | ||
| options: Map[String, String], | ||
| hadoopConf: Configuration): PartitionedFile => Iterator[InternalRow] = { | ||
| val outputSchema = StructType(requiredSchema.fields ++ partitionSchema.fields) | ||
| val isCount = requiredSchema.isEmpty && !isMOR && !isIncremental | ||
| val augmentedStorageConf = new HadoopStorageConfiguration(hadoopConf).getInline | ||
| setSchemaEvolutionConfigs(augmentedStorageConf) | ||
| val (remainingPartitionSchemaArr, fixedPartitionIndexesArr) = partitionSchema.fields.toSeq.zipWithIndex.filter(p => !mandatoryFields.contains(p._1.name)).unzip | ||
|
|
||
| // The schema of the partition cols we want to append the value instead of reading from the file | ||
| val remainingPartitionSchema = StructType(remainingPartitionSchemaArr) | ||
|
|
||
| // index positions of the remainingPartitionSchema fields in partitionSchema | ||
| val fixedPartitionIndexes = fixedPartitionIndexesArr.toSet | ||
|
|
||
| // schema that we want fg reader to output to us | ||
| val requestedSchema = StructType(requiredSchema.fields ++ partitionSchema.fields.filter(f => mandatoryFields.contains(f.name))) | ||
| val requestedAvroSchema = AvroConversionUtils.convertStructTypeToAvroSchema(requestedSchema, sanitizedTableName) | ||
|
Contributor
Author
There was a problem hiding this comment. filter out partition columns that we will read from the parquet. Also get the indexes in the partition schema so that we can prune the partition column values to only the remaining partition schema columns |
||
| val dataAvroSchema = AvroConversionUtils.convertStructTypeToAvroSchema(dataSchema, sanitizedTableName) | ||
| val parquetFileReader = spark.sparkContext.broadcast(sparkAdapter.createParquetFileReader(supportBatchResult, | ||
| spark.sessionState.conf, options, augmentedStorageConf.unwrap())) | ||
|
|
@@ -154,18 +189,23 @@ class HoodieFileGroupReaderBasedParquetFileFormat(tableState: HoodieTableState, | |
| // Append partition values to rows and project to output schema | ||
| appendPartitionAndProject( | ||
| reader.getClosableIterator, | ||
| requestedSchema, | ||
| remainingPartitionSchema, | ||
jonvex marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| outputSchema, | ||
| fileSliceMapping.getPartitionValues, | ||
| fixedPartitionIndexes) | ||
|
|
||
| case _ => | ||
| readBaseFile(file, parquetFileReader.value, requestedSchema, remainingPartitionSchema, fixedPartitionIndexes, | ||
| requiredSchema, partitionSchema, outputSchema, filters, storageConf) | ||
| } | ||
| // CDC queries. | ||
| case hoodiePartitionCDCFileGroupSliceMapping: HoodiePartitionCDCFileGroupMapping => | ||
| buildCDCRecordIterator(hoodiePartitionCDCFileGroupSliceMapping, parquetFileReader.value, storageConf, fileIndexProps, requiredSchema) | ||
|
|
||
| case _ => | ||
| readBaseFile(file, parquetFileReader.value, requestedSchema, remainingPartitionSchema, fixedPartitionIndexes, | ||
| requiredSchema, partitionSchema, outputSchema, filters, storageConf) | ||
| } | ||
| } | ||
| } | ||
|
|
@@ -177,11 +217,11 @@ class HoodieFileGroupReaderBasedParquetFileFormat(tableState: HoodieTableState, | |
| } | ||
| } | ||
|
|
||
| private def buildCDCRecordIterator(hoodiePartitionCDCFileGroupSliceMapping: HoodiePartitionCDCFileGroupMapping, | ||
| parquetFileReader: SparkParquetReader, | ||
| storageConf: StorageConfiguration[Configuration], | ||
| props: TypedProperties, | ||
| requiredSchema: StructType): Iterator[InternalRow] = { | ||
| val fileSplits = hoodiePartitionCDCFileGroupSliceMapping.getFileSplits().toArray | ||
| val cdcFileGroupSplit: HoodieCDCFileGroupSplit = HoodieCDCFileGroupSplit(fileSplits) | ||
| props.setProperty(HoodieTableConfig.HOODIE_TABLE_NAME_KEY, tableName) | ||
|
|
@@ -204,13 +244,22 @@ class HoodieFileGroupReaderBasedParquetFileFormat(tableState: HoodieTableState, | |
| inputSchema: StructType, | ||
| partitionSchema: StructType, | ||
| to: StructType, | ||
| partitionValues: InternalRow, | ||
| fixedPartitionIndexes: Set[Int]): Iterator[InternalRow] = { | ||
| if (partitionSchema.isEmpty) { | ||
| //'inputSchema' and 'to' should be the same so the projection will just be an identity func | ||
| projectSchema(iter, inputSchema, to) | ||
| } else { | ||
| val fixedPartitionValues = if (partitionSchema.length == partitionValues.numFields) { | ||
| //need to append all of the partition fields | ||
| partitionValues | ||
| } else { | ||
| //some partition fields read from file, some were not | ||
| getFixedPartitionValues(partitionValues, partitionSchema, fixedPartitionIndexes) | ||
| } | ||
| val unsafeProjection = generateUnsafeProjection(StructType(inputSchema.fields ++ partitionSchema.fields), to) | ||
| val joinedRow = new JoinedRow() | ||
| makeCloseableFileGroupMappingRecordIterator(iter, d => unsafeProjection(joinedRow(d, fixedPartitionValues))) | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -221,8 +270,8 @@ class HoodieFileGroupReaderBasedParquetFileFormat(tableState: HoodieTableState, | |
| makeCloseableFileGroupMappingRecordIterator(iter, d => unsafeProjection(d)) | ||
| } | ||
|
|
||
| private def makeCloseableFileGroupMappingRecordIterator(closeableFileGroupRecordIterator: HoodieFileGroupReader.HoodieFileGroupReaderIterator[InternalRow], | ||
| mappingFunction: Function[InternalRow, InternalRow]): Iterator[InternalRow] = { | ||
| new Iterator[InternalRow] with Closeable { | ||
| override def hasNext: Boolean = closeableFileGroupRecordIterator.hasNext | ||
|
|
||
|
|
@@ -231,4 +280,43 @@ class HoodieFileGroupReaderBasedParquetFileFormat(tableState: HoodieTableState, | |
| override def close(): Unit = closeableFileGroupRecordIterator.close() | ||
| } | ||
| } | ||
|
|
||
| private def readBaseFile(file: PartitionedFile, parquetFileReader: SparkParquetReader, requestedSchema: StructType, | ||
| remainingPartitionSchema: StructType, fixedPartitionIndexes: Set[Int], requiredSchema: StructType, | ||
| partitionSchema: StructType, outputSchema: StructType, filters: Seq[Filter], | ||
| storageConf: StorageConfiguration[Configuration]): Iterator[InternalRow] = { | ||
| if (remainingPartitionSchema.fields.length == partitionSchema.fields.length) { | ||
|
Contributor
There was a problem hiding this comment. Are all of these different branches covered by the tests?
Contributor
Author
There was a problem hiding this comment. I tested with TestSparkSqlWithCustomKeyGenerator by changing partitionColumnsToRead in HoodieHadoopFsRelationFactory by doing: to fake what [HUDI-8036] + [HUDI-8098] will do. This exposed a case that I didn't test. For MOR with log files where we read some, but not all of the partition columns, I was not doing appending correctly. I have updated HoodieFileGroupReaderBasedParquetFileFormat.appendPartitionAndProject to do this correctly now. I feel like appendPartitionAndProject and readBaseFile have overlapping logic, but can't think of a better way to do this for now.
Contributor
|
||
| //none of partition fields are read from the file, so the reader will do the appending for us | ||
| parquetFileReader.read(file, requiredSchema, partitionSchema, internalSchemaOpt, filters, storageConf) | ||
| } else if (remainingPartitionSchema.fields.length == 0) { | ||
| //we read all of the partition fields from the file | ||
| val pfileUtils = sparkAdapter.getSparkPartitionedFileUtils | ||
| //we need to modify the partitioned file so that the partition values are empty | ||
| val modifiedFile = pfileUtils.createPartitionedFile(InternalRow.empty, pfileUtils.getPathFromPartitionedFile(file), file.start, file.length) | ||
| //and we pass an empty schema for the partition schema | ||
| parquetFileReader.read(modifiedFile, outputSchema, new StructType(), internalSchemaOpt, filters, storageConf) | ||
| } else { | ||
| //need to do an additional projection here. The case in mind is that partition schema is "a,b,c" mandatoryFields is "a,c", | ||
| //then we will read (dataSchema + a + c) and append b. So the final schema will be (data schema + a + c +b) | ||
| //but expected output is (data schema + a + b + c) | ||
| val pfileUtils = sparkAdapter.getSparkPartitionedFileUtils | ||
| val partitionValues = getFixedPartitionValues(file.partitionValues, partitionSchema, fixedPartitionIndexes) | ||
| val modifiedFile = pfileUtils.createPartitionedFile(partitionValues, pfileUtils.getPathFromPartitionedFile(file), file.start, file.length) | ||
| val iter = parquetFileReader.read(modifiedFile, requestedSchema, remainingPartitionSchema, internalSchemaOpt, filters, storageConf) | ||
| projectIter(iter, StructType(requestedSchema.fields ++ remainingPartitionSchema.fields), outputSchema) | ||
| } | ||
| } | ||
|
|
||
| private def projectIter(iter: Iterator[Any], from: StructType, to: StructType): Iterator[InternalRow] = { | ||
| val unsafeProjection = generateUnsafeProjection(from, to) | ||
| val batchProjection = ColumnarBatchUtils.generateProjection(from, to) | ||
| iter.map { | ||
| case ir: InternalRow => unsafeProjection(ir) | ||
| case cb: ColumnarBatch => batchProjection(cb) | ||
| }.asInstanceOf[Iterator[InternalRow]] | ||
| } | ||
|
|
||
| private def getFixedPartitionValues(allPartitionValues: InternalRow, partitionSchema: StructType, fixedPartitionIndexes: Set[Int]): InternalRow = { | ||
| InternalRow.fromSeq(allPartitionValues.toSeq(partitionSchema).zipWithIndex.filter(p => fixedPartitionIndexes.contains(p._2)).map(p => p._1)) | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.