Add a SchemaPruningTest trait for running tests with schema pruning

turned on

Author: mallman

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/SchemaPruningTest.scala

@@ -0,0 +1,52 @@
+/*
+ * 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.catalyst
+
+import org.scalatest.BeforeAndAfterAll
+
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.internal.SQLConf.NESTED_SCHEMA_PRUNING_ENABLED
+
+/**
+ * A PlanTest that ensures that all tests in this suite are run with nested schema pruning enabled.
+ * Remove this trait once the default value of SQLConf.NESTED_SCHEMA_PRUNING_ENABLED is set to true.
+ */
+private[sql] trait SchemaPruningTest extends PlanTest with BeforeAndAfterAll {
+  private var originalConfSchemaPruningEnabled = false
+
+  override protected def beforeAll(): Unit = {
+    // Call `withSQLConf` eagerly because some subtypes of `PlanTest` (I'm looking at you,
+    // `SQLTestUtils`) override `withSQLConf` to reset the existing `SQLConf` with a new one without
+    // copying existing settings first. This here is an awful, ugly way to get around that behavior
+    // by initializing the "real" `SQLConf` with an noop call to `withSQLConf`. I don't want to risk
+    // "fixing" the downstream behavior, breaking everything else that's expecting these semantics.
+    // Oh well...
+    withSQLConf()(())
+    originalConfSchemaPruningEnabled = conf.nestedSchemaPruningEnabled
+    conf.setConf(NESTED_SCHEMA_PRUNING_ENABLED, true)
+    super.beforeAll()
+  }
+
+  override protected def afterAll(): Unit = {
+    try {
+      super.afterAll()
+    } finally {
+      conf.setConf(NESTED_SCHEMA_PRUNING_ENABLED, originalConfSchemaPruningEnabled)
+    }
+  }
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/AggregateFieldExtractionPushdownSuite.scala

@@ -17,15 +17,15 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
+import org.apache.spark.sql.catalyst.SchemaPruningTest
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions.aggregate.Count
-import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
 import org.apache.spark.sql.types._
 
-class AggregateFieldExtractionPushdownSuite extends PlanTest {
+class AggregateFieldExtractionPushdownSuite extends SchemaPruningTest {
   private val testRelation =
     LocalRelation(
       StructField("a", StructType(

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinFieldExtractionPushdownSuite.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
+import org.apache.spark.sql.catalyst.SchemaPruningTest
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions._
@@ -25,7 +26,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
 import org.apache.spark.sql.types._
 
-class JoinFieldExtractionPushdownSuite extends PlanTest {
+class JoinFieldExtractionPushdownSuite extends SchemaPruningTest {
   private val leftRelation =
     LocalRelation(
       StructField("la", StructType(

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala

@@ -2267,15 +2267,17 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
   }
 
   test("SPARK-4502: Nested column pruning shouldn't fail filter") {
-    withTempPath { dir =>
-      val path = dir.getCanonicalPath
-      val data =
-        """{"a":{"b":1,"c":2}}
-          |{}""".stripMargin
-      Seq(data).toDF().repartition(1).write.text(path)
-      checkAnswer(
-        spark.read.json(path).filter($"a.b" > 1).select($"a.b"),
-        Seq.empty)
+    withSQLConf(SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key -> "true") {
+      withTempPath { dir =>
+        val path = dir.getCanonicalPath
+        val data =
+          """{"a":{"b":1,"c":2}}
+            |{}""".stripMargin
+        Seq(data).toDF().repartition(1).write.text(path)
+        checkAnswer(
+          spark.read.json(path).filter($"a.b" > 1).select($"a.b"),
+          Seq.empty)
+      }
     }
   }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/FileSchemaPruningTest.scala

@@ -21,10 +21,11 @@ import org.scalactic.Equality
 import org.scalatest.Assertions
 
 import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.catalyst.SchemaPruningTest
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
 import org.apache.spark.sql.types.StructType
 
-private[sql] trait FileSchemaPruningTest {
+private[sql] trait FileSchemaPruningTest extends SchemaPruningTest {
   _: Assertions =>
 
   private val schemaEquality = new Equality[StructType] {