andygrove commented on code in PR #4507:
URL: https://github.com/apache/datafusion-comet/pull/4507#discussion_r3343600116
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -624,21 +668,102 @@ object CometShuffleExchangeExec
}
}
+ /**
+ * Build a Comet native shuffle dependency around an existing
`RDD[ColumnarBatch]` of real
+ * batches. Used by [[org.apache.spark.sql.comet.CometCollectLimitExec]] and
+ * [[org.apache.spark.sql.comet.CometTakeOrderedAndProjectExec]] where the
input is the result
+ * of a local-limit / topK transform and there is no separate child native
subtree to inline.
+ *
+ * Implemented as a thin wrapper around [[prepareNativeShuffleDependency]]:
synthesizes a
+ * `Scan("ShuffleWriterInput")` as the child native op (so the writer's plan
is still
+ * `ShuffleWriter -> Scan`, consuming JVM batches via Arrow C Stream), wraps
`rdd` as the single
+ * leaf input of a thin scheduling RDD, and supplies a minimal
[[NativeExecContext]]. Lets the
+ * writer use one code path for both this case and the
[[CometShuffleExchangeExec]] case.
+ */
def prepareShuffleDependency(
rdd: RDD[ColumnarBatch],
outputAttributes: Seq[Attribute],
outputPartitioning: Partitioning,
serializer: Serializer,
metrics: Map[String, SQLMetric]): ShuffleDependency[Int, ColumnarBatch,
ColumnarBatch] = {
- val numParts = rdd.getNumPartitions
+
+ val scanBuilder =
OperatorOuterClass.Scan.newBuilder().setSource("ShuffleWriterInput")
+ val scanTypes = outputAttributes.flatten { attr =>
+ QueryPlanSerde.serializeDataType(attr.dataType)
Review Comment:
nit: This might read more clearly as a `flatMap`
--
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]
