[SPARK-53401][SQL] Enable Direct Passthrough Partitioning in the DataFrame API

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

@@ -2045,6 +2045,19 @@ object functions {
    */
   def spark_partition_id(): Column = Column.fn("spark_partition_id")
 
+  /**
+   * Returns the partition ID specified by the given column expression for direct shuffle
+   * partitioning. The input expression must evaluate to an integral type and must not be null.
+   *
+   * This function is used with DataFrame.repartitionByExpr to allow users to directly specify
+   * target partition IDs instead of using hash-based partitioning.
+   *
+   * @group misc_funcs
+   * @since 4.1.0
+   */
+  def direct_shuffle_partition_id(e: Column): Column =
+    Column.fn("direct_shuffle_partition_id", e)
+
   /**
    * Computes the square root of the specified float value.
    *

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -773,6 +773,7 @@ object FunctionRegistry {
     expression[AesEncrypt]("aes_encrypt"),
     expression[AesDecrypt]("aes_decrypt"),
     expression[SparkPartitionID]("spark_partition_id"),
+    expression[DirectShufflePartitionID]("direct_shuffle_partition_id"),
     expression[InputFileName]("input_file_name"),
     expression[InputFileBlockStart]("input_file_block_start"),
     expression[InputFileBlockLength]("input_file_block_length"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/DirectShufflePartitionID.scala

@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode, FalseLiteral}
+import org.apache.spark.sql.catalyst.expressions.codegen.Block._
+import org.apache.spark.sql.types.{AbstractDataType, DataType, LongType}
+
+/**
+ * Expression that takes a partition ID value and passes it through directly for use in
+ * shuffle partitioning. This is used with RepartitionByExpression to allow users to
+ * directly specify target partition IDs.
+ *
+ * The child expression must evaluate to an integral type and must not be null.
+ * The resulting partition ID must be in the range [0, numPartitions).
+ */
+@ExpressionDescription(
+  usage = "_FUNC_(expr) - Returns the partition ID specified by expr for direct shuffle " +
+    "partitioning.",
+  arguments = """
+    Arguments:
+      * expr - an integral expression that specifies the target partition ID
+  """,
+  examples = """
+    Examples:
+      > df.repartition(10, direct_shuffle_partition_id($"partition_id"))
+      > df.repartition(10, expr("direct_shuffle_partition_id(id % 5)"))
+  """,
+  since = "4.1.0",
+  group = "misc_funcs"
+)
+case class DirectShufflePartitionID(child: Expression)
+    extends UnaryExpression
+    with ExpectsInputTypes {
+
+  override def dataType: DataType = child.dataType
+
+  override def inputTypes: Seq[AbstractDataType] = LongType :: Nil
+
+  override def nullable: Boolean = false
+
+  override val prettyName: String = "direct_shuffle_partition_id"
+
+  override def eval(input: InternalRow): Any = {
+    val result = child.eval(input)
+    if (result == null) {
+      throw new IllegalArgumentException(
+        "The partition ID expression must not be null.")
+    }
+    nullSafeEval(result)
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val childGen = child.genCode(ctx)
+    val resultCode =
+      s"""
+         |${childGen.code}
+         |if (${childGen.isNull}) {
+         |  throw new IllegalArgumentException(
+         |    "The partition ID expression must not be null.");
+         |}
+         |""".stripMargin
+
+    ev.copy(code = code"$resultCode", isNull = FalseLiteral, value = childGen.value)
+  }
+
+  override protected def withNewChildInternal(newChild: Expression): DirectShufflePartitionID =
+    copy(child = newChild)
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.catalog.CatalogTable.VIEW_STORING_ANALYZED_
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, TypedImperativeAggregate}
 import org.apache.spark.sql.catalyst.plans._
-import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition}
+import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning, ShufflePartitionIdPassThrough, SinglePartition}
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 import org.apache.spark.sql.catalyst.trees.TreePattern._
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
@@ -1871,19 +1871,30 @@ trait HasPartitionExpressions extends SQLConfHelper {
   protected def partitioning: Partitioning = if (partitionExpressions.isEmpty) {
     RoundRobinPartitioning(numPartitions)
   } else {
-    val (sortOrder, nonSortOrder) = partitionExpressions.partition(_.isInstanceOf[SortOrder])
-    require(sortOrder.isEmpty || nonSortOrder.isEmpty,
-      s"${getClass.getSimpleName} expects that either all its `partitionExpressions` are of type " +
-        "`SortOrder`, which means `RangePartitioning`, or none of them are `SortOrder`, which " +
-        "means `HashPartitioning`. In this case we have:" +
-        s"""
-           |SortOrder: $sortOrder
-           |NonSortOrder: $nonSortOrder
-       """.stripMargin)
-    if (sortOrder.nonEmpty) {
-      RangePartitioning(sortOrder.map(_.asInstanceOf[SortOrder]), numPartitions)
+    // Check if we have DirectShufflePartitionID expressions
+    val directShuffleExprs = partitionExpressions.filter(_.isInstanceOf[DirectShufflePartitionID])
+    if (directShuffleExprs.nonEmpty) {
+      require(directShuffleExprs.length == 1 && partitionExpressions.length == 1,
+        s"DirectShufflePartitionID can only be used as a single partition expression, " +
+          s"but found ${directShuffleExprs.length} DirectShufflePartitionID expressions " +
+          s"out of ${partitionExpressions.length} total expressions")
+      ShufflePartitionIdPassThrough(
+        partitionExpressions.head.asInstanceOf[DirectShufflePartitionID], numPartitions)
     } else {
-      HashPartitioning(partitionExpressions, numPartitions)
+      val (sortOrder, nonSortOrder) = partitionExpressions.partition(_.isInstanceOf[SortOrder])
+      require(sortOrder.isEmpty || nonSortOrder.isEmpty,
+        s"${getClass.getSimpleName} expects that either all its `partitionExpressions` are of " +
+          "`SortOrder`, which means `RangePartitioning`, or none of them are `SortOrder`, which " +
+          "means `HashPartitioning`. In this case we have:" +
+          s"""
+             |SortOrder: $sortOrder
+             |NonSortOrder: $nonSortOrder
+         """.stripMargin)
+      if (sortOrder.nonEmpty) {
+        RangePartitioning(sortOrder.map(_.asInstanceOf[SortOrder]), numPartitions)
+      } else {
+        HashPartitioning(partitionExpressions, numPartitions)
+      }
     }
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

@@ -946,3 +946,24 @@ case class ShuffleSpecCollection(specs: Seq[ShuffleSpec]) extends ShuffleSpec {
     specs.head.numPartitions
   }
 }
+
+/**
+ * Represents a partitioning where partition IDs are passed through directly from the
+ * DirectShufflePartitionID expression. This partitioning scheme is used when users
+ * want to directly control partition placement rather than using hash-based partitioning.
+ *
+ * This partitioning maps directly to the PartitionIdPassthrough RDD partitioner.
+ */
+case class ShufflePartitionIdPassThrough(
+    expr: DirectShufflePartitionID,
+    numPartitions: Int) extends HashPartitioningLike {
+
+  // TODO(SPARK-53401): Support Shuffle Spec in Direct Partition ID Pass Through
+  def partitionIdExpression: Expression = Pmod(expr, Literal(numPartitions))
+
+  override def expressions: Seq[Expression] = expr :: Nil
+
+  override protected def withNewChildrenInternal(
+      newChildren: IndexedSeq[Expression]): ShufflePartitionIdPassThrough =
+    copy(expr = newChildren.head.asInstanceOf[DirectShufflePartitionID])
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala

@@ -344,6 +344,10 @@ object ShuffleExchangeExec {
         // For HashPartitioning, the partitioning key is already a valid partition ID, as we use
         // `HashPartitioning.partitionIdExpression` to produce partitioning key.
         new PartitionIdPassthrough(n)
+      case ShufflePartitionIdPassThrough(_, n) =>
+        // For ShufflePartitionIdPassThrough, the DirectShufflePartitionID expression directly
+        // produces partition IDs, so we use PartitionIdPassthrough to pass them through directly.
+        new PartitionIdPassthrough(n)
       case RangePartitioning(sortingExpressions, numPartitions) =>
         // Extract only fields used for sorting to avoid collecting large fields that does not
         // affect sorting result when deciding partition bounds in RangePartitioner
@@ -399,6 +403,10 @@ object ShuffleExchangeExec {
       case SinglePartition => identity
       case KeyGroupedPartitioning(expressions, _, _, _) =>
         row => bindReferences(expressions, outputAttributes).map(_.eval(row))
+      case s: ShufflePartitionIdPassThrough =>
+        // For ShufflePartitionIdPassThrough, the expression directly evaluates to the partition ID
+        val projection = UnsafeProjection.create(s.expressions, outputAttributes)
+        row => projection(row).getInt(0)
       case _ => throw SparkException.internalError(s"Exchange not implemented for $newPartitioning")
     }
 

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala

@@ -2785,6 +2785,51 @@ class DataFrameSuite extends QueryTest
     val df1 = df.select("a").orderBy("b").orderBy("all")
     checkAnswer(df1, Seq(Row(1), Row(4)))
   }
+
+  test("SPARK-53401: direct_shuffle_partition_id - should partition rows to the specified " +
+    "partition ID") {
+    val numPartitions = 10
+    val df = spark.range(100).withColumn("p_id", col("id") % numPartitions)
+
+    val repartitioned = df.repartition(numPartitions, direct_shuffle_partition_id($"p_id"))
+    val result = repartitioned.withColumn("actual_p_id", spark_partition_id())
+
+    assert(result.filter(col("p_id") =!= col("actual_p_id")).count() == 0)
+
+    assert(result.rdd.getNumPartitions == numPartitions)
+  }
+
+  test("SPARK-53401: direct_shuffle_partition_id - should work with expr()") {
+    val numPartitions = 5
+    val df = spark.range(50).withColumn("p_id", col("id") % numPartitions)
+
+    val repartitioned = df.repartition(numPartitions, expr("direct_shuffle_partition_id(p_id)"))
+    val result = repartitioned.withColumn("actual_p_id", spark_partition_id())
+
+    assert(result.filter(col("p_id") =!= col("actual_p_id")).count() == 0)
+  }
+
+  test("SPARK-53401: direct_shuffle_partition_id - should fail when partition ID is null") {
+    val df = spark.range(10).withColumn("p_id",
+      when(col("id") < 5, col("id")).otherwise(lit(null).cast("long"))
+    )
+    val repartitioned = df.repartition(5, direct_shuffle_partition_id($"p_id"))
+
+    val e = intercept[SparkException] {
+      repartitioned.collect()
+    }
+    assert(e.getCause.isInstanceOf[IllegalArgumentException])
+    assert(e.getCause.getMessage.contains("The partition ID expression must not be null."))
+  }
+
+  test("SPARK-53401: direct_shuffle_partition_id - should fail analysis for non-integral types") {
+    val df = spark.range(5).withColumn("s", lit("a"))
+    val e = intercept[AnalysisException] {
+      df.repartition(5, direct_shuffle_partition_id($"s")).collect()
+    }
+    // Should fail with type error from DirectShufflePartitionID expression
+    assert(e.getMessage.contains("requires an integral type"))
+  }
 }
 
 case class GroupByKey(a: Int, b: Int)