This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 909090f4825e [SPARK-56504][SQL] Extend join pushdown to accept pushed
samples
909090f4825e is described below
commit 909090f4825e4a20cb61c51a949c832d90a9c6ce
Author: Adithya Ajith <[email protected]>
AuthorDate: Fri Jul 3 21:50:42 2026 +0800
[SPARK-56504][SQL] Extend join pushdown to accept pushed samples
### What changes were proposed in this pull request?
This PR allows DataSource V2 join pushdown to compose with pushed table
samples when the connector can preserve the previously pushed sample state
([SPARK-56504](https://issues.apache.org/jira/browse/SPARK-56504)).
Before this change, Spark blocked join pushdown when either side of the
join had a real pushed table sample. That guard prevented incorrect results,
because a connector that pushed the join without preserving the already-pushed
sample could silently drop the sample.
This PR now removes that hard Spark-side block and relies on the existing
stateful `ScanBuilder` model:
1. Table sample pushdown is attempted before join pushdown.
2. If a connector accepts table sample pushdown, its scan builder already
stores that pushed sample state.
3. When `pushDownJoin(...)` is called later, the connector can decide from
its own stored state whether it can preserve the sample.
4. If it cannot preserve previously pushed state such as a table sample, it
must return `false`.
The PR updates the existing `SupportsPushDownJoin.pushDownJoin(...)`
documentation to make that contract explicit.
For JDBC, this PR preserves pushed table sample clauses when building the
left and right subqueries used for join pushdown. After the joined SQL query is
built, JDBC clears the stored table sample clause so it is not applied again.
This PR does not add a new public join pushdown API.
### Why are the changes needed?
SPARK-55978 added table sample pushdown and intentionally blocked composing
pushed table samples with join pushdown. That was safe, but too restrictive for
connectors that can correctly preserve both operations.
JDBC is one such connector. It already stores accepted table sample
pushdown state internally as a table sample clause. During join pushdown, JDBC
can include that clause in each pushed join-side subquery, preserving the
correct semantics.
This change allows Spark to push down joins on top of pushed table samples
when the connector accepts the composition through the existing
`pushDownJoin(...)` contract.
If a connector cannot preserve previously pushed state, it must reject join
pushdown by returning `false`.
### Does this PR introduce _any_ user-facing change?
Yes, for DataSource V2 JDBC queries with both table sample pushdown and
join pushdown enabled.
Previously, Spark did not push down joins when either side had a real
pushed table sample. After this change, JDBC can push down such joins while
preserving the table sample clause in the generated SQL.
There is no new public API. The existing `SupportsPushDownJoin` contract is
clarified for connectors that hold previously pushed state.
### How was this patch tested?
Added and updated tests in:
- `DataSourceV2TableSampleSuite`
- `JDBCSuite`
- in-memory test connector/catalog fixtures for sample and join pushdown
behavior
Ran:
```bash
build/sbt 'sql/Test/compile'
build/sbt 'sql/testOnly *DataSourceV2TableSampleSuite -- -z SPARK-56504'
'sql/testOnly org.apache.spark.sql.jdbc.JDBCSuite -- -z SPARK-56504'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex GPT 5.5
Closes #56486 from XdithyX/SPARK-56504.
Authored-by: Adithya Ajith <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit 273474b4cbc9213f733e4298e9cff3853531df5c)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../sql/connector/read/SupportsPushDownJoin.java | 6 +-
.../catalog/InMemoryTableWithTableSample.scala | 125 +++++++++++++++++++--
.../InMemoryTableWithTableSampleCatalog.scala | 27 +++++
.../datasources/v2/V2ScanRelationPushDown.scala | 13 +--
.../datasources/v2/jdbc/JDBCScanBuilder.scala | 8 +-
.../connector/DataSourceV2TableSampleSuite.scala | 87 ++++++++++++--
.../org/apache/spark/sql/jdbc/JDBCSuite.scala | 25 +++++
7 files changed, 256 insertions(+), 35 deletions(-)
diff --git
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownJoin.java
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownJoin.java
index a48a78671922..36e9d31a6a46 100644
---
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownJoin.java
+++
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownJoin.java
@@ -43,6 +43,9 @@ public interface SupportsPushDownJoin extends ScanBuilder {
/**
* Pushes down the join of the current {@code SupportsPushDownJoin} and the
other side of join
* {@code SupportsPushDownJoin}.
+ * <p>
+ * If this scan builder holds previously pushed state, such as a pushed
table sample, and cannot
+ * preserve it when pushing down the join, this method must return false.
*
* @param other {@code SupportsPushDownJoin} that this {@code
SupportsPushDownJoin}
* gets joined with.
@@ -52,7 +55,8 @@ public interface SupportsPushDownJoin extends ScanBuilder {
* @param rightSideRequiredColumnsWithAliases required output of the
* right side {@code
SupportsPushDownJoin}
* @param condition join condition. Columns are named after the specified
aliases in
- * {@code leftSideRequiredColumnWithAliases} and {@code
rightSideRequiredColumnWithAliases}
+ * {@code leftSideRequiredColumnsWithAliases} and {@code
rightSideRequiredColumnsWithAliases}
+ *
* @return True if join has been successfully pushed down.
*/
boolean pushDownJoin(
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithTableSample.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithTableSample.scala
index 514a7f3beda4..4c4ad71f3c68 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithTableSample.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithTableSample.scala
@@ -66,11 +66,11 @@ class InMemoryTableWithTableSample(
options: CaseInsensitiveStringMap)
extends InMemoryScanBuilder(tableSchema, options) with
SupportsPushDownTableSample {
- private var sampleFraction: Double = 1.0
- private var sampleSeed: Long = 0L
- private var sampleMethod: SampleMethod = SampleMethod.BERNOULLI
- private var sampleWithReplacement: Boolean = false
- private var samplePushed: Boolean = false
+ protected var sampleFraction: Double = 1.0
+ protected var sampleSeed: Long = 0L
+ protected var sampleMethod: SampleMethod = SampleMethod.BERNOULLI
+ protected var sampleWithReplacement: Boolean = false
+ protected var samplePushed: Boolean = false
override def pushTableSample(
lowerBound: Double,
@@ -118,6 +118,20 @@ class InMemoryTableWithTableSample(
InMemoryBatchScan(filteredPartitions, schema, tableSchema, options)
}
}
+
+ protected def hasResultAffectingSample: Boolean = {
+ samplePushed && (sampleWithReplacement || sampleFraction < 1.0)
+ }
+
+ protected def sampleDescription(side: String): Option[String] = {
+ if (samplePushed) {
+ val pct = sampleFraction * 100
+ val method = sampleMethod.toString.toUpperCase(Locale.ROOT)
+ Some(s"$side SAMPLE: $method SAMPLE ($pct)")
+ } else {
+ None
+ }
+ }
}
private class InMemoryBatchScanWithSample(
@@ -141,7 +155,7 @@ class InMemoryTableWithTableSample(
/**
* An in-memory table that supports both TABLESAMPLE pushdown and JOIN
pushdown.
- * Used to test the guard that prevents join pushdown when a side has a pushed
sample.
+ * Used to test join pushdown when a side has a pushed sample.
*/
class InMemoryTableWithJoinAndSample(
name: String,
@@ -163,6 +177,7 @@ class InMemoryTableWithJoinAndSample(
private[catalog] val ownSchema: StructType = tableSchema
private var pushed: Array[Predicate] = Array.empty
private var joinedSchema: Option[StructType] = None
+ private var joinedSampleDescription: Option[String] = None
override def pushPredicates(predicates: Array[Predicate]):
Array[Predicate] = {
pushed = predicates
@@ -189,7 +204,103 @@ class InMemoryTableWithJoinAndSample(
leftSideRequiredColumnsWithAliases: Array[ColumnWithAlias],
rightSideRequiredColumnsWithAliases: Array[ColumnWithAlias],
condition: Predicate): Boolean = {
- val otherSchema =
other.asInstanceOf[InMemoryJoinAndSampleScanBuilder].ownSchema
+ val otherScanBuilder =
other.asInstanceOf[InMemoryJoinAndSampleScanBuilder]
+ if (sampleWithReplacement || otherScanBuilder.sampleWithReplacement) {
+ return false
+ }
+
+ val otherSchema = otherScanBuilder.ownSchema
+ val leftFields = leftSideRequiredColumnsWithAliases.map { col =>
+ val name = if (col.alias() != null) col.alias() else col.colName()
+ tableSchema(col.colName()).copy(name = name)
+ }
+ val rightFields = rightSideRequiredColumnsWithAliases.map { col =>
+ val name = if (col.alias() != null) col.alias() else col.colName()
+ otherSchema(col.colName()).copy(name = name)
+ }
+ joinedSchema = Some(StructType(leftFields ++ rightFields))
+ joinedSampleDescription = Some(
+ Seq(sampleDescription("LEFT"),
otherScanBuilder.sampleDescription("RIGHT"))
+ .flatten
+ .mkString(" "))
+ true
+ }
+
+ override def build: Scan = {
+ joinedSchema match {
+ case Some(js) =>
+ new InMemoryBatchScanWithJoinSample(
+ data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq,
+ js, tableSchema, options, joinedSampleDescription.getOrElse(""))
+ case None => super.build
+ }
+ }
+
+ private class InMemoryBatchScanWithJoinSample(
+ data: Seq[InputPartition],
+ readSchema: StructType,
+ tableSchema: StructType,
+ options: CaseInsensitiveStringMap,
+ sampleDescription: String)
+ extends InMemoryBatchScan(data, readSchema, tableSchema, options) {
+
+ override def description(): String = {
+ Seq(super.description(),
sampleDescription).filter(_.nonEmpty).mkString(" ")
+ }
+ }
+ }
+}
+
+/**
+ * An in-memory table that supports TABLESAMPLE pushdown and only the original
+ * JOIN pushdown API. Used to test a connector rejecting pushed samples from
pushDownJoin.
+ */
+class InMemoryTableWithLegacyJoinAndSample(
+ name: String,
+ columns: Array[Column],
+ partitioning: Array[Transform],
+ properties: util.Map[String, String])
+ extends InMemoryTableWithTableSample(name, columns, partitioning,
properties) {
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder
= {
+ new InMemoryLegacyJoinAndSampleScanBuilder(schema, options)
+ }
+
+ class InMemoryLegacyJoinAndSampleScanBuilder(
+ tableSchema: StructType,
+ options: CaseInsensitiveStringMap)
+ extends InMemoryTableSampleScanBuilder(tableSchema, options)
+ with SupportsPushDownJoin with SupportsPushDownV2Filters {
+
+ private[catalog] val ownSchema: StructType = tableSchema
+ private var pushed: Array[Predicate] = Array.empty
+ private var joinedSchema: Option[StructType] = None
+
+ override def pushPredicates(predicates: Array[Predicate]):
Array[Predicate] = {
+ pushed = predicates
+ Array.empty
+ }
+
+ override def pushFilters(filters: Array[Filter]): Array[Filter] = {
+ Array.empty
+ }
+
+ override def pushedPredicates(): Array[Predicate] = pushed
+
+ override def isOtherSideCompatibleForJoin(other: SupportsPushDownJoin):
Boolean = true
+
+ override def pushDownJoin(
+ other: SupportsPushDownJoin,
+ joinType: JoinType,
+ leftSideRequiredColumnsWithAliases: Array[ColumnWithAlias],
+ rightSideRequiredColumnsWithAliases: Array[ColumnWithAlias],
+ condition: Predicate): Boolean = {
+ val otherScanBuilder =
other.asInstanceOf[InMemoryLegacyJoinAndSampleScanBuilder]
+ if (hasResultAffectingSample ||
otherScanBuilder.hasResultAffectingSample) {
+ return false
+ }
+
+ val otherSchema = otherScanBuilder.ownSchema
val leftFields = leftSideRequiredColumnsWithAliases.map { col =>
val name = if (col.alias() != null) col.alias() else col.colName()
tableSchema(col.colName()).copy(name = name)
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithTableSampleCatalog.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithTableSampleCatalog.scala
index 12da978ea11a..30cb00eab6fe 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithTableSampleCatalog.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithTableSampleCatalog.scala
@@ -74,6 +74,33 @@ class InMemoryTableWithJoinAndSampleCatalog extends
InMemoryTableCatalog {
}
}
+class InMemoryTableWithLegacyJoinAndSampleCatalog extends InMemoryTableCatalog
{
+ import CatalogV2Implicits._
+
+ override def createTable(
+ ident: Identifier,
+ columns: Array[Column],
+ partitions: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ if (tables.containsKey(ident)) {
+ throw new TableAlreadyExistsException(ident.asMultipartIdentifier)
+ }
+
+ InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties)
+
+ val tableName = s"$name.${ident.quoted}"
+ val table = new InMemoryTableWithLegacyJoinAndSample(
+ tableName, columns, partitions, properties)
+ tables.put(ident, table)
+ namespaces.putIfAbsent(ident.namespace.toList, Map())
+ table
+ }
+
+ override def createTable(ident: Identifier, tableInfo: TableInfo): Table = {
+ createTable(ident, tableInfo.columns(), tableInfo.partitions(),
tableInfo.properties)
+ }
+}
+
class InMemoryTableWithLegacyTableSampleCatalog extends InMemoryTableCatalog {
import CatalogV2Implicits._
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
index 42afb513f406..14d9f754f95c 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
@@ -165,18 +165,6 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
rightProjections.forall(_.isInstanceOf[AttributeReference]) &&
// Cross joins are not supported because they increase the amount of
data.
condition.isDefined &&
- // Do not push down join if either side has a pushed sample with
- // fraction < 1, because the merged scan builder would silently
- // discard it and change the result. At fraction = 1 without
- // replacement the sample is a no-op on the result set, so dropping
- // it is safe. With replacement (Poisson sampling), even fraction 1
- // can emit each row 0, 1, 2, ... times, so it is not a no-op.
- // TODO(SPARK-56504): Extend SupportsPushDownJoin to accept pushed
- // samples so sources supporting both can handle the composition.
- leftHolder.pushedSample.forall(s =>
- !s.withReplacement && s.upperBound - s.lowerBound >= 1.0) &&
- rightHolder.pushedSample.forall(s =>
- !s.withReplacement && s.upperBound - s.lowerBound >= 1.0) &&
lBuilder.isOtherSideCompatibleForJoin(rBuilder) =>
// Process left and right columns in original order
val (leftSideRequiredColumnsWithAliases,
rightSideRequiredColumnsWithAliases) =
@@ -1172,6 +1160,7 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
sHolder.pushedLimit, sHolder.pushedOffset, sHolder.sortOrders,
sHolder.pushedPredicates,
sHolder.joinedRelationsPushedDownOperators, optRelationName)
}
+
}
case class ScanBuilderHolder(
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala
index 45d5f920b9be..c6eec4487f5b 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala
@@ -263,13 +263,13 @@ case class JDBCScanBuilder(
val quotedAliases = columnsWithAliases
.map(col => Option(col.alias()).map(dialect.quoteIdentifier))
- // Only filters can be pushed down before join pushdown, so we need to
craft SQL query
- // that contains filters as well.
- // Joins on top of samples are not supported so we don't need to provide
tableSample here.
- dialect
+ val builder = dialect
.getJdbcSQLQueryBuilder(jdbcOptions)
.withPredicates(pushedPredicate, JDBCPartition(whereClause = null, idx =
1))
.withAliasedColumns(quotedColumns, quotedAliases)
+
+ tableSampleClause.foreach(builder.withTableSampleClause)
+ builder
}
override def pushTableSample(
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2TableSampleSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2TableSampleSuite.scala
index 164c098e95e8..6e1670ce2cb8 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2TableSampleSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2TableSampleSuite.scala
@@ -18,7 +18,11 @@
package org.apache.spark.sql.connector
import org.apache.spark.sql.AnalysisException
-import
org.apache.spark.sql.connector.catalog.{InMemoryTableWithJoinAndSampleCatalog,
InMemoryTableWithLegacyTableSampleCatalog, InMemoryTableWithTableSampleCatalog}
+import org.apache.spark.sql.connector.catalog.{
+ InMemoryTableWithJoinAndSampleCatalog,
+ InMemoryTableWithLegacyJoinAndSampleCatalog,
+ InMemoryTableWithLegacyTableSampleCatalog,
+ InMemoryTableWithTableSampleCatalog}
import org.apache.spark.sql.internal.SQLConf
class DataSourceV2TableSampleSuite extends DatasourceV2SQLBase
@@ -156,7 +160,7 @@ class DataSourceV2TableSampleSuite extends
DatasourceV2SQLBase
}
}
- test("SPARK-55978: join pushdown is skipped when a side has a pushed
sample") {
+ test("SPARK-56504: join pushdown includes left side pushed sample") {
val joinSampleCatalog = "testjoinsample"
registerCatalog(joinSampleCatalog,
classOf[InMemoryTableWithJoinAndSampleCatalog])
val t1 = s"$joinSampleCatalog.ns.t1"
@@ -171,15 +175,47 @@ class DataSourceV2TableSampleSuite extends
DatasourceV2SQLBase
val dfNoSample = sql(s"SELECT * FROM $t1 JOIN $t2 ON $t1.id = $t2.id")
checkJoinPushed(dfNoSample)
- // With a SYSTEM sample (fraction < 1) on one side: join pushdown
- // should be skipped because the merged scan builder would silently
- // discard the sample.
+ // The connector preserves its previously pushed sample when pushing
down the join.
val dfWithSample = sql(
s"SELECT * FROM $t1 TABLESAMPLE SYSTEM (50 PERCENT) " +
s"JOIN $t2 ON $t1.id = $t2.id")
- checkJoinNotPushed(dfWithSample)
- // The sample should still be pushed down though
+ checkJoinPushed(dfWithSample)
checkSamplePushed(dfWithSample, pushed = true)
+ checkPushedInfo(dfWithSample, "LEFT SAMPLE: SYSTEM SAMPLE (50.0)")
+ }
+ } finally {
+ sql(s"DROP TABLE IF EXISTS $t1")
+ sql(s"DROP TABLE IF EXISTS $t2")
+ }
+ }
+
+ test("SPARK-56504: join pushdown includes right and both side pushed
samples") {
+ val joinSampleCatalog = "testjoinsampleboth"
+ registerCatalog(joinSampleCatalog,
classOf[InMemoryTableWithJoinAndSampleCatalog])
+ val t1 = s"$joinSampleCatalog.ns.t1"
+ val t2 = s"$joinSampleCatalog.ns.t2"
+ sql(s"CREATE TABLE $t1 (id bigint, data string) USING _")
+ sql(s"CREATE TABLE $t2 (id bigint, data string) USING _")
+ try {
+ sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b'), (3, 'c')")
+ sql(s"INSERT INTO $t2 VALUES (2, 'x'), (3, 'y'), (4, 'z')")
+ withSQLConf(SQLConf.DATA_SOURCE_V2_JOIN_PUSHDOWN.key -> "true") {
+ val rightSample = sql(
+ s"SELECT * FROM $t1 JOIN $t2 TABLESAMPLE BERNOULLI (50 PERCENT) " +
+ s"ON $t1.id = $t2.id")
+ checkJoinPushed(rightSample)
+ checkSamplePushed(rightSample, pushed = true)
+ checkPushedInfo(rightSample, "RIGHT SAMPLE: BERNOULLI SAMPLE (50.0)")
+
+ val bothSamples = sql(
+ s"SELECT * FROM $t1 TABLESAMPLE SYSTEM (50 PERCENT) " +
+ s"JOIN $t2 TABLESAMPLE BERNOULLI (50 PERCENT) ON $t1.id = $t2.id")
+ checkJoinPushed(bothSamples)
+ checkSamplePushed(bothSamples, pushed = true)
+ checkPushedInfo(
+ bothSamples,
+ "LEFT SAMPLE: SYSTEM SAMPLE (50.0)",
+ "RIGHT SAMPLE: BERNOULLI SAMPLE (50.0)")
}
} finally {
sql(s"DROP TABLE IF EXISTS $t1")
@@ -198,10 +234,8 @@ class DataSourceV2TableSampleSuite extends
DatasourceV2SQLBase
sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b'), (3, 'c')")
sql(s"INSERT INTO $t2 VALUES (2, 'x'), (3, 'y'), (4, 'z')")
withSQLConf(SQLConf.DATA_SOURCE_V2_JOIN_PUSHDOWN.key -> "true") {
- // At fraction = 1 the sample is a no-op on the result set, so
- // dropping it inside the merged scan builder is safe. The guard
- // in V2ScanRelationPushDown short-circuits and join pushdown
- // proceeds.
+ // At fraction = 1 the sample is a no-op on the result set, so this
connector
+ // can safely preserve it while pushing down the join.
val dfWithSample = sql(
s"SELECT * FROM $t1 TABLESAMPLE SYSTEM (100 PERCENT) " +
s"JOIN $t2 ON $t1.id = $t2.id")
@@ -213,6 +247,37 @@ class DataSourceV2TableSampleSuite extends
DatasourceV2SQLBase
}
}
+ test("SPARK-56504: connector rejects join pushdown when it cannot preserve
pushed samples") {
+ val joinSampleCatalog = "testlegjoinandsample"
+ registerCatalog(joinSampleCatalog,
classOf[InMemoryTableWithLegacyJoinAndSampleCatalog])
+ val t1 = s"$joinSampleCatalog.ns.t1"
+ val t2 = s"$joinSampleCatalog.ns.t2"
+ sql(s"CREATE TABLE $t1 (id bigint, data string) USING _")
+ sql(s"CREATE TABLE $t2 (id bigint, data string) USING _")
+ try {
+ sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b'), (3, 'c')")
+ sql(s"INSERT INTO $t2 VALUES (2, 'x'), (3, 'y'), (4, 'z')")
+ withSQLConf(SQLConf.DATA_SOURCE_V2_JOIN_PUSHDOWN.key -> "true") {
+ val noSample = sql(s"SELECT * FROM $t1 JOIN $t2 ON $t1.id = $t2.id")
+ checkJoinPushed(noSample)
+
+ val noOpSample = sql(
+ s"SELECT * FROM $t1 TABLESAMPLE SYSTEM (100 PERCENT) " +
+ s"JOIN $t2 ON $t1.id = $t2.id")
+ checkJoinPushed(noOpSample)
+
+ val realSample = sql(
+ s"SELECT * FROM $t1 TABLESAMPLE SYSTEM (50 PERCENT) " +
+ s"JOIN $t2 ON $t1.id = $t2.id")
+ checkJoinNotPushed(realSample)
+ checkSamplePushed(realSample, pushed = true)
+ }
+ } finally {
+ sql(s"DROP TABLE IF EXISTS $t1")
+ sql(s"DROP TABLE IF EXISTS $t2")
+ }
+ }
+
test("SPARK-55978: with-replacement sample blocks join pushdown even at
fraction 1") {
val joinSampleCatalog = "testjoinsamplerepl"
registerCatalog(joinSampleCatalog,
classOf[InMemoryTableWithJoinAndSampleCatalog])
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala
index f353f6f2a699..0c9a053c59e6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala
@@ -1289,6 +1289,31 @@ class JDBCSuite extends SharedSparkSession {
"SELECT tab.* FROM (SELECT a,b FROM test ) tab WHERE rownum <= 123")
}
+ test("SPARK-56504: JdbcSQLQueryBuilder preserves table sample in pushed join
sides") {
+ // JDBC url is a required option but is not used in this test.
+ val options = new JDBCOptions(Map("url" -> "jdbc:h2://host:port",
"dbtable" -> "test"))
+ val dialect = JdbcDialects.get("jdbc:h2://host:port")
+ val left = dialect
+ .getJdbcSQLQueryBuilder(options)
+ .withColumns(Array("a"))
+ .withTableSampleClause("TABLESAMPLE SYSTEM (50)")
+ val right = dialect
+ .getJdbcSQLQueryBuilder(options)
+ .withColumns(Array("b"))
+ .withTableSampleClause("TABLESAMPLE BERNOULLI (25)")
+
+ val query = dialect
+ .getJdbcSQLQueryBuilder(options)
+ .withJoin(left, right, "L", "R", Array("a", "b"), "INNER JOIN", "L.a =
R.b")
+ .build()
+ .replaceAll("\\s+", " ")
+
+ assert(query.contains("SELECT a FROM test TABLESAMPLE SYSTEM (50)"))
+ assert(query.contains("SELECT b FROM test TABLESAMPLE BERNOULLI (25)"))
+ assert(query.contains("INNER JOIN"))
+ assert(query.contains("ON L.a = R.b"))
+ }
+
test("MsSqlServerDialect jdbc type mapping") {
val msSqlServerDialect = JdbcDialects.get("jdbc:sqlserver")
assert(msSqlServerDialect.getJDBCType(TimestampType).map(_.databaseTypeDefinition).get
==
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
