fix: preserve ordering through retained double negation

Treat an exact deterministic -(-x) key as determining x within its tie group, and make permutation analysis use the shared satisfied-key path.

Carry the same relationship through projections that evaluate the retained expression, so grouped aggregates, windows, and aliases preserve the full suffix ordering without cancelling checked temporal negation.

Signed-off-by: Amogh Ramesh <ramogh2404@gmail.com>

Author: Amogh-2404

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use std::sync::Arc;
+use std::time::Duration;
 
 use crate::memory_limit::DummyStreamPartition;
 use crate::physical_optimizer::test_utils::{
@@ -31,9 +32,12 @@ use crate::physical_optimizer::test_utils::{
 
 use arrow::compute::{SortOptions};
 use arrow::datatypes::{DataType, SchemaRef};
+use arrow::error::ArrowError;
 use datafusion_common::config::{ConfigOptions, CsvOptions};
 use datafusion_common::tree_node::{TreeNode, TransformedResult};
-use datafusion_common::{create_array, DataFusionError, NullEquality, Result, TableReference};
+use datafusion_common::{
+    DataFusionError, NullEquality, Result, ScalarValue, TableReference, create_array,
+};
 use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
 use datafusion_datasource::source::DataSourceExec;
 use datafusion_expr_common::operator::Operator;
@@ -49,7 +53,12 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
-use datafusion_physical_plan::{displayable, get_plan_string, ExecutionPlan, ExecutionPlanProperties};
+use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_physical_plan::streaming::PartitionStream;
+use datafusion_physical_plan::{
+    ExecutionPlan, ExecutionPlanProperties, SendableRecordBatchStream, displayable,
+    get_plan_string,
+};
 use datafusion::datasource::physical_plan::CsvSource;
 use datafusion::datasource::listing::PartitionedFile;
 use datafusion_physical_optimizer::enforce_sorting::{PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, parallelize_sorts, ensure_sorting};
@@ -2944,6 +2953,334 @@ async fn test_sort_with_streaming_table() -> Result<()> {
     Ok(())
 }
 
+#[derive(Debug)]
+struct PendingOrderedPartition {
+    schema: SchemaRef,
+    batches: Vec<RecordBatch>,
+}
+
+impl PartitionStream for PendingOrderedPartition {
+    fn schema(&self) -> &SchemaRef {
+        &self.schema
+    }
+
+    fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
+        let stream = futures::stream::iter(
+            self.batches
+                .clone()
+                .into_iter()
+                .map(Ok::<_, DataFusionError>),
+        )
+        .chain(futures::stream::pending());
+        Box::pin(RecordBatchStreamAdapter::new(
+            Arc::clone(&self.schema),
+            stream,
+        ))
+    }
+}
+
+fn pending_ordered_context(
+    value: &ScalarValue,
+    nullable: bool,
+) -> Result<SessionContext> {
+    let schema: SchemaRef = Arc::new(Schema::new(vec![
+        Field::new("i", value.data_type(), nullable),
+        Field::new("j", DataType::Int64, false),
+    ]));
+    let batches = [0_i64, 1_i64]
+        .into_iter()
+        .map(|j| -> Result<RecordBatch> {
+            Ok(RecordBatch::try_new(
+                Arc::clone(&schema),
+                vec![
+                    value.to_array_of_size(1)?,
+                    ScalarValue::Int64(Some(j)).to_array_of_size(1)?,
+                ],
+            )?)
+        })
+        .collect::<Result<Vec<_>>>()?;
+    let partition: Arc<dyn PartitionStream> = Arc::new(PendingOrderedPartition {
+        schema: Arc::clone(&schema),
+        batches,
+    });
+    let provider = StreamingTable::try_new(Arc::clone(&schema), vec![partition])?
+        .with_sort_order(vec![
+            datafusion_expr::col("i").sort(true, false),
+            datafusion_expr::col("j").sort(true, false),
+        ])
+        .with_infinite_table(true);
+    let ctx = SessionContext::new_with_config(
+        SessionConfig::new()
+            .with_target_partitions(1)
+            .with_batch_size(1),
+    );
+    ctx.register_table("s", Arc::new(provider))?;
+    Ok(ctx)
+}
+
+async fn assert_double_negation_group_streams(
+    case: &str,
+    value: ScalarValue,
+    nullable: bool,
+) -> Result<()> {
+    let ctx = pending_ordered_context(&value, nullable)?;
+    let plan = ctx
+        .sql(
+            "SELECT -(-i) AS k, j, COUNT(*) AS n \
+             FROM s GROUP BY -(-i), j LIMIT 1",
+        )
+        .await?
+        .create_physical_plan()
+        .await?;
+    let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+    let aggregate = plan_text
+        .lines()
+        .find(|line| line.contains("AggregateExec:"))
+        .expect("aggregate plan");
+    assert!(
+        aggregate.contains("ordering_mode=Sorted"),
+        "{case}: expected fully sorted aggregation:\n{plan_text}"
+    );
+    assert!(
+        !plan_text.contains("SortExec:"),
+        "{case}: existing input order should avoid a sort:\n{plan_text}"
+    );
+
+    let mut stream = plan.execute(0, ctx.task_ctx())?;
+    let output = tokio::time::timeout(Duration::from_secs(5), stream.next())
+        .await
+        .expect("sorted aggregation withheld a completed group")
+        .expect("stream ended before producing a group")?;
+    assert_eq!(output.num_rows(), 1, "{case}");
+    assert_eq!(
+        ScalarValue::try_from_array(output.column(0), 0)?,
+        value,
+        "{case}"
+    );
+    assert_eq!(
+        ScalarValue::try_from_array(output.column(1), 0)?,
+        ScalarValue::Int64(Some(0)),
+        "{case}"
+    );
+    assert_eq!(
+        ScalarValue::try_from_array(output.column(2), 0)?,
+        ScalarValue::Int64(Some(1)),
+        "{case}"
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn double_negation_group_by_streams_on_full_input_order() -> Result<()> {
+    for (case, value, nullable) in [
+        (
+            "timestamp",
+            ScalarValue::TimestampNanosecond(Some(1), None),
+            false,
+        ),
+        (
+            "nullable timestamp",
+            ScalarValue::TimestampNanosecond(None, None),
+            true,
+        ),
+        (
+            "year-month interval",
+            ScalarValue::IntervalYearMonth(Some(1)),
+            false,
+        ),
+    ] {
+        assert_double_negation_group_streams(case, value, nullable).await?;
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn double_negation_projection_preserves_suffix_ordering() -> Result<()> {
+    for (case, sql) in [
+        (
+            "projection",
+            "SELECT k, j FROM (SELECT -(-i) AS k, j FROM s) \
+             ORDER BY k, j LIMIT 1",
+        ),
+        (
+            "aggregate output",
+            "SELECT k, j FROM ( \
+               SELECT -(-i) AS k, j, COUNT(*) AS n \
+               FROM s GROUP BY -(-i), j \
+             ) ORDER BY k, j LIMIT 1",
+        ),
+    ] {
+        let ctx = pending_ordered_context(
+            &ScalarValue::TimestampNanosecond(Some(1), None),
+            false,
+        )?;
+        let plan = ctx.sql(sql).await?.create_physical_plan().await?;
+        let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+        assert!(
+            !plan_text.contains("SortExec:"),
+            "{case}: projected ordering should include the suffix:\n{plan_text}"
+        );
+
+        let mut stream = plan.execute(0, ctx.task_ctx())?;
+        let output = tokio::time::timeout(Duration::from_secs(5), stream.next())
+            .await
+            .unwrap_or_else(|_| panic!("{case}: ordered projection did not stream"))
+            .expect("stream ended before producing a row")?;
+        assert_eq!(output.num_rows(), 1, "{case}");
+        assert_eq!(
+            ScalarValue::try_from_array(output.column(0), 0)?,
+            ScalarValue::TimestampNanosecond(Some(1), None),
+            "{case}"
+        );
+        assert_eq!(
+            ScalarValue::try_from_array(output.column(1), 0)?,
+            ScalarValue::Int64(Some(0)),
+            "{case}"
+        );
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn double_negation_group_by_preserves_timestamp_overflow() -> Result<()> {
+    let ctx = pending_ordered_context(
+        &ScalarValue::TimestampNanosecond(Some(i64::MIN), None),
+        false,
+    )?;
+    let plan = ctx
+        .sql(
+            "SELECT -(-i) AS k, j, COUNT(*) AS n \
+             FROM s GROUP BY -(-i), j LIMIT 1",
+        )
+        .await?
+        .create_physical_plan()
+        .await?;
+
+    let mut stream = plan.execute(0, ctx.task_ctx())?;
+    let err = tokio::time::timeout(Duration::from_secs(5), stream.next())
+        .await
+        .expect("timestamp overflow was not reported")
+        .expect("stream ended before reporting timestamp overflow")
+        .expect_err("double negation must still evaluate its checked inner negation");
+    assert!(
+        matches!(
+            err.find_root(),
+            DataFusionError::ArrowError(error, _)
+                if matches!(error.as_ref(), ArrowError::ArithmeticOverflow(_))
+        ),
+        "unexpected error: {err}"
+    );
+    Ok(())
+}
+
+async fn assert_double_negation_lead_streams(
+    case: &str,
+    value: ScalarValue,
+    nullable: bool,
+) -> Result<()> {
+    let ctx = pending_ordered_context(&value, nullable)?;
+    let plan = ctx
+        .sql(
+            "SELECT i, j, LEAD(j, 1) OVER (PARTITION BY -(-i), j) AS n \
+             FROM s LIMIT 1",
+        )
+        .await?
+        .create_physical_plan()
+        .await?;
+    let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+    let window = plan_text
+        .lines()
+        .find(|line| line.contains("BoundedWindowAggExec:"))
+        .expect("bounded window plan");
+    assert!(
+        window.contains("mode=[Sorted]"),
+        "{case}: expected fully sorted window partitions:\n{plan_text}"
+    );
+    assert!(
+        !plan_text.contains("SortExec:"),
+        "{case}: existing input order should avoid a sort:\n{plan_text}"
+    );
+
+    let mut stream = plan.execute(0, ctx.task_ctx())?;
+    let output = tokio::time::timeout(Duration::from_secs(5), stream.next())
+        .await
+        .unwrap_or_else(|_| panic!("{case}: LEAD withheld a completed partition"))
+        .expect("stream ended before producing a row")?;
+    assert_eq!(output.num_rows(), 1, "{case}");
+    assert_eq!(
+        ScalarValue::try_from_array(output.column(0), 0)?,
+        value,
+        "{case}"
+    );
+    assert_eq!(
+        ScalarValue::try_from_array(output.column(1), 0)?,
+        ScalarValue::Int64(Some(0)),
+        "{case}"
+    );
+    assert_eq!(
+        ScalarValue::try_from_array(output.column(2), 0)?,
+        ScalarValue::Int64(None),
+        "{case}"
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn double_negation_lead_partition_streams_on_full_input_order() -> Result<()> {
+    for (case, value, nullable) in [
+        (
+            "timestamp",
+            ScalarValue::TimestampNanosecond(Some(1), None),
+            false,
+        ),
+        (
+            "nullable timestamp",
+            ScalarValue::TimestampNanosecond(None, None),
+            true,
+        ),
+        (
+            "year-month interval",
+            ScalarValue::IntervalYearMonth(Some(1)),
+            false,
+        ),
+    ] {
+        assert_double_negation_lead_streams(case, value, nullable).await?;
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn double_negation_window_partition_preserves_timestamp_overflow() -> Result<()> {
+    let ctx = pending_ordered_context(
+        &ScalarValue::TimestampNanosecond(Some(i64::MIN), None),
+        false,
+    )?;
+    let plan = ctx
+        .sql(
+            "SELECT LEAD(j, 1) OVER (PARTITION BY -(-i), j) AS n \
+             FROM s LIMIT 1",
+        )
+        .await?
+        .create_physical_plan()
+        .await?;
+
+    let mut stream = plan.execute(0, ctx.task_ctx())?;
+    let err = tokio::time::timeout(Duration::from_secs(5), stream.next())
+        .await
+        .expect("timestamp overflow was not reported")
+        .expect("stream ended before reporting timestamp overflow")
+        .expect_err("the original PARTITION BY expression must still be evaluated");
+    assert!(
+        matches!(
+            err.find_root(),
+            DataFusionError::ArrowError(error, _)
+                if matches!(error.as_ref(), ArrowError::ArithmeticOverflow(_))
+        ),
+        "unexpected error: {err}"
+    );
+    Ok(())
+}
+
 /// Regression: `parallelize_sorts` must not relocate a per-partition `SortExec`
 /// below an order-preserving `ProjectionExec` that *reorders* columns without
 /// remapping the sort-key column indices.