read from parsedStats in loadActions

Author: dhruvarya-db

spark/src/main/scala/org/apache/spark/sql/delta/CheckpointProvider.scala

@@ -81,6 +81,15 @@ trait CheckpointProvider extends UninitializedCheckpointProvider {
    * This is only intended to be used for logging and metrics.
    */
   def checkpointPolicy: Option[CheckpointPolicy.Policy]
+
+  /**
+   * List of different file indexes and corresponding schemas which could help derive full
+   * state-reconstruction for the checkpoint.
+   * Different FileIndexes could have different schemas depending on `stats_parsed` / `stats`
+   * columns in the underlying file(s).
+   */
+  def allActionsFileIndexesAndSchemas(
+    spark: SparkSession, deltaLog: DeltaLog): Seq[(DeltaLogFileIndex, StructType)]
 }
 
 object CheckpointProvider extends DeltaLogging {
@@ -112,7 +121,11 @@ object CheckpointProvider extends DeltaLogging {
               s"has no CheckpointMetadata action")
           }
           require(isV2CheckpointEnabled(snapshotDescriptor.protocol))
-          V2CheckpointProvider(uninitializedV2CheckpointProvider, checkpointMetadata, sidecarFiles)
+          V2CheckpointProvider(
+            uninitializedV2CheckpointProvider,
+            checkpointMetadata,
+            sidecarFiles,
+            snapshotDescriptor.deltaLog)
         }
       }
     case provider: UninitializedV1OrV2ParquetCheckpointProvider
@@ -136,7 +149,7 @@ object CheckpointProvider extends DeltaLogging {
           checkpointMetadataOpt match {
             case Some(cm) =>
               require(isV2CheckpointEnabled(snapshotDescriptor))
-              V2CheckpointProvider(provider, cm, sidecarFiles)
+              V2CheckpointProvider(provider, cm, sidecarFiles, snapshotDescriptor.deltaLog)
             case None =>
               PreloadedCheckpointProvider(provider.topLevelFiles, provider.lastCheckpointInfoOpt)
           }
@@ -166,6 +179,24 @@ object CheckpointProvider extends DeltaLogging {
     checksumOpt.flatMap(checksum => Option(checksum.protocol)).map(isV2CheckpointEnabled)
   }
 
+  private[delta] def getParquetSchema(
+      spark: SparkSession,
+      deltaLog: DeltaLog,
+      parquetFile: FileStatus,
+      schemaFromLastCheckpoint: Option[StructType]): StructType = {
+    // Try to get the checkpoint schema from the last_checkpoint.
+    // If it is not there then get it from filesystem by doing I/O.
+    val fetchChkSchemaFromLastCheckpoint = spark.sessionState.conf.getConf(
+      DeltaSQLConf.USE_CHECKPOINT_SCHEMA_FROM_CHECKPOINT_METADATA)
+    schemaFromLastCheckpoint match {
+      case Some(schema) if fetchChkSchemaFromLastCheckpoint => schema
+      case _ =>
+        recordDeltaOperation(deltaLog, "snapshot.checkpointSchema.fromFileSystem") {
+          Snapshot.getParquetFileSchemaAndRowCount(spark, deltaLog, parquetFile)._1
+        }
+    }
+  }
+
   private def sendEventForV2CheckpointRead(
       startTimeMs: Long,
       fileStatus: FileStatus,
@@ -294,6 +325,20 @@ case class PreloadedCheckpointProvider(
   override lazy val topLevelFileIndex: Option[DeltaLogFileIndex] = Some(fileIndex)
 
   override def checkpointPolicy: Option[CheckpointPolicy.Policy] = Some(CheckpointPolicy.Classic)
+
+  override def allActionsFileIndexesAndSchemas(
+      spark: SparkSession, deltaLog: DeltaLog): Seq[(DeltaLogFileIndex, StructType)] = {
+    Seq((fileIndex, checkpointSchema(spark, deltaLog)))
+  }
+
+  private val checkpointSchemaWithCaching = new LazyCheckpointSchemaGetter {
+    override def fileStatus: FileStatus = topLevelFiles.head
+    override def schemaFromLastCheckpoint: Option[StructType] =
+      lastCheckpointInfoOpt.flatMap(_.checkpointSchema)
+  }
+  private def checkpointSchema(spark: SparkSession, deltaLog: DeltaLog): StructType =
+    checkpointSchemaWithCaching.get(spark, deltaLog)
+
 }
 
 /**
@@ -312,6 +357,8 @@ object EmptyCheckpointProvider extends CheckpointProvider {
   override def allActionsFileIndexes(): Seq[DeltaLogFileIndex] = Nil
   override def topLevelFileIndex: Option[DeltaLogFileIndex] = None
   override def checkpointPolicy: Option[CheckpointPolicy.Policy] = None
+  override def allActionsFileIndexesAndSchemas(
+    spark: SparkSession, deltaLog: DeltaLog): Seq[(DeltaLogFileIndex, StructType)] = Nil
 }
 
 /** A trait representing a v2 [[UninitializedCheckpointProvider]] */
@@ -426,6 +473,11 @@ abstract class LazyCompleteCheckpointProvider(
 
   override def checkpointPolicy: Option[CheckpointPolicy.Policy] =
     underlyingCheckpointProvider.checkpointPolicy
+
+  override def allActionsFileIndexesAndSchemas(
+      spark: SparkSession, deltaLog: DeltaLog): Seq[(DeltaLogFileIndex, StructType)] = {
+    underlyingCheckpointProvider.allActionsFileIndexesAndSchemas(spark, deltaLog)
+  }
 }
 
 /**
@@ -438,6 +490,8 @@ abstract class LazyCompleteCheckpointProvider(
  * @param sidecarFiles          seq of [[SidecarFile]] for the v2 checkpoint
  * @param lastCheckpointInfoOpt optional last checkpoint info for the v2 checkpoint
  * @param logPath               delta log path for the underlying delta table
+ * @param sidecarSchemaFetcher     function to fetch sidecar schema.
+ *                              Returns None if there are no sidecar files.
  */
 case class V2CheckpointProvider(
     override val version: Long,
@@ -446,7 +500,8 @@ case class V2CheckpointProvider(
     checkpointMetadata: CheckpointMetadata,
     sidecarFiles: Seq[SidecarFile],
     lastCheckpointInfoOpt: Option[LastCheckpointInfo],
-    logPath: Path
+    logPath: Path,
+    sidecarSchemaFetcher: () => Option[StructType]
   ) extends CheckpointProvider with DeltaLogging {
 
   private[delta] def sidecarFileStatuses: Seq[FileStatus] =
@@ -473,22 +528,85 @@ case class V2CheckpointProvider(
 
   override def checkpointPolicy: Option[CheckpointPolicy.Policy] = Some(CheckpointPolicy.V2)
 
+  private val v2SchemaWithCaching = new LazyCheckpointSchemaGetter {
+    override def fileStatus: FileStatus = v2CheckpointFile
+    override def schemaFromLastCheckpoint: Option[StructType] =
+      lastCheckpointInfoOpt.flatMap(_.checkpointSchema)
+  }
+
+  protected def schemaForV2Checkpoint(
+      spark: SparkSession, deltaLog: DeltaLog): StructType = {
+    if (v2CheckpointFormat != V2Checkpoint.Format.PARQUET) {
+      return Action.logSchema
+    }
+    v2SchemaWithCaching.get(spark, deltaLog)
+  }
+
+  protected def schemaForSidecarFile(spark: SparkSession, deltaLog: DeltaLog): StructType = {
+    sidecarSchemaFetcher()
+      .getOrElse {
+        throw DeltaErrors.assertionFailedError("Sidecar schema asked without any sidecar files")
+      }
+  }
+
+  override def allActionsFileIndexesAndSchemas(
+      spark: SparkSession, deltaLog: DeltaLog): Seq[(DeltaLogFileIndex, StructType)] = {
+    (fileIndexForV2Checkpoint, schemaForV2Checkpoint(spark, deltaLog)) +:
+      fileIndexesForSidecarFiles.map((_, schemaForSidecarFile(spark, deltaLog)))
+  }
 }
 
 object V2CheckpointProvider {
-
   /** Alternate constructor which uses [[UninitializedV2LikeCheckpointProvider]] */
   def apply(
       uninitializedV2LikeCheckpointProvider: UninitializedV2LikeCheckpointProvider,
       checkpointMetadata: CheckpointMetadata,
-      sidecarFiles: Seq[SidecarFile]): V2CheckpointProvider = {
+      sidecarFiles: Seq[SidecarFile],
+      deltaLog: DeltaLog): V2CheckpointProvider = {
+    def getSidecarSchemaFetcher: () => Option[StructType] = {
+      val nonFateSharingSidecarSchemaFuture: NonFateSharingFuture[Option[StructType]] = {
+        checkpointV2ThreadPool.submitNonFateSharing { spark: SparkSession =>
+          sidecarFiles.headOption.map { sidecarFile =>
+            val sidecarFileStatus =
+              sidecarFile.toFileStatus(uninitializedV2LikeCheckpointProvider.logPath)
+            CheckpointProvider.getParquetSchema(
+              spark, deltaLog, sidecarFileStatus, schemaFromLastCheckpoint = None)
+          }
+        }
+      }
+      () => nonFateSharingSidecarSchemaFuture.get(Duration.Inf)
+    }
     V2CheckpointProvider(
       uninitializedV2LikeCheckpointProvider.version,
       uninitializedV2LikeCheckpointProvider.fileStatus,
       uninitializedV2LikeCheckpointProvider.v2CheckpointFormat,
       checkpointMetadata,
       sidecarFiles,
       uninitializedV2LikeCheckpointProvider.lastCheckpointInfoOpt,
-      uninitializedV2LikeCheckpointProvider.logPath)
+      uninitializedV2LikeCheckpointProvider.logPath,
+      getSidecarSchemaFetcher
+    )
   }
 }
+
+abstract class LazyCheckpointSchemaGetter {
+  protected def fileStatus: FileStatus
+  protected def schemaFromLastCheckpoint: Option[StructType]
+
+  private var lazySchema = Option.empty[StructType]
+
+  def get(spark: SparkSession, deltaLog: DeltaLog): StructType = {
+    lazySchema.getOrElse {
+      this.synchronized {
+        // re-check with lock held, in case of races with other initializers
+        if (lazySchema.isEmpty) {
+          lazySchema = Some(CheckpointProvider.getParquetSchema(
+            spark, deltaLog, fileStatus, schemaFromLastCheckpoint))
+        }
+        lazySchema.get
+      }
+    }
+  }
+
+  def getIfKnown: Option[StructType] = lazySchema
+}

spark/src/main/scala/org/apache/spark/sql/delta/Checkpoints.scala

@@ -35,6 +35,7 @@ import org.apache.spark.sql.delta.util.{DeltaFileOperations, DeltaLogGroupingIte
 import org.apache.spark.sql.delta.util.{Utils => DeltaUtils}
 import org.apache.spark.sql.delta.util.FileNames._
 import org.apache.spark.sql.delta.util.JsonUtils
+import org.apache.spark.sql.util.ScalaExtensions._
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileStatus, FileSystem, Path}
 import org.apache.hadoop.mapred.{JobConf, TaskAttemptContextImpl, TaskAttemptID}
@@ -1196,6 +1197,7 @@ object Checkpoints
       val partitionValues = Checkpoints.extractPartitionValues(
         snapshot.metadata.partitionSchema, "add.partitionValues")
       additionalCols ++= partitionValues
+      additionalCols ++= Checkpoints.extractStats(snapshot.statsSchema, "add.stats")
     }
     state.withColumn("add",
       when(col("add").isNotNull, struct(Seq(
@@ -1215,7 +1217,8 @@ object Checkpoints
   }
 
   def shouldWriteStatsAsStruct(conf: SQLConf, snapshot: Snapshot): Boolean = {
-    DeltaConfigs.CHECKPOINT_WRITE_STATS_AS_STRUCT.fromMetaData(snapshot.metadata)
+    DeltaConfigs.CHECKPOINT_WRITE_STATS_AS_STRUCT.fromMetaData(snapshot.metadata) &&
+      !conf.getConf(DeltaSQLConf.STATS_AS_STRUCT_IN_CHECKPOINT_FORCE_DISABLED).getOrElse(false)
   }
 
   def shouldWriteStatsAsJson(snapshot: Snapshot): Boolean = {
@@ -1247,6 +1250,17 @@ object Checkpoints
       None
     } else Some(struct(partitionValues: _*).as(STRUCT_PARTITIONS_COL_NAME))
   }
+  // This method can be overridden in tests to create a checkpoint with parsed stats.
+  def includeStatsParsedInCheckpoint(): Boolean = true
+
+  /** Parse the stats from JSON and keep as a struct field when available. */
+  def extractStats(statsSchema: StructType, statsColName: String): Option[Column] = {
+    import org.apache.spark.sql.functions.from_json
+    Option.when(includeStatsParsedInCheckpoint() && statsSchema.nonEmpty) {
+      from_json(col(statsColName), statsSchema, DeltaFileProviderUtils.jsonStatsParseOption)
+        .as(Checkpoints.STRUCT_STATS_COL_NAME)
+    }
+  }
 }
 
 object V2Checkpoint {

spark/src/main/scala/org/apache/spark/sql/delta/Snapshot.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.delta
 // scalastyle:off import.ordering.noEmptyLine
 import java.util.{Locale, TimeZone}
 
+import scala.collection.JavaConverters._
 import scala.collection.mutable
 
 import org.apache.spark.sql.delta.actions._
@@ -38,10 +39,14 @@ import org.apache.spark.sql.delta.util.StateCache
 import org.apache.spark.sql.util.ScalaExtensions._
 import io.delta.storage.commit.CommitCoordinatorClient
 import org.apache.hadoop.fs.{FileStatus, Path}
+import org.apache.parquet.format.converter.ParquetMetadataConverter.NO_FILTER
+import org.apache.parquet.hadoop.Footer
+import org.apache.parquet.hadoop.ParquetFileReader
 
 import org.apache.spark.internal.{MDC, MessageWithContext}
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.catalog.CatalogTable
+import org.apache.spark.sql.execution.datasources.parquet.{ParquetFileFormat, ParquetToSparkSchemaConverter}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.Utils
@@ -528,11 +533,51 @@ class Snapshot(
    * when sorted in ascending order, will order older actions before newer ones, as required by
    * [[InMemoryLogReplay]]); and [[ADD_STATS_TO_USE_COL_NAME]] (to handle certain combinations of
    * config settings for delta.checkpoint.writeStatsAsJson and delta.checkpoint.writeStatsAsStruct).
+   * When we see a V2 checkpoint without the old stats column, but the stats_parsed column, we
+   * json encode the stats_parsed column back as "stats" again. This is a temporary correctness
+   * hack.
    */
   protected def loadActions: DataFrame = {
-    fileIndices.map(deltaLog.loadIndex(_))
-      .reduceOption(_.union(_)).getOrElse(emptyDF)
-      .withColumn(ADD_STATS_TO_USE_COL_NAME, col("add.stats"))
+    if (fileIndices.isEmpty) return emptyDF
+
+    // Augment the schema with a NullType add.stats_parsed column, as a place-holder for
+    // compatibility with the checkpoint parquet. Both deltas and checkpoints generally use this
+    // schema. HOWEVER, IF (and only if) a checkpoint actually exists, AND it provides an
+    // add.stats_parsed column AND it lacks an add.stats column, THEN (and only then) the checkpoint
+    // DF includes the actual add.stats_parsed column -- not a NullType placeholder -- from which we
+    // generate the add_stats_to_use column (add.stats is unused in that case). Meanwhile, JSON
+    // deltas always map add.stats to add_stats_to_use, and always use the placeholder.
+    val logSchemaToUse = Action.logSchema
+    val jsonStatsCol = col("add.stats")
+    val deltas = deltaFileIndexOpt.map(deltaLog.loadIndex(_, logSchemaToUse))
+      .map(_.withColumn(ADD_STATS_TO_USE_COL_NAME, jsonStatsCol))
+
+    val checkpointDataframes = checkpointProvider
+      .allActionsFileIndexesAndSchemas(spark, deltaLog)
+      .map { case (index, schema) =>
+        val addSchema = schema("add").dataType.asInstanceOf[StructType]
+        val (checkpointSchemaToUse, checkpointStatsColToUse) =
+          if (addSchema.exists(_.name == "stats_parsed") && !addSchema.exists(_.name == "stats")) {
+            val checkpointSchemaToUse =
+              Action.logSchemaWithAddStatsParsed(addSchema("stats_parsed"))
+            (
+              checkpointSchemaToUse,
+              to_json(
+                col("add.stats_parsed")
+              )
+            )
+          } else {
+            // Normal (JSON-like) schema suffices
+            (logSchemaToUse, jsonStatsCol)
+          }
+
+        // For schema compat, make sure to discard add.stats_parsed (if present)
+        deltaLog.loadIndex(index, checkpointSchemaToUse)
+          .withColumn(COMMIT_VERSION_COLUMN, lit(checkpointProvider.version))
+          .withColumn(ADD_STATS_TO_USE_COL_NAME, checkpointStatsColToUse)
+          .withColumn("add", col("add").dropFields("stats_parsed"))
+      }
+      (checkpointDataframes ++ deltas).reduce(_.union(_))
   }
 
   /**
@@ -817,6 +862,30 @@ object Snapshot extends DeltaLogging {
       base
     }
   }
+
+  /**
+   * Gets the schema of a single parquet file by reading its footer. Code here is copied from
+   * ParquetFileFormat.
+   */
+  private[delta] def getParquetFileSchemaAndRowCount(
+      spark: SparkSession,
+      deltaLog: DeltaLog,
+      file: FileStatus): (StructType, Long) = {
+    // Converter used to convert Parquet `MessageType` to Spark SQL `StructType`
+    val converter = new ParquetToSparkSchemaConverter(
+      assumeBinaryIsString = spark.sessionState.conf.isParquetBinaryAsString,
+      assumeInt96IsTimestamp = spark.sessionState.conf.isParquetINT96AsTimestamp)
+
+    val conf = deltaLog.newDeltaHadoopConf()
+
+    val parquetMetadata = {
+      ParquetFileReader.readFooter(deltaLog.newDeltaHadoopConf(), file.getPath)
+    }
+    val rowCount = parquetMetadata.getBlocks.asScala.map(_.getRowCount).sum
+
+    val footer = new Footer(file.getPath(), parquetMetadata)
+    (ParquetFileFormat.readSchemaFromFooter(footer, converter), rowCount)
+  }
 }
 
 /**

spark/src/main/scala/org/apache/spark/sql/delta/actions/actions.scala

@@ -115,6 +115,18 @@ object Action {
 
   lazy val logSchema = ExpressionEncoder[SingleAction]().schema
   lazy val addFileSchema = logSchema("add").dataType.asInstanceOf[StructType]
+
+  /**
+   * Same as [[logSchema]], but with a user-specified add.stats_parsed column. This is useful for
+   * reading parquet checkpoint files that provide add.stats_parsed instead of add.stats.
+   */
+  def logSchemaWithAddStatsParsed(statsParsed: StructField): StructType = {
+    val logAddSchema = logSchema("add").dataType.asInstanceOf[StructType]
+    val fields = logSchema.map { f =>
+      if (f.name == "add") f.copy(dataType = logAddSchema.add(statsParsed)) else f
+    }
+    StructType(fields)
+  }
 }
 
 /**