From be04cd3eaafe3c5ccca5d6b020e7f60535ec76ed Mon Sep 17 00:00:00 2001 From: Tam Nguyen Duc <1218621+tamnd@users.noreply.github.com> Date: Sat, 10 Oct 2026 08:35:05 +0700 Subject: [PATCH] Write COPY TO with PARTITION_BY as hive partitioned directories --- crates/rudb-bind/src/lib.rs | 2 + crates/rudb-bind/src/partition.rs | 278 ++++++++++++++++++++++++++++++ crates/rudb-bind/src/statement.rs | 79 ++++----- crates/rudb/src/database.rs | 9 +- crates/rudb/src/export.rs | 249 ++++++++++++++++++++++++-- crates/rudb/tests/copy_to.rs | 144 ++++++++++++++++ 6 files changed, 687 insertions(+), 74 deletions(-) create mode 100644 crates/rudb-bind/src/partition.rs diff --git a/crates/rudb-bind/src/lib.rs b/crates/rudb-bind/src/lib.rs index 89bf64db0..060f98156 100644 --- a/crates/rudb-bind/src/lib.rs +++ b/crates/rudb-bind/src/lib.rs @@ -38,6 +38,7 @@ mod orderedset; mod ordinality; mod overcall; mod parameters; +mod partition; mod pgcalls; mod pivot; mod query_text; @@ -58,6 +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 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 new file mode 100644 index 000000000..69bdace4d --- /dev/null +++ b/crates/rudb-bind/src/partition.rs @@ -0,0 +1,278 @@ +//! The options of a `COPY ... TO` that say how the rows are split over files, read before the +//! format sees the rest. +//! +//! `PARTITION_BY` names the columns whose values pick a directory, one level a column, the way +//! the pin's hive writer lays them out. The options around it say what to do with what is +//! already there, what the files are called and whether the partition columns are written too. +//! The checks are made in the pin's order: a column the query does not have and a column named +//! twice as the options are read, then the options that cannot go together, then a partition +//! that leaves nothing to write. + +use rudb_common::{Error, Result}; +use rudb_parse::ast; + +/// How a `COPY ... TO` splits its rows over files. With no columns it writes the one file. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct Partitioned { + /// The columns of the query that pick the directory, in the order of the levels. + pub columns: Vec, + /// Whether the partition columns are written into the files as well. + pub write_columns: bool, + /// Whether the files go side by side in the one directory, numbered across all of them, + /// rather than one directory a partition, which is `HIVE_FILE_PATTERN false`. + pub flat: bool, + /// What is done about files already in the directory. + pub existing: Existing, + /// What a file is called, before its extension. + pub pattern: Vec, + /// The extension of a file, after the dot. + pub extension: String, +} + +/// What a partitioned `COPY ... TO` does about the files a directory already holds. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum Existing { + /// Refuses a directory that holds any file. + #[default] + Refuse, + /// Removes every file under the directory first, which is `OVERWRITE`. + Overwrite, + /// Writes over a file of the same name and leaves the others, which is + /// `OVERWRITE_OR_IGNORE`. + Ignore, + /// Adds files of new names, which is `APPEND`. + Append, +} + +/// One piece of a `FILENAME_PATTERN`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum NamePiece { + /// Text written as it is. + Text(String), + /// The number of the file, which is `{i}`. + Offset, + /// A random UUID, which is `{uuid}` or `{uuidv4}`. + Uuid, + /// A time ordered UUID, which is `{uuidv7}`. + Uuid7, +} + +/// The options read here, which the format never sees. +const READ_HERE: [&str; 9] = [ + "partition_by", + "write_partition_columns", + "overwrite", + "overwrite_or_ignore", + "append", + "filename_pattern", + "file_extension", + "hive_file_pattern", + "use_tmp_file", +]; + +/// Reads the partitioning options of `copy` against the columns the query has, and answers +/// them with the copy left holding only the options the format reads. +/// +/// Without `PARTITION_BY`, the options about the files of a directory have nothing to act on +/// and are taken and dropped, as the pin drops them for one file. +pub(crate) fn partitioning( + copy: &ast::CopyTo, + names: &[String], + format: &str, +) -> Result<(Partitioned, ast::CopyTo)> { + let partitioned = copy.options.iter().any(|(name, _)| name == "partition_by"); + let mut out = Partitioned { extension: format.to_string(), ..Partitioned::default() }; + let mut rest = copy.clone(); + rest.options.clear(); + rest.values.clear(); + let (mut overwrite, mut ignore, mut append) = (false, false, false); + let (mut pattern, mut tmp, mut per_thread) = (None, false, false); + for (index, (name, value)) in copy.options.iter().enumerate() { + let here = READ_HERE.contains(&name.as_str()) || partitioned && name == "per_thread_output"; + if !here { + rest.options.push((name.clone(), value.clone())); + if let Some(&expr) = copy.values.get(index) { + rest.values.push(expr); + } + continue; + } + match name.as_str() { + "partition_by" => out.columns = columns(value.as_deref(), names)?, + "write_partition_columns" => out.write_columns = boolean(name, value.as_deref())?, + "overwrite" => overwrite = boolean(name, value.as_deref())?, + "overwrite_or_ignore" => ignore = boolean(name, value.as_deref())?, + "append" => append = boolean(name, value.as_deref())?, + "filename_pattern" => pattern = Some(value.clone().unwrap_or_default()), + "file_extension" => out.extension = value.clone().unwrap_or_default(), + "hive_file_pattern" => out.flat = !boolean(name, value.as_deref())?, + "use_tmp_file" => tmp = boolean(name, value.as_deref())?, + _ => per_thread = boolean(name, value.as_deref())?, + } + } + if [overwrite, ignore, append].iter().filter(|&&set| set).count() > 1 { + return Err(Error::binder("Can only set one of OVERWRITE_OR_IGNORE, OVERWRITE or APPEND")); + } + out.existing = if overwrite { + Existing::Overwrite + } else if ignore { + Existing::Ignore + } else if append { + Existing::Append + } else { + Existing::Refuse + }; + out.pattern = match pattern { + Some(pattern) => pieces(&pattern), + None if append => vec![NamePiece::Uuid], + None => vec![NamePiece::Text("data_".to_string()), NamePiece::Offset], + }; + if append + && !out.pattern.iter().any(|piece| matches!(piece, NamePiece::Uuid | NamePiece::Uuid7)) + { + return Err(Error::binder("APPEND mode requires a {uuid} label in filename_pattern")); + } + if !partitioned { + return Ok((out, rest)); + } + if tmp { + return Err(Error::not_implemented( + "Can't combine USE_TMP_FILE and PARTITIONED BY for COPY", + )); + } + if per_thread { + return Err(Error::not_implemented( + "Can't combine PER_THREAD_OUTPUT and PARTITIONED BY for COPY", + )); + } + if !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. \ + WRITE_PARTITION_COLUMNS option can be used to write partition columns.", + )); + } + Ok((out, rest)) +} + +/// The columns a `PARTITION_BY` names, as a list in parentheses or a single name, each bare or +/// quoted either way, matched without regard to case, where `*` is every column. +fn columns(written: Option<&str>, names: &[String]) -> Result> { + let written = written.unwrap_or_default().trim(); + let list = written.strip_prefix('(').and_then(|rest| rest.strip_suffix(')')).unwrap_or(written); + let mut out: Vec = Vec::new(); + for column in split(list) { + let column = column.trim(); + if column == "*" { + return Ok((0..names.len()).collect()); + } + let unquoted = ['"', '\''] + .iter() + .find_map(|"e| column.strip_prefix(quote).and_then(|rest| rest.strip_suffix(quote))) + .unwrap_or(column); + let Some(index) = names.iter().position(|name| name.eq_ignore_ascii_case(unquoted)) else { + return Err(Error::binder(format!( + "\"partition_by\" expected to find {unquoted}, but it was not found in the table" + ))); + }; + if out.contains(&index) { + return Err(Error::binder(format!( + "\"partition_by\" does now allow duplicate columns (found: {unquoted})" + ))); + } + out.push(index); + } + Ok(out) +} + +/// Splits a list on the commas that are not inside quotes. +fn split(list: &str) -> Vec<&str> { + let mut out = Vec::new(); + let (mut start, mut quote) = (0, None); + for (at, character) in list.char_indices() { + match (character, quote) { + ('"' | '\'', None) => quote = Some(character), + (_, Some(open)) if character == open => quote = None, + (',', None) => { + out.push(&list[start..at]); + start = at + 1; + } + _ => {} + } + } + out.push(&list[start..]); + out +} + +/// An option that is on or off, where an option with no value is on. +fn boolean(name: &str, written: Option<&str>) -> Result { + let Some(written) = written else { return Ok(true) }; + match written.trim_matches('\'').to_ascii_lowercase().as_str() { + "true" | "t" | "1" | "on" | "y" | "yes" => Ok(true), + "false" | "f" | "0" | "off" | "n" | "no" => Ok(false), + _ => Err(Error::binder(format!("\"{name}\" expects a boolean value, not {written}"))), + } +} + +/// The pieces of a `FILENAME_PATTERN`, where a pattern of text alone has the number of the file +/// put after it. +fn pieces(pattern: &str) -> Vec { + const LABELS: [(&str, NamePiece); 4] = [ + ("{i}", NamePiece::Offset), + ("{uuid}", NamePiece::Uuid), + ("{uuidv4}", NamePiece::Uuid), + ("{uuidv7}", NamePiece::Uuid7), + ]; + let mut out = Vec::new(); + let (mut text, mut rest) = (String::new(), pattern); + while let Some(character) = rest.chars().next() { + if let Some((label, piece)) = LABELS.iter().find(|(label, _)| rest.starts_with(label)) { + if !text.is_empty() { + out.push(NamePiece::Text(std::mem::take(&mut text))); + } + out.push(piece.clone()); + rest = &rest[label.len()..]; + } else { + text.push(character); + rest = &rest[character.len_utf8()..]; + } + } + if !text.is_empty() { + out.push(NamePiece::Text(text)); + } + if matches!(out.as_slice(), [NamePiece::Text(_)]) { + out.push(NamePiece::Offset); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_pattern_of_text_alone_is_numbered_and_the_labels_are_read() { + assert_eq!(pieces("foo"), [NamePiece::Text("foo".into()), NamePiece::Offset]); + assert_eq!( + pieces("a{uuid}_{i}{x}"), + [ + NamePiece::Text("a".into()), + NamePiece::Uuid, + NamePiece::Text("_".into()), + NamePiece::Offset, + NamePiece::Text("{x}".into()), + ] + ); + assert_eq!(pieces("{uuidv7}"), [NamePiece::Uuid7]); + } + + #[test] + fn the_columns_are_found_in_any_case_and_any_quotes() { + let names = ["id".to_string(), "K".to_string(), "a b".to_string()]; + assert_eq!(columns(Some("(k, \"a b\")"), &names).unwrap(), [1, 2]); + assert_eq!(columns(Some("('id','k')"), &names).unwrap(), [0, 1]); + assert_eq!(columns(Some("*"), &names).unwrap(), [0, 1, 2]); + let missing = columns(Some("zz"), &names).unwrap_err().to_string(); + assert!(missing.contains("expected to find zz"), "{missing}"); + let twice = columns(Some("(k, K)"), &names).unwrap_err().to_string(); + assert!(twice.contains("(found: K)"), "{twice}"); + } +} diff --git a/crates/rudb-bind/src/statement.rs b/crates/rudb-bind/src/statement.rs index 022f0918e..7fb973eee 100644 --- a/crates/rudb-bind/src/statement.rs +++ b/crates/rudb-bind/src/statement.rs @@ -169,6 +169,8 @@ pub struct CopyTo { pub date_format: Option, /// For JSON, the `strftime` format a timestamp of any precision or zone is written with. pub timestamp_format: Option, + /// How the rows are split over files, which by default they are not. + pub partitioned: crate::Partitioned, } /// A bound `SET` or `RESET`. @@ -850,10 +852,11 @@ pub(crate) fn bind_one( ast::Statement::CopyTo(index) => { let copy = &ast.copies[index as usize]; let mut binder = Binder::with(catalog, parameters, session); - let (root, _) = binder.bind_query(ast, copy.query)?; + let (root, scope) = binder.bind_query(ast, copy.query)?; let plan = finish(binder, root)?; let typed = copy_values(ast, copy, catalog, parameters, session); - copy_to(copy, &typed, plan).map(Bound::CopyTo) + let names = scope.columns.into_iter().map(|column| column.name).collect::>(); + copy_to(copy, &typed, &names, plan).map(Bound::CopyTo) } ast::Statement::Vacuum(index) => { crate::vacuum::vacuum(ast, &ast.vacuums[index as usize], catalog, session) @@ -952,7 +955,8 @@ fn refuse_written(copy: &ast::CopyTo, typed: &[Option], format: &str) - let Some(written) = written else { continue }; // A bare NULL for `HEADER` is refused here for every format, before JSON says it does // not know the option. - if written.null() && (!json || !bare(written) || name == "header") { + if written.null() && (!json || !bare(written) || name == "header") && name != "partition_by" + { return Err(Error::binder(format!( "NULL is not supported as a valid option for COPY option \"{name}\"" ))); @@ -1000,7 +1004,12 @@ fn refuse_written(copy: &ast::CopyTo, typed: &[Option], format: &str) - /// CSV, JSON and Parquet are written, the format picked by the file's extension unless `FORMAT` /// names one. An option the pin takes and this does not is /// refused by name, and one the pin does not take either gets the first line of its refusal. -fn copy_to(copy: &ast::CopyTo, typed: &[Option], plan: Plan) -> Result { +fn copy_to( + copy: &ast::CopyTo, + typed: &[Option], + names: &[String], + plan: Plan, +) -> Result { let lowered = copy.path.to_ascii_lowercase(); let mut format = if lowered.ends_with(".parquet") { "parquet" @@ -1017,11 +1026,18 @@ fn copy_to(copy: &ast::CopyTo, typed: &[Option], plan: Plan) -> Result< return Err(Error::catalog(format!("Copy Function with name {format} does not exist!"))); } refuse_written(copy, typed, &format)?; - match format.as_str() { - "json" => return json_to(copy, plan), - "parquet" => return parquet_to(copy, plan), - _ => {} - } + let (partitioned, rest) = crate::partition::partitioning(copy, names, &format)?; + let mut out = match format.as_str() { + "json" => json_to(&rest, plan)?, + "parquet" => parquet_to(&rest, plan)?, + _ => csv_to(&rest, plan)?, + }; + out.partitioned = partitioned; + Ok(out) +} + +/// Reads the options of a `COPY ... TO` a CSV file. +fn csv_to(copy: &ast::CopyTo, plan: Plan) -> Result { let mut out = CopyTo { plan, path: copy.path.clone(), @@ -1039,6 +1055,7 @@ fn copy_to(copy: &ast::CopyTo, typed: &[Option], plan: Plan) -> Result< row_group_size: 0, date_format: None, timestamp_format: None, + partitioned: crate::Partitioned::default(), }; for (name, value) in ©.options { let text = || { @@ -1082,26 +1099,9 @@ fn copy_to(copy: &ast::CopyTo, typed: &[Option], plan: Plan) -> Result< } } } - "compression" - | "dateformat" - | "date_format" - | "timestampformat" - | "timestamp_format" - | "new_line" - | "prefix" - | "suffix" - | "per_thread_output" - | "file_size_bytes" - | "partition_by" - | "overwrite" - | "overwrite_or_ignore" - | "filename_pattern" - | "file_extension" - | "use_tmp_file" - | "return_files" - | "write_partition_columns" - | "preserve_order" - | "force_not_null" + "compression" | "dateformat" | "date_format" | "timestampformat" + | "timestamp_format" | "new_line" | "prefix" | "suffix" | "per_thread_output" + | "file_size_bytes" | "return_files" | "preserve_order" | "force_not_null" | "encoding" => { return Err(Error::not_implemented(format!( "COPY TO with the option {name} is not supported yet" @@ -1140,6 +1140,7 @@ fn json_to(copy: &ast::CopyTo, plan: Plan) -> Result { row_group_size: 0, date_format: None, timestamp_format: None, + partitioned: crate::Partitioned::default(), }; for (name, value) in ©.options { match name.as_str() { @@ -1165,17 +1166,7 @@ fn json_to(copy: &ast::CopyTo, plan: Plan) -> Result { "Option \"encoding\" is not supported for writing - only for reading", )); } - "compression" - | "per_thread_output" - | "file_size_bytes" - | "partition_by" - | "overwrite" - | "overwrite_or_ignore" - | "filename_pattern" - | "file_extension" - | "use_tmp_file" - | "return_files" - | "write_partition_columns" + "compression" | "per_thread_output" | "file_size_bytes" | "return_files" | "preserve_order" => { return Err(Error::not_implemented(format!( "COPY TO with the option {name} is not supported yet" @@ -1234,6 +1225,7 @@ fn parquet_to(copy: &ast::CopyTo, plan: Plan) -> Result { row_group_size: 122_880, date_format: None, timestamp_format: None, + partitioned: crate::Partitioned::default(), }; for (name, value) in ©.options { let written = value.as_deref().unwrap_or_default().trim_matches('\''); @@ -1278,14 +1270,7 @@ fn parquet_to(copy: &ast::CopyTo, plan: Plan) -> Result { | "geoparquet_version" | "per_thread_output" | "file_size_bytes" - | "partition_by" - | "overwrite" - | "overwrite_or_ignore" - | "filename_pattern" - | "file_extension" - | "use_tmp_file" | "return_files" - | "write_partition_columns" | "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 6e96f68aa..2ecf3a687 100644 --- a/crates/rudb/src/database.rs +++ b/crates/rudb/src/database.rs @@ -7149,14 +7149,7 @@ impl Shared { .after(Planning { parse_ns, bind_ns, rewrite_ns, optimize_ns }); let result = run(sql, ©.plan, &catalog, cancel, under)?; let zone = session.session_time_zone(); - let rows = if copy.parquet { - crate::export::write_parquet(©, &result)? - } else if copy.json { - crate::export::write_json(©, &result, zone)? - } else { - crate::export::write_csv(©, &result, zone)? - }; - QueryResult::changed(rows) + QueryResult::changed(crate::export::write(©, &result, zone)?) } 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 fbb52e19e..4256a2221 100644 --- a/crates/rudb/src/export.rs +++ b/crates/rudb/src/export.rs @@ -15,27 +15,239 @@ //! //! A Parquet file is written by `rudb-parquet`, a row group at a time, with the columns cast first //! to the types that crate stores: an enum as its text, a coarse timestamp in microseconds. +//! +//! With `PARTITION_BY`, the rows are split by the values of the partition columns into a directory +//! a value, laid out the way Hive lays them out, with a file of the chosen format in each. +use std::cmp::Ordering; +use std::collections::HashMap; use std::fs::File; use std::io::{BufWriter, Write}; +use std::path::Path; -use rudb_bind::CopyTo; -use rudb_common::{Error, Field, LogicalType, Result, SessionTimeZone, Value}; +use rudb_bind::{CopyTo, Existing, NamePiece}; +use rudb_common::{Error, Field, LogicalType, Memory, Result, SessionTimeZone, Value}; use rudb_compress::Codec; use rudb_kernels::cast::{cast_in_time_zone, cast_value}; +use rudb_kernels::compare::order_with_nulls; use rudb_kernels::strftime::Format; use rudb_vector::{Chunk, Vector}; use crate::QueryResult; -/// Writes the rows of `result` to the file `copy` names and answers how many there were. -pub(crate) fn write_csv( +/// Writes the rows of `result` the way `copy` asks, to the one file or split over a directory, +/// and answers how many there were. +pub(crate) fn write(copy: &CopyTo, result: &QueryResult, zone: SessionTimeZone) -> Result { + if copy.partitioned.columns.is_empty() { + write_file(copy, ©.path, result, zone) + } else { + write_partitioned(copy, result, zone) + } +} + +/// Writes the rows of `result` to the one file at `path`, in the format `copy` names. +fn write_file( copy: &CopyTo, + path: &str, result: &QueryResult, zone: SessionTimeZone, ) -> Result { - let file = File::create(©.path) - .map_err(|error| Error::io(format!("Cannot open file \"{}\": {error}", copy.path)))?; + if copy.parquet { + write_parquet(copy, path, result) + } else if copy.json { + write_json(copy, path, result, zone) + } else { + write_csv(copy, path, result, zone) + } +} + +/// What a directory is called for a NULL partition value, which is the name Hive gave it. +const NULL_PARTITION: &str = "__HIVE_DEFAULT_PARTITION__"; + +/// Writes the rows of `result` split by the values of the partition columns, a directory a +/// column named `column=value`, both percent encoded, with a file in each of the deepest ones. +/// +/// The rows of a partition keep their order. A partition's file is numbered 0 in its own +/// directory, or, with the files side by side, by the order of the partition values, a column at +/// a time with a NULL last, which is how the pin numbers them. +fn write_partitioned(copy: &CopyTo, result: &QueryResult, zone: SessionTimeZone) -> Result { + let split = ©.partitioned; + let root = Path::new(©.path); + match split.existing { + Existing::Refuse if holds_files(root) => { + return Err(Error::io(format!( + "Directory \"{}\" is not empty! Enable OVERWRITE option to overwrite files", + copy.path + ))); + } + Existing::Overwrite => remove_files(root)?, + _ => {} + } + let made = |path: &Path| { + std::fs::create_dir_all(path).map_err(|error| { + Error::io(format!("Failed to create directory \"{}\": {error}", path.display())) + }) + }; + made(root)?; + let names = result.names(); + let chunks = result.chunks(); + let mut found: HashMap = HashMap::new(); + let mut partitions: Vec<(String, Vec, Vec>)> = Vec::new(); + for (at, chunk) in chunks.iter().enumerate() { + let mut keys = Vec::with_capacity(split.columns.len()); + for &column in &split.columns { + let vector = chunk.column(column)?; + keys.push(cast_in_time_zone(vector, &LogicalType::Varchar, false, Some(zone))?); + } + for row in 0..chunk.len() { + let mut directory = String::new(); + for (key, &column) in keys.iter().zip(&split.columns) { + if !directory.is_empty() { + directory.push('/'); + } + encode(&mut directory, &names[column]); + directory.push('='); + match key.try_value_at(row)? { + Value::Null => directory.push_str(NULL_PARTITION), + Value::Varchar(text) => encode_value(&mut directory, &text), + other => encode_value(&mut directory, &other.to_string()), + } + } + let partition = match found.get(&directory) { + Some(&partition) => partition, + None => { + let values = split + .columns + .iter() + .map(|&column| chunk.column(column)?.try_value_at(row)) + .collect::>>()?; + found.insert(directory.clone(), partitions.len()); + partitions.push((directory, values, vec![Vec::new(); chunks.len()])); + partitions.len() - 1 + } + }; + partitions[partition].2[at].push(u32::try_from(row).unwrap_or(u32::MAX)); + } + } + let kept = (0..names.len()) + .filter(|column| split.write_columns || !split.columns.contains(column)) + .collect::>(); + let kept_names = kept.iter().map(|&column| names[column].clone()).collect::>(); + let kept_types = kept.iter().map(|&column| result.types()[column].clone()).collect::>(); + if split.flat { + partitions.sort_by(|(_, left, _), (_, right, _)| { + left.iter() + .zip(right) + .map(|(left, right)| { + order_with_nulls(left, right, false).unwrap_or(Ordering::Equal) + }) + .find(|&order| order != Ordering::Equal) + .unwrap_or(Ordering::Equal) + }); + } + let mut rows = 0; + for (number, (directory, _, picked)) in partitions.iter().enumerate() { + let mut parts = Vec::new(); + for (chunk, rows) in chunks.iter().zip(picked) { + if rows.is_empty() { + continue; + } + let columns = kept + .iter() + .map(|&column| chunk.column(column)?.gather(rows)) + .collect::>>()?; + parts.push(Chunk::new(columns)?); + } + let part = QueryResult::new( + kept_names.clone(), + kept_types.clone(), + parts, + Memory::unlimited().reservation(), + ); + let (place, number) = + if split.flat { (root.to_path_buf(), number) } else { (root.join(directory), 0) }; + made(&place)?; + let name = format!("{}.{}", file_name(&split.pattern, number, zone)?, split.extension); + rows += write_file(copy, &place.join(name).to_string_lossy(), &part, zone)?; + } + Ok(rows) +} + +/// Whether a directory holds a file at any depth. +fn holds_files(directory: &Path) -> bool { + let Ok(entries) = std::fs::read_dir(directory) else { return false }; + entries.flatten().any(|entry| { + entry.file_type().is_ok_and(|kind| !kind.is_dir()) || holds_files(&entry.path()) + }) +} + +/// Removes every file under a directory and leaves the directories, which is what the pin's +/// `OVERWRITE` does. +fn remove_files(directory: &Path) -> Result<()> { + let Ok(entries) = std::fs::read_dir(directory) else { return Ok(()) }; + for entry in entries.flatten() { + let path = entry.path(); + if entry.file_type().is_ok_and(|kind| kind.is_dir()) { + remove_files(&path)?; + } else { + std::fs::remove_file(&path).map_err(|error| { + Error::io(format!("Could not remove file \"{}\": {error}", path.display())) + })?; + } + } + Ok(()) +} + +/// Appends `text` percent encoded, where everything but a letter, a digit and `_-~.` is a `%` +/// and two upper case hex digits a byte. +fn encode(out: &mut String, text: &str) { + for byte in text.bytes() { + if byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'~' | b'.') { + out.push(char::from(byte)); + } else { + out.push_str(&format!("%{byte:02X}")); + } + } +} + +/// Appends a partition value percent encoded, where a value that would read as the NULL +/// directory has its first character encoded as well so that it gets a directory of its own. +fn encode_value(out: &mut String, text: &str) { + let start = out.len(); + encode(out, text); + if out[start..].eq_ignore_ascii_case(NULL_PARTITION) { + let first = out.as_bytes()[start]; + out.replace_range(start..=start, &format!("%{first:02X}")); + } +} + +/// The name of a file before its extension, from the pieces of a `FILENAME_PATTERN`. +fn file_name(pattern: &[NamePiece], number: usize, zone: SessionTimeZone) -> Result { + let mut out = String::new(); + for piece in pattern { + match piece { + NamePiece::Text(text) => out.push_str(text), + NamePiece::Offset => out.push_str(&number.to_string()), + NamePiece::Uuid | NamePiece::Uuid7 => { + let kind = if *piece == NamePiece::Uuid { "uuidv4" } else { "uuidv7" }; + let drawn = rudb_kernels::drawn(kind, 1)?; + let value = drawn.try_value_at(0)?; + out.push_str(&text(&value, &LogicalType::Uuid, zone)?.unwrap_or_default()); + } + } + } + Ok(out) +} + +/// Writes the rows of `result` to the CSV file at `path` and answers how many there were. +fn write_csv( + copy: &CopyTo, + path: &str, + result: &QueryResult, + zone: SessionTimeZone, +) -> Result { + let file = File::create(path) + .map_err(|error| Error::io(format!("Cannot open file \"{path}\": {error}")))?; let mut out = BufWriter::new(file); let names = result.names(); let forced = names @@ -52,9 +264,8 @@ pub(crate) fn write_csv( ))); } } - let written = |error: std::io::Error| { - Error::io(format!("Could not write file \"{}\": {error}", copy.path)) - }; + let written = + |error: std::io::Error| Error::io(format!("Could not write file \"{path}\": {error}")); let mut line = String::new(); if copy.header { for (column, name) in names.iter().enumerate() { @@ -136,18 +347,18 @@ fn field(line: &mut String, text: &str, forced: bool, copy: &CopyTo) { line.push_str(©.quote); } -/// Writes the rows of `result` to the JSON file `copy` names and answers how many there were. -pub(crate) fn write_json( +/// Writes the rows of `result` to the JSON file at `path` and answers how many there were. +fn write_json( copy: &CopyTo, + path: &str, result: &QueryResult, zone: SessionTimeZone, ) -> Result { - let file = File::create(©.path) - .map_err(|error| Error::io(format!("Cannot open file \"{}\": {error}", copy.path)))?; + let file = File::create(path) + .map_err(|error| Error::io(format!("Cannot open file \"{path}\": {error}")))?; let mut out = BufWriter::new(file); - let written = |error: std::io::Error| { - Error::io(format!("Could not write file \"{}\": {error}", copy.path)) - }; + let written = + |error: std::io::Error| Error::io(format!("Could not write file \"{path}\": {error}")); let moments = Moments { zone, date: copy.date_format.as_deref().map(Format::parse).transpose()?, @@ -201,9 +412,9 @@ pub(crate) fn write_json( /// Writes a result as a Parquet file, a row group every `ROW_GROUP_SIZE` rows or so, since a /// group ends at the end of the chunk that fills it. -pub(crate) fn write_parquet(copy: &CopyTo, result: &QueryResult) -> Result { - let file = File::create(©.path) - .map_err(|error| Error::io(format!("Cannot open file \"{}\": {error}", copy.path)))?; +fn write_parquet(copy: &CopyTo, path: &str, result: &QueryResult) -> Result { + let file = File::create(path) + .map_err(|error| Error::io(format!("Cannot open file \"{path}\": {error}")))?; let types = result.types(); let mut stored = Vec::with_capacity(types.len()); let mut fields = Vec::with_capacity(types.len()); diff --git a/crates/rudb/tests/copy_to.rs b/crates/rudb/tests/copy_to.rs index 6f6296e2c..0bbb2573c 100644 --- a/crates/rudb/tests/copy_to.rs +++ b/crates/rudb/tests/copy_to.rs @@ -469,3 +469,147 @@ fn an_exported_state_is_written_as_its_layout_and_reads_back_into_a_state() { assert_eq!(got, vec![vec![text("42"), text("42")]]); let _ = std::fs::remove_file(&path); } + +/// Every file under `root`, as its path below it and what it holds, in path order. +fn files(root: &std::path::Path) -> String { + fn walk(directory: &std::path::Path, out: &mut Vec) { + for entry in std::fs::read_dir(directory).expect("lists the directory").flatten() { + if entry.file_type().expect("has a type").is_dir() { + walk(&entry.path(), out); + } else { + out.push(entry.path()); + } + } + } + let mut found = Vec::new(); + walk(root, &mut found); + found.sort(); + let mut out = String::new(); + for path in found { + let name = path.strip_prefix(root).expect("is below the root").display().to_string(); + let text = std::fs::read_to_string(&path).expect("reads the file"); + out.push_str(&format!("== {name}\n{text}")); + } + out +} + +#[test] +fn partition_by_writes_a_directory_a_value_the_way_the_pin_lays_them_out() { + let root = std::env::temp_dir().join(format!("rudb-copy-to-hive-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + let dir = root.display().to_string(); + let db = Database::new(); + db.execute( + "CREATE TABLE p AS SELECT * FROM (VALUES (1, 'x', DATE '2020-01-01'), (2, 'a b', NULL), \ + (3, 'x', DATE '2020-01-01'), (4, '__hive_default_partition__', DATE '2020-01-02')) \ + v(id, k, d)", + ) + .expect("creates"); + let run = |sql: &str| db.execute(&sql.replace("DIR", &dir)); + let copied = run("COPY p TO 'DIR/a' (FORMAT csv, PARTITION_BY (k, d))").expect("copies"); + assert_eq!(copied.changes(), Some(4)); + assert_eq!( + files(&root.join("a")), + "== k=%5F_hive_default_partition__/d=2020-01-02/data_0.csv\nid\n4\n\ + == k=a%20b/d=__HIVE_DEFAULT_PARTITION__/data_0.csv\nid\n2\n\ + == k=x/d=2020-01-01/data_0.csv\nid\n1\n3\n" + ); + let again = run("COPY p TO 'DIR/a' (FORMAT csv, PARTITION_BY (k, d))").unwrap_err(); + assert_eq!( + again.to_string(), + format!( + "IO Error: Directory \"{dir}/a\" is not empty! Enable OVERWRITE option to overwrite \ + files" + ) + ); + run("COPY (SELECT 9 AS id, 'x' AS k, DATE '2020-01-01' AS d) TO 'DIR/a' \ + (FORMAT csv, PARTITION_BY (k, d), OVERWRITE_OR_IGNORE)") + .expect("writes over the one file"); + let read = "SELECT string_agg(id || ':' || k || ':' || coalesce(d::VARCHAR, '-'), ' ' \ + ORDER BY id) FROM read_csv('DIR/a/*/*/*.csv', hive_partitioning = true)"; + let read = db.query(&read.replace("DIR", &dir)).expect("reads back"); + assert_eq!( + read.value_at(0, 0), + rudb_common::Value::Varchar( + "2:a b:- 4:__hive_default_partition__:2020-01-02 9:x:2020-01-01".into() + ) + ); + + // Side by side, the files are numbered by the order of the values. + run("COPY p TO 'DIR/b' (FORMAT csv, PARTITION_BY k, HIVE_FILE_PATTERN false, \ + FILENAME_PATTERN 'out_{i}', FILE_EXTENSION 'txt', WRITE_PARTITION_COLUMNS)") + .expect("copies"); + assert_eq!( + files(&root.join("b")), + "== out_0.txt\nid,k,d\n4,__hive_default_partition__,2020-01-02\n\ + == out_1.txt\nid,k,d\n2,a b,\n\ + == out_2.txt\nid,k,d\n1,x,2020-01-01\n3,x,2020-01-01\n" + ); + run("COPY (SELECT * FROM (VALUES (1, 30), (2, 10), (3, 20), (4, 5)) v(id, k)) TO 'DIR/n' \ + (FORMAT csv, PARTITION_BY 'k', HIVE_FILE_PATTERN false)") + .expect("copies"); + assert_eq!( + files(&root.join("n")), + "== data_0.csv\nid\n4\n== data_1.csv\nid\n2\n== data_2.csv\nid\n3\n== data_3.csv\nid\n1\n" + ); + + // OVERWRITE takes away every file that was there. + run("COPY p TO 'DIR/c' (FORMAT json, PARTITION_BY (d))").expect("copies"); + run("COPY (SELECT 1 AS id, 'q' AS k) TO 'DIR/c' (FORMAT json, PARTITION_BY (k), OVERWRITE)") + .expect("copies"); + assert_eq!(files(&root.join("c")), "== k=q/data_0.json\n{\"id\":1}\n"); + + // APPEND names its files with a UUID, so a second copy adds to the first. + for _ in 0..2 { + run("COPY p TO 'DIR/e' (FORMAT parquet, PARTITION_BY (k), APPEND)").expect("appends"); + } + let count = db + .query(&format!("SELECT count(*) FROM read_parquet('{dir}/e/*/*.parquet')")) + .expect("reads back"); + assert_eq!(count.value_at(0, 0), rudb_common::Value::BigInt(8)); + let _ = std::fs::remove_dir_all(&root); +} + +#[test] +fn a_partitioned_copy_refuses_what_the_pin_refuses() { + let db = Database::new(); + db.execute("CREATE TABLE p AS SELECT 1 AS id, 'x' AS k").expect("creates"); + for (sql, want) in [ + ( + "COPY p TO 'e' (PARTITION_BY (k), APPEND, FILENAME_PATTERN 'f')", + "Binder Error: APPEND mode requires a {uuid} label in filename_pattern", + ), + ( + "COPY p TO 'e.csv' (APPEND, FILENAME_PATTERN 'x')", + "Binder Error: APPEND mode requires a {uuid} label in filename_pattern", + ), + ( + "COPY p TO 'e' (PARTITION_BY (k), APPEND, OVERWRITE)", + "Binder Error: Can only set one of OVERWRITE_OR_IGNORE, OVERWRITE or APPEND", + ), + ( + "COPY p TO 'e' (PARTITION_BY (k), USE_TMP_FILE)", + "Not implemented Error: Can't combine USE_TMP_FILE and PARTITIONED BY for COPY", + ), + ( + "COPY p TO 'e' (PARTITION_BY (k), PER_THREAD_OUTPUT)", + "Not implemented Error: Can't combine PER_THREAD_OUTPUT and PARTITIONED BY for COPY", + ), + ( + "COPY p TO 'e' (PARTITION_BY (zz))", + "Binder Error: \"partition_by\" expected to find zz, but it was not found in the table", + ), + ( + "COPY p TO 'e' (PARTITION_BY (k, K))", + "Binder Error: \"partition_by\" does now allow duplicate columns (found: K)", + ), + ( + "COPY p TO 'e' (PARTITION_BY (k, id))", + "Not implemented Error: No column to write as all columns are specified as partition \ + columns. WRITE_PARTITION_COLUMNS option can be used to write partition columns.", + ), + ] { + let error = db.execute(sql).expect_err(sql); + assert_eq!(error.to_string(), want, "{sql}"); + } +}