neilconway commented on code in PR #23804:
URL: https://github.com/apache/datafusion/pull/23804#discussion_r3631229964
##########
datafusion/physical-expr/src/equivalence/properties/dependency.rs:
##########
@@ -1617,4 +1607,47 @@ mod tests {
Ok(())
}
+
+ #[test]
+ fn test_concat_projection_does_not_preserve_lex_ordering() -> Result<()> {
+ let input_schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Utf8, false),
+ Field::new("b", DataType::Utf8, false),
+ ]));
+ let col_a = col("a", &input_schema)?;
+ let col_b = col("b", &input_schema)?;
+
+ let mut input_properties =
EquivalenceProperties::new(Arc::clone(&input_schema));
+ input_properties.add_ordering([
+ PhysicalSortExpr::new_default(Arc::clone(&col_a)),
+ PhysicalSortExpr::new_default(Arc::clone(&col_b)),
+ ]);
+
+ let a_concat_b = Arc::new(ScalarFunctionExpr::new(
+ "concat",
+ concat(),
+ vec![col_a, Arc::clone(&col_b)],
+ Field::new("ab", DataType::Utf8, true).into(),
+ Arc::new(ConfigOptions::default()),
+ )) as Arc<dyn PhysicalExpr>;
+ let projection_mapping = ProjectionMapping::try_new(
+ vec![(a_concat_b, "ab".to_string()), (col_b, "b".to_string())],
+ &input_schema,
+ )?;
+ let output_schema = output_schema(&projection_mapping, &input_schema)?;
+ let projected =
+ input_properties.project(&projection_mapping,
Arc::clone(&output_schema));
+ let unsafe_projected_ordering: LexOrdering = [
+ PhysicalSortExpr::new_default(col("ab", &output_schema)?),
+ PhysicalSortExpr::new_default(col("b", &output_schema)?),
+ ]
+ .into();
+
+ assert!(
+ !projected.ordering_satisfy(unsafe_projected_ordering)?,
+ "concat(a, b) must not replace a in the projected [a, b] ordering"
+ );
+
+ Ok(())
+ }
Review Comment:
This test passes with the bugfix reverted. It is probably redundant with the
other tests anyway -- maybe just remove it?
##########
datafusion/core/tests/physical_optimizer/enforce_sorting.rs:
##########
@@ -3255,3 +3255,137 @@ async fn
test_does_not_push_fetch_sort_through_projection_over_union() -> Result
Ok(())
}
+
+#[tokio::test]
+async fn test_concat_lex_sort_in_project_with_optimization() -> Result<()> {
+ test_concat_lex_ordering_with_project(true).await
+}
+
+#[tokio::test]
+async fn test_concat_lex_sort_in_project_without_optimization() -> Result<()> {
+ test_concat_lex_ordering_with_project(false).await
+}
+
+/// The input plan is:
+/// ```text
+/// SortExec: expr=[c@0 ASC], preserve_partitioning=[false]
+/// ProjectionExec: expr=[concat(a@0, b@1) as c]
+/// DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0
ASC
+/// ```
+///
+/// The output must be sorted whether or not sort enforcement is optimized.
+async fn test_concat_lex_ordering_with_project(optimize: bool) -> Result<()> {
+ // Sorted by a ASC ("a" < "a0") and by b ASC ("1" == "1"), but
+ // concat(a, b) is NOT sorted: "a1" > "a01".
+ let batch = record_batch!(("a", Utf8, ["a", "a0"]), ("b", Utf8, ["1",
"1"]))?;
+ 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)]).unwrap(),
+ LexOrdering::new([sort_expr("b", &schema)]).unwrap(),
+ ])?,
+ ))) as Arc<dyn ExecutionPlan>;
+
+ let concat =
Arc::new(datafusion_physical_expr::ScalarFunctionExpr::try_new(
+ datafusion::functions::string::concat(),
+ vec![col("a", &schema)?, col("b", &schema)?],
+ &schema,
+ Arc::new(ConfigOptions::default()),
+ )?) as Arc<dyn PhysicalExpr>;
+ let projection = projection_exec(vec![(concat, "c".to_string())], source)?;
+ let plan = sort_exec(
+ LexOrdering::new([sort_expr("c", &projection.schema())]).unwrap(),
+ projection,
+ );
+
+ assert_plan_output_eq(plan, optimize, vec!["a01", "a1"]).await
+}
+
+#[tokio::test]
+async fn test_concat_lex_sort_in_filter_with_optimization() -> Result<()> {
+ test_concat_lex_ordering_with_filter(true).await
+}
+
+#[tokio::test]
+async fn test_concat_lex_sort_in_filter_without_optimization() -> Result<()> {
+ test_concat_lex_ordering_with_filter(false).await
+}
+
+/// The input plan is:
+/// ```text
+/// SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false]
+/// FilterExec: c@2 = concat(a@0, b@1)
+/// SortExec: expr=[c@2 ASC, a@0 ASC, b@1 ASC],
preserve_partitioning=[false]
+/// DataSourceExec: partitions=1, partition_sizes=[1]
+/// ```
+///
+/// The output must be sorted whether or not sort enforcement is optimized.
+async fn test_concat_lex_ordering_with_filter(optimize: bool) -> Result<()> {
+ // Rows are sorted by (c, a, b) and satisfy c = concat(a, b), but are
+ // NOT sorted by (a, b): ("a0", "1") > ("a", "1").
+ let batch = record_batch!(
+ ("a", Utf8, ["a0", "a"]),
+ ("b", Utf8, ["1", "1"]),
+ ("c", Utf8, ["a01", "a1"])
+ )?;
+ let schema = batch.schema();
+ let source =
datafusion::datasource::memory::MemorySourceConfig::try_new_exec(
+ &[vec![batch]],
+ schema.clone(),
+ None,
+ )?;
+ let sort_cab = sort_exec(
+ LexOrdering::new([
+ sort_expr("c", &schema),
+ sort_expr("a", &schema),
+ sort_expr("b", &schema),
+ ])
+ .unwrap(),
+ source,
+ );
+
+ let concat =
Arc::new(datafusion_physical_expr::ScalarFunctionExpr::try_new(
+ datafusion::functions::string::concat(),
+ vec![col("a", &schema)?, col("b", &schema)?],
+ &schema,
+ Arc::new(ConfigOptions::default()),
+ )?) as Arc<dyn PhysicalExpr>;
+ let predicate =
+ Arc::new(BinaryExpr::new(col("c", &schema)?, Operator::Eq, concat)) as
_;
+ let filter = filter_exec(predicate, sort_cab);
+ let plan = sort_exec(
+ LexOrdering::new([sort_expr("a", &schema), sort_expr("b",
&schema)]).unwrap(),
+ filter,
+ );
+
+ assert_plan_output_eq(plan, optimize, vec!["a", "a0"]).await
+}
+
+async fn assert_plan_output_eq(
+ plan: Arc<dyn ExecutionPlan>,
+ optimize: bool,
+ expected_result: Vec<&str>,
+) -> Result<()> {
+ let plan = if optimize {
+ EnsureRequirements::new().optimize(plan, &ConfigOptions::default())?
+ } else {
+ plan
+ };
+ let plan_str = displayable(plan.as_ref()).indent(true).to_string();
+
+ let results =
+ datafusion_physical_plan::collect(plan,
Arc::new(TaskContext::default())).await?;
+ let all = arrow::compute::concat_batches(&results[0].schema(), &results)?;
+ let expected = create_array!(Utf8, expected_result) as ArrayRef;
+ assert_eq!(
+ all.column(0),
+ &expected,
+ "output must satisfy the requested ordering (optimize={optimize});
plan:\n{plan_str}"
+ );
+ Ok(())
+}
Review Comment:
Not critical, but this seems like a _lot_ of test machinery to me that
doesn't necessarily justify itself. We could adequately test that the
end-to-end behavior is correct with an SLT over the plan shape, no?
##########
datafusion/physical-expr/src/equivalence/properties/dependency.rs:
##########
@@ -1102,7 +1097,7 @@ mod tests {
}
#[test]
- fn test_ordering_equivalence_with_concat_equality() -> Result<()> {
+ fn test_ordering_equivalence_with_non_lex_monotonic_concat_equality() ->
Result<()> {
Review Comment:
This test doesn't seem that valuable now that the `concat` behavior has
changed; consider removing it?
--
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]
