ProjectExec does not keep sorting information for passthrough expression #23798

/// A pass-through wrapper around a column: just assert that column does not contain any nulls
#[derive(Debug, Eq)]
struct AssertNotNull {
    inner: Arc<dyn PhysicalExpr>,
}

impl AssertNotNull {
    fn new(inner: Arc<dyn PhysicalExpr>) -> Arc<Self> {
        Arc::new(Self { inner })
    }
}

impl PartialEq for AssertNotNull {
    fn eq(&self, other: &Self) -> bool {
        self.inner.eq(&other.inner)
    }
}

impl std::hash::Hash for AssertNotNull {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.inner.hash(state);
    }
}

impl std::fmt::Display for AssertNotNull {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "assert_not_null({})", self.inner)
    }
}

impl PhysicalExpr for AssertNotNull {
    fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
        self.inner.data_type(input_schema)
    }

    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
        Ok(false)
    }

    fn evaluate(
        &self,
        batch: &RecordBatch,
    ) -> Result<datafusion_expr::ColumnarValue> {
        let child = self.inner.evaluate(batch)?;
        match child {
            ColumnarValue::Array(a) if a.logical_null_count() > 0 => {
                return Err(DataFusionError::Internal(
                    "AssertNotNull evaluated to null".to_string(),
                ))
            } 
            ColumnarValue::Scalar(s) if s.is_null() => {
                return Err(DataFusionError::Internal(
                    "AssertNotNull evaluated to null".to_string(),
                ))
            } 
            child => Ok(child)
        }
    }

    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
        vec![&self.inner]
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn PhysicalExpr>>,
    ) -> Result<Arc<dyn PhysicalExpr>> {
        Ok(Arc::new(AssertNotNull {
            inner: Arc::clone(&children[0]),
        }))
    }

    fn get_properties(
        &self,
        children: &[datafusion_expr::sort_properties::ExprProperties],
    ) -> Result<datafusion_expr::sort_properties::ExprProperties> {
        Ok(children[0].clone())
    }

    fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "assert_not_null({})", self.inner)
    }
}

#[tokio::test]
async fn test_passthrough_wrapper_projection_keeps_ordering() -> Result<()> {
    fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr {
        PhysicalSortExpr {
            expr: col(name, schema).unwrap(),
            options: Default::default(),
        }
    }

    pub fn projection_exec(
        expr: Vec<(Arc<dyn PhysicalExpr>, String)>,
        input: Arc<dyn ExecutionPlan>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        let proj_exprs: Vec<ProjectionExpr> = expr
          .into_iter()
          .map(|(expr, alias)| ProjectionExpr { expr, alias })
          .collect();
        Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?))
    }

    let batch = record_batch!(
        ("a", Utf8, ["x", "y"]),
        ("b", Utf8, ["1", "2"]),
        ("c", Utf8, ["1", "2"])
    )?;
    let schema = batch.schema();
    let source = Arc::new(DataSourceExec::new(Arc::new(
        datafusion::datasource::memory::MemorySourceConfig::try_new(
            &[vec![batch]],
            schema.clone(),
            None,
        )?
        .try_with_sort_information(vec![
            LexOrdering::new([
                sort_expr("a", &schema),
                sort_expr("b", &schema),
                sort_expr("c", &schema),
            ])
            .unwrap(),
        ])?,
    ))) as Arc<dyn ExecutionPlan>;

    let projection = projection_exec(
        vec![
            (
                AssertNotNull::new(
                    col("a", &schema)?,
                ),
                "a".to_string(),
            ),
            (
                AssertNotNull::new(
                    col("b", &schema)?,
                ),
                "b".to_string(),
            ),
            (
                AssertNotNull::new(
                    col("c", &schema)?,
                ),
                "c".to_string(),
            ),
        ],
        source,
    )?;

    let plan = SortExec::new(
        LexOrdering::new([
            sort_expr("a", &projection.schema()),
            sort_expr("b", &projection.schema()),
            sort_expr("c", &projection.schema()),
        ])
        .unwrap(),
        projection,
    );

    let sort_satisfied = plan
      .input()
      .equivalence_properties()
      .ordering_satisfy(plan.expr().clone())?;

    assert!(sort_satisfied, "sort should be satisfied");

    let optimized =
        EnsureRequirements::new().optimize(Arc::new(plan), &ConfigOptions::default())?;
    let plan_str = displayable(optimized.as_ref()).indent(true).to_string();
    assert!(
        !plan_str.contains("SortExec"),
        "sort should be elided; plan:\n{plan_str}"
    );

    Ok(())
}