rluvaton opened a new issue, #23798:
URL: https://github.com/apache/datafusion/issues/23798
### 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:
```rust
let sort_satisfied = project_plan
.equivalence_properties()
.ordering_satisfy(lex_ordering)?;
```
**will return `false`**
### To Reproduce
<details><summary>Full test</summary>
```rust
/// 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(())
}
```
</details>
### Expected behavior
The sort properties should be kept if the expression keep the sorting
### Additional context
_No response_
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
