Fix: Preserve sorting for the COPY statement

Author: Bert Vermeiren

datafusion/core/src/physical_planner.rs

@@ -97,6 +97,7 @@ use datafusion_sql::TableReference;
 use sqlparser::ast::NullTreatment;
 
 use async_trait::async_trait;
+use datafusion_physical_expr_common::sort_expr::LexRequirement;
 use datafusion_physical_plan::async_func::{AsyncFuncExec, AsyncMapper};
 use futures::{StreamExt, TryStreamExt};
 use itertools::{multiunzip, Itertools};
@@ -556,8 +557,17 @@ impl DefaultPhysicalPlanner {
                     file_extension,
                 };
 
+                let ordering = input_exec.properties().output_ordering().cloned();
+                let sort_req: Option<LexRequirement> =
+                    ordering.map(|o| LexRequirement::from(o));
+
                 sink_format
-                    .create_writer_physical_plan(input_exec, session_state, config, None)
+                    .create_writer_physical_plan(
+                        input_exec,
+                        session_state,
+                        config,
+                        sort_req,
+                    )
                     .await?
             }
             LogicalPlan::Dml(DmlStatement {

datafusion/sqllogictest/test_files/copy.slt

@@ -427,6 +427,70 @@ select * from validate_parquet_single;
 1 Foo
 2 Bar
 
+
+# copy with sorting requirement
+
+statement ok
+create table to_order_by_table(col1 integer) as values (1), (4), (3), (2);
+
+query TT
+EXPLAIN COPY (SELECT * FROM to_order_by_table ORDER BY col1 ASC NULLS LAST) to 'test_files/scratch/copy/sorted_table.parquet/';
+----
+logical_plan
+01)CopyTo: format=parquet output_url=test_files/scratch/copy/sorted_table.parquet/ options: ()
+02)--Sort: to_order_by_table.col1 ASC NULLS LAST
+03)----TableScan: to_order_by_table projection=[col1]
+physical_plan
+01)DataSinkExec: sink=ParquetSink(file_groups=[])
+02)--SortExec: expr=[col1@0 ASC NULLS LAST], preserve_partitioning=[false]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+
+query I
+COPY (SELECT * FROM to_order_by_table ORDER BY col1 ASC NULLS LAST) to 'test_files/scratch/copy/sorted_table.parquet';
+----
+4
+
+statement ok
+drop table to_order_by_table;
+
+statement ok
+CREATE EXTERNAL TABLE validate_sorted_table STORED AS PARQUET LOCATION 'test_files/scratch/copy/sorted_table.parquet';
+
+# verify col1 is sorted ASC within the just written parquet file
+query I
+WITH
+    numbered_rows AS (
+        SELECT
+            ROW_NUMBER() OVER () as row_num,
+            col1
+        FROM validate_sorted_table
+    ),
+    adjacent_check AS (
+        SELECT
+            a.row_num,
+            a.col1 as current_value,
+            b.col1 as next_value
+        FROM numbered_rows a
+                 JOIN numbered_rows b ON a.row_num + 1 = b.row_num
+    ),
+    violations AS (
+        SELECT
+            row_num,
+            current_value,
+            next_value,
+            'VIOLATION: ' || current_value || ' > ' || next_value as issue
+        FROM adjacent_check
+        WHERE current_value > next_value
+    )
+    SELECT COUNT(*) FROM violations;
+----
+0
+
+statement ok
+drop table validate_sorted_table;
+
+
 # copy from table to folder of compressed json files
 query I
 COPY source_table  to 'test_files/scratch/copy/table_json_gz' STORED AS JSON OPTIONS ('format.compression' gzip);