[SPARK-36516][SQL] Support File Metadata Cache for ORC

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -967,6 +967,28 @@ object SQLConf {
     .booleanConf
     .createWithDefault(false)
 
+  val FILE_META_CACHE_ORC_ENABLED = buildConf("spark.sql.fileMetaCache.orc.enabled")
+    .doc("To indicate if enable orc file meta cache, it is recommended to enabled " +
+      "this config when multiple queries are performed on the same dataset, default is false.")
+    .version("3.3.0")
+    .booleanConf
+    .createWithDefault(false)
+
+  val FILE_META_CACHE_TTL_SINCE_LAST_ACCESS =
+    buildConf("spark.sql.fileMetaCache.ttlSinceLastAccess")
+      .version("3.3.0")
+      .doc("Time-to-live for file metadata cache entry after last access, the unit is seconds.")
+      .timeConf(TimeUnit.SECONDS)
+      .createWithDefault(3600L)
+
+  val FILE_META_CACHE_MAXIMUM_SIZE =
+    buildConf("spark.sql.fileMetaCache.maximumSize")
+      .version("3.3.0")
+      .doc("Maximum number of file meta entries the file meta cache contains.")
+      .intConf
+      .checkValue(_ > 0, "The value of fileMetaCache maximumSize must be positive")
+      .createWithDefault(1000)
+
   val HIVE_VERIFY_PARTITION_PATH = buildConf("spark.sql.hive.verifyPartitionPath")
     .doc("When true, check all the partition paths under the table\'s root directory " +
          "when reading data stored in HDFS. This configuration will be deprecated in the future " +
@@ -3600,6 +3622,8 @@ class SQLConf extends Serializable with Logging {
 
   def parquetVectorizedReaderBatchSize: Int = getConf(PARQUET_VECTORIZED_READER_BATCH_SIZE)
 
+  def fileMetaCacheOrcEnabled: Boolean = getConf(FILE_META_CACHE_ORC_ENABLED)
+
   def columnBatchSize: Int = getConf(COLUMN_BATCH_SIZE)
 
   def cacheVectorizedReaderEnabled: Boolean = getConf(CACHE_VECTORIZED_READER_ENABLED)

sql/core/src/main/java/org/apache/spark/sql/execution/datasources/orc/OrcColumnarBatchReader.java

@@ -29,6 +29,7 @@
 import org.apache.orc.OrcFile;
 import org.apache.orc.Reader;
 import org.apache.orc.TypeDescription;
+import org.apache.orc.impl.OrcTail;
 import org.apache.orc.mapred.OrcInputFormat;
 
 import org.apache.spark.sql.catalyst.InternalRow;
@@ -74,6 +75,8 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
   // The wrapped ORC column vectors.
   private org.apache.spark.sql.vectorized.ColumnVector[] orcVectorWrappers;
 
+  private OrcTail cachedTail;
+
   public OrcColumnarBatchReader(int capacity) {
     this.capacity = capacity;
   }
@@ -124,7 +127,8 @@ public void initialize(
       fileSplit.getPath(),
       OrcFile.readerOptions(conf)
         .maxLength(OrcConf.MAX_FILE_LENGTH.getLong(conf))
-        .filesystem(fileSplit.getPath().getFileSystem(conf)));
+        .filesystem(fileSplit.getPath().getFileSystem(conf))
+        .orcTail(cachedTail));
     Reader.Options options =
       OrcInputFormat.buildOptions(conf, reader, fileSplit.getStart(), fileSplit.getLength());
     recordReader = reader.rows(options);
@@ -189,6 +193,10 @@ public void initBatch(
     columnarBatch = new ColumnarBatch(orcVectorWrappers);
   }
 
+  public void setCachedTail(OrcTail cachedTail) {
+    this.cachedTail = cachedTail;
+  }
+
   /**
    * Return true if there exists more data in the next batch. If exists, prepare the next batch
    * by copying from ORC VectorizedRowBatch columns to Spark ColumnarBatch columns.

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileMetaCacheManager.scala

@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources
+
+import java.util.concurrent.TimeUnit
+
+import com.github.benmanes.caffeine.cache.{CacheLoader, Caffeine}
+import com.github.benmanes.caffeine.cache.stats.CacheStats
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.SparkEnv
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * A singleton Cache Manager to caching file meta. We cache these file metas in order to speed up
+ * iterated queries over the same dataset. Otherwise, each query would have to hit remote storage
+ * in order to fetch file meta before read files.
+ *
+ * We should implement the corresponding `FileMetaKey` for a specific file format, for example
+ * `ParquetFileMetaKey` or `OrcFileMetaKey`. By default, the file path is used as the identification
+ * of the `FileMetaKey` and the `getFileMeta` method of `FileMetaKey` is used to return the file
+ * meta of the corresponding file format.
+ */
+object FileMetaCacheManager extends Logging {
+
+  private lazy val cacheLoader = new CacheLoader[FileMetaKey, FileMeta]() {
+    override def load(fileMetaKey: FileMetaKey): FileMeta = {
+      logDebug(s"Loading Data File Meta ${fileMetaKey.path}")
+      fileMetaKey.getFileMeta
+    }
+  }
+
+  private lazy val ttlTime =
+    SparkEnv.get.conf.get(SQLConf.FILE_META_CACHE_TTL_SINCE_LAST_ACCESS)
+
+  private lazy val maximumSize =
+    SparkEnv.get.conf.get(SQLConf.FILE_META_CACHE_MAXIMUM_SIZE)
+
+  private lazy val cache = Caffeine
+    .newBuilder()
+    .expireAfterAccess(ttlTime, TimeUnit.SECONDS)
+    .maximumSize(maximumSize)
+    .recordStats()
+    .build[FileMetaKey, FileMeta](cacheLoader)
+
+  /**
+   * Returns the `FileMeta` associated with the `FileMetaKey` in the `FileMetaCacheManager`,
+   * obtaining that the `FileMeta` from `cacheLoader.load(FileMetaKey)` if necessary.
+   */
+  def get(fileMeteKey: FileMetaKey): FileMeta = cache.get(fileMeteKey)
+
+  /**
+   * Return current snapshot of FileMeta Cache's cumulative statistics
+   * include cache hitCount, missCount and so on.
+   * This method is only called when testing now.
+   */
+  private def cacheStats: CacheStats = cache.stats()
+
+  /**
+   * Use to cleanUp entries in the FileMeta Cache.
+   * This method is only called when testing now.
+   */
+  private def cleanUp(): Unit = cache.cleanUp()
+}
+
+abstract class FileMetaKey {
+  def path: Path
+  def configuration: Configuration
+  def getFileMeta: FileMeta
+  override def hashCode(): Int = path.hashCode
+  override def equals(other: Any): Boolean = other match {
+    case df: FileMetaKey => path.equals(df.path)
+    case _ => false
+  }
+}
+
+trait FileMeta

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFileFormat.scala

@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection
 import org.apache.spark.sql.execution.datasources._
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.sources._
 import org.apache.spark.sql.types._
 import org.apache.spark.util.{SerializableConfiguration, Utils}
@@ -154,11 +155,18 @@ class OrcFileFormat
 
     (file: PartitionedFile) => {
       val conf = broadcastedConf.value.value
+      val metaCacheEnabled =
+        conf.getBoolean(SQLConf.FILE_META_CACHE_ORC_ENABLED.key, false)
 
       val filePath = new Path(new URI(file.filePath))
 
       val fs = filePath.getFileSystem(conf)
-      val readerOptions = OrcFile.readerOptions(conf).filesystem(fs)
+      val readerOptions = if (metaCacheEnabled) {
+        val tail = OrcFileMeta.readTailFromCache(filePath, conf)
+        OrcFile.readerOptions(conf).filesystem(fs).orcTail(tail)
+      } else {
+        OrcFile.readerOptions(conf).filesystem(fs)
+      }
       val resultedColPruneInfo =
         Utils.tryWithResource(OrcFile.createReader(filePath, readerOptions)) { reader =>
           OrcUtils.requestedColumnIds(
@@ -200,6 +208,10 @@ class OrcFileFormat
           val requestedDataColIds = requestedColIds ++ Array.fill(partitionSchema.length)(-1)
           val requestedPartitionColIds =
             Array.fill(requiredSchema.length)(-1) ++ Range(0, partitionSchema.length)
+          if (metaCacheEnabled) {
+            val tail = OrcFileMeta.readTailFromCache(filePath, conf)
+            batchReader.setCachedTail(tail)
+          }
           batchReader.initialize(fileSplit, taskAttemptContext)
           batchReader.initBatch(
             TypeDescription.fromString(resultSchemaString),

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFileMeta.scala

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.orc
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.Path
+import org.apache.orc.OrcFile
+import org.apache.orc.impl.{OrcTail, ReaderImpl}
+
+import org.apache.spark.sql.execution.datasources.{FileMeta, FileMetaCacheManager, FileMetaKey}
+import org.apache.spark.util.Utils
+
+case class OrcFileMetaKey(path: Path, configuration: Configuration)
+  extends FileMetaKey {
+  override def getFileMeta: OrcFileMeta = OrcFileMeta(path, configuration)
+}
+
+case class OrcFileMeta(tail: OrcTail) extends FileMeta
+
+object OrcFileMeta {
+  def apply(path: Path, conf: Configuration): OrcFileMeta = {
+    val fs = path.getFileSystem(conf)
+    val readerOptions = OrcFile.readerOptions(conf).filesystem(fs)
+    Utils.tryWithResource(new ReaderImpl(path, readerOptions)) { fileReader =>
+      new OrcFileMeta(new OrcTail(fileReader.getFileTail, fileReader.getSerializedFileFooter))
+    }
+  }
+
+  def readTailFromCache(path: Path, conf: Configuration): OrcTail =
+    FileMetaCacheManager.get(OrcFileMetaKey(path, conf)).asInstanceOf[OrcFileMeta].tail
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/orc/OrcPartitionReaderFactory.scala

@@ -31,7 +31,7 @@ import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader}
 import org.apache.spark.sql.execution.datasources.PartitionedFile
-import org.apache.spark.sql.execution.datasources.orc.{OrcColumnarBatchReader, OrcDeserializer, OrcFilters, OrcUtils}
+import org.apache.spark.sql.execution.datasources.orc.{OrcColumnarBatchReader, OrcDeserializer, OrcFileMeta, OrcFilters, OrcUtils}
 import org.apache.spark.sql.execution.datasources.v2._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.sources.Filter
@@ -60,6 +60,7 @@ case class OrcPartitionReaderFactory(
   private val capacity = sqlConf.orcVectorizedReaderBatchSize
   private val orcFilterPushDown = sqlConf.orcFilterPushDown
   private val ignoreCorruptFiles = sqlConf.ignoreCorruptFiles
+  private val orcMetaCacheEnabled = sqlConf.fileMetaCacheOrcEnabled
 
   override def supportColumnarReads(partition: InputPartition): Boolean = {
     sqlConf.orcVectorizedReaderEnabled && sqlConf.wholeStageEnabled &&
@@ -88,7 +89,12 @@ case class OrcPartitionReaderFactory(
     pushDownPredicates(filePath, conf)
 
     val fs = filePath.getFileSystem(conf)
-    val readerOptions = OrcFile.readerOptions(conf).filesystem(fs)
+    val readerOptions = if (orcMetaCacheEnabled) {
+      val tail = OrcFileMeta.readTailFromCache(filePath, conf)
+      OrcFile.readerOptions(conf).filesystem(fs).orcTail(tail)
+    } else {
+      OrcFile.readerOptions(conf).filesystem(fs)
+    }
     val resultedColPruneInfo =
       Utils.tryWithResource(OrcFile.createReader(filePath, readerOptions)) { reader =>
         OrcUtils.requestedColumnIds(
@@ -135,7 +141,12 @@ case class OrcPartitionReaderFactory(
     pushDownPredicates(filePath, conf)
 
     val fs = filePath.getFileSystem(conf)
-    val readerOptions = OrcFile.readerOptions(conf).filesystem(fs)
+    val readerOptions = if (orcMetaCacheEnabled) {
+      val tail = OrcFileMeta.readTailFromCache(filePath, conf)
+      OrcFile.readerOptions(conf).filesystem(fs).orcTail(tail)
+    } else {
+      OrcFile.readerOptions(conf).filesystem(fs)
+    }
     val resultedColPruneInfo =
       Utils.tryWithResource(OrcFile.createReader(filePath, readerOptions)) { reader =>
         OrcUtils.requestedColumnIds(
@@ -158,6 +169,10 @@ case class OrcPartitionReaderFactory(
       val taskAttemptContext = new TaskAttemptContextImpl(taskConf, attemptId)
 
       val batchReader = new OrcColumnarBatchReader(capacity)
+      if (orcMetaCacheEnabled) {
+        val tail = OrcFileMeta.readTailFromCache(filePath, conf)
+        batchReader.setCachedTail(tail)
+      }
       batchReader.initialize(fileSplit, taskAttemptContext)
       val requestedPartitionColIds =
         Array.fill(readDataSchema.length)(-1) ++ Range(0, partitionSchema.length)

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcQuerySuite.scala

@@ -35,7 +35,7 @@ import org.apache.spark.{SparkConf, SparkException}
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.execution.FileSourceScanExec
-import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation, RecordReaderIterator}
+import org.apache.spark.sql.execution.datasources.{FileMetaCacheManager, HadoopFsRelation, LogicalRelation, RecordReaderIterator}
 import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
@@ -742,6 +742,43 @@ abstract class OrcQuerySuite extends OrcQueryTest with SharedSparkSession {
       }
     }
   }
+
+  test("SPARK-36516: simple select queries with file meta cache") {
+    withSQLConf(SQLConf.FILE_META_CACHE_ORC_ENABLED.key -> "true") {
+      import org.scalatest.PrivateMethodTester._
+      import com.github.benmanes.caffeine.cache.stats.CacheStats
+      val cacheStatsMethod = PrivateMethod[CacheStats](Symbol("cacheStats"))
+      val cleanUpMethod = PrivateMethod[Unit](Symbol("cleanUp"))
+      val tableName = "orc_use_meta_cache"
+      withTable(tableName) {
+        (0 until 10).map(i => (i, i.toString)).toDF("id", "value")
+          .write.format("orc").saveAsTable(tableName)
+        try {
+          val statsBeforeQuery = FileMetaCacheManager.invokePrivate(cacheStatsMethod())
+          checkAnswer(sql(s"SELECT id FROM $tableName where id > 5"),
+            (6 until 10).map(Row.apply(_)))
+          val statsAfterQuery1 = FileMetaCacheManager.invokePrivate(cacheStatsMethod())
+          // The 1st query triggers 4 times file meta read: 2 times related to
+          // push down filter and 2 times related to file read. The 1st query
+          // run twice: df.collect() and df.rdd.count(), so it triggers 8 times
+          // file meta read in total. missCount is 2 because cache is empty and
+          // 2 meta files need load, other 6 times will read meta from cache.
+          assert(statsAfterQuery1.missCount() - statsBeforeQuery.missCount() == 2)
+          assert(statsAfterQuery1.hitCount() - statsBeforeQuery.hitCount() == 6)
+          checkAnswer(sql(s"SELECT id FROM $tableName where id < 5"),
+            (0 until 5).map(Row.apply(_)))
+          val statsAfterQuery2 = FileMetaCacheManager.invokePrivate(cacheStatsMethod())
+          // The 2nd query also triggers 8 times file meta read in total and
+          // all read from meta cache, so missCount no growth and hitCount
+          // increase 8 times.
+          assert(statsAfterQuery2.missCount() - statsAfterQuery1.missCount() == 0)
+          assert(statsAfterQuery2.hitCount() - statsAfterQuery1.hitCount() == 8)
+        } finally {
+          FileMetaCacheManager.invokePrivate(cleanUpMethod())
+        }
+      }
+    }
+  }
 }
 
 class OrcV1QuerySuite extends OrcQuerySuite {