comphead commented on code in PR #5048: URL: https://github.com/apache/datafusion-comet/pull/5048#discussion_r4277251387
########## docs/source/contributor-guide/timezones.md: ########## @@ -188,7 +188,9 @@ Two Spark settings are not handled natively: - Comet disables itself for a session with `spark.sql.parquet.int96TimestampConversion=true`, which shifts `INT96` values written by Impala. - Comet does not rebase dates and timestamps that were written with the legacy hybrid calendar. - See [#5010](https://github.com/apache/datafusion-comet/issues/5010). Spark's rebase of legacy + It falls back to Spark for these scans, unless row-group statistics show that no value needs a + rebase. Native rebasing is tracked in + [#5047](https://github.com/apache/datafusion-comet/pull/5047). Spark's rebase of legacy Review Comment: Since this branch last merged `main`, #5365 added `native/core/src/parquet/datetime_rebase.rs`, which resolves Spark's per-file rebase specs natively. Only the Delta arm enables it (`rebase_from_file_metadata = true` in `delta_spark_scan.rs`), and `NativeScan` still passes `false`. After a rebase this sentence would be wrong for Delta tables, which are rebased natively instead of falling back. Could it be scoped to the built-in Parquet scan? The tree would also hold two implementations of `DataSourceUtils.datetimeRebaseSpec` and `int96RebaseSpec`, `CometScanUtils.rebaseReason` here and `resolve_file_rebase_policies` natively. Is the plan to keep both, or should #5047 build on `datetime_rebase.rs` so `NativeScan` can opt in? The branch also conflicts with `main` in `CometDataWritingCommand.scala` and `ParquetReadSuite.scala`. ########## spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala: ########## @@ -1540,6 +1595,16 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { // With 1000 rows and default parallelism, we should get multiple partitions assert(partFiles.length > 1, "Expected multiple part files to be created") + val conf = spark.sparkContext.hadoopConfiguration + partFiles.foreach { partFile => + val inputFile = HadoopInputFile.fromPath(new Path(partFile.getAbsolutePath), conf) + Using.resource(ParquetFileReader.open(inputFile)) { reader => + val metadata = reader.getFooter.getFileMetaData.getKeyValueMetaData + assert(metadata.get("org.apache.spark.version") == SPARK_VERSION_SHORT) + assert(!metadata.containsKey("org.apache.comet.datetimeRebaseMode")) Review Comment: I could not find `org.apache.comet.datetimeRebaseMode` anywhere else at this head. It looks left over from an earlier revision that stamped a Comet-specific key, so this assertion cannot fail. Could it go? ########## spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala: ########## @@ -19,11 +19,423 @@ package org.apache.spark.sql.comet +import java.util.concurrent.{Callable, ExecutorCompletionService} + +import scala.collection.mutable.ListBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.{FileStatus, Path} +import org.apache.parquet.HadoopReadOptions +import org.apache.parquet.column.statistics.Statistics +import org.apache.parquet.format.converter.ParquetMetadataConverter.{NO_FILTER, SKIP_ROW_GROUPS} +import org.apache.parquet.hadoop.ParquetFileReader +import org.apache.parquet.hadoop.metadata.{ColumnChunkMetaData, ParquetMetadata} +import org.apache.parquet.hadoop.util.HadoopInputFile +import org.apache.parquet.schema.{PrimitiveType, Type} +import org.apache.parquet.schema.LogicalTypeAnnotation.{DateLogicalTypeAnnotation, TimestampLogicalTypeAnnotation, TimeUnit} +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, Literal} +import org.apache.spark.sql.catalyst.util.RebaseDateTime import org.apache.spark.sql.execution.{InSubqueryExec, SubqueryAdaptiveBroadcastExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, DateType, StructType, TimestampNTZType, TimestampType} +import org.apache.spark.util.ThreadUtils + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.serde.SupportLevel object CometScanUtils { + /** + * One Parquet file of the scan's listing, as the datetime rebase check sees it: its cache + * identity and the status its footer is read with. The length and modification time must be the + * listing's. + */ + case class ParquetFileInfo(path: Path, length: Long, modificationTime: Long) + + private val LegacyDateTimeKey = "org.apache.spark.legacyDateTime" + private val LegacyInt96Key = "org.apache.spark.legacyINT96" + + /** How a top-level Parquet column stores datetime values. */ + private sealed trait DatetimeEncoding + private case object DateDays extends DatetimeEncoding + private case object TimestampMicros extends DatetimeEncoding + private case object TimestampMillis extends DatetimeEncoding + private case object TimestampInt96 extends DatetimeEncoding + + /** + * A top-level datetime column of one file. `min` is its smallest value in all row groups, in + * the unit of the encoding. It is None when a row group has no usable statistics. + */ + private case class DatetimeColumn(encoding: DatetimeEncoding, min: Option[Long]) + + /** + * Datetime-relevant footer metadata of one Parquet file, independent of any read mode. + * `columns` is None until a footer read with row groups, and empty for a file that Spark never + * rebases. + */ + private case class DatetimeFooterFacts( + sparkVersion: Option[String], + legacyKeys: Set[String], + columns: Option[Map[String, DatetimeColumn]]) + + private type FooterCacheKey = (String, Long, Long) + + // The cache bound most recently applied by applyFooterFactsCacheMaxSize. + @volatile private var footerFactsCacheMaxSize: Int = + CometConf.COMET_SCAN_PARQUET_CHECK_DATETIME_REBASE_MAX_CACHED_FILES.defaultValue.get + + // Bounded LRU cache of per-file footer facts, keyed by (path, length, modificationTime) so a + // rewritten file is re-read. The cached facts are independent of the read modes and requested + // types, so one entry answers the rebase question for any query. This keeps AQE re-planning + // and repeated queries over the same files from re-reading footers on the driver. + private val footerFactsCache = + java.util.Collections.synchronizedMap( + new java.util.LinkedHashMap[FooterCacheKey, DatetimeFooterFacts](64, 0.75f, true) { + override def removeEldestEntry( + eldest: java.util.Map.Entry[FooterCacheKey, DatetimeFooterFacts]): Boolean = + size() > footerFactsCacheMaxSize + }) + + /** + * Applies the configured bound to the cache. The cache outlives any one session, so the bound + * in effect is the one from the session that last ran the check. A lower bound evicts the least + * recently used entries at once, since `removeEldestEntry` only evicts one entry per insert. + */ + private def applyFooterFactsCacheMaxSize(maxSize: Int): Unit = { + if (maxSize != footerFactsCacheMaxSize) { + footerFactsCache.synchronized { + footerFactsCacheMaxSize = maxSize + // The map is in access order, so iteration starts at the least recently used entry. + val keys = footerFactsCache.keySet.iterator + while (footerFactsCache.size() > maxSize && keys.hasNext) { + keys.next() + keys.remove() + } + } + } + } + + // Spark reads a value at or after these cutoffs as stored, in every read mode. + private val DateCutoffDays: Long = RebaseDateTime.lastSwitchJulianDay.toLong + private val TimestampCutoffMicros: Long = RebaseDateTime.lastSwitchJulianTs + // Spark compares millis after it scales them to micros, so round the cutoff up. + private val TimestampCutoffMillis: Long = Math.floorDiv(TimestampCutoffMicros + 999L, 1000L) + + /** One of the two rebase specs of Spark, with the read mode of the scan. */ + private case class RebaseSpec(name: String, mode: String, minVersion: String, legacyKey: String) + + /** A rebase spec that applies to a column, and the values that Spark rebases under it. */ + private case class RebaseCheck(spec: RebaseSpec, values: String) + + /** + * A requested field with datetime values. Only a top-level DATE, TIMESTAMP or TIMESTAMP_NTZ + * uses statistics. + */ + private case class RequestedColumn( + name: String, + dataType: DataType, + usesStatistics: Boolean, + checks: Seq[RebaseCheck]) + + private sealed trait Verdict + private case object NoRebase extends Verdict + private case object NeedsStatistics extends Verdict + private case class Fallback(reason: String) extends Verdict + + // Spark 4.0+ reads a DATE column as TIMESTAMP_NTZ and rebases the days. It rebases no other + // NTZ read. + private def readsDates(dataType: DataType): Boolean = + SupportLevel.containsType(dataType, classOf[DateType]) || + (isSpark40Plus && SupportLevel.containsType(dataType, classOf[TimestampNTZType])) + + private def readsTimestamps(dataType: DataType): Boolean = + SupportLevel.containsType(dataType, classOf[TimestampType]) + + /** Whether a scan with this required schema reads values that Spark can rebase. */ + def readsRebasableDatetimes(requiredSchema: StructType): Boolean = + requiredSchema.fields.exists(f => readsDates(f.dataType) || readsTimestamps(f.dataType)) + + /** + * The fallback reason when Spark rebases `check` values of `column` in a file, or None when + * Spark reads them as stored. + */ + private def rebaseReason( + facts: DatetimeFooterFacts, + column: String, + check: RebaseCheck): Option[String] = { + val spec = check.spec + val statistics = + s", and row-group statistics do not rule out ${check.values} in column `$column`" + def legacyFile(detail: String): String = + s"Native Parquet scan does not rebase datetime values in a Parquet file $detail$statistics" + // Mirrors Spark's DataSourceUtils.datetimeRebaseSpec/int96RebaseSpec: when the file has no + // Spark version key the configured mode decides (EXCEPTION must also fall back, because + // Spark would raise on ancient values while Comet would not); when a version is present the + // mode is ignored and only the version and the legacy markers matter. The `v < minVersion` + // comparison is deliberately the same lexicographic string comparison Spark uses. + facts.sparkVersion match { + case None if spec.mode != "CORRECTED" => + Some( + s"Native Parquet scan does not support the ${spec.mode} ${spec.name} read mode for a " + + s"Parquet file without org.apache.spark.version$statistics. Set " + + s"${SQLConf.PARQUET_REBASE_MODE_IN_READ.key} and " + + s"${SQLConf.PARQUET_INT96_REBASE_MODE_IN_READ.key} to CORRECTED if the files use " + + "the proleptic Gregorian calendar") + case Some(v) if v < spec.minVersion => Some(legacyFile(s"written by Spark $v")) + case Some(_) if facts.legacyKeys.contains(spec.legacyKey) => + Some(legacyFile(s"with ${spec.legacyKey} metadata")) + case _ => None + } + } + + /** How a top-level column stores datetime values, or None for another type. */ + private def encodingOf(field: PrimitiveType): Option[DatetimeEncoding] = + (field.getPrimitiveTypeName, field.getLogicalTypeAnnotation) match { + case (PrimitiveTypeName.INT32, _: DateLogicalTypeAnnotation) => Some(DateDays) + case (PrimitiveTypeName.INT64, t: TimestampLogicalTypeAnnotation) + if t.getUnit == TimeUnit.MICROS => + Some(TimestampMicros) + case (PrimitiveTypeName.INT64, t: TimestampLogicalTypeAnnotation) + if t.getUnit == TimeUnit.MILLIS => + Some(TimestampMillis) + case (PrimitiveTypeName.INT96, _) => Some(TimestampInt96) + case _ => None + } + + /** The smallest value of a column chunk, Long.MaxValue without values, or None if unknown. */ + private def chunkMin(chunk: ColumnChunkMetaData): Option[Long] = + if (chunk.getValueCount == 0) { + Some(Long.MaxValue) + } else { + statisticsMin(chunk.getStatistics, chunk.getValueCount) + } + + private def statisticsMin(stats: Statistics[_], valueCount: Long): Option[Long] = + if (stats == null) { + None + } else if (stats.hasNonNullValue) { + val min: Any = stats.genericGetMin + min match { + case v: java.lang.Integer => Some(v.longValue) + case v: java.lang.Long => Some(v.longValue) + case _ => None + } + } else if (stats.isNumNullsSet && stats.getNumNulls == valueCount) { + // Every value of the chunk is null. + Some(Long.MaxValue) + } else { + None + } + + /** The top-level datetime columns of a file, with their smallest value in all row groups. */ + private def datetimeColumns(footer: ParquetMetadata): Map[String, DatetimeColumn] = { + val encodings = footer.getFileMetaData.getSchema.getFields.asScala.flatMap { + case field: PrimitiveType if !field.isRepetition(Type.Repetition.REPEATED) => + encodingOf(field).map(field.getName -> _) + case _ => None + } + val blocks = footer.getBlocks.asScala + val chunks = blocks + .flatMap(_.getColumns.asScala) + .filter(_.getPath.size == 1) + .groupBy(_.getPath.toArray.head) + encodings.map { case (name, encoding) => + val columnChunks = chunks.get(name).map(_.toList).getOrElse(Nil) + // INT96 has no reliable statistics. A valid file has one chunk per row group. + val min = + if (encoding == TimestampInt96 || columnChunks.size != blocks.size) { + None + } else { + columnChunks.foldLeft(Option(Long.MaxValue)) { (acc, chunk) => + acc.flatMap(a => chunkMin(chunk).map(math.min(a, _))) + } + } + name -> DatetimeColumn(encoding, min) + }.toMap + } + + /** Reads the facts of one file. `withStatistics` also reads row groups for column minimums. */ + private def readFacts( + file: ParquetFileInfo, + conf: Configuration, + withStatistics: Boolean): DatetimeFooterFacts = { + // Open the file with the status the listing already returned: HadoopInputFile.fromPath + // would ask the file system for it again, a HEAD request per file on object stores. The + // footer is located from the end of the file, so the length must be the listed one. + val status = new FileStatus(file.length, false, 0, 0, file.modificationTime, file.path) + val inputFile = HadoopInputFile.fromStatus(status, conf) + val readOptions = HadoopReadOptions + .builder(conf, file.path) + .withMetadataFilter(if (withStatistics) NO_FILTER else SKIP_ROW_GROUPS) + .build() + val reader = ParquetFileReader.open(inputFile, readOptions) + try { + val footer = reader.getFooter + val metadata = footer.getFileMetaData.getKeyValueMetaData + val sparkVersion = Option(metadata.get("org.apache.spark.version")) + val legacyKeys = Set(LegacyDateTimeKey, LegacyInt96Key).filter(metadata.containsKey) + // Spark never rebases a file from Spark 3.1.0 or later without legacy keys. + val rebasable = sparkVersion.forall(_ < "3.1.0") || legacyKeys.nonEmpty + val columns = + if (!withStatistics) None + else if (rebasable) Some(datetimeColumns(footer)) + else Some(Map.empty[String, DatetimeColumn]) + DatetimeFooterFacts(sparkVersion, legacyKeys, columns) + } finally { + reader.close() + } + } + + /** + * The reason to read `files` with Spark, or None when Spark reads their datetime values as + * stored. With `useStatistics`, row-group minimums can prove that Spark reads a column as + * stored. + */ + def datetimeRebaseFallbackReason( + files: Seq[ParquetFileInfo], + conf: Configuration, + datetimeMode: String, + int96Mode: String, + requiredSchema: StructType, + useStatistics: Boolean): Option[String] = { + applyFooterFactsCacheMaxSize( + CometConf.COMET_SCAN_PARQUET_CHECK_DATETIME_REBASE_MAX_CACHED_FILES.get()) + + val datetimeSpec = RebaseSpec("datetime", datetimeMode, "3.0.0", LegacyDateTimeKey) + val int96Spec = RebaseSpec("INT96", int96Mode, "3.1.0", LegacyInt96Key) + val dateCheck = RebaseCheck(datetimeSpec, "dates before 1582-10-15") + val timestampCheck = RebaseCheck(datetimeSpec, "timestamps before 1900-01-01T00:00:00Z") + val int96Check = RebaseCheck(int96Spec, "timestamps before 1900-01-01T00:00:00Z") + + val requested = requiredSchema.fields.toSeq.flatMap { field => + val dates = if (readsDates(field.dataType)) Seq(dateCheck) else Nil + val timestamps = + if (readsTimestamps(field.dataType)) Seq(timestampCheck, int96Check) else Nil + val usesStatistics = useStatistics && (field.dataType == DateType || + field.dataType == TimestampType || field.dataType == TimestampNTZType) + Some(RequestedColumn(field.name, field.dataType, usesStatistics, dates ++ timestamps)) + .filter(_.checks.nonEmpty) + } + if (requested.isEmpty) { + return None + } + + // The checks left for a top-level column once its encoding and minimum in the file are known. + def checksFor(column: RequestedColumn, stored: Option[DatetimeColumn]): Seq[RebaseCheck] = { + def unlessFrom(min: Option[Long], cutoff: Long, check: RebaseCheck): Seq[RebaseCheck] = + if (min.exists(_ >= cutoff)) Nil else Seq(check) + (column.dataType, stored) match { + case (DateType, Some(DatetimeColumn(DateDays, min))) => + unlessFrom(min, DateCutoffDays, dateCheck) + case (TimestampType, Some(DatetimeColumn(TimestampMicros, min))) => + unlessFrom(min, TimestampCutoffMicros, timestampCheck) + case (TimestampType, Some(DatetimeColumn(TimestampMillis, min))) => + unlessFrom(min, TimestampCutoffMillis, timestampCheck) + case (TimestampType, Some(DatetimeColumn(TimestampInt96, _))) => Seq(int96Check) + case (TimestampNTZType, Some(DatetimeColumn(DateDays, min))) => + unlessFrom(min, DateCutoffDays, dateCheck) + case (TimestampNTZType, Some(_)) => Nil + case _ => column.checks + } + } + + def verdict(facts: DatetimeFooterFacts): Verdict = { + val outcomes = requested.map { column => + val checks = facts.columns match { + case Some(columns) if column.usesStatistics => + checksFor(column, columns.get(column.name)) + case _ => column.checks + } + val reason = checks.flatMap(check => rebaseReason(facts, column.name, check)).headOption + (reason, column.usesStatistics && facts.columns.isEmpty) + } + outcomes + .collectFirst { case (Some(reason), false) => Fallback(reason) } + .getOrElse(if (outcomes.exists(_._1.isDefined)) NeedsStatistics else NoRebase) + } + + // Answer from the cache where possible. Only misses and files that need statistics pay a read. + var cachedReason: Option[String] = None + val pending = new ListBuffer[(FooterCacheKey, ParquetFileInfo, Option[DatetimeFooterFacts])] + // Spark never opens a zero-length file, and a footer read fails on one. + files.iterator.filter(_.length > 0).foreach { file => + val key = (file.path.toString, file.length, file.modificationTime) + val cached = Option(footerFactsCache.get(key)) + cached.map(verdict) match { + case Some(NoRebase) => + case Some(Fallback(reason)) => cachedReason = cachedReason.orElse(Some(reason)) + case _ => pending += ((key, file, cached)) + } + } + if (cachedReason.isDefined || pending.isEmpty) { + return cachedReason + } + + // A file without a Spark version needs statistics under this datetime mode, so read them now. + val statisticsFirst = datetimeMode != "CORRECTED" && requested.exists(_.usesStatistics) Review Comment: This looks only at the datetime mode, but a requested `TIMESTAMP` column also carries the `int96Check`. By my reading, on Spark 3.x with `datetimeRebaseModeInRead=CORRECTED` and the default `int96RebaseModeInRead=EXCEPTION`, an unversioned file is first read with `SKIP_ROW_GROUPS`. The INT96 check then makes `verdict` return `NeedsStatistics`, and `resolve` opens the footer a second time, so each such file pays two footer reads on the driver. Would deriving it from the checks work, for example `requested.exists(c => c.usesStatistics && c.checks.exists(_.spec.mode != "CORRECTED"))`? I traced this by reading and did not run it. -- 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. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
