[SPARK-19149][SQL] Follow-up: simplify cache implementation.

## What changes were proposed in this pull request?
This patch simplifies slightly the logical plan statistics cache implementation, as discussed in https://github.com/apache/spark/pull/16529

## How was this patch tested?
N/A - this has no behavior change.

Author: Reynold Xin <rxin@databricks.com>

Closes #16544 from rxin/SPARK-19149.

1 files changed, 13 insertions, 8 deletions

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
index 9e5ba9ca8f..0587a59214 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
@@ -82,17 +82,22 @@ abstract class LogicalPlan extends QueryPlan[LogicalPlan] with Logging {
   }
 
   /** A cache for the estimated statistics, such that it will only be computed once. */
-  private val statsCache = new ThreadLocal[Option[Statistics]] {
-    override protected def initialValue: Option[Statistics] = None
-  }
+  private var statsCache: Option[Statistics] = None
 
-  def stats(conf: CatalystConf): Statistics = statsCache.get.getOrElse {
-    statsCache.set(Some(computeStats(conf)))
-    statsCache.get.get
+  /**
+   * Returns the estimated statistics for the current logical plan node. Under the hood, this
+   * method caches the return value, which is computed based on the configuration passed in the
+   * first time. If the configuration changes, the cache can be invalidated by calling
+   * [[invalidateStatsCache()]].
+   */
+  final def stats(conf: CatalystConf): Statistics = statsCache.getOrElse {
+    statsCache = Some(computeStats(conf))
+    statsCache.get
   }
 
-  def invalidateStatsCache(): Unit = {
-    statsCache.set(None)
+  /** Invalidates the stats cache. See [[stats]] for more information. */
+  final def invalidateStatsCache(): Unit = {
+    statsCache = None
     children.foreach(_.invalidateStatsCache())
   }