LinSimon-901101 opened a new issue, #6570: URL: https://github.com/apache/datafusion-comet/issues/6570
### Describe the bug ### Describe the bug As identified in the review of #5875, the JVM codegen dispatcher initializes each kernel using `TaskContext.partitionId()`, while Spark and Comet's native execution path use the index of the partition being computed. These indexes can differ under `CometUnionExec` and `CometCoalesceExec`. Consequently, partition-dependent expressions inside a dispatched subtree, such as `rand`, `uuid`, `monotonically_increasing_id`, and `spark_partition_id`, can behave differently from Spark. The reviewer noted that this already affects `branch-1.1`, for example when evaluating `map(1, spark_partition_id())` in a later `UNION ALL` branch. #5875 extends dispatch to additional map lookups and exposes them to the same underlying issue. ### Steps to reproduce The following scenario was reported in the linked review. A standalone executable reproducer and execution logs are not included yet. 1. Enable Comet and the JVM codegen dispatcher. 2. Evaluate `map(1, spark_partition_id())` within a later branch of a `UNION ALL` query. 3. Confirm that the expression is dispatched within that branch and that the physical plan uses `CometUnionExec`. 4. Compare the returned partition IDs with the equivalent query using Spark execution. The review also identifies expressions evaluated below `CometCoalesceExec` as affected when the task partition ID differs from the partition being computed. ### Expected behavior Dispatched kernels should receive the index of the partition being computed, consistent with Spark and Comet's native execution path. Regression coverage should exercise both union and coalesce, verify the dispatched execution route, and check partition-dependent behavior against Spark. ### Additional context - Follow-up requested in [[Andy's review of #5875](https://github.com/apache/datafusion-comet/pull/5875#pullrequestreview-5401526275)](https://github.com/apache/datafusion-comet/pull/5875#pullrequestreview-5401526275), section “Partition index in dispatched lookups”. - The fix belongs in the shared dispatcher and is tracked separately from the map-lookup changes in #5875. - Relevant code: `CometScalaUDFCodegen.currentPartitionIndex()` and the call to `kernel.init(...)` in `spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala`. ### Steps to reproduce _No response_ ### Expected behavior _No response_ ### Additional context _No response_ -- 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]
