diff --git a/crates/rudb-bind/src/lib.rs b/crates/rudb-bind/src/lib.rs index 060f9815..e594a208 100644 --- a/crates/rudb-bind/src/lib.rs +++ b/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, diff --git a/crates/rudb-bind/src/partition.rs b/crates/rudb-bind/src/partition.rs index 6ae84d36..5a246e63 100644 --- a/crates/rudb-bind/src/partition.rs +++ b/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 { + 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 { match rudb_kernels::cast_value(value, &LogicalType::UBigInt, true) { diff --git a/crates/rudb-bind/src/statement.rs b/crates/rudb-bind/src/statement.rs index 859aadb7..ed981036 100644 --- a/crates/rudb-bind/src/statement.rs +++ b/crates/rudb-bind/src/statement.rs @@ -1100,8 +1100,8 @@ fn csv_to(copy: &ast::CopyTo, plan: Plan) -> Result { } } "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 { "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 { | "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" diff --git a/crates/rudb/src/database.rs b/crates/rudb/src/database.rs index 991668c3..33b0d112 100644 --- a/crates/rudb/src/database.rs +++ b/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, ©.plan, &catalog, cancel, under)?; let zone = session.session_time_zone(); let threads = self.inner.pool.threads(); - QueryResult::changed(crate::export::write(©, &result, zone, threads)?) + let (rows, files) = crate::export::write(©, &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") => diff --git a/crates/rudb/src/export.rs b/crates/rudb/src/export.rs index d65a707f..9b38084d 100644 --- a/crates/rudb/src/export.rs +++ b/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 { +) -> Result<(usize, Vec)> { let split = ©.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, ©.path, result, zone)?; sink.write(result.chunks())?; - return sink.finish(); + return Ok((sink.finish(&mut files)?, files)); } let root = Path::new(©.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, +) -> Result { let split = ©.partitioned; let path = |number: usize| -> Result { 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, +) -> Result { 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, ) -> Result { let split = ©.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 { + /// 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) -> Result { 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) } } diff --git a/crates/rudb/src/result.rs b/crates/rudb/src/result.rs index bf3324bc..397dd62f 100644 --- a/crates/rudb/src/result.rs +++ b/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) -> Result { + 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] diff --git a/crates/rudb/tests/copy_to.rs b/crates/rudb/tests/copy_to.rs index 2c3ae93b..332535e6 100644 --- a/crates/rudb/tests/copy_to.rs +++ b/crates/rudb/tests/copy_to.rs @@ -796,3 +796,70 @@ fn a_rotating_copy_refuses_what_the_pin_refuses() { assert_eq!(error.to_string(), want, "{sql}"); } } + +#[test] +fn return_files_answers_the_count_and_the_files_written_as_the_pin_does() { + use rudb_common::Value; + let root = std::env::temp_dir().join(format!("rudb-copy-to-files-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + std::fs::create_dir_all(&root).expect("makes the directory"); + let dir = root.display().to_string(); + let db = Database::new(); + db.execute("CREATE TABLE t AS SELECT range i, range % 2 k FROM range(5)").expect("creates"); + let files = |sql: &str| { + let sql = sql.replace("DIR", &dir); + let result = db.execute(&sql).unwrap_or_else(|error| panic!("{sql}: {error}")); + assert_eq!(result.names(), ["Count", "Files"], "{sql}"); + assert_eq!(result.changes(), None, "{sql}"); + let Value::List { values, .. } = result.value_at(0, 1) else { panic!("{sql}: a list") }; + let listed = values.iter().map(|value| value.to_string().replace(&dir, "DIR")); + (result.value_at(0, 0), listed.collect::>()) + }; + let count = |rows| Value::BigInt(rows); + + assert_eq!(files("COPY t TO 'DIR/a.csv' (RETURN_FILES)"), (count(5), vec!["DIR/a.csv".into()])); + assert_eq!( + files("COPY t TO 'DIR/p' (FORMAT csv, PARTITION_BY k, RETURN_FILES TRUE)"), + (count(5), vec!["DIR/p/k=0/data_0.csv".into(), "DIR/p/k=1/data_0.csv".into()]) + ); + // No rows still write the one file, unless WRITE_EMPTY_FILE says not to. + assert_eq!( + files("COPY (FROM t LIMIT 0) TO 'DIR/e.parquet' (RETURN_FILES 1)"), + (count(0), vec!["DIR/e.parquet".into()]) + ); + assert_eq!( + files("COPY (FROM t LIMIT 0) TO 'DIR/f.parquet' (RETURN_FILES, WRITE_EMPTY_FILE false)"), + (count(0), vec![]) + ); + assert_eq!( + files("COPY (FROM range(5000)) TO 'DIR/r' (FORMAT csv, FILE_SIZE_BYTES 1, RETURN_FILES)"), + (count(5000), ["0", "1", "2"].map(|i| format!("DIR/r/data_{i}.csv")).to_vec()) + ); + let off = db.execute(&format!("COPY t TO '{dir}/g.csv' (RETURN_FILES false)")).expect("copies"); + assert_eq!(off.changes(), Some(5)); + + for (sql, message) in [ + ( + "COPY t TO 'DIR/x.csv' (RETURN_FILES, RETURN_STATS)", + "Binder Error: Can only set one of RETURN_FILES or RETURN_STATS for COPY", + ), + ( + "COPY t TO 'DIR/x.csv' (RETURN_FILES 'x')", + "Invalid Input Error: Copy option \"return_files\" expected an argument of type \ + BOOLEAN - the argument \"x\" of type VARCHAR could not be cast as this type", + ), + ( + "COPY t TO 'DIR/x.csv' (RETURN_STATS)", + "Not implemented Error: RETURN_STATS is not supported for the \"csv\" copy format", + ), + ( + "COPY t TO 'DIR/x.json' (RETURN_STATS)", + "Not implemented Error: RETURN_STATS is not supported for the \"csv\" copy format", + ), + ] { + let sql = sql.replace("DIR", &dir); + let error = db.execute(&sql).expect_err(&sql).to_string(); + assert_eq!(error, message, "{sql}"); + } + let _ = std::fs::remove_dir_all(&root); +}