@@ -46,6 +46,7 @@ use datafusion::logical_expr::{
     ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
     Volatility, create_udf,
 };
+use datafusion::physical_optimizer::optimizer::PhysicalOptimizer;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::*;
 use datafusion::{
@@ -114,6 +115,7 @@ impl TestContext {
     /// If `None` is returned (e.g. because some needed feature is not
     /// enabled), the file should be skipped
     pub async fn try_new_for_test_file(relative_path: &Path) -> Option<Self> {
+        let file_name = relative_path.file_name().unwrap().to_str().unwrap();
         let config = SessionConfig::new()
             // hardcode target partitions so plans are deterministic
             .with_target_partitions(4);
@@ -177,11 +179,23 @@ impl TestContext {
             state_builder = state_builder.with_statistics_registry(registry);
         }
 
+        if file_name == "sort_pushdown.slt" {
+            let mut rules = PhysicalOptimizer::new().rules;
+            let ensure_requirements = rules
+                .iter()
+                .position(|rule| rule.name() == "EnsureRequirements")
+                .expect("default optimizer includes EnsureRequirements");
+            rules.insert(
+                ensure_requirements,
+                Arc::new(custom_sort_pushdown::CustomLimitRule),
+            );
+            state_builder = state_builder.with_physical_optimizer_rules(rules);
+        }
+
         let state = state_builder.build();
 
         let mut test_ctx = TestContext::new(SessionContext::new_with_state(state));
 
-        let file_name = relative_path.file_name().unwrap().to_str().unwrap();
         match file_name {
             "sort_pushdown.slt" => {
                 custom_sort_pushdown::register_custom_sort_pushdown(
 
@@ -23,19 +23,22 @@ use arrow::datatypes::{self as arrow_schema, SchemaRef};
 use async_trait::async_trait;
 use datafusion::catalog::{Session, TableProvider};
 use datafusion::common::Result;
-use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::config::ConfigOptions;
+use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::logical_expr::{Expr, Operator, TableType};
 use datafusion::physical_expr::expressions::{binary, col, lit, negative};
 use datafusion::physical_expr::projection::ProjectionExprs;
 use datafusion::physical_expr::{PhysicalExpr, PhysicalSortExpr};
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
 use datafusion::physical_plan::filter::FilterExec;
 use datafusion::physical_plan::limit::LocalLimitExec;
 use datafusion::physical_plan::projection::ProjectionExec;
+use datafusion::physical_plan::sorts::sort::SortExec;
 use datafusion::physical_plan::test::TestMemoryExec;
 use datafusion::physical_plan::{
-    ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
-    ReplaceChildrenOptions,
+    ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan,
+    ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions,
 };
 use datafusion::prelude::SessionContext;
 
@@ -79,18 +82,65 @@ pub(super) fn register_custom_sort_pushdown(ctx: &SessionContext) -> Result<()>
 
     let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Int32, [30, 20, 10]))?;
     let schema = batch.schema();
-    let ordering = [PhysicalSortExpr::new_default(col("a", &schema)?)].into();
+    let ordering =
+        [PhysicalSortExpr::new_default(col("a", &schema)?).nulls_last()].into();
     let input = Arc::new(
         TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)?
             .try_with_sort_information(vec![ordering])?,
     );
-    ctx.register_table(
-        "custom_limited_input",
-        Arc::new(CustomExec(Arc::new(LocalLimitExec::new(input, 2)))),
-    )?;
+    ctx.register_table("custom_ordered_input", Arc::new(CustomExec(input)))?;
     Ok(())
 }
 
+/// Exposes the SQL fixture's ordered selection as a custom limit before sort pushdown.
+#[derive(Debug)]
+pub(super) struct CustomLimitRule;
+
+impl PhysicalOptimizerRule for CustomLimitRule {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        _: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        plan.transform_up(|plan| {
+            let Some(sort) = plan.downcast_ref::<SortExec>() else {
+                return Ok(Transformed::no(plan));
+            };
+            let Some(fetch) = sort.fetch() else {
+                return Ok(Transformed::no(plan));
+            };
+            let input = sort.input();
+            let source = input
+                .downcast_ref::<ProjectionExec>()
+                .map_or(input, |projection| projection.input());
+            let Some(source) = source.downcast_ref::<CustomExec>() else {
+                return Ok(Transformed::no(plan));
+            };
+            // Only the dedicated ordered source is a wrapped memory scan.
+            if !source.0.is::<TestMemoryExec>()
+                || input.output_partitioning().partition_count() != 1
+                || !input
+                    .equivalence_properties()
+                    .ordering_satisfy(sort.expr().clone())?
+            {
+                return Ok(Transformed::no(plan));
+            }
+            let limit = LocalLimitExec::new(Arc::clone(input), fetch);
+            let plan: Arc<dyn ExecutionPlan> = Arc::new(CustomExec(Arc::new(limit)));
+            Ok(Transformed::yes(plan))
+        })
+        .map(|transformed| transformed.data)
+    }
+
+    fn name(&self) -> &str {
+        "CustomLimitRule"
+    }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
+}
+
 /// Exposes a native plan's children and properties under a distinct plan type,
 /// exercising the generic custom-plan sort pushdown path.
 #[derive(Debug, Clone)]
 
@@ -73,21 +73,26 @@ physical_plan
 03)----CooperativeExec
 04)------DataSourceExec: partitions=1, partition_sizes=[1]
 
-# The input is ordered by a. Sorting by b must keep the first two selected rows.
+# The inner ORDER BY determines which rows LIMIT selects; the outer sort only
+# changes their display order. The fixture plans that selection as a custom limit.
 query II
-SELECT a, b FROM custom_limited_input ORDER BY b;
+SELECT a, b
+FROM (SELECT a, b FROM custom_ordered_input ORDER BY a LIMIT 2) s
+ORDER BY b;
 ----
 2 20
 1 30
 
 query TT
-EXPLAIN SELECT a, b FROM custom_limited_input ORDER BY b;
+EXPLAIN SELECT a, b
+FROM (SELECT a, b FROM custom_ordered_input ORDER BY a LIMIT 2) s
+ORDER BY b;
 ----
 physical_plan
 01)SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
 02)--CustomLocalLimitExec: fetch=2
 03)----CooperativeExec
-04)------DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 ASC
+04)------CustomDataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 ASC NULLS LAST
 
 statement ok
 SET datafusion.explain.physical_plan_only = false;