+ * Examples: + *
+ * For example, to return reducer for reducing f_source = bucket(4, x) on f_target = bucket(2, x) + *
+ * Example of reducing f_source = days(x) on f_target = hours(x) + *
+ * A [[Reducer]] exists for a transform expression function if it is + * 'reducible' on the other expression function. + *
+ * @return reducer function or None if not reducible on the other transform expression + */ + def reducers(other: TransformExpression): Option[Reducer[_, _]] = { + (function, other.function) match { + case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) => + reducer(e1, numBucketsOpt, e2, other.numBucketsOpt) + case _ => None + } + } + + // Return a Reducer for a reducible function on another reducible function + private def reducer( + thisFunction: ReducibleFunction[_, _], + thisNumBucketsOpt: Option[Int], + otherFunction: ReducibleFunction[_, _], + otherNumBucketsOpt: Option[Int]): Option[Reducer[_, _]] = { + val res = (thisNumBucketsOpt, otherNumBucketsOpt) match { + case (Some(numBuckets), Some(otherNumBuckets)) => + thisFunction.reducer(numBuckets, otherFunction, otherNumBuckets) + case _ => thisFunction.reducer(otherFunction) + } + Option(res) + } + override def dataType: DataType = function.resultType() override protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala index dfe07a443a230..63a000f50489d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala @@ -33,6 +33,8 @@ import org.codehaus.janino.util.ClassFile import org.apache.spark.{SparkException, SparkIllegalArgumentException, TaskContext, TaskKilledException} import org.apache.spark.executor.InputMetrics import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKey._ +import org.apache.spark.internal.MDC import org.apache.spark.metrics.source.CodegenMetrics import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.encoders.HashableWeakReference @@ -1593,7 +1595,8 @@ object CodeGenerator extends Logging { if (byteCodeSize > DEFAULT_JVM_HUGE_METHOD_LIMIT) { logInfo("Generated method too long to be JIT compiled: " + - s"${cf.getThisClassName}.${method.getName} is $byteCodeSize bytes") + log"${MDC(CLASS_NAME, cf.getThisClassName)}.${MDC(METHOD_NAME, method.getName)} " + + log"is ${MDC(BYTECODE_SIZE, byteCodeSize)} bytes") } byteCodeSize @@ -1638,7 +1641,7 @@ object CodeGenerator extends Logging { val timeMs: Double = duration.toDouble / NANOS_PER_MILLIS CodegenMetrics.METRIC_SOURCE_CODE_SIZE.update(code.body.length) CodegenMetrics.METRIC_COMPILATION_TIME.update(timeMs.toLong) - logInfo(s"Code generated in $timeMs ms") + logInfo(log"Code generated in ${MDC(TIME_UNITS, timeMs)} ms") _compileTime.add(duration) result } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 122cba5b74f8d..0e31595f919d2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -20,6 +20,8 @@ package org.apache.spark.sql.catalyst.optimizer import scala.collection.mutable import org.apache.spark.SparkException +import org.apache.spark.internal.LogKey._ +import org.apache.spark.internal.MDC import org.apache.spark.sql.catalyst.SQLConfHelper import org.apache.spark.sql.catalyst.analysis._ import org.apache.spark.sql.catalyst.catalog.{InMemoryCatalog, SessionCatalog} @@ -455,7 +457,8 @@ abstract class Optimizer(catalogManager: CatalogManager) val filteredRules = batch.rules.filter { rule => val exclude = excludedRules.contains(rule.ruleName) if (exclude) { - logInfo(s"Optimization rule '${rule.ruleName}' is excluded from the optimizer.") + logInfo(log"Optimization rule '${MDC(RULE_NAME, rule.ruleName)}' " + + log"is excluded from the optimizer.") } !exclude } @@ -464,8 +467,8 @@ abstract class Optimizer(catalogManager: CatalogManager) } else if (filteredRules.nonEmpty) { Some(Batch(batch.name, batch.strategy, filteredRules: _*)) } else { - logInfo(s"Optimization batch '${batch.name}' is excluded from the optimizer " + - s"as all enclosed rules have been excluded.") + logInfo(log"Optimization batch '${MDC(RULE_BATCH_NAME, batch.name)}' " + + log"is excluded from the optimizer as all enclosed rules have been excluded.") None } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index c98a2a92a3abb..2364130f79e4c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala @@ -24,6 +24,7 @@ import org.apache.spark.{SparkException, SparkUnsupportedOperationException} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper +import org.apache.spark.sql.connector.catalog.functions.Reducer import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, IntegerType} @@ -833,10 +834,42 @@ case class KeyGroupedShuffleSpec( (left, right) match { case (_: LeafExpression, _: LeafExpression) => true case (left: TransformExpression, right: TransformExpression) => - left.isSameFunction(right) + if (SQLConf.get.v2BucketingPushPartValuesEnabled && + !SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled && + SQLConf.get.v2BucketingAllowCompatibleTransforms) { + left.isCompatible(right) + } else { + left.isSameFunction(right) + } case _ => false } + /** + * Return a set of [[Reducer]] for the partition expressions of this shuffle spec, + * on the partition expressions of another shuffle spec. + *
+ * A [[Reducer]] exists for a partition expression function of this shuffle spec if it is + * 'reducible' on the corresponding partition expression function of the other shuffle spec. + *
+ * If a value is returned, there must be one [[Reducer]] per partition expression. + * A None value in the set indicates that the particular partition expression is not reducible + * on the corresponding expression on the other shuffle spec. + *
+ * Returning none also indicates that none of the partition expressions can be reduced on the
+ * corresponding expression on the other shuffle spec.
+ *
+ * @param other other key-grouped shuffle spec
+ */
+ def reducers(other: KeyGroupedShuffleSpec): Option[Seq[Option[Reducer[_, _]]]] = {
+ val results = partitioning.expressions.zip(other.partitioning.expressions).map {
+ case (e1: TransformExpression, e2: TransformExpression) => e1.reducers(e2)
+ case (_, _) => None
+ }
+
+ // optimize to not return a value, if none of the partition expressions are reducible
+ if (results.forall(p => p.isEmpty)) None else Some(results)
+ }
+
override def canCreatePartitioning: Boolean = SQLConf.get.v2BucketingShuffleEnabled &&
// Only support partition expressions are AttributeReference for now
partitioning.expressions.forall(_.isInstanceOf[AttributeReference])
@@ -846,6 +879,21 @@ case class KeyGroupedShuffleSpec(
}
}
+object KeyGroupedShuffleSpec {
+ def reducePartitionValue(
+ row: InternalRow,
+ expressions: Seq[Expression],
+ reducers: Seq[Option[Reducer[_, _]]]):
+ InternalRowComparableWrapper = {
+ val partitionVals = row.toSeq(expressions.map(_.dataType))
+ val reducedRow = partitionVals.zip(reducers).map{
+ case (v, Some(reducer: Reducer[Any, Any])) => reducer.reduce(v)
+ case (v, _) => v
+ }.toArray
+ InternalRowComparableWrapper(new GenericInternalRow(reducedRow), expressions)
+ }
+}
+
case class ShuffleSpecCollection(specs: Seq[ShuffleSpec]) extends ShuffleSpec {
override def isCompatibleWith(other: ShuffleSpec): Boolean = {
specs.exists(_.isCompatibleWith(other))
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleExecutor.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleExecutor.scala
index d5cd5a90e3382..ff2d3fd631127 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleExecutor.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleExecutor.scala
@@ -18,7 +18,9 @@
package org.apache.spark.sql.catalyst.rules
import org.apache.spark.SparkException
-import org.apache.spark.internal.Logging
+import org.apache.spark.internal.{Logging, MessageWithContext}
+import org.apache.spark.internal.LogKey._
+import org.apache.spark.internal.MDC
import org.apache.spark.sql.catalyst.QueryPlanningTracker
import org.apache.spark.sql.catalyst.trees.TreeNode
import org.apache.spark.sql.catalyst.util.DateTimeConstants.NANOS_PER_SECOND
@@ -56,10 +58,10 @@ class PlanChangeLogger[TreeType <: TreeNode[_]] extends Logging {
def logRule(ruleName: String, oldPlan: TreeType, newPlan: TreeType): Unit = {
if (!newPlan.fastEquals(oldPlan)) {
if (logRules.isEmpty || logRules.get.contains(ruleName)) {
- def message(): String = {
- s"""
- |=== Applying Rule $ruleName ===
- |${sideBySide(oldPlan.treeString, newPlan.treeString).mkString("\n")}
+ def message(): MessageWithContext = {
+ log"""
+ |=== Applying Rule ${MDC(RULE_NAME, ruleName)} ===
+ |${MDC(QUERY_PLAN, sideBySide(oldPlan.treeString, newPlan.treeString).mkString("\n"))}
""".stripMargin
}
@@ -70,14 +72,14 @@ class PlanChangeLogger[TreeType <: TreeNode[_]] extends Logging {
def logBatch(batchName: String, oldPlan: TreeType, newPlan: TreeType): Unit = {
if (logBatches.isEmpty || logBatches.get.contains(batchName)) {
- def message(): String = {
+ def message(): MessageWithContext = {
if (!oldPlan.fastEquals(newPlan)) {
- s"""
- |=== Result of Batch $batchName ===
- |${sideBySide(oldPlan.treeString, newPlan.treeString).mkString("\n")}
+ log"""
+ |=== Result of Batch ${MDC(RULE_BATCH_NAME, batchName)} ===
+ |${MDC(QUERY_PLAN, sideBySide(oldPlan.treeString, newPlan.treeString).mkString("\n"))}
""".stripMargin
} else {
- s"Batch $batchName has no effect."
+ log"Batch ${MDC(RULE_BATCH_NAME, batchName)} has no effect."
}
}
@@ -88,26 +90,26 @@ class PlanChangeLogger[TreeType <: TreeNode[_]] extends Logging {
def logMetrics(metrics: QueryExecutionMetrics): Unit = {
val totalTime = metrics.time / NANOS_PER_SECOND.toDouble
val totalTimeEffective = metrics.timeEffective / NANOS_PER_SECOND.toDouble
- val message =
- s"""
+ val message: MessageWithContext =
+ log"""
|=== Metrics of Executed Rules ===
- |Total number of runs: ${metrics.numRuns}
- |Total time: $totalTime seconds
- |Total number of effective runs: ${metrics.numEffectiveRuns}
- |Total time of effective runs: $totalTimeEffective seconds
+ |Total number of runs: ${MDC(RULE_NUMBER_OF_RUNS, metrics.numRuns)}
+ |Total time: ${MDC(TIME_UNITS, totalTime)} seconds
+ |Total number of effective runs: ${MDC(RULE_NUMBER_OF_RUNS, metrics.numEffectiveRuns)}
+ |Total time of effective runs: ${MDC(TIME_UNITS, totalTimeEffective)} seconds
""".stripMargin
logBasedOnLevel(message)
}
- private def logBasedOnLevel(f: => String): Unit = {
+ private def logBasedOnLevel(f: => MessageWithContext): Unit = {
logLevel match {
- case "TRACE" => logTrace(f)
- case "DEBUG" => logDebug(f)
+ case "TRACE" => logTrace(f.message)
+ case "DEBUG" => logDebug(f.message)
case "INFO" => logInfo(f)
case "WARN" => logWarning(f)
case "ERROR" => logError(f)
- case _ => logTrace(f)
+ case _ => logTrace(f.message)
}
}
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/ValidatorUtil.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/ValidatorUtil.scala
index 51140d0101939..3d93c4e8742ab 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/ValidatorUtil.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/ValidatorUtil.scala
@@ -27,6 +27,8 @@ import org.apache.hadoop.fs.Path
import org.apache.spark.SparkFiles
import org.apache.spark.deploy.SparkHadoopUtil
import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKey._
+import org.apache.spark.internal.MDC
/**
* Utilities for working with XSD validation.
@@ -56,7 +58,8 @@ object ValidatorUtil extends Logging {
case e: Throwable =>
// Handle case where it was added with sc.addFile
// When they are added via sc.addFile, they are always downloaded to local file system
- logInfo(s"$xsdPath was not found, falling back to look up files added by Spark")
+ logInfo(log"${MDC(XSD_PATH, xsdPath)} was not found, " +
+ log"falling back to look up files added by Spark")
val f = new File(SparkFiles.get(xsdPath.toString))
if (f.exists()) {
new FileInputStream(f)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 9f07722528e88..73cb4fba8637d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -1558,6 +1558,18 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS =
+ buildConf("spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled")
+ .doc("Whether to allow storage-partition join in the case where the partition transforms " +
+ "are compatible but not identical. This config requires both " +
+ s"${V2_BUCKETING_ENABLED.key} and ${V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key} to be " +
+ s"enabled and ${V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key} " +
+ "to be disabled."
+ )
+ .version("4.0.0")
+ .booleanConf
+ .createWithDefault(false)
+
val BUCKETING_MAX_BUCKETS = buildConf("spark.sql.sources.bucketing.maxBuckets")
.doc("The maximum number of buckets allowed.")
.version("2.4.0")
@@ -5323,6 +5335,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
def v2BucketingAllowJoinKeysSubsetOfPartitionKeys: Boolean =
getConf(SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS)
+ def v2BucketingAllowCompatibleTransforms: Boolean =
+ getConf(SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS)
+
def dataFrameSelfJoinAutoResolveAmbiguity: Boolean =
getConf(DATAFRAME_SELF_JOIN_AUTO_RESOLVE_AMBIGUITY)
diff --git a/sql/core/benchmarks/StateStoreBasicOperationsBenchmark-jdk21-results.txt b/sql/core/benchmarks/StateStoreBasicOperationsBenchmark-jdk21-results.txt
index d3b3aafc21e57..0317e61163752 100644
--- a/sql/core/benchmarks/StateStoreBasicOperationsBenchmark-jdk21-results.txt
+++ b/sql/core/benchmarks/StateStoreBasicOperationsBenchmark-jdk21-results.txt
@@ -6,33 +6,66 @@ OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
putting 10000 rows (10000 rows to overwrite - rate 100): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------------
-In-memory 9 10 1 1.1 894.7 1.0X
-RocksDB (trackTotalNumberOfRows: true) 41 42 2 0.2 4064.6 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 15 1 0.7 1466.8 0.6X
+In-memory 9 10 1 1.1 936.2 1.0X
+RocksDB (trackTotalNumberOfRows: true) 41 42 1 0.2 4068.9 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1500.4 0.6X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
putting 10000 rows (5000 rows to overwrite - rate 50): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------
-In-memory 9 10 0 1.1 893.1 1.0X
-RocksDB (trackTotalNumberOfRows: true) 40 40 1 0.3 3959.7 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1510.8 0.6X
+In-memory 9 11 1 1.1 929.8 1.0X
+RocksDB (trackTotalNumberOfRows: true) 40 41 1 0.3 3955.7 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1497.3 0.6X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
putting 10000 rows (1000 rows to overwrite - rate 10): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------
-In-memory 9 9 0 1.1 872.0 1.0X
-RocksDB (trackTotalNumberOfRows: true) 39 40 1 0.3 3887.2 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.7 1532.3 0.6X
+In-memory 9 10 1 1.1 907.5 1.0X
+RocksDB (trackTotalNumberOfRows: true) 39 40 1 0.3 3886.5 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1497.2 0.6X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
putting 10000 rows (0 rows to overwrite - rate 0): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------
-In-memory 9 10 1 1.1 874.5 1.0X
-RocksDB (trackTotalNumberOfRows: true) 40 41 1 0.3 3967.1 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.7 1526.2 0.6X
+In-memory 9 10 1 1.1 904.0 1.0X
+RocksDB (trackTotalNumberOfRows: true) 39 40 1 0.3 3859.8 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.7 1497.2 0.6X
+
+
+================================================================================================
+merge rows
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
+AMD EPYC 7763 64-Core Processor
+merging 10000 rows with 10 values per key (10000 rows to overwrite - rate 100): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+--------------------------------------------------------------------------------------------------------------------------------------------------------------
+RocksDB (trackTotalNumberOfRows: true) 519 533 7 0.0 51916.6 1.0X
+RocksDB (trackTotalNumberOfRows: false) 171 177 3 0.1 17083.9 3.0X
+
+OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
+AMD EPYC 7763 64-Core Processor
+merging 10000 rows with 10 values per key (5000 rows to overwrite - rate 50): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------------------------------
+RocksDB (trackTotalNumberOfRows: true) 506 521 7 0.0 50644.0 1.0X
+RocksDB (trackTotalNumberOfRows: false) 170 176 3 0.1 17022.0 3.0X
+
+OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
+AMD EPYC 7763 64-Core Processor
+merging 10000 rows with 10 values per key (1000 rows to overwrite - rate 10): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------------------------------
+RocksDB (trackTotalNumberOfRows: true) 493 508 6 0.0 49319.3 1.0X
+RocksDB (trackTotalNumberOfRows: false) 169 175 3 0.1 16897.6 2.9X
+
+OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
+AMD EPYC 7763 64-Core Processor
+merging 10000 rows with 10 values per key (0 rows to overwrite - rate 0): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+--------------------------------------------------------------------------------------------------------------------------------------------------------
+RocksDB (trackTotalNumberOfRows: true) 495 508 6 0.0 49462.5 1.0X
+RocksDB (trackTotalNumberOfRows: false) 169 175 3 0.1 16896.6 2.9X
================================================================================================
@@ -43,33 +76,33 @@ OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
trying to delete 10000 rows from 10000 rows(10000 rows are non-existing - rate 100): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 0 1 0 20.9 47.9 1.0X
-RocksDB (trackTotalNumberOfRows: true) 40 40 1 0.3 3956.8 0.0X
-RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.6 1541.9 0.0X
+In-memory 0 1 0 26.3 38.0 1.0X
+RocksDB (trackTotalNumberOfRows: true) 39 41 1 0.3 3942.0 0.0X
+RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1529.2 0.0X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
trying to delete 10000 rows from 10000 rows(5000 rows are non-existing - rate 50): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 10 1 1.3 773.4 1.0X
-RocksDB (trackTotalNumberOfRows: true) 40 41 1 0.2 4024.1 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1537.8 0.5X
+In-memory 8 9 1 1.3 790.4 1.0X
+RocksDB (trackTotalNumberOfRows: true) 40 41 1 0.2 4036.7 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.7 1536.9 0.5X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
trying to delete 10000 rows from 10000 rows(1000 rows are non-existing - rate 10): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 10 1 1.2 817.7 1.0X
-RocksDB (trackTotalNumberOfRows: true) 41 42 1 0.2 4111.7 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.6 1540.5 0.5X
+In-memory 8 10 1 1.2 847.0 1.0X
+RocksDB (trackTotalNumberOfRows: true) 41 42 1 0.2 4099.8 0.2X
+RocksDB (trackTotalNumberOfRows: false) 16 16 0 0.6 1563.3 0.5X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
trying to delete 10000 rows from 10000 rows(0 rows are non-existing - rate 0): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 10 1 1.2 820.0 1.0X
-RocksDB (trackTotalNumberOfRows: true) 41 42 1 0.2 4133.0 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.7 1526.2 0.5X
+In-memory 9 10 1 1.2 859.4 1.0X
+RocksDB (trackTotalNumberOfRows: true) 41 42 1 0.2 4118.9 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1507.8 0.6X
================================================================================================
@@ -80,32 +113,30 @@ OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
evicting 10000 rows (maxTimestampToEvictInMillis: 9999) from 10000 rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 9 0 1.2 805.5 1.0X
-RocksDB (trackTotalNumberOfRows: true) 39 40 1 0.3 3888.6 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.7 1538.2 0.5X
+In-memory 8 9 1 1.2 831.0 1.0X
+RocksDB (trackTotalNumberOfRows: true) 40 40 1 0.3 3956.6 0.2X
+RocksDB (trackTotalNumberOfRows: false) 16 16 0 0.6 1571.3 0.5X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
evicting 5000 rows (maxTimestampToEvictInMillis: 4999) from 10000 rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 8 0 1.3 754.7 1.0X
-RocksDB (trackTotalNumberOfRows: true) 21 22 0 0.5 2091.7 0.4X
-RocksDB (trackTotalNumberOfRows: false) 9 9 0 1.1 916.1 0.8X
+In-memory 8 8 1 1.3 787.6 1.0X
+RocksDB (trackTotalNumberOfRows: true) 21 22 0 0.5 2112.6 0.4X
+RocksDB (trackTotalNumberOfRows: false) 9 9 0 1.1 932.9 0.8X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
evicting 1000 rows (maxTimestampToEvictInMillis: 999) from 10000 rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 7 8 1 1.4 692.8 1.0X
-RocksDB (trackTotalNumberOfRows: true) 7 7 0 1.5 654.6 1.1X
-RocksDB (trackTotalNumberOfRows: false) 4 4 0 2.4 423.8 1.6X
+In-memory 7 8 0 1.4 715.7 1.0X
+RocksDB (trackTotalNumberOfRows: true) 7 7 0 1.5 676.3 1.1X
+RocksDB (trackTotalNumberOfRows: false) 4 5 0 2.3 442.3 1.6X
OpenJDK 64-Bit Server VM 21.0.2+13-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
evicting 0 rows (maxTimestampToEvictInMillis: -1) from 10000 rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 0 0 0 24.2 41.2 1.0X
-RocksDB (trackTotalNumberOfRows: true) 3 3 0 3.4 290.1 0.1X
-RocksDB (trackTotalNumberOfRows: false) 3 3 0 3.4 290.6 0.1X
-
-
+In-memory 0 0 0 23.8 41.9 1.0X
+RocksDB (trackTotalNumberOfRows: true) 3 3 0 3.2 309.5 0.1X
+RocksDB (trackTotalNumberOfRows: false) 3 3 0 3.2 309.9 0.1X
diff --git a/sql/core/benchmarks/StateStoreBasicOperationsBenchmark-results.txt b/sql/core/benchmarks/StateStoreBasicOperationsBenchmark-results.txt
index 86d3d44003316..d2aa646d5ec1d 100644
--- a/sql/core/benchmarks/StateStoreBasicOperationsBenchmark-results.txt
+++ b/sql/core/benchmarks/StateStoreBasicOperationsBenchmark-results.txt
@@ -6,33 +6,66 @@ OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
putting 10000 rows (10000 rows to overwrite - rate 100): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------------
-In-memory 9 10 1 1.1 907.3 1.0X
-RocksDB (trackTotalNumberOfRows: true) 40 42 2 0.2 4048.4 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1508.4 0.6X
+In-memory 10 12 1 1.0 960.1 1.0X
+RocksDB (trackTotalNumberOfRows: true) 42 43 2 0.2 4173.9 0.2X
+RocksDB (trackTotalNumberOfRows: false) 16 16 1 0.6 1551.6 0.6X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
putting 10000 rows (5000 rows to overwrite - rate 50): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------
-In-memory 9 10 1 1.1 901.9 1.0X
-RocksDB (trackTotalNumberOfRows: true) 39 41 1 0.3 3922.0 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1513.5 0.6X
+In-memory 10 12 1 1.0 970.1 1.0X
+RocksDB (trackTotalNumberOfRows: true) 41 42 1 0.2 4095.8 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 17 1 0.6 1544.6 0.6X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
putting 10000 rows (1000 rows to overwrite - rate 10): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------
-In-memory 9 9 0 1.1 880.2 1.0X
-RocksDB (trackTotalNumberOfRows: true) 38 39 1 0.3 3829.7 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1506.0 0.6X
+In-memory 9 11 1 1.1 933.3 1.0X
+RocksDB (trackTotalNumberOfRows: true) 40 41 1 0.3 3966.2 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.6 1540.2 0.6X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
putting 10000 rows (0 rows to overwrite - rate 0): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------
-In-memory 9 9 0 1.1 878.6 1.0X
-RocksDB (trackTotalNumberOfRows: true) 38 39 1 0.3 3803.4 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1500.3 0.6X
+In-memory 9 11 1 1.1 936.1 1.0X
+RocksDB (trackTotalNumberOfRows: true) 39 41 1 0.3 3942.4 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.7 1530.1 0.6X
+
+
+================================================================================================
+merge rows
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
+AMD EPYC 7763 64-Core Processor
+merging 10000 rows with 10 values per key (10000 rows to overwrite - rate 100): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+--------------------------------------------------------------------------------------------------------------------------------------------------------------
+RocksDB (trackTotalNumberOfRows: true) 525 538 6 0.0 52516.4 1.0X
+RocksDB (trackTotalNumberOfRows: false) 170 177 4 0.1 16960.4 3.1X
+
+OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
+AMD EPYC 7763 64-Core Processor
+merging 10000 rows with 10 values per key (5000 rows to overwrite - rate 50): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------------------------------
+RocksDB (trackTotalNumberOfRows: true) 514 528 6 0.0 51351.9 1.0X
+RocksDB (trackTotalNumberOfRows: false) 168 174 4 0.1 16794.0 3.1X
+
+OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
+AMD EPYC 7763 64-Core Processor
+merging 10000 rows with 10 values per key (1000 rows to overwrite - rate 10): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------------------------------
+RocksDB (trackTotalNumberOfRows: true) 500 513 6 0.0 49955.1 1.0X
+RocksDB (trackTotalNumberOfRows: false) 169 174 2 0.1 16867.1 3.0X
+
+OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
+AMD EPYC 7763 64-Core Processor
+merging 10000 rows with 10 values per key (0 rows to overwrite - rate 0): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+--------------------------------------------------------------------------------------------------------------------------------------------------------
+RocksDB (trackTotalNumberOfRows: true) 492 508 8 0.0 49225.8 1.0X
+RocksDB (trackTotalNumberOfRows: false) 168 173 3 0.1 16757.2 2.9X
================================================================================================
@@ -43,33 +76,33 @@ OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
trying to delete 10000 rows from 10000 rows(10000 rows are non-existing - rate 100): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 0 1 0 26.5 37.8 1.0X
-RocksDB (trackTotalNumberOfRows: true) 38 39 1 0.3 3779.3 0.0X
-RocksDB (trackTotalNumberOfRows: false) 15 15 1 0.7 1462.9 0.0X
+In-memory 0 1 0 26.1 38.3 1.0X
+RocksDB (trackTotalNumberOfRows: true) 38 40 1 0.3 3835.6 0.0X
+RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1455.7 0.0X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
trying to delete 10000 rows from 10000 rows(5000 rows are non-existing - rate 50): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 7 8 0 1.3 742.4 1.0X
-RocksDB (trackTotalNumberOfRows: true) 39 40 1 0.3 3913.8 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 15 0 0.7 1461.1 0.5X
+In-memory 8 9 1 1.3 793.9 1.0X
+RocksDB (trackTotalNumberOfRows: true) 40 41 1 0.2 4018.1 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 0 0.7 1505.6 0.5X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
trying to delete 10000 rows from 10000 rows(1000 rows are non-existing - rate 10): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 8 0 1.3 794.8 1.0X
-RocksDB (trackTotalNumberOfRows: true) 40 41 1 0.2 4012.8 0.2X
-RocksDB (trackTotalNumberOfRows: false) 15 15 0 0.7 1461.5 0.5X
+In-memory 8 10 1 1.2 837.2 1.0X
+RocksDB (trackTotalNumberOfRows: true) 41 42 1 0.2 4073.9 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 16 1 0.7 1470.6 0.6X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
trying to delete 10000 rows from 10000 rows(0 rows are non-existing - rate 0): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 9 0 1.2 809.7 1.0X
-RocksDB (trackTotalNumberOfRows: true) 40 42 1 0.2 4043.4 0.2X
-RocksDB (trackTotalNumberOfRows: false) 14 15 1 0.7 1445.2 0.6X
+In-memory 8 9 0 1.2 843.6 1.0X
+RocksDB (trackTotalNumberOfRows: true) 41 42 1 0.2 4088.7 0.2X
+RocksDB (trackTotalNumberOfRows: false) 15 15 0 0.7 1466.1 0.6X
================================================================================================
@@ -80,32 +113,30 @@ OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
evicting 10000 rows (maxTimestampToEvictInMillis: 9999) from 10000 rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 9 0 1.2 806.2 1.0X
-RocksDB (trackTotalNumberOfRows: true) 40 41 1 0.3 3980.0 0.2X
-RocksDB (trackTotalNumberOfRows: false) 16 16 0 0.6 1599.4 0.5X
+In-memory 8 9 0 1.2 833.5 1.0X
+RocksDB (trackTotalNumberOfRows: true) 40 41 0 0.3 3976.5 0.2X
+RocksDB (trackTotalNumberOfRows: false) 16 16 0 0.6 1588.1 0.5X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
evicting 5000 rows (maxTimestampToEvictInMillis: 4999) from 10000 rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 8 8 0 1.3 752.2 1.0X
-RocksDB (trackTotalNumberOfRows: true) 22 22 0 0.5 2170.8 0.3X
-RocksDB (trackTotalNumberOfRows: false) 10 10 0 1.0 967.7 0.8X
+In-memory 8 8 0 1.3 784.3 1.0X
+RocksDB (trackTotalNumberOfRows: true) 22 22 0 0.5 2155.1 0.4X
+RocksDB (trackTotalNumberOfRows: false) 10 10 0 1.0 986.9 0.8X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
evicting 1000 rows (maxTimestampToEvictInMillis: 999) from 10000 rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-----------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 7 7 0 1.4 694.9 1.0X
-RocksDB (trackTotalNumberOfRows: true) 7 7 0 1.4 700.1 1.0X
-RocksDB (trackTotalNumberOfRows: false) 5 5 0 2.2 465.0 1.5X
+In-memory 7 8 0 1.4 722.3 1.0X
+RocksDB (trackTotalNumberOfRows: true) 7 7 0 1.4 718.8 1.0X
+RocksDB (trackTotalNumberOfRows: false) 5 5 0 2.0 488.7 1.5X
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Linux 6.5.0-1016-azure
AMD EPYC 7763 64-Core Processor
evicting 0 rows (maxTimestampToEvictInMillis: -1) from 10000 rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------------------
-In-memory 1 1 0 19.7 50.7 1.0X
-RocksDB (trackTotalNumberOfRows: true) 3 3 0 3.0 332.1 0.2X
-RocksDB (trackTotalNumberOfRows: false) 3 3 0 3.0 331.7 0.2X
-
-
+In-memory 0 1 0 21.3 46.9 1.0X
+RocksDB (trackTotalNumberOfRows: true) 4 4 0 2.8 358.9 0.1X
+RocksDB (trackTotalNumberOfRows: false) 4 4 0 2.8 358.7 0.1X
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala
index 7cce599040189..f949dbf71a371 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala
@@ -24,9 +24,10 @@ import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.QueryPlan
-import org.apache.spark.sql.catalyst.plans.physical.{KeyGroupedPartitioning, Partitioning, SinglePartition}
+import org.apache.spark.sql.catalyst.plans.physical.{KeyGroupedPartitioning, KeyGroupedShuffleSpec, Partitioning, SinglePartition}
import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper}
import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.connector.catalog.functions.Reducer
import org.apache.spark.sql.connector.read._
import org.apache.spark.util.ArrayImplicits._
@@ -164,6 +165,18 @@ case class BatchScanExec(
(groupedParts, expressions)
}
+ // Also re-group the partitions if we are reducing compatible partition expressions
+ val finalGroupedPartitions = spjParams.reducers match {
+ case Some(reducers) =>
+ val result = groupedPartitions.groupBy { case (row, _) =>
+ KeyGroupedShuffleSpec.reducePartitionValue(row, partExpressions, reducers)
+ }.map { case (wrapper, splits) => (wrapper.row, splits.flatMap(_._2)) }.toSeq
+ val rowOrdering = RowOrdering.createNaturalAscendingOrdering(
+ partExpressions.map(_.dataType))
+ result.sorted(rowOrdering.on((t: (InternalRow, _)) => t._1))
+ case _ => groupedPartitions
+ }
+
// When partially clustered, the input partitions are not grouped by partition
// values. Here we'll need to check `commonPartitionValues` and decide how to group
// and replicate splits within a partition.
@@ -174,7 +187,7 @@ case class BatchScanExec(
.get
.map(t => (InternalRowComparableWrapper(t._1, partExpressions), t._2))
.toMap
- val nestGroupedPartitions = groupedPartitions.map { case (partValue, splits) =>
+ val nestGroupedPartitions = finalGroupedPartitions.map { case (partValue, splits) =>
// `commonPartValuesMap` should contain the part value since it's the super set.
val numSplits = commonPartValuesMap
.get(InternalRowComparableWrapper(partValue, partExpressions))
@@ -207,7 +220,7 @@ case class BatchScanExec(
} else {
// either `commonPartitionValues` is not defined, or it is defined but
// `applyPartialClustering` is false.
- val partitionMapping = groupedPartitions.map { case (partValue, splits) =>
+ val partitionMapping = finalGroupedPartitions.map { case (partValue, splits) =>
InternalRowComparableWrapper(partValue, partExpressions) -> splits
}.toMap
@@ -259,6 +272,7 @@ case class StoragePartitionJoinParams(
keyGroupedPartitioning: Option[Seq[Expression]] = None,
joinKeyPositions: Option[Seq[Int]] = None,
commonPartitionValues: Option[Seq[(InternalRow, Int)]] = None,
+ reducers: Option[Seq[Option[Reducer[_, _]]]] = None,
applyPartialClustering: Boolean = false,
replicatePartitions: Boolean = false) {
override def equals(other: Any): Boolean = other match {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
index 2a7c1206bb410..a0f74ef6c3d02 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
@@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.plans._
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
+import org.apache.spark.sql.connector.catalog.functions.Reducer
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec}
@@ -505,11 +506,28 @@ case class EnsureRequirements(
}
}
- // Now we need to push-down the common partition key to the scan in each child
- newLeft = populatePartitionValues(left, mergedPartValues, leftSpec.joinKeyPositions,
- applyPartialClustering, replicateLeftSide)
- newRight = populatePartitionValues(right, mergedPartValues, rightSpec.joinKeyPositions,
- applyPartialClustering, replicateRightSide)
+ // in case of compatible but not identical partition expressions, we apply 'reduce'
+ // transforms to group one side's partitions as well as the common partition values
+ val leftReducers = leftSpec.reducers(rightSpec)
+ val rightReducers = rightSpec.reducers(leftSpec)
+
+ if (leftReducers.isDefined || rightReducers.isDefined) {
+ mergedPartValues = reduceCommonPartValues(mergedPartValues,
+ leftSpec.partitioning.expressions,
+ leftReducers)
+ mergedPartValues = reduceCommonPartValues(mergedPartValues,
+ rightSpec.partitioning.expressions,
+ rightReducers)
+ val rowOrdering = RowOrdering
+ .createNaturalAscendingOrdering(partitionExprs.map(_.dataType))
+ mergedPartValues = mergedPartValues.sorted(rowOrdering.on((t: (InternalRow, _)) => t._1))
+ }
+
+ // Now we need to push-down the common partition information to the scan in each child
+ newLeft = populateCommonPartitionInfo(left, mergedPartValues, leftSpec.joinKeyPositions,
+ leftReducers, applyPartialClustering, replicateLeftSide)
+ newRight = populateCommonPartitionInfo(right, mergedPartValues, rightSpec.joinKeyPositions,
+ rightReducers, applyPartialClustering, replicateRightSide)
}
}
@@ -527,11 +545,12 @@ case class EnsureRequirements(
joinType == LeftAnti || joinType == LeftOuter
}
- // Populate the common partition values down to the scan nodes
- private def populatePartitionValues(
+ // Populate the common partition information down to the scan nodes
+ private def populateCommonPartitionInfo(
plan: SparkPlan,
values: Seq[(InternalRow, Int)],
joinKeyPositions: Option[Seq[Int]],
+ reducers: Option[Seq[Option[Reducer[_, _]]]],
applyPartialClustering: Boolean,
replicatePartitions: Boolean): SparkPlan = plan match {
case scan: BatchScanExec =>
@@ -539,13 +558,26 @@ case class EnsureRequirements(
spjParams = scan.spjParams.copy(
commonPartitionValues = Some(values),
joinKeyPositions = joinKeyPositions,
+ reducers = reducers,
applyPartialClustering = applyPartialClustering,
replicatePartitions = replicatePartitions
)
)
case node =>
- node.mapChildren(child => populatePartitionValues(
- child, values, joinKeyPositions, applyPartialClustering, replicatePartitions))
+ node.mapChildren(child => populateCommonPartitionInfo(
+ child, values, joinKeyPositions, reducers, applyPartialClustering, replicatePartitions))
+ }
+
+ private def reduceCommonPartValues(
+ commonPartValues: Seq[(InternalRow, Int)],
+ expressions: Seq[Expression],
+ reducers: Option[Seq[Option[Reducer[_, _]]]]) = {
+ reducers match {
+ case Some(reducers) => commonPartValues.groupBy { case (row, _) =>
+ KeyGroupedShuffleSpec.reducePartitionValue(row, expressions, reducers)
+ }.map{ case(wrapper, splits) => (wrapper.row, splits.map(_._2).sum) }.toSeq
+ case _ => commonPartValues
+ }
}
/**
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index 7fdc703007c2c..ec275fe101fd6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -63,11 +63,17 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
Collections.emptyMap[String, String]
}
private val table: String = "tbl"
+
private val columns: Array[Column] = Array(
Column.create("id", IntegerType),
Column.create("data", StringType),
Column.create("ts", TimestampType))
+ private val columns2: Array[Column] = Array(
+ Column.create("store_id", IntegerType),
+ Column.create("dept_id", IntegerType),
+ Column.create("data", StringType))
+
test("clustered distribution: output partitioning should be KeyGroupedPartitioning") {
val partitions: Array[Transform] = Array(Expressions.years("ts"))
@@ -1309,6 +1315,474 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
}
}
+ test("SPARK-47094: Support compatible buckets") {
+ val table1 = "tab1e1"
+ val table2 = "table2"
+
+ Seq(
+ ((2, 4), (4, 2)),
+ ((4, 2), (2, 4)),
+ ((2, 2), (4, 6)),
+ ((6, 2), (2, 2))).foreach {
+ case ((table1buckets1, table1buckets2), (table2buckets1, table2buckets2)) =>
+ catalog.clearTables()
+
+ val partition1 = Array(bucket(table1buckets1, "store_id"),
+ bucket(table1buckets2, "dept_id"))
+ val partition2 = Array(bucket(table2buckets1, "store_id"),
+ bucket(table2buckets2, "dept_id"))
+
+ Seq((table1, partition1), (table2, partition2)).foreach { case (tab, part) =>
+ createTable(tab, columns2, part)
+ val insertStr = s"INSERT INTO testcat.ns.$tab VALUES " +
+ "(0, 0, 'aa'), " +
+ "(0, 0, 'ab'), " + // duplicate partition key
+ "(0, 1, 'ac'), " +
+ "(0, 2, 'ad'), " +
+ "(0, 3, 'ae'), " +
+ "(0, 4, 'af'), " +
+ "(0, 5, 'ag'), " +
+ "(1, 0, 'ah'), " +
+ "(1, 0, 'ai'), " + // duplicate partition key
+ "(1, 1, 'aj'), " +
+ "(1, 2, 'ak'), " +
+ "(1, 3, 'al'), " +
+ "(1, 4, 'am'), " +
+ "(1, 5, 'an'), " +
+ "(2, 0, 'ao'), " +
+ "(2, 0, 'ap'), " + // duplicate partition key
+ "(2, 1, 'aq'), " +
+ "(2, 2, 'ar'), " +
+ "(2, 3, 'as'), " +
+ "(2, 4, 'at'), " +
+ "(2, 5, 'au'), " +
+ "(3, 0, 'av'), " +
+ "(3, 0, 'aw'), " + // duplicate partition key
+ "(3, 1, 'ax'), " +
+ "(3, 2, 'ay'), " +
+ "(3, 3, 'az'), " +
+ "(3, 4, 'ba'), " +
+ "(3, 5, 'bb'), " +
+ "(4, 0, 'bc'), " +
+ "(4, 0, 'bd'), " + // duplicate partition key
+ "(4, 1, 'be'), " +
+ "(4, 2, 'bf'), " +
+ "(4, 3, 'bg'), " +
+ "(4, 4, 'bh'), " +
+ "(4, 5, 'bi'), " +
+ "(5, 0, 'bj'), " +
+ "(5, 0, 'bk'), " + // duplicate partition key
+ "(5, 1, 'bl'), " +
+ "(5, 2, 'bm'), " +
+ "(5, 3, 'bn'), " +
+ "(5, 4, 'bo'), " +
+ "(5, 5, 'bp')"
+
+ // additional unmatched partitions to test push down
+ val finalStr = if (tab == table1) {
+ insertStr ++ ", (8, 0, 'xa'), (8, 8, 'xx')"
+ } else {
+ insertStr ++ ", (9, 0, 'ya'), (9, 9, 'yy')"
+ }
+
+ sql(finalStr)
+ }
+
+ Seq(true, false).foreach { allowJoinKeysSubsetOfPartitionKeys =>
+ withSQLConf(
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
+ allowJoinKeysSubsetOfPartitionKeys.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("t1", "t2")}
+ |t1.store_id, t1.dept_id, t1.data, t2.data
+ |FROM testcat.ns.$table1 t1 JOIN testcat.ns.$table2 t2
+ |ON t1.store_id = t2.store_id AND t1.dept_id = t2.dept_id
+ |ORDER BY t1.store_id, t1.dept_id, t1.data, t2.data
+ |""".stripMargin)
+
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.isEmpty, "SPJ should be triggered")
+
+ val scans = collectScans(df.queryExecution.executedPlan).map(_.inputRDD.
+ partitions.length)
+ val expectedBuckets = Math.min(table1buckets1, table2buckets1) *
+ Math.min(table1buckets2, table2buckets2)
+ assert(scans == Seq(expectedBuckets, expectedBuckets))
+
+ checkAnswer(df, Seq(
+ Row(0, 0, "aa", "aa"),
+ Row(0, 0, "aa", "ab"),
+ Row(0, 0, "ab", "aa"),
+ Row(0, 0, "ab", "ab"),
+ Row(0, 1, "ac", "ac"),
+ Row(0, 2, "ad", "ad"),
+ Row(0, 3, "ae", "ae"),
+ Row(0, 4, "af", "af"),
+ Row(0, 5, "ag", "ag"),
+ Row(1, 0, "ah", "ah"),
+ Row(1, 0, "ah", "ai"),
+ Row(1, 0, "ai", "ah"),
+ Row(1, 0, "ai", "ai"),
+ Row(1, 1, "aj", "aj"),
+ Row(1, 2, "ak", "ak"),
+ Row(1, 3, "al", "al"),
+ Row(1, 4, "am", "am"),
+ Row(1, 5, "an", "an"),
+ Row(2, 0, "ao", "ao"),
+ Row(2, 0, "ao", "ap"),
+ Row(2, 0, "ap", "ao"),
+ Row(2, 0, "ap", "ap"),
+ Row(2, 1, "aq", "aq"),
+ Row(2, 2, "ar", "ar"),
+ Row(2, 3, "as", "as"),
+ Row(2, 4, "at", "at"),
+ Row(2, 5, "au", "au"),
+ Row(3, 0, "av", "av"),
+ Row(3, 0, "av", "aw"),
+ Row(3, 0, "aw", "av"),
+ Row(3, 0, "aw", "aw"),
+ Row(3, 1, "ax", "ax"),
+ Row(3, 2, "ay", "ay"),
+ Row(3, 3, "az", "az"),
+ Row(3, 4, "ba", "ba"),
+ Row(3, 5, "bb", "bb"),
+ Row(4, 0, "bc", "bc"),
+ Row(4, 0, "bc", "bd"),
+ Row(4, 0, "bd", "bc"),
+ Row(4, 0, "bd", "bd"),
+ Row(4, 1, "be", "be"),
+ Row(4, 2, "bf", "bf"),
+ Row(4, 3, "bg", "bg"),
+ Row(4, 4, "bh", "bh"),
+ Row(4, 5, "bi", "bi"),
+ Row(5, 0, "bj", "bj"),
+ Row(5, 0, "bj", "bk"),
+ Row(5, 0, "bk", "bj"),
+ Row(5, 0, "bk", "bk"),
+ Row(5, 1, "bl", "bl"),
+ Row(5, 2, "bm", "bm"),
+ Row(5, 3, "bn", "bn"),
+ Row(5, 4, "bo", "bo"),
+ Row(5, 5, "bp", "bp")
+ ))
+ }
+ }
+ }
+ }
+
+ test("SPARK-47094: Support compatible buckets with common divisor") {
+ val table1 = "tab1e1"
+ val table2 = "table2"
+
+ Seq(
+ ((6, 4), (4, 6)),
+ ((6, 6), (4, 4)),
+ ((4, 4), (6, 6)),
+ ((4, 6), (6, 4))).foreach {
+ case ((table1buckets1, table1buckets2), (table2buckets1, table2buckets2)) =>
+ catalog.clearTables()
+
+ val partition1 = Array(bucket(table1buckets1, "store_id"),
+ bucket(table1buckets2, "dept_id"))
+ val partition2 = Array(bucket(table2buckets1, "store_id"),
+ bucket(table2buckets2, "dept_id"))
+
+ Seq((table1, partition1), (table2, partition2)).foreach { case (tab, part) =>
+ createTable(tab, columns2, part)
+ val insertStr = s"INSERT INTO testcat.ns.$tab VALUES " +
+ "(0, 0, 'aa'), " +
+ "(0, 0, 'ab'), " + // duplicate partition key
+ "(0, 1, 'ac'), " +
+ "(0, 2, 'ad'), " +
+ "(0, 3, 'ae'), " +
+ "(0, 4, 'af'), " +
+ "(0, 5, 'ag'), " +
+ "(1, 0, 'ah'), " +
+ "(1, 0, 'ai'), " + // duplicate partition key
+ "(1, 1, 'aj'), " +
+ "(1, 2, 'ak'), " +
+ "(1, 3, 'al'), " +
+ "(1, 4, 'am'), " +
+ "(1, 5, 'an'), " +
+ "(2, 0, 'ao'), " +
+ "(2, 0, 'ap'), " + // duplicate partition key
+ "(2, 1, 'aq'), " +
+ "(2, 2, 'ar'), " +
+ "(2, 3, 'as'), " +
+ "(2, 4, 'at'), " +
+ "(2, 5, 'au'), " +
+ "(3, 0, 'av'), " +
+ "(3, 0, 'aw'), " + // duplicate partition key
+ "(3, 1, 'ax'), " +
+ "(3, 2, 'ay'), " +
+ "(3, 3, 'az'), " +
+ "(3, 4, 'ba'), " +
+ "(3, 5, 'bb'), " +
+ "(4, 0, 'bc'), " +
+ "(4, 0, 'bd'), " + // duplicate partition key
+ "(4, 1, 'be'), " +
+ "(4, 2, 'bf'), " +
+ "(4, 3, 'bg'), " +
+ "(4, 4, 'bh'), " +
+ "(4, 5, 'bi'), " +
+ "(5, 0, 'bj'), " +
+ "(5, 0, 'bk'), " + // duplicate partition key
+ "(5, 1, 'bl'), " +
+ "(5, 2, 'bm'), " +
+ "(5, 3, 'bn'), " +
+ "(5, 4, 'bo'), " +
+ "(5, 5, 'bp')"
+
+ // additional unmatched partitions to test push down
+ val finalStr = if (tab == table1) {
+ insertStr ++ ", (8, 0, 'xa'), (8, 8, 'xx')"
+ } else {
+ insertStr ++ ", (9, 0, 'ya'), (9, 9, 'yy')"
+ }
+
+ sql(finalStr)
+ }
+
+ Seq(true, false).foreach { allowJoinKeysSubsetOfPartitionKeys =>
+ withSQLConf(
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
+ allowJoinKeysSubsetOfPartitionKeys.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("t1", "t2")}
+ |t1.store_id, t1.dept_id, t1.data, t2.data
+ |FROM testcat.ns.$table1 t1 JOIN testcat.ns.$table2 t2
+ |ON t1.store_id = t2.store_id AND t1.dept_id = t2.dept_id
+ |ORDER BY t1.store_id, t1.dept_id, t1.data, t2.data
+ |""".stripMargin)
+
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.isEmpty, "SPJ should be triggered")
+
+ val scans = collectScans(df.queryExecution.executedPlan).map(_.inputRDD.
+ partitions.length)
+
+ def gcd(a: Int, b: Int): Int = BigInt(a).gcd(BigInt(b)).toInt
+ val expectedBuckets = gcd(table1buckets1, table2buckets1) *
+ gcd(table1buckets2, table2buckets2)
+ assert(scans == Seq(expectedBuckets, expectedBuckets))
+
+ checkAnswer(df, Seq(
+ Row(0, 0, "aa", "aa"),
+ Row(0, 0, "aa", "ab"),
+ Row(0, 0, "ab", "aa"),
+ Row(0, 0, "ab", "ab"),
+ Row(0, 1, "ac", "ac"),
+ Row(0, 2, "ad", "ad"),
+ Row(0, 3, "ae", "ae"),
+ Row(0, 4, "af", "af"),
+ Row(0, 5, "ag", "ag"),
+ Row(1, 0, "ah", "ah"),
+ Row(1, 0, "ah", "ai"),
+ Row(1, 0, "ai", "ah"),
+ Row(1, 0, "ai", "ai"),
+ Row(1, 1, "aj", "aj"),
+ Row(1, 2, "ak", "ak"),
+ Row(1, 3, "al", "al"),
+ Row(1, 4, "am", "am"),
+ Row(1, 5, "an", "an"),
+ Row(2, 0, "ao", "ao"),
+ Row(2, 0, "ao", "ap"),
+ Row(2, 0, "ap", "ao"),
+ Row(2, 0, "ap", "ap"),
+ Row(2, 1, "aq", "aq"),
+ Row(2, 2, "ar", "ar"),
+ Row(2, 3, "as", "as"),
+ Row(2, 4, "at", "at"),
+ Row(2, 5, "au", "au"),
+ Row(3, 0, "av", "av"),
+ Row(3, 0, "av", "aw"),
+ Row(3, 0, "aw", "av"),
+ Row(3, 0, "aw", "aw"),
+ Row(3, 1, "ax", "ax"),
+ Row(3, 2, "ay", "ay"),
+ Row(3, 3, "az", "az"),
+ Row(3, 4, "ba", "ba"),
+ Row(3, 5, "bb", "bb"),
+ Row(4, 0, "bc", "bc"),
+ Row(4, 0, "bc", "bd"),
+ Row(4, 0, "bd", "bc"),
+ Row(4, 0, "bd", "bd"),
+ Row(4, 1, "be", "be"),
+ Row(4, 2, "bf", "bf"),
+ Row(4, 3, "bg", "bg"),
+ Row(4, 4, "bh", "bh"),
+ Row(4, 5, "bi", "bi"),
+ Row(5, 0, "bj", "bj"),
+ Row(5, 0, "bj", "bk"),
+ Row(5, 0, "bk", "bj"),
+ Row(5, 0, "bk", "bk"),
+ Row(5, 1, "bl", "bl"),
+ Row(5, 2, "bm", "bm"),
+ Row(5, 3, "bn", "bn"),
+ Row(5, 4, "bo", "bo"),
+ Row(5, 5, "bp", "bp")
+ ))
+ }
+ }
+ }
+ }
+
+ test("SPARK-47094: Support compatible buckets with less join keys than partition keys") {
+ val table1 = "tab1e1"
+ val table2 = "table2"
+
+ Seq((2, 4), (4, 2), (2, 6), (6, 2)).foreach {
+ case (table1buckets, table2buckets) =>
+ catalog.clearTables()
+
+ val partition1 = Array(identity("data"),
+ bucket(table1buckets, "dept_id"))
+ val partition2 = Array(bucket(3, "store_id"),
+ bucket(table2buckets, "dept_id"))
+
+ createTable(table1, columns2, partition1)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(0, 0, 'aa'), " +
+ "(1, 0, 'ab'), " +
+ "(2, 1, 'ac'), " +
+ "(3, 2, 'ad'), " +
+ "(4, 3, 'ae'), " +
+ "(5, 4, 'af'), " +
+ "(6, 5, 'ag'), " +
+
+ // value without other side match
+ "(6, 6, 'xx')"
+ )
+
+ createTable(table2, columns2, partition2)
+ sql(s"INSERT INTO testcat.ns.$table2 VALUES " +
+ "(6, 0, '01'), " +
+ "(5, 1, '02'), " + // duplicate partition key
+ "(5, 1, '03'), " +
+ "(4, 2, '04'), " +
+ "(3, 3, '05'), " +
+ "(2, 4, '06'), " +
+ "(1, 5, '07'), " +
+
+ // value without other side match
+ "(7, 7, '99')"
+ )
+
+
+ withSQLConf(
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("t1", "t2")}
+ |t1.store_id, t2.store_id, t1.dept_id, t2.dept_id, t1.data, t2.data
+ |FROM testcat.ns.$table1 t1 JOIN testcat.ns.$table2 t2
+ |ON t1.dept_id = t2.dept_id
+ |ORDER BY t1.store_id, t1.dept_id, t1.data, t2.data
+ |""".stripMargin)
+
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.isEmpty, "SPJ should be triggered")
+
+ val scans = collectScans(df.queryExecution.executedPlan).map(_.inputRDD.
+ partitions.length)
+
+ val expectedBuckets = Math.min(table1buckets, table2buckets)
+
+ assert(scans == Seq(expectedBuckets, expectedBuckets))
+
+ checkAnswer(df, Seq(
+ Row(0, 6, 0, 0, "aa", "01"),
+ Row(1, 6, 0, 0, "ab", "01"),
+ Row(2, 5, 1, 1, "ac", "02"),
+ Row(2, 5, 1, 1, "ac", "03"),
+ Row(3, 4, 2, 2, "ad", "04"),
+ Row(4, 3, 3, 3, "ae", "05"),
+ Row(5, 2, 4, 4, "af", "06"),
+ Row(6, 1, 5, 5, "ag", "07")
+ ))
+ }
+ }
+ }
+
+ test("SPARK-47094: Compatible buckets does not support SPJ with " +
+ "push-down values or partially-clustered") {
+ val table1 = "tab1e1"
+ val table2 = "table2"
+
+ val partition1 = Array(bucket(4, "store_id"),
+ bucket(2, "dept_id"))
+ val partition2 = Array(bucket(2, "store_id"),
+ bucket(2, "dept_id"))
+
+ createTable(table1, columns2, partition1)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(0, 0, 'aa'), " +
+ "(1, 1, 'bb'), " +
+ "(2, 2, 'cc')"
+ )
+
+ createTable(table2, columns2, partition2)
+ sql(s"INSERT INTO testcat.ns.$table2 VALUES " +
+ "(0, 0, 'aa'), " +
+ "(1, 1, 'bb'), " +
+ "(2, 2, 'cc')"
+ )
+
+ Seq(true, false).foreach{ allowPushDown =>
+ Seq(true, false).foreach{ partiallyClustered =>
+ withSQLConf(
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> allowPushDown.toString,
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
+ partiallyClustered.toString,
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("t1", "t2")}
+ |t1.store_id, t1.store_id, t1.dept_id, t2.dept_id, t1.data, t2.data
+ |FROM testcat.ns.$table1 t1 JOIN testcat.ns.$table2 t2
+ |ON t1.store_id = t2.store_id AND t1.dept_id = t2.dept_id
+ |ORDER BY t1.store_id, t1.dept_id, t1.data, t2.data
+ |""".stripMargin)
+
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ val scans = collectScans(df.queryExecution.executedPlan).map(_.inputRDD.
+ partitions.length)
+
+ (allowPushDown, partiallyClustered) match {
+ case (true, false) =>
+ assert(shuffles.isEmpty, "SPJ should be triggered")
+ assert(scans == Seq(2, 2))
+ case (_, _) =>
+ assert(shuffles.nonEmpty, "SPJ should not be triggered")
+ assert(scans == Seq(3, 2))
+ }
+
+ checkAnswer(df, Seq(
+ Row(0, 0, 0, 0, "aa", "aa"),
+ Row(1, 1, 1, 1, "bb", "bb"),
+ Row(2, 2, 2, 2, "cc", "cc")
+ ))
+ }
+ }
+ }
+ }
+
test("SPARK-44647: test join key is the second cluster key") {
val table1 = "tab1e1"
val table2 = "table2"
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala
index 61895d49c4a2a..5cdb900901056 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala
@@ -76,7 +76,7 @@ object UnboundBucketFunction extends UnboundFunction {
override def name(): String = "bucket"
}
-object BucketFunction extends ScalarFunction[Int] {
+object BucketFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] {
override def inputTypes(): Array[DataType] = Array(IntegerType, LongType)
override def resultType(): DataType = IntegerType
override def name(): String = "bucket"
@@ -85,6 +85,26 @@ object BucketFunction extends ScalarFunction[Int] {
override def produceResult(input: InternalRow): Int = {
(input.getLong(1) % input.getInt(0)).toInt
}
+
+ override def reducer(
+ thisNumBuckets: Int,
+ otherFunc: ReducibleFunction[_, _],
+ otherNumBuckets: Int): Reducer[Int, Int] = {
+
+ if (otherFunc == BucketFunction) {
+ val gcd = this.gcd(thisNumBuckets, otherNumBuckets)
+ if (gcd != thisNumBuckets) {
+ return BucketReducer(thisNumBuckets, gcd)
+ }
+ }
+ null
+ }
+
+ private def gcd(a: Int, b: Int): Int = BigInt(a).gcd(BigInt(b)).toInt
+}
+
+case class BucketReducer(thisNumBuckets: Int, divisor: Int) extends Reducer[Int, Int] {
+ override def reduce(bucket: Int): Int = bucket % divisor
}
object UnboundStringSelfFunction extends UnboundFunction {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/StateStoreBasicOperationsBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/StateStoreBasicOperationsBenchmark.scala
index a5c393ac0567f..36035e35ee258 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/StateStoreBasicOperationsBenchmark.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/StateStoreBasicOperationsBenchmark.scala
@@ -38,7 +38,8 @@ import org.apache.spark.util.Utils
* 2. build/sbt "sql/Test/runMain