/// 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(())
}
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:
here the Sort should already be satisfied, but calling:
will return
falseTo Reproduce
Full test
Expected behavior
The sort properties should be kept if the expression keep the sorting
Additional context
No response