Skip to content

ProjectExec does not keep sorting information for passthrough expression #23798

Description

@rluvaton

Describe the bug

Lets say I have an physical expression/udf that just passthrough the child expression (for example, an expression that does assert_not_null)

then we have a plan like this:

SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC], preserve_partitioning=[false]
  ProjectionExec: expr=[assert_not_null(a@0) as a, assert_not_null(b@1) as b, assert_not_null(c@2) as c]
    DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 ASC, b@1 ASC, c@2 ASC

here the Sort should already be satisfied, but calling:

let sort_satisfied = project_plan
  .equivalence_properties()
  .ordering_satisfy(lex_ordering)?;

will return false

To Reproduce

Full test
/// 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(())
}

Expected behavior

The sort properties should be kept if the expression keep the sorting

Additional context

No response

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions