buraksenn commented on code in PR #23804:
URL: https://github.com/apache/datafusion/pull/23804#discussion_r3631734323
##########
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:
I've actually copy-pasted this part from the issue but I'll adjust all tests
accordingly
--
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]
