feat: unload preserve order if the source is ordered. (#17864)

* feat: unload preserve order if the source is ordered.

* fix

Author: youngsofun

src/query/catalog/src/plan/datasource/datasource_info/stage.rs

@@ -30,10 +30,14 @@ use databend_common_storage::StageFilesInfo;
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, Eq)]
 pub struct StageTableInfo {
+    // common
+    pub stage_root: String,
+    pub stage_info: StageInfo,
+
+    // copy into table only
     pub schema: TableSchemaRef,
     pub default_exprs: Option<Vec<RemoteDefaultExpr>>,
     pub files_info: StageFilesInfo,
-    pub stage_info: StageInfo,
     pub files_to_copy: Option<Vec<StageFileInfo>>,
     // files that
     // - are listed as candidates to be copied
@@ -42,9 +46,11 @@ pub struct StageTableInfo {
     // - may need to be purged as well (depends on the copy options)
     pub duplicated_files_detected: Vec<String>,
     pub is_select: bool,
-    pub copy_into_location_options: CopyIntoLocationOptions,
     pub copy_into_table_options: CopyIntoTableOptions,
-    pub stage_root: String,
+
+    // copy into location only
+    pub copy_into_location_ordered: bool,
+    pub copy_into_location_options: CopyIntoLocationOptions,
 }
 
 impl StageTableInfo {

src/query/service/src/interpreters/interpreter_copy_into_location.rs

@@ -114,6 +114,7 @@ impl CopyIntoLocationInterpreter {
                 copy_into_location_options: options.clone(),
                 copy_into_table_options: Default::default(),
                 stage_root: "".to_string(),
+                copy_into_location_ordered: self.plan.is_ordered,
             },
         }));
 

src/query/service/src/sessions/query_ctx.rs

@@ -1636,6 +1636,7 @@ impl TableContext for QueryContext {
                     copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
                     stage_root,
+                    copy_into_location_ordered: false,
                 };
                 OrcTable::try_create(info).await
             }
@@ -1655,6 +1656,7 @@ impl TableContext for QueryContext {
                     copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
                     stage_root,
+                    copy_into_location_ordered: false,
                 };
                 StageTable::try_create(info)
             }
@@ -1692,6 +1694,7 @@ impl TableContext for QueryContext {
                     copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
                     stage_root,
+                    copy_into_location_ordered: false,
                 };
                 StageTable::try_create(info)
             }

src/query/sql/src/planner/binder/copy_into_location.rs

@@ -109,6 +109,13 @@ impl Binder {
                     .await
             }
         }?;
+        let mut is_ordered = false;
+        if let Plan::Query { s_expr, .. } = &query {
+            let p = s_expr.derive_relational_prop()?;
+            if !p.orderings.is_empty() {
+                is_ordered = true;
+            }
+        }
 
         let (mut stage_info, path) = resolve_file_location(self.ctx.as_ref(), &stmt.dst).await?;
 
@@ -140,6 +147,7 @@ impl Binder {
             path,
             from: Box::new(query),
             options: stmt.options.clone(),
+            is_ordered,
         }))
     }
 }

src/query/sql/src/planner/binder/copy_into_table.rs

@@ -224,6 +224,7 @@ impl Binder {
                 copy_into_location_options: Default::default(),
                 copy_into_table_options: stmt.options.clone(),
                 stage_root: "".to_string(),
+                copy_into_location_ordered: false,
             },
             values_consts: vec![],
             required_source_schema: required_values_schema.clone(),
@@ -404,6 +405,7 @@ impl Binder {
                 copy_into_location_options: Default::default(),
                 copy_into_table_options: options,
                 stage_root: "".to_string(),
+                copy_into_location_ordered: false,
             },
             write_mode,
             query: None,

src/query/sql/src/planner/optimizer/optimizer.rs

@@ -149,11 +149,13 @@ pub async fn optimize(opt_ctx: Arc<OptimizerContext>, plan: Plan) -> Result<Plan
             path,
             from,
             options,
+            is_ordered,
         }) => Ok(Plan::CopyIntoLocation(CopyIntoLocationPlan {
             stage,
             path,
             from: Box::new(Box::pin(optimize(opt_ctx, *from)).await?),
             options,
+            is_ordered,
         })),
         Plan::CopyIntoTable(mut plan) if !plan.no_file_to_copy => {
             plan.enable_distributed = opt_ctx.get_enable_distributed_optimization()

src/query/sql/src/planner/plans/copy_into_location.rs

@@ -31,6 +31,7 @@ pub struct CopyIntoLocationPlan {
     pub path: String,
     pub from: Box<Plan>,
     pub options: CopyIntoLocationOptions,
+    pub is_ordered: bool,
 }
 
 impl CopyIntoLocationPlan {

src/query/storages/stage/src/append/do_append.rs

@@ -35,7 +35,11 @@ impl StageTable {
 
         let fmt = self.table_info.stage_info.file_format_params.clone();
         let mem_limit = settings.get_max_memory_usage()? as usize;
-        let max_threads = settings.get_max_threads()? as usize;
+        let mut max_threads = settings.get_max_threads()? as usize;
+        if self.table_info.copy_into_location_ordered {
+            max_threads = 1;
+            pipeline.try_resize(1)?;
+        }
 
         let op = StageTable::get_op(&self.table_info.stage_info)?;
         let query_id = ctx.get_id();

src/query/storages/stage/src/append/parquet_file/writer_processor.rs

@@ -52,7 +52,7 @@ pub struct ParquetFileWriter {
     arrow_schema: Arc<Schema>,
     compression: Compression,
 
-    input_data: Vec<DataBlock>,
+    input_data: VecDeque<DataBlock>,
 
     input_bytes: usize,
     row_counts: usize,
@@ -150,7 +150,7 @@ impl ParquetFileWriter {
             unload_output,
             unload_output_blocks: None,
             writer,
-            input_data: Vec::new(),
+            input_data: VecDeque::new(),
             input_bytes: 0,
             file_to_write: None,
             data_accessor,
@@ -231,11 +231,11 @@ impl Processor for ParquetFileWriter {
         } else if self.input.has_data() {
             let block = self.input.pull_data().unwrap()?;
             if self.targe_file_size.is_none() {
-                self.input_data.push(block);
+                self.input_data.push_back(block);
             } else {
                 let block_meta = block.get_owned_meta().unwrap();
                 let blocks = BlockBatch::downcast_from(block_meta).unwrap();
-                self.input_data.extend_from_slice(&blocks.blocks);
+                self.input_data.extend(blocks.blocks);
             }
 
             self.input.set_not_need_data();
@@ -247,7 +247,7 @@ impl Processor for ParquetFileWriter {
     }
 
     fn process(&mut self) -> Result<()> {
-        while let Some(b) = self.input_data.pop() {
+        while let Some(b) = self.input_data.pop_front() {
             self.input_bytes += b.memory_size();
             self.row_counts += b.num_rows();
             let batch = b.to_record_batch(&self.table_info.schema)?;

tests/sqllogictests/suites/stage/ordered_unload.test

@@ -0,0 +1,59 @@
+statement ok
+create or replace stage s1;
+
+# not ordered csv
+statement ok
+copy /*+ set_var(max_block_size=10) */ INTO @s1 from (select * from numbers(10000)) file_format=(type=csv) max_file_size=100;
+
+query
+SELECT COUNT(*)<10000 FROM (SELECT $1 AS a, rank() OVER (ORDER BY metadata$filename, metadata$file_row_number, $1) AS rank FROM '@s1' ( FILE_FORMAT => 'csv', )) WHERE a + 1 = rank
+----
+1
+
+statement ok
+remove @s1;
+
+# ordered csv
+
+statement ok
+copy /*+ set_var(max_block_size=10) */ INTO @s1 from (select * from numbers(10000) order by number) file_format=(type=csv) max_file_size=100;
+
+query
+SELECT COUNT(*) FROM (SELECT $1 AS a, rank() OVER (ORDER BY metadata$filename, metadata$file_row_number, $1) AS rank FROM '@s1' ( FILE_FORMAT => 'csv', )) WHERE a + 1 = rank
+----
+10000
+
+# data_af2ab6dc-8725-46e5-a601-3dad9c512769_0000_00000770.csv
+query
+SELECT * from list_stage(location => '@s1') where substr(name, 43, 4) != '0000'
+----
+
+statement ok
+remove @s1;
+
+
+# ordered parquet
+
+statement ok
+copy /*+ set_var(max_block_size=10) */ INTO @s1 from (select * from numbers(10000) order by number) file_format=(type=parquet) max_file_size=100;
+
+query
+SELECT COUNT(*) FROM (SELECT $1 AS a, rank() OVER (ORDER BY metadata$filename, metadata$file_row_number, $1) AS rank FROM '@s1' ( FILE_FORMAT => 'parquet', )) WHERE a + 1 = rank
+----
+10000
+
+statement ok
+remove @s1;
+
+# ordered ndjson
+
+statement ok
+copy /*+ set_var(max_block_size=10) */ INTO @s1 from (select * from numbers(10000) order by number) file_format=(type=ndjson) max_file_size=100;
+
+query
+SELECT COUNT(*) FROM (SELECT $1:number::int AS a, rank() OVER (ORDER BY metadata$filename, metadata$file_row_number, a) AS rank FROM '@s1' ( FILE_FORMAT => 'ndjson', )) WHERE a + 1 = rank
+----
+10000
+
+statement ok
+drop stage s1;