[SPARK-37038][SQL] DSV2 Sample Push Down

docs/sql-data-sources-jdbc.md

@@ -255,6 +255,15 @@ logging into the data sources.
     <td>read</td>
   </tr>
 
+  <tr>
+    <td><code>pushDownTableSample</code></td>
+    <td><code>false</code></td>
+    <td>
+     The option to enable or disable TABLESAMPLE push-down into the JDBC data source. The default value is false, in which case Spark does not push down TABLESAMPLE to the JDBC data source. Otherwise, if value sets to true, TABLESAMPLE is pushed down to the JDBC data source.
+    </td>
+    <td>read</td>
+  </tr>
+
   <tr>
     <td><code>keytab</code></td>
     <td>(none)</td>

external/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala

@@ -49,6 +49,9 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite with V2JDBCTes
   override def sparkConf: SparkConf = super.sparkConf
     .set("spark.sql.catalog.postgresql", classOf[JDBCTableCatalog].getName)
     .set("spark.sql.catalog.postgresql.url", db.getJdbcUrl(dockerIp, externalPort))
+    .set("spark.sql.catalog.postgresql.pushDownTableSample", "true")
+    .set("spark.sql.catalog.postgresql.pushDownLimit", "true")
+
   override def dataPreparation(conn: Connection): Unit = {}
 
   override def testUpdateColumnType(tbl: String): Unit = {
@@ -75,4 +78,6 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite with V2JDBCTes
     val expectedSchema = new StructType().add("ID", IntegerType, true, defaultMetadata)
     assert(t.schema === expectedSchema)
   }
+
+  override def supportsTableSample: Boolean = true
 }

external/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala

@@ -23,16 +23,20 @@ import org.apache.log4j.Level
 
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.analysis.{IndexAlreadyExistsException, NoSuchIndexException}
+import org.apache.spark.sql.catalyst.plans.logical.Sample
 import org.apache.spark.sql.connector.catalog.{Catalogs, Identifier, TableCatalog}
 import org.apache.spark.sql.connector.catalog.index.SupportsIndex
 import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2ScanRelation, V1ScanWrapper}
 import org.apache.spark.sql.jdbc.DockerIntegrationFunSuite
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types._
 import org.apache.spark.tags.DockerTest
 
 @DockerTest
 private[v2] trait V2JDBCTest extends SharedSparkSession with DockerIntegrationFunSuite {
+  import testImplicits._
+
   val catalogName: String
   // dialect specific update column type test
   def testUpdateColumnType(tbl: String): Unit
@@ -284,4 +288,83 @@ private[v2] trait V2JDBCTest extends SharedSparkSession with DockerIntegrationFu
       testIndexUsingSQL(s"$catalogName.new_table")
     }
   }
+
+  def supportsTableSample: Boolean = false
+
+  test("Test TABLESAMPLE") {
+    require(supportsTableSample)
+    withTable(s"$catalogName.new_table") {
+      sql(s"CREATE TABLE $catalogName.new_table (col1 INT, col2 INT)")
+      spark.range(10).select($"id" * 2, $"id" * 2 + 1).write.insertInto(s"$catalogName.new_table")
+
+      val df1 = sql(s"SELECT col1 FROM $catalogName.new_table TABLESAMPLE (BUCKET 6 OUT OF 10)" +
+        s" REPEATABLE (12345)")
+      val scan1 = df1.queryExecution.optimizedPlan.collectFirst {
+        case s: DataSourceV2ScanRelation => s
+      }.get
+      assert(scan1.schema.names.sameElements(Seq("col1")))
+
+      val sample1 = df1.queryExecution.optimizedPlan.collect {
+        case s: Sample => s
+      }
+      assert(sample1.isEmpty)
+      assert(df1.collect().length <= 7)
+
+      val df2 = sql(s"SELECT * FROM $catalogName.new_table TABLESAMPLE (50 PERCENT)" +
+        s" REPEATABLE (12345)")
+      val sample2 = df2.queryExecution.optimizedPlan.collect {
+        case s: Sample => s
+      }
+      assert(sample2.isEmpty)
+      assert(df2.collect().length <= 7)
+
+      val df3 = sql(s"SELECT col1 FROM $catalogName.new_table TABLESAMPLE (BUCKET 6 OUT OF 10)" +
+        s" LIMIT 2")
+      val sample3 = df3.queryExecution.optimizedPlan.collect {
+        case s: Sample => s
+      }
+      assert(sample3.isEmpty)
+      df3.queryExecution.optimizedPlan.collectFirst {
+        case s@DataSourceV2ScanRelation(_, scan, _) => scan match {
+          case v1: V1ScanWrapper =>
+            assert(v1.pushedDownOperators.limit.nonEmpty &&
+              v1.pushedDownOperators.limit.get === 2)
+            s.schema.names.sameElements(Seq("col1"))
+        }
+      }
+      assert(df3.collect().length == 2)
+
+      val df4 = sql(s"SELECT col1 FROM $catalogName.new_table" +
+        s" TABLESAMPLE (50 PERCENT) REPEATABLE (12345) LIMIT 2")
+      val sample4 = df4.queryExecution.optimizedPlan.collect {
+        case s: Sample => s
+      }
+      assert(sample4.isEmpty)
+      df4.queryExecution.optimizedPlan.collect {
+        case s@DataSourceV2ScanRelation(_, scan, _) => scan match {
+          case v1: V1ScanWrapper =>
+            assert(v1.pushedDownOperators.limit.nonEmpty &&
+              v1.pushedDownOperators.limit.get === 2)
+            s.schema.names.sameElements(Seq("col1"))
+        }
+      }
+      assert(df4.collect().length == 2)
+
+      // Push down order is filter -> sample -> limit
+      // in this test only limit is pushed down because sample is after limit
+      // Filter in combination with sample is not allowed so no need to test
+      val df5 = spark.read.table(s"$catalogName.new_table").limit(2).sample(0.5)
+      val sample5 = df5.queryExecution.optimizedPlan.collect {
+        case s: Sample => s
+      }
+      assert(sample5.nonEmpty)
+      df5.queryExecution.optimizedPlan.collect {
+        case DataSourceV2ScanRelation(_, scan, _) => scan match {
+          case v1: V1ScanWrapper =>
+            assert(v1.pushedDownOperators.limit.nonEmpty &&
+              v1.pushedDownOperators.limit.get === 2)
+        }
+      }
+    }
+  }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/ScanBuilder.java

@@ -23,7 +23,7 @@
  * An interface for building the {@link Scan}. Implementations can mixin SupportsPushDownXYZ
  * interfaces to do operator push down, and keep the operator push down result in the returned
  * {@link Scan}. When pushing down operators, the push down order is:
- * filter -> aggregate -> limit -> column pruning.
+ * sample -> filter -> aggregate -> limit -> column pruning.
  *
  * @since 3.0.0
  */

sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownTableSample.java

@@ -0,0 +1,39 @@
+/*
+ * 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.connector.read;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * A mix-in interface for {@link Scan}. Data sources can implement this interface to
+ * push down SAMPLE.
+ *
+ * @since 3.3.0
+ */
+@Evolving
+public interface SupportsPushDownTableSample extends ScanBuilder {
+
+  /**
+   * Pushes down SAMPLE to the data source.
+   */
+  boolean pushTableSample(
+      double lowerBound,
+      double upperBound,
+      boolean withReplacement,
+      long seed);
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala

@@ -31,9 +31,9 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, UnknownPartitioning}
 import org.apache.spark.sql.catalyst.util.truncatedString
-import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
 import org.apache.spark.sql.execution.datasources._
 import org.apache.spark.sql.execution.datasources.parquet.{ParquetFileFormat => ParquetSource}
+import org.apache.spark.sql.execution.datasources.v2.PushedDownOperators
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.sources.{BaseRelation, Filter}
@@ -103,8 +103,7 @@ case class RowDataSourceScanExec(
     requiredSchema: StructType,
     filters: Set[Filter],
     handledFilters: Set[Filter],
-    aggregation: Option[Aggregation],
-    limit: Option[Int],
+    pushedDownOperators: PushedDownOperators,
     rdd: RDD[InternalRow],
     @transient relation: BaseRelation,
     tableIdentifier: Option[TableIdentifier])
@@ -135,9 +134,9 @@ case class RowDataSourceScanExec(
 
     def seqToString(seq: Seq[Any]): String = seq.mkString("[", ", ", "]")
 
-    val (aggString, groupByString) = if (aggregation.nonEmpty) {
-      (seqToString(aggregation.get.aggregateExpressions),
-        seqToString(aggregation.get.groupByColumns))
+    val (aggString, groupByString) = if (pushedDownOperators.aggregation.nonEmpty) {
+      (seqToString(pushedDownOperators.aggregation.get.aggregateExpressions),
+        seqToString(pushedDownOperators.aggregation.get.groupByColumns))
     } else {
       ("[]", "[]")
     }
@@ -155,7 +154,10 @@ case class RowDataSourceScanExec(
       "PushedFilters" -> seqToString(markedFilters.toSeq),
       "PushedAggregates" -> aggString,
       "PushedGroupby" -> groupByString) ++
-      limit.map(value => "PushedLimit" -> s"LIMIT $value")
+      pushedDownOperators.limit.map(value => "PushedLimit" -> s"LIMIT $value") ++
+      pushedDownOperators.sample.map(v => "PushedSample" ->
+        s"SAMPLE ${v.lowerBound} ${v.upperBound} ${v.withReplacement} ${v.seed}"
+      )
   }
 
   // Don't care about `rdd` and `tableIdentifier` when canonicalizing.

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

@@ -45,6 +45,7 @@ import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Coun
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.{InSubqueryExec, RowDataSourceScanExec, SparkPlan}
 import org.apache.spark.sql.execution.command._
+import org.apache.spark.sql.execution.datasources.v2.PushedDownOperators
 import org.apache.spark.sql.execution.streaming.StreamingRelation
 import org.apache.spark.sql.internal.SQLConf.StoreAssignmentPolicy
 import org.apache.spark.sql.sources._
@@ -335,8 +336,7 @@ object DataSourceStrategy
         l.output.toStructType,
         Set.empty,
         Set.empty,
-        None,
-        None,
+        PushedDownOperators(None, None, None),
         toCatalystRDD(l, baseRelation.buildScan()),
         baseRelation,
         None) :: Nil
@@ -410,8 +410,7 @@ object DataSourceStrategy
         requestedColumns.toStructType,
         pushedFilters.toSet,
         handledFilters,
-        None,
-        None,
+        PushedDownOperators(None, None, None),
         scanBuilder(requestedColumns, candidatePredicates, pushedFilters),
         relation.relation,
         relation.catalogTable.map(_.identifier))
@@ -434,8 +433,7 @@ object DataSourceStrategy
         requestedColumns.toStructType,
         pushedFilters.toSet,
         handledFilters,
-        None,
-        None,
+        PushedDownOperators(None, None, None),
         scanBuilder(requestedColumns, candidatePredicates, pushedFilters),
         relation.relation,
         relation.catalogTable.map(_.identifier))

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

@@ -194,6 +194,9 @@ class JDBCOptions(
   // An option to allow/disallow pushing down LIMIT into JDBC data source
   val pushDownLimit = parameters.getOrElse(JDBC_PUSHDOWN_LIMIT, "false").toBoolean
 
+  // An option to allow/disallow pushing down TABLESAMPLE into JDBC data source
+  val pushDownTableSample = parameters.getOrElse(JDBC_PUSHDOWN_TABLESAMPLE, "false").toBoolean
+
   // The local path of user's keytab file, which is assumed to be pre-uploaded to all nodes either
   // by --files option of spark-submit or manually
   val keytab = {
@@ -270,6 +273,7 @@ object JDBCOptions {
   val JDBC_PUSHDOWN_PREDICATE = newOption("pushDownPredicate")
   val JDBC_PUSHDOWN_AGGREGATE = newOption("pushDownAggregate")
   val JDBC_PUSHDOWN_LIMIT = newOption("pushDownLimit")
+  val JDBC_PUSHDOWN_TABLESAMPLE = newOption("pushDownTableSample")
   val JDBC_KEYTAB = newOption("keytab")
   val JDBC_PRINCIPAL = newOption("principal")
   val JDBC_TABLE_COMMENT = newOption("tableComment")

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

@@ -26,6 +26,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Count, CountStar, Max, Min, Sum}
+import org.apache.spark.sql.execution.datasources.v2.TableSample
 import org.apache.spark.sql.jdbc.{JdbcDialect, JdbcDialects}
 import org.apache.spark.sql.sources._
 import org.apache.spark.sql.types._
@@ -181,6 +182,7 @@ object JDBCRDD extends Logging {
    * @param groupByColumns - The pushed down group by columns.
    * @param limit - The pushed down limit. If the value is 0, it means no limit or limit
    *                is not pushed down.
+   * @param sample - The pushed down tableSample.
    *
    * @return An RDD representing "SELECT requiredColumns FROM fqTable".
    */
@@ -193,6 +195,7 @@ object JDBCRDD extends Logging {
       options: JDBCOptions,
       outputSchema: Option[StructType] = None,
       groupByColumns: Option[Array[String]] = None,
+      sample: Option[TableSample] = None,
       limit: Int = 0): RDD[InternalRow] = {
     val url = options.url
     val dialect = JdbcDialects.get(url)
@@ -212,6 +215,7 @@ object JDBCRDD extends Logging {
       url,
       options,
       groupByColumns,
+      sample,
       limit)
   }
 }
@@ -231,6 +235,7 @@ private[jdbc] class JDBCRDD(
     url: String,
     options: JDBCOptions,
     groupByColumns: Option[Array[String]],
+    sample: Option[TableSample],
     limit: Int)
   extends RDD[InternalRow](sc, Nil) {
 
@@ -354,10 +359,12 @@ private[jdbc] class JDBCRDD(
 
     val myWhereClause = getWhereClause(part)
 
+    val myTableSampleClause: String = JdbcDialects.get(url).getTableSample(sample)
+
     val myLimitClause: String = dialect.getLimitClause(limit)
 
     val sqlText = s"SELECT $columnList FROM ${options.tableOrQuery} $myWhereClause" +
-      s" $getGroupByClause $myLimitClause"
+      s" $getGroupByClause $myTableSampleClause $myLimitClause"
     stmt = conn.prepareStatement(sqlText,
         ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)
     stmt.setFetchSize(options.fetchSize)

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

@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.util.{DateFormatter, DateTimeUtils, TimestampFormatter}
 import org.apache.spark.sql.catalyst.util.DateTimeUtils.{getZoneId, stringToDate, stringToTimestamp}
 import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.execution.datasources.v2.TableSample
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.jdbc.JdbcDialects
 import org.apache.spark.sql.sources._
@@ -299,7 +300,9 @@ private[sql] case class JDBCRelation(
       finalSchema: StructType,
       filters: Array[Filter],
       groupByColumns: Option[Array[String]],
-      limit: Int): RDD[Row] = {
+      tableSample: Option[TableSample],
+      limit: Int
+      ): RDD[Row] = {
     // Rely on a type erasure hack to pass RDD[InternalRow] back as RDD[Row]
     JDBCRDD.scanTable(
       sparkSession.sparkContext,
@@ -310,6 +313,7 @@ private[sql] case class JDBCRelation(
       jdbcOptions,
       Some(finalSchema),
       groupByColumns,
+      tableSample,
       limit).asInstanceOf[RDD[Row]]
   }
 

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

@@ -93,8 +93,8 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
   }
 
   override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
-    case PhysicalOperation(project, filters,
-        DataSourceV2ScanRelation(_, V1ScanWrapper(scan, pushed, aggregate, limit), output)) =>
+    case PhysicalOperation(project, filters, DataSourceV2ScanRelation(
+      _, V1ScanWrapper(scan, pushed, pushedDownOperators), output)) =>
       val v1Relation = scan.toV1TableScan[BaseRelation with TableScan](session.sqlContext)
       if (v1Relation.schema != scan.readSchema()) {
         throw QueryExecutionErrors.fallbackV1RelationReportsInconsistentSchemaError(
@@ -108,8 +108,7 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
         output.toStructType,
         Set.empty,
         pushed.toSet,
-        aggregate,
-        limit,
+        pushedDownOperators,
         unsafeRowRDD,
         v1Relation,
         tableIdentifier = None)