diff --git a/.gitignore b/.gitignore index 2bcc0950d01b3..0c71255e6ac34 100644 --- a/.gitignore +++ b/.gitignore @@ -78,3 +78,5 @@ profile.json.gz # Claude Code personal settings .claude/settings.local.json + +.codex/config.toml \ No newline at end of file diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index 96b911e8db130..afdaaa9dd7ebb 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -18,6 +18,7 @@ use std::collections::HashMap; use super::*; +use datafusion::assert_batches_eq; use datafusion_common::{ParamValues, ScalarValue, metadata::ScalarAndMetadata}; use insta::assert_snapshot; @@ -493,3 +494,39 @@ async fn test_recursive_cte_batch_schema_stable_with_order_by_limit() -> Result< } Ok(()) } + +#[tokio::test] +async fn ordered_offset_subquery_preserves_selected_row() -> Result<()> { + let ctx = SessionContext::new_with_config( + SessionConfig::new().set_usize("datafusion.optimizer.max_passes", 1), + ); + let results = ctx + .sql( + "SELECT grp, COUNT(*) AS n + FROM ( + SELECT grp, sort_key + FROM ( + VALUES ('physical_second', 2), ('sorted_first', 1) + ) AS t(grp, sort_key) + ORDER BY sort_key ASC NULLS LAST + OFFSET 1 + ) q + GROUP BY grp + ORDER BY grp", + ) + .await? + .collect() + .await?; + + assert_batches_eq!( + [ + "+-----------------+---+", + "| grp | n |", + "+-----------------+---+", + "| physical_second | 1 |", + "+-----------------+---+", + ], + &results + ); + Ok(()) +} diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs index 6efaf76457919..e5e52a9573722 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -444,7 +444,11 @@ pub fn ensure_sorting( ); child = update_sort_ctx_children_data(child, true)?; } - } else if physical_ordering.is_none() || !plan.maintains_input_order()[idx] { + } else if !is_limit(plan) + && (physical_ordering.is_none() || !plan.maintains_input_order()[idx]) + { + // Limit is excluded because it consumes the input sequence, so + // removing a linked sort can change which rows it skips or returns. // We have a `SortExec` whose effect may be neutralized by another // order-imposing operator, remove this sort: child = update_child_to_remove_unnecessary_sort(idx, child, plan)?;