sunchao commented on a change in pull request #35657:
URL: https://github.com/apache/spark/pull/35657#discussion_r829317532
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala
##########
@@ -44,20 +52,108 @@ object V2ExpressionUtils extends SQLConfHelper {
refs.map(ref => resolveRef[T](ref, plan))
}
- def toCatalyst(expr: V2Expression, query: LogicalPlan): Expression = {
+ /**
+ * Converts the array of input V2 [[V2SortOrder]] into their counterparts in
catalyst.
+ */
+ def toCatalystOrdering(ordering: Array[V2SortOrder], query: LogicalPlan):
Seq[SortOrder] = {
+ sequenceToOption(ordering.map(toCatalyst(_,
query))).asInstanceOf[Option[Seq[SortOrder]]]
+ .getOrElse(Seq.empty)
+ }
+
+ /**
+ * Converts the V2 [[V2Distribution]] into its counterpart in catalyst
[[Distribution]].
+ *
+ * If `funCatalogOpt` is set, it will be used to lookup any V2 function
referenced
+ * by the input distribution. For instance, a bucket transform in the
distribution requires a
+ * corresponding function to exist in the catalog, in order for Spark to
leverage bucket join
+ * for the V2 data sources.
+ *
+ * This returns [[UnspecifiedDistribution]] if any non-identity transform is
used in the
+ * distribution, AND the `funCatalogOpt` is not set OR the corresponding
function is not
+ * defined in the catalog.
+ */
+ def toCatalystDistribution(
+ distribution: V2Distribution,
+ query: LogicalPlan,
+ funCatalogOpt: Option[FunctionCatalog] = None): Distribution =
distribution match {
+ case d: V2OrderedDistribution =>
+ val resolvedOrdering = toCatalystOrdering(d.ordering(), query)
+ OrderedDistribution(resolvedOrdering)
+ case d: V2ClusteredDistribution =>
+ sequenceToOption(d.clustering.map(toCatalyst(_, query, funCatalogOpt)))
+ .map(ClusteredDistribution(_))
+ .getOrElse(UnspecifiedDistribution)
+ case _: V2UnspecifiedDistribution =>
+ UnspecifiedDistribution
+ }
+
+ def toCatalyst(
+ expr: V2Expression,
+ query: LogicalPlan,
+ funCatalogOpt: Option[FunctionCatalog] = None): Option[Expression] = {
expr match {
+ case t: Transform =>
+ toCatalystTransform(t, query, funCatalogOpt)
case SortValue(child, direction, nullOrdering) =>
- val catalystChild = toCatalyst(child, query)
- SortOrder(catalystChild, toCatalyst(direction),
toCatalyst(nullOrdering), Seq.empty)
- case IdentityTransform(ref) =>
- resolveRef[NamedExpression](ref, query)
+ toCatalyst(child, query, funCatalogOpt).map { catalystChild =>
+ SortOrder(catalystChild, toCatalyst(direction),
toCatalyst(nullOrdering), Seq.empty)
+ }
case ref: FieldReference =>
- resolveRef[NamedExpression](ref, query)
+ Some(resolveRef[NamedExpression](ref, query))
case _ =>
throw new AnalysisException(s"$expr is not currently supported")
}
}
+ def toCatalystTransform(
+ trans: Transform,
+ query: LogicalPlan,
+ funCatalogOpt: Option[FunctionCatalog] = None): Option[Expression] =
trans match {
+ case IdentityTransform(ref) =>
+ Some(resolveRef[NamedExpression](ref, query))
+ case BucketTransform(numBuckets, refs, sorted) if sorted.isEmpty &&
refs.length == 1 =>
+ val resolvedRefs = refs.map(r => resolveRef[NamedExpression](r, query))
+ funCatalogOpt.flatMap { catalog =>
+ loadV2Function(catalog, "bucket", resolvedRefs).map { bound =>
+ DataSourceBucketTransformExpression(numBuckets, bound, resolvedRefs)
Review comment:
It's a function for the bucket transform: `bucket(num_buckets, c)` (it
could be `bucket(num_buckets, c1, c2, ..)` in future).
The issue here is `canonicalName` for the bucket `BoundFunction`, for
obvious reason, doesn't consider the value of `numBuckets`. However, to check
of two bucket transforms are compatible, we need to take that into account.
That's why we need the extra `DataSourceBucketTransformExpression`
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
##########
@@ -305,6 +306,63 @@ case class HashPartitioning(expressions: Seq[Expression],
numPartitions: Int)
newChildren: IndexedSeq[Expression]): HashPartitioning = copy(expressions
= newChildren)
}
+/**
+ * Represents a partitioning where rows are split across partitions based on
transforms defined
+ * by `expressions`. `partitionValuesOpt`, if defined, should contain value of
partition key(s) in
+ * ascending order, after evaluated by the transforms in `expressions`, for
each input partition.
+ * In addition, its length must be the same as the number of input partitions
(and thus is a 1-1
+ * mapping), and each row in `partitionValuesOpt` must be unique.
+ *
+ * For example, if `expressions` is `[years(ts_col)]`, then a valid value of
`partitionValuesOpt` is
+ * `[0, 1, 2]`, which represents 3 input partitions with distinct partition
values. All rows
+ * in each partition have the same value for column `ts_col` (which is of
timestamp type), after
+ * being applied by the `years` transform.
Review comment:
Hmm I'm thinking whether we can still use `DataSourcePartitioning` for
the V2 interface, following similar naming conventions such as `DataSourceRDD`,
`DataSourceScanExec`, etc. As you mentioned above, `KeyGroupedPartitioning` is
not limited to V2 sources but rather more general and can be applied to other
things like Hive bucketing in future.
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
##########
@@ -305,6 +306,63 @@ case class HashPartitioning(expressions: Seq[Expression],
numPartitions: Int)
newChildren: IndexedSeq[Expression]): HashPartitioning = copy(expressions
= newChildren)
}
+/**
+ * Represents a partitioning where rows are split across partitions based on
transforms defined
+ * by `expressions`. `partitionValuesOpt`, if defined, should contain value of
partition key(s) in
+ * ascending order, after evaluated by the transforms in `expressions`, for
each input partition.
+ * In addition, its length must be the same as the number of input partitions
(and thus is a 1-1
+ * mapping), and each row in `partitionValuesOpt` must be unique.
+ *
+ * For example, if `expressions` is `[years(ts_col)]`, then a valid value of
`partitionValuesOpt` is
+ * `[0, 1, 2]`, which represents 3 input partitions with distinct partition
values. All rows
+ * in each partition have the same value for column `ts_col` (which is of
timestamp type), after
+ * being applied by the `years` transform.
+ *
+ * On the other hand, `[0, 0, 1]` is not a valid value for
`partitionValuesOpt` since `0` is
+ * duplicated twice.
+ *
+ * @param expressions partition expressions for the partitioning.
+ * @param numPartitions the number of partitions
+ * @param partitionValuesOpt if set, the values for the cluster keys of the
distribution, must be
+ * in ascending order.
+ */
+case class KeyGroupedPartitioning(
+ expressions: Seq[Expression],
+ numPartitions: Int,
+ partitionValuesOpt: Option[Seq[InternalRow]] = None) extends Partitioning {
+
+ override def satisfies0(required: Distribution): Boolean = {
+ super.satisfies0(required) || {
+ required match {
+ case c @ ClusteredDistribution(requiredClustering,
requireAllClusterKeys, _) =>
+ if (requireAllClusterKeys) {
+ // Checks whether this partitioning is partitioned on exactly same
clustering keys of
+ // `ClusteredDistribution`.
+ c.areAllClusterKeysMatched(expressions)
+ } else {
+ // We'll need to find leaf attributes from the partition
expressions first.
+ val attributes = expressions.flatMap(_.collectLeaves())
+ attributes.forall(x =>
requiredClustering.exists(_.semanticEquals(x)))
Review comment:
For simplicity this PR only support transforms with a single argument,
so the check here is very similar to `HashPartitioning`. I plan to work on the
support of multi-arguments as a separate PR.
> I think the theory is, as long as the transform function is deterministic,
if two rows have the same value of [a, b, c], the calculated key values [f1(a,
b), f2(b, c)] are also the same.
Yea agree with this. It may be more complicated if you have duplicated keys
in distribution or/and partitioning though. We also need to think how to
implement `EnsureRequirements.reorderJoinPredicates` for this case.
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
##########
@@ -137,8 +138,16 @@ case class EnsureRequirements(
Some(finalCandidateSpecs.values.maxBy(_.numPartitions))
}
+ // Check if 1) all children are of `KeyGroupedPartitioning` and 2) they
are all compatible
+ // with each other. If both are true, skip shuffle.
+ val allCompatible = childrenIndexes.sliding(2).map {
+ case Seq(a, b) =>
+ checkKeyGroupedSpec(specs(a)) && checkKeyGroupedSpec(specs(b)) &&
Review comment:
> ... It looks like we can always eliminate shuffles if all the children
are compatible.
When children are compatible, it still doesn't mean each child's partition
keys **fully match** the distribution keys. Since in
`KeyGroupedPartitioning.satisfies`, we only check if partition keys are subset
of distribution keys.
That's why we need an additional check here to make sure that when
`spark.sql.requireAllClusterKeysForCoPartition` is true, the child partition
keys fully match the distribution keys.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
