Answer COPY TO with RETURN_FILES as the count and the list of files written

crates/rudb-bind/src/lib.rs

@@ -59,7 +59,7 @@ pub use collation::check_collation;
 pub use context::{is_session_word, micros_now};
 pub use macros::kept_macro;
 pub use parameters::{Capture, Caught, Described, Parameters, Placeholders, Written};
-pub use partition::{Existing, NamePiece, Partitioned};
+pub use partition::{Existing, NamePiece, Partitioned, Returns};
 pub use statement::{
     Attach, Bound, Call, Checks, Conflict, ConflictAction, CopyTo, CreateTable, CreateView,
     DropTable, Insert, MacroChange, SchemaChange, SequenceChange, TriggerChange, TypeChange, Write,

crates/rudb-bind/src/partition.rs

@@ -10,6 +10,7 @@
 //!
 //! `FILE_SIZE_BYTES`, `ROW_GROUPS_PER_FILE` and `PER_THREAD_OUTPUT` are read here too, since they
 //! also turn the path into a directory of numbered files, which is what the pin calls rotating.
+//! So are `RETURN_FILES` and `RETURN_STATS`, which say what the statement answers with.
 
 use rudb_common::{Error, LogicalType, Result, Value, parse_size};
 use rudb_parse::ast;
@@ -43,6 +44,8 @@ pub struct Partitioned {
     pub per_thread: bool,
     /// Whether no file at all is written for no rows, which is `WRITE_EMPTY_FILE false`.
     pub skip_empty: bool,
+    /// What the statement answers with.
+    pub returns: Returns,
 }
 
 impl Partitioned {
@@ -74,6 +77,19 @@ pub enum Existing {
     Append,
 }
 
+/// What a `COPY ... TO` answers with.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub enum Returns {
+    /// How many rows it wrote.
+    #[default]
+    Count,
+    /// How many rows it wrote and the list of the files, which is `RETURN_FILES`.
+    Files,
+    /// A row for each file it wrote, with the statistics of its columns, which is
+    /// `RETURN_STATS`.
+    Stats,
+}
+
 /// One piece of a `FILENAME_PATTERN`.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum NamePiece {
@@ -88,7 +104,7 @@ pub enum NamePiece {
 }
 
 /// The options read here, which the format never sees.
-const READ_HERE: [&str; 12] = [
+const READ_HERE: [&str; 14] = [
     "partition_by",
     "write_partition_columns",
     "overwrite",
@@ -101,6 +117,8 @@ const READ_HERE: [&str; 12] = [
     "per_thread_output",
     "file_size_bytes",
     "write_empty_file",
+    "return_files",
+    "return_stats",
 ];
 
 /// The options that count the batches of a file, which JSON does not take.
@@ -164,6 +182,18 @@ pub(crate) fn partitioning(
             "per_thread_output" => out.per_thread = boolean(name, value.as_deref())?,
             "file_size_bytes" => out.file_size = Some(bytes(value.as_deref(), written)?),
             "write_empty_file" => out.skip_empty = !boolean(name, value.as_deref())?,
+            // Only an option that is on counts as set.
+            "return_files" | "return_stats" => {
+                if flag(name, value.as_deref(), written)? {
+                    if out.returns != Returns::Count {
+                        return Err(Error::binder(
+                            "Can only set one of RETURN_FILES or RETURN_STATS for COPY",
+                        ));
+                    }
+                    out.returns =
+                        if name == "return_files" { Returns::Files } else { Returns::Stats };
+                }
+            }
             _ => out.batches = Some(unsigned(name, value.as_deref(), written)?),
         }
     }
@@ -187,6 +217,17 @@ pub(crate) fn partitioning(
         return Err(Error::binder("APPEND mode requires a {uuid} label in filename_pattern"));
     }
     refuse_together(&out, tmp, partitioned)?;
+    // The pin writes JSON through its CSV writer, which is the name it gives.
+    if out.returns == Returns::Stats && format != "parquet" {
+        return Err(Error::not_implemented(
+            "RETURN_STATS is not supported for the \"csv\" copy format",
+        ));
+    }
+    if out.returns == Returns::Stats {
+        return Err(Error::not_implemented(
+            "COPY TO with the option return_stats is not supported yet",
+        ));
+    }
     if partitioned && !out.write_columns && out.columns.len() == names.len() {
         return Err(Error::not_implemented(
             "No column to write as all columns are specified as partition columns. \
@@ -265,6 +306,25 @@ fn unsigned(name: &str, text: Option<&str>, written: Option<&Written>) -> Result
     })
 }
 
+/// An option that is on or off the way the pin reads a `BOOLEAN` copy option, where an option with
+/// no value is on and anything else is cast, in the pin's words when it cannot be.
+fn flag(name: &str, text: Option<&str>, written: Option<&Written>) -> Result<bool> {
+    let Some(text) = text else { return Ok(true) };
+    let (ty, value) = match written
+        .and_then(|written| written.value.as_ref().map(|value| (&written.ty, value)))
+    {
+        Some((ty, value)) if *ty != LogicalType::Varchar => (ty.clone(), value.clone()),
+        _ => (LogicalType::Varchar, Value::Varchar(text.trim_matches('\'').into())),
+    };
+    match rudb_kernels::cast_value(&value, &LogicalType::Boolean, true) {
+        Ok(Value::Boolean(on)) => Ok(on),
+        _ => Err(Error::invalid_input(format!(
+            "Copy option \"{name}\" expected an argument of type BOOLEAN - the argument \
+             \"{value}\" of type {ty} could not be cast as this type"
+        ))),
+    }
+}
+
 /// A value cast to an unsigned 64 bit number, if it is one.
 fn cast_unsigned(value: &Value) -> Option<u64> {
     match rudb_kernels::cast_value(value, &LogicalType::UBigInt, true) {

crates/rudb-bind/src/statement.rs

@@ -1100,8 +1100,8 @@ fn csv_to(copy: &ast::CopyTo, plan: Plan) -> Result<CopyTo> {
                 }
             }
             "compression" | "dateformat" | "date_format" | "timestampformat"
-            | "timestamp_format" | "new_line" | "prefix" | "suffix" | "return_files"
-            | "preserve_order" | "force_not_null" | "encoding" => {
+            | "timestamp_format" | "new_line" | "prefix" | "suffix" | "preserve_order"
+            | "force_not_null" | "encoding" => {
                 return Err(Error::not_implemented(format!(
                     "COPY TO with the option {name} is not supported yet"
                 )));
@@ -1165,7 +1165,7 @@ fn json_to(copy: &ast::CopyTo, plan: Plan) -> Result<CopyTo> {
                     "Option \"encoding\" is not supported for writing - only for reading",
                 ));
             }
-            "compression" | "return_files" | "preserve_order" => {
+            "compression" | "preserve_order" => {
                 return Err(Error::not_implemented(format!(
                     "COPY TO with the option {name} is not supported yet"
                 )));
@@ -1265,7 +1265,6 @@ fn parquet_to(copy: &ast::CopyTo, plan: Plan) -> Result<CopyTo> {
             | "bloom_filter_false_positive_ratio"
             | "parquet_version"
             | "geoparquet_version"
-            | "return_files"
             | "preserve_order" => {
                 return Err(Error::not_implemented(format!(
                     "COPY TO with the option {name} is not supported yet"

crates/rudb/src/database.rs

@@ -7,7 +7,7 @@ use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex, MutexGuard, PoisonError, Weak};
 use std::time::{Duration, Instant};
 
-use rudb_bind::{Bound, Described, Parameters, Placeholders, Write};
+use rudb_bind::{Bound, Described, Parameters, Placeholders, Returns, Write};
 use rudb_catalog::{
     Alteration, Catalog, DEFAULT_CATALOG, DETACHED, Entry, Key, KeyLog, QualifiedName,
     TEMP_CATALOG, View, same_name,
@@ -7150,7 +7150,12 @@ impl Shared {
                 let result = run(sql, &copy.plan, &catalog, cancel, under)?;
                 let zone = session.session_time_zone();
                 let threads = self.inner.pool.threads();
-                QueryResult::changed(crate::export::write(&copy, &result, zone, threads)?)
+                let (rows, files) = crate::export::write(&copy, &result, zone, threads)?;
+                if copy.partitioned.returns == Returns::Files {
+                    QueryResult::copied_files(rows, files)
+                } else {
+                    QueryResult::changed(rows)
+                }
             }
             Bound::Setting(setting)
                 if setting.pragma && setting.name.eq_ignore_ascii_case("device_card_refresh") =>

crates/rudb/src/export.rs

@@ -46,21 +46,22 @@ use crate::QueryResult;
 const BATCH: usize = 2048;
 
 /// Writes the rows of `result` the way `copy` asks, to the one file or to a directory of them,
-/// and answers how many there were.
+/// and answers how many there were and the paths of the files, in the order they were written.
 pub(crate) fn write(
     copy: &CopyTo,
     result: &QueryResult,
     zone: SessionTimeZone,
     threads: usize,
-) -> Result<usize> {
+) -> Result<(usize, Vec<String>)> {
     let split = &copy.partitioned;
+    let mut files = Vec::new();
     if split.skip_empty && result.is_empty() {
-        return Ok(0);
+        return Ok((0, files));
     }
     if !split.directory() {
         let mut sink = Sink::open(copy, &copy.path, result, zone)?;
         sink.write(result.chunks())?;
-        return sink.finish();
+        return Ok((sink.finish(&mut files)?, files));
     }
     let root = Path::new(&copy.path);
     match split.existing {
@@ -74,10 +75,12 @@ pub(crate) fn write(
         _ => {}
     }
     make_directory(root)?;
-    if split.columns.is_empty() {
-        return Ok(write_files(copy, root, 0, result, zone, threads)?.0);
-    }
-    write_partitioned(copy, root, result, zone)
+    let rows = if split.columns.is_empty() {
+        write_files(copy, root, 0, result, zone, threads, &mut files)?
+    } else {
+        write_partitioned(copy, root, result, zone, &mut files)?
+    };
+    Ok((rows, files))
 }
 
 /// Makes a directory and the ones above it.
@@ -87,8 +90,8 @@ fn make_directory(path: &Path) -> Result<()> {
     })
 }
 
-/// Writes the rows of `result` into `directory`, the files numbered from `first`, and answers how
-/// many rows and how many files there were.
+/// Writes the rows of `result` into `directory`, the files numbered from `first`, adds their paths
+/// to `files` and answers how many rows there were.
 ///
 /// When the `COPY` rotates, a file is closed before the next batch once it holds one and has
 /// reached the size or the count of batches asked for, so a file is never left empty, where the
@@ -101,41 +104,42 @@ fn write_files(
     result: &QueryResult,
     zone: SessionTimeZone,
     threads: usize,
-) -> Result<(usize, usize)> {
+    files: &mut Vec<String>,
+) -> Result<usize> {
     let split = &copy.partitioned;
     let path = |number: usize| -> Result<String> {
         let name = format!("{}.{}", file_name(&split.pattern, number, zone)?, split.extension);
         Ok(directory.join(name).to_string_lossy().into_owned())
     };
     if split.per_thread && !split.rotates() {
-        return write_threads(copy, &path, first, result, zone, threads);
+        return write_threads(copy, &path, first, result, zone, threads, files);
     }
     let mut sink = Sink::open(copy, &path(first)?, result, zone)?;
     if !split.rotates() {
         sink.write(result.chunks())?;
-        return Ok((sink.finish()?, 1));
+        return sink.finish(files);
     }
     let least = if copy.parquet { copy.row_group_size } else { 0 };
-    let (mut rows, mut files, mut held) = (0, 1, 0);
+    let (mut rows, mut number, mut held) = (0, first, 0);
     for batch in batches(result.chunks(), least)? {
         let full = held > 0
             && (split.file_size.is_some_and(|most| sink.size() >= most)
                 || split.batches.is_some_and(|most| held >= most));
         if full {
-            rows += sink.finish()?;
-            sink = Sink::open(copy, &path(first + files)?, result, zone)?;
-            files += 1;
+            rows += sink.finish(files)?;
+            number += 1;
+            sink = Sink::open(copy, &path(number)?, result, zone)?;
             held = 0;
         }
         sink.write(&batch)?;
         held += 1;
     }
-    Ok((rows + sink.finish()?, files))
+    Ok(rows + sink.finish(files)?)
 }
 
 /// Writes the rows of `result` the way `PER_THREAD_OUTPUT` does, a file for each of `threads`
-/// that gets rows, each taking the next run of batches, and answers how many rows and files there
-/// were.
+/// that gets rows, each taking the next run of batches, adds their paths to `files` and answers
+/// how many rows there were.
 ///
 /// The pin writes a file for each thread that ends up with rows, so how many there are depends on
 /// how its scan was split, and a source that one thread reads, such as `range`, gives it one file.
@@ -148,24 +152,24 @@ fn write_threads(
     result: &QueryResult,
     zone: SessionTimeZone,
     threads: usize,
-) -> Result<(usize, usize)> {
+    files: &mut Vec<String>,
+) -> Result<usize> {
     let batches = batches(result.chunks(), 0)?;
     let each = batches.len().div_ceil(threads.max(1)).max(1);
     let mut runs = batches.chunks(each).peekable();
     if runs.peek().is_none() {
         let sink = Sink::open(copy, &path(first)?, result, zone)?;
-        return Ok((sink.finish()?, 1));
+        return sink.finish(files);
     }
-    let (mut rows, mut files) = (0, 0);
-    for run in runs {
-        let mut sink = Sink::open(copy, &path(first + files)?, result, zone)?;
+    let mut rows = 0;
+    for (number, run) in (first..).zip(runs) {
+        let mut sink = Sink::open(copy, &path(number)?, result, zone)?;
         for batch in run {
             sink.write(batch)?;
         }
-        rows += sink.finish()?;
-        files += 1;
+        rows += sink.finish(files)?;
     }
-    Ok((rows, files))
+    Ok(rows)
 }
 
 /// The rows cut into the batches the pin closes a file between: [`BATCH`] rows each, or as many of
@@ -211,6 +215,7 @@ fn write_partitioned(
     root: &Path,
     result: &QueryResult,
     zone: SessionTimeZone,
+    files: &mut Vec<String>,
 ) -> Result<usize> {
     let split = &copy.partitioned;
     let names = result.names();
@@ -291,9 +296,9 @@ fn write_partitioned(
         let (place, first) =
             if split.flat { (root.to_path_buf(), next) } else { (root.join(directory), 0) };
         make_directory(&place)?;
-        let (written, files) = write_files(copy, &place, first, &part, zone, 1)?;
-        rows += written;
-        next += files;
+        let before = files.len();
+        rows += write_files(copy, &place, first, &part, zone, 1, files)?;
+        next += files.len() - before;
     }
     Ok(rows)
 }
@@ -594,9 +599,9 @@ impl<'a> Sink<'a> {
         Ok(())
     }
 
-    /// Writes what comes after the last row and closes the file, and answers how many rows it
-    /// holds.
-    fn finish(self) -> Result<usize> {
+    /// Writes what comes after the last row and closes the file, adds its path to `files` and
+    /// answers how many rows it holds.
+    fn finish(self, files: &mut Vec<String>) -> Result<usize> {
         let written = failed(&self.path);
         match self.body {
             Body::Csv { mut out, .. } => out.flush().map_err(written)?,
@@ -612,6 +617,7 @@ impl<'a> Sink<'a> {
                 writer.finish()?;
             }
         }
+        files.push(self.path);
         Ok(self.rows)
     }
 }

crates/rudb/src/result.rs

@@ -323,6 +323,27 @@ impl QueryResult {
         Ok(result)
     }
 
+    /// What a `COPY ... TO` with `RETURN_FILES` answers: how many rows it wrote and the list of
+    /// the files it wrote them to, as a query's one row rather than a count of changes.
+    pub(crate) fn copied_files(rows: usize, files: Vec<String>) -> Result<Self> {
+        let count = i64::try_from(rows).unwrap_or(i64::MAX);
+        let list = LogicalType::list(LogicalType::Varchar);
+        let files = Value::List {
+            element: LogicalType::Varchar,
+            values: files.into_iter().map(Value::Varchar).collect(),
+        };
+        let chunk = Chunk::new(vec![
+            Vector::from_values(LogicalType::BigInt, &[Value::BigInt(count)])?,
+            Vector::from_values(list.clone(), &[files])?,
+        ])?;
+        Ok(Self::new(
+            vec!["Count".to_owned(), "Files".to_owned()],
+            vec![LogicalType::BigInt, list],
+            vec![chunk],
+            Memory::unlimited().reservation(),
+        ))
+    }
+
     /// How many rows the statement wrote, when it was an `INSERT`, `UPDATE`, `DELETE` or `CREATE
     /// TABLE`, and `None` when the rows are the answer to a query.
     #[must_use]