Skip to content

Conversation

@ShreyeshArangath
Copy link
Contributor

Which issue does this PR close?

Closes ##1888

Rationale for this change

This is part of the effort to add non-deterministic expressions to Auron.

What changes are included in this PR?

This PR adds support for spark_partition_id using native

Are there any user-facing changes?

N/A

How was this patch tested?

Unit tests

@ShreyeshArangath ShreyeshArangath marked this pull request as ready for review January 18, 2026 20:32
@Tartarus0zm Tartarus0zm requested a review from Copilot January 19, 2026 03:34
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR implements support for the spark_partition_id() function using native execution, enabling Auron to handle this non-deterministic expression natively rather than falling back to Spark's implementation.

Changes:

  • Added Rust implementation of SparkPartitionIdExpr that returns the partition ID as a constant Int32 array
  • Integrated the expression into the planner, protobuf schema, and Scala conversion layers
  • Included comprehensive unit tests for the Rust implementation

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.

Show a summary per file
File Description
native-engine/datafusion-ext-exprs/src/spark_partition_id.rs New Rust implementation of SparkPartitionIdExpr with unit tests
native-engine/datafusion-ext-exprs/src/lib.rs Registers the new spark_partition_id module
native-engine/auron-planner/src/planner.rs Integrates SparkPartitionIdExpr into the planner, passing partition_id from PhysicalPlanner
native-engine/auron-planner/proto/auron.proto Adds SparkPartitionIdExprNode protobuf message definition
spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala Adds conversion for StubExpr("SparkPartitionID") to protobuf
spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala Adds conversion for SparkPartitionID expression to protobuf and imports the catalyst expression

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant