dbtsai commented on a change in pull request #28319: URL: https://github.com/apache/spark/pull/28319#discussion_r414269670
########## File path: sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ParquetNestedPredicatePushDownBenchmark.scala ########## @@ -0,0 +1,120 @@ +/* + * 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.benchmark + +import org.apache.spark.SparkConf +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession} +import org.apache.spark.sql.internal.SQLConf + +/** + * Synthetic benchmark for nested fields predicate push down performance for Parquet datasource. + * To run this benchmark: + * {{{ + * 1. without sbt: + * bin/spark-submit --class <this class> --jars <spark core test jar> <sql core test jar> + * 2. build/sbt "sql/test:runMain <this class>" + * 3. generate result: + * SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/test:runMain <this class>" + * Results will be written to "benchmarks/ParquetNestedPredicatePushDownBenchmark-results.txt". + * }}} + */ +object ParquetNestedPredicatePushDownBenchmark extends SqlBasedBenchmark { + + private val N = 100 * 1024 * 1024 + private val NUMBER_OF_ITER = 10 + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName(this.getClass.getSimpleName) + // Since `spark.master` always exists, overrides this value + .set("spark.master", "local[1]") + + SparkSession.builder().config(conf).getOrCreate() + } + + private val df: DataFrame = spark.range(1, N, 1, 4).toDF("id") + .selectExpr("id", "STRUCT(id x, STRUCT(CAST(id AS STRING) z) y) nested") + .sort("id") + + private def addCase(benchmark: Benchmark, inputPath: String, + enableNestedPD: Boolean, + name: String, filterFn: DataFrame => DataFrame): Unit = { + val loadDF = spark.read.parquet(inputPath) + benchmark.addCase(name) { + _ => + withSQLConf((SQLConf.NESTED_PREDICATE_PUSHDOWN_ENABLED.key, + enableNestedPD.toString)) { + filterFn(loadDF).noop() + } + } + } + + private def createAndRunBenchmark(name: String, filterFn: DataFrame => DataFrame): Unit = { + withTempPath { + tempDir => + val outputPath = tempDir.getCanonicalPath + df.write.mode(SaveMode.Overwrite).parquet(tempDir.getCanonicalPath) + val benchmark = new Benchmark(name, N, NUMBER_OF_ITER, output = output) + addCase(benchmark, outputPath, enableNestedPD = false, + "NestedFieldsPredicatePushDownDisabled", filterFn) + addCase(benchmark, outputPath, enableNestedPD = true, + "NestedFieldsPredicatePushDownEnabled", filterFn) + benchmark.run() + } + } + + /** + * Benchmark for sorted data with a filter which allows to filter out all the row groups + * when nested fields predicate push down enabled + */ + def runLoadNoRowGroupWhenPredicatePushedDown(): Unit = { + // no row group will be loaded when predicate pushed down + val filterFn: DataFrame => DataFrame = df => df.filter("nested.x < 0") Review comment: nit, you might call `filterFn` as `withFilter`. Typically, the type is not required for non-public variable. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
