[CARMEL-7516][CARMEL-5431] Limit the cached FileStatus size in driver memory (apache#163)

Author: wangyum

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala

@@ -26,7 +26,7 @@ import javax.annotation.concurrent.GuardedBy
 import scala.collection.mutable
 import scala.util.{Failure, Success, Try}
 
-import com.google.common.cache.{Cache, CacheBuilder}
+import com.google.common.cache.{Cache, CacheBuilder, RemovalCause, RemovalListener, RemovalNotification, Weigher}
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
@@ -72,6 +72,7 @@ class SessionCatalog(
     functionResourceLoader: FunctionResourceLoader,
     functionExpressionBuilder: FunctionExpressionBuilder,
     cacheSize: Int = SQLConf.get.tableRelationCacheSize,
+    cacheMemorySizeInBytes: Int = SQLConf.get.tableRelationCacheMemorySizeInBytes,
     cacheTTL: Long = SQLConf.get.metadataCacheTTL,
     defaultDatabase: String = SQLConf.get.defaultDatabase,
     scratchSessionPath: Option[Path] = None,
@@ -95,6 +96,7 @@ class SessionCatalog(
       DummyFunctionResourceLoader,
       DummyFunctionExpressionBuilder,
       conf.tableRelationCacheSize,
+      conf.tableRelationCacheMemorySizeInBytes,
       conf.metadataCacheTTL,
       conf.defaultDatabase)
   }
@@ -152,6 +154,8 @@ class SessionCatalog(
 
   private val validNameFormat = "([\\w_]+)".r
 
+  def defaultPlanCacheSize: Int = cacheMemorySizeInBytes / cacheSize
+
   /**
    * Checks if the given name conforms the Hive standard ("[a-zA-Z_0-9]+"),
    * i.e. if this name only contains characters, numbers, and _.
@@ -210,33 +214,61 @@ class SessionCatalog(
     }
   }
 
-  private val tableRelationCache: Cache[QualifiedTableName, LogicalPlan] = {
+  private val tableRelationCache: Cache[QualifiedTableName, (LogicalPlan, Int)] = {
     var builder = CacheBuilder.newBuilder()
-      .maximumSize(cacheSize)
+
+    if (cacheMemorySizeInBytes > 0) {
+      val weigher = new Weigher[QualifiedTableName, (LogicalPlan, Int)] {
+        override def weigh(key: QualifiedTableName, value: (LogicalPlan, Int)): Int = {
+          value._2
+        }
+      }
+
+      val removeListener = new RemovalListener[QualifiedTableName, (LogicalPlan, Int)]() {
+        override def onRemoval(
+            removed: RemovalNotification[QualifiedTableName, (LogicalPlan, Int)]): Unit = {
+          if (removed.getCause == RemovalCause.SIZE) {
+            logWarning(
+              "Evicting session cache relation from memory due to size constraints " +
+                "(spark.sql.hive.filesourcePartitionFileCacheSize = "
+                + cacheMemorySizeInBytes + " bytes)." +
+                " This may impact following planning performance.")
+          }
+        }
+      }
+      builder.weigher(weigher)
+        .removalListener(removeListener)
+        .maximumWeight(cacheMemorySizeInBytes)
+    } else {
+      builder.maximumSize(cacheSize)
+    }
 
     if (cacheTTL > 0) {
       builder = builder.expireAfterWrite(cacheTTL, TimeUnit.SECONDS)
     }
 
-    builder.build[QualifiedTableName, LogicalPlan]()
+    builder.build[QualifiedTableName, (LogicalPlan, Int)]()
   }
 
   private val listener = new ExternalCatalogListener(this)
   liveBus.map(_.addToExternalCatalogQueue(listener))
 
   /** This method provides a way to get a cached plan. */
-  def getCachedPlan(t: QualifiedTableName, c: Callable[LogicalPlan]): LogicalPlan = {
-    tableRelationCache.get(t, c)
+  def getCachedPlan(t: QualifiedTableName, c: Callable[(LogicalPlan, Int)]): LogicalPlan = {
+    tableRelationCache.get(t, c)._1
   }
 
   /** This method provides a way to get a cached plan if the key exists. */
   def getCachedTable(key: QualifiedTableName): LogicalPlan = {
-    tableRelationCache.getIfPresent(key)
+    Option(tableRelationCache.getIfPresent(key)).map(_._1).orNull
   }
 
   /** This method provides a way to cache a plan. */
-  def cacheTable(t: QualifiedTableName, l: LogicalPlan): Unit = {
-    tableRelationCache.put(t, l)
+  def cacheTable(
+      t: QualifiedTableName,
+      l: LogicalPlan,
+      estimatePlanSize: Int = defaultPlanCacheSize): Unit = {
+    tableRelationCache.put(t, (l, estimatePlanSize))
   }
 
   /** This method provides a way to invalidate a cached plan. */

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

@@ -5150,6 +5150,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
   def tableRelationCacheSize: Int =
     getConf(StaticSQLConf.FILESOURCE_TABLE_RELATION_CACHE_SIZE)
 
+  def tableRelationCacheMemorySizeInBytes: Int =
+    getConf(StaticSQLConf.FILESOURCE_TABLE_RELATION_CACHE_MEM_SIZE_IN_BYTES)
+
   def codegenCacheMaxEntries: Int = getConf(StaticSQLConf.CODEGEN_CACHE_MAX_ENTRIES)
 
   def exchangeReuseEnabled: Boolean = getConf(EXCHANGE_REUSE_ENABLED)

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

@@ -92,6 +92,18 @@ object StaticSQLConf {
       .checkValue(cacheSize => cacheSize >= 0, "The maximum size of the cache must not be negative")
       .createWithDefault(1000)
 
+  val FILESOURCE_TABLE_RELATION_CACHE_MEM_SIZE_IN_BYTES =
+    buildStaticConf("spark.sql.filesourceTableRelationCache.memory.size")
+      .internal()
+      .doc("The maximum memory size of the cache that maps qualified table names to" +
+        " table relation plans. Compare to the conf: spark.sql.filesourceTableRelationCacheSize," +
+        " this conf is for the memory size limit, not for cache entry number, to avoid the" +
+        " cached relation consume too much driver memory.")
+      .version("3.2.0")
+      .intConf
+      .checkValue(cacheSize => cacheSize >= 0, "The maximum size of the cache must not be negative")
+      .createWithDefault(200 * 1024 * 1024)
+
   val CODEGEN_CACHE_MAX_ENTRIES = buildStaticConf("spark.sql.codegen.cache.maxEntries")
       .internal()
       .doc("When nonzero, enable caching of generated classes for operators and expressions. " +

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalogSuite.scala

@@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.{AliasIdentifier, FunctionIdentifier, Quali
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
-import org.apache.spark.sql.catalyst.plans.logical.{LeafCommand, LogicalPlan, Project, Range, SubqueryAlias, View}
+import org.apache.spark.sql.catalyst.plans.logical.{Command, LeafCommand, LogicalPlan, Project, Range, SubqueryAlias, View}
 import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns
 import org.apache.spark.sql.connector.catalog.CatalogManager
 import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME
@@ -1922,4 +1922,40 @@ abstract class SessionCatalogSuite extends AnalysisTest with Eventually {
         "use temporary table")
     }
   }
+
+  test("Remove table relation cache if size exceed") {
+    case class TestCommand() extends Command {
+      override def children: Seq[LogicalPlan] = Nil
+
+      override protected def withNewChildrenInternal(
+          newChildren: IndexedSeq[LogicalPlan]): LogicalPlan = { this }
+    }
+
+    val conf = new SQLConf()
+    conf.setConf(StaticSQLConf.FILESOURCE_TABLE_RELATION_CACHE_MEM_SIZE_IN_BYTES, 1000000)
+
+    withConfAndEmptyCatalog(conf) { catalog =>
+      val table1 = QualifiedTableName(catalog.getCurrentDatabase, "test1")
+      val table2 = QualifiedTableName(catalog.getCurrentDatabase, "test2")
+      val table3 = QualifiedTableName(catalog.getCurrentDatabase, "test3")
+
+      // First, make sure the test table is not cached.
+      assert(catalog.getCachedTable(table1) === null)
+      assert(catalog.getCachedTable(table2) === null)
+      assert(catalog.getCachedTable(table3) === null)
+
+      catalog.cacheTable(table1, TestCommand(), 100000)
+      assert(catalog.getCachedTable(table1) !== null)
+
+      var plan = catalog.getCachedPlan(table2, () => {
+        (TestCommand(), 100000)
+      })
+      assert(catalog.getCachedTable(table2) !== null)
+
+      catalog.getCachedPlan(table3, () => {
+        (TestCommand(), 1000000)
+      })
+      assert(catalog.getCachedTable(table3) === null)
+    }
+  }
 }

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

@@ -48,13 +48,15 @@ import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Aggr
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.{RowDataSourceScanExec, SparkPlan}
 import org.apache.spark.sql.execution.command._
+import org.apache.spark.sql.execution.datasources.DataSourceAnalysis.estimateRelationPlanSize
 import org.apache.spark.sql.execution.datasources.v2.PushedDownOperators
 import org.apache.spark.sql.execution.streaming.StreamingRelation
 import org.apache.spark.sql.sources._
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.util.{PartitioningUtils => CatalystPartitioningUtils}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.util.SizeEstimator
 
 /**
  * Replaces generic operations with specific variants that are designed to work with Spark
@@ -230,6 +232,19 @@ object DataSourceAnalysis extends Rule[LogicalPlan] {
       }
       insertCommand
   }
+
+  // Only consider the file status size in the hadoopFsRelation
+  def estimateRelationPlanSize(logicalRelation: LogicalRelation, defaultPlanSize: Int): Int = {
+    var size = logicalRelation match {
+      case LogicalRelation(HadoopFsRelation(l: InMemoryFileIndex, _, _, _, _, _), _, _, _) =>
+        SizeEstimator.estimate(l.leafFiles)
+      case _ => defaultPlanSize
+    }
+    if (size > Int.MaxValue) {
+      size = Int.MaxValue
+    }
+    size.toInt
+  }
 }
 
 
@@ -257,7 +272,9 @@ class FindDataSourceTable(sparkSession: SparkSession) extends Rule[LogicalPlan]
           className = table.provider.get,
           options = dsOptions,
           catalogTable = Some(table))
-      LogicalRelation(dataSource.resolveRelation(checkFilesExist = false), table)
+      val rel = LogicalRelation(dataSource.resolveRelation(checkFilesExist = false), table)
+      val estimateSize = estimateRelationPlanSize(rel, catalog.defaultPlanCacheSize)
+      (rel, estimateSize)
     })
   }
 

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala

@@ -32,6 +32,7 @@ import org.apache.spark.sql.catalyst.catalog._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
 import org.apache.spark.sql.execution.datasources._
+import org.apache.spark.sql.execution.datasources.DataSourceAnalysis.estimateRelationPlanSize
 import org.apache.spark.sql.execution.datasources.parquet.{ParquetFileFormat, ParquetOptions}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.SQLConf.HiveCaseSensitiveInferenceMode._
@@ -260,7 +261,8 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log
             fileFormat = fileFormat,
             options = enableDynamicPartition)(sparkSession = sparkSession)
           val created = LogicalRelation(fsRelation, updatedTable)
-          catalogProxy.cacheTable(tableIdentifier, created)
+          val estimateSize = estimateRelationPlanSize(created, catalogProxy.defaultPlanCacheSize)
+          catalogProxy.cacheTable(tableIdentifier, created, estimateSize)
           created
         }
 
@@ -289,8 +291,8 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log
                 options = hiveOptions.filter { case (k, _) => !k.equalsIgnoreCase("path") },
                 className = fileType).resolveRelation(),
               table = updatedTable)
-
-          catalogProxy.cacheTable(tableIdentifier, created)
+          val estimateSize = estimateRelationPlanSize(created, catalogProxy.defaultPlanCacheSize)
+          catalogProxy.cacheTable(tableIdentifier, created, estimateSize)
           created
         }