From 347a1950612c51c564326ac970021a18d1a1b0c0 Mon Sep 17 00:00:00 2001 From: Tam Nguyen Duc <1218621+tamnd@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:30:03 +0700 Subject: [PATCH] Answer COPY TO with RETURN_STATS as the rows, sizes and column statistics of each Parquet file --- crates/rudb-bind/src/partition.rs | 5 - crates/rudb-parquet/src/lib.rs | 2 +- crates/rudb-parquet/src/write.rs | 73 +++++++++- crates/rudb/src/database.rs | 11 +- crates/rudb/src/export.rs | 214 +++++++++++++++++++++++++++--- crates/rudb/tests/copy_to.rs | 96 ++++++++++++++ 6 files changed, 373 insertions(+), 28 deletions(-) diff --git a/crates/rudb-bind/src/partition.rs b/crates/rudb-bind/src/partition.rs index 5a246e630..7cd2a331e 100644 --- a/crates/rudb-bind/src/partition.rs +++ b/crates/rudb-bind/src/partition.rs @@ -223,11 +223,6 @@ pub(crate) fn partitioning( "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. \ diff --git a/crates/rudb-parquet/src/lib.rs b/crates/rudb-parquet/src/lib.rs index 5983cb127..e9c76a72e 100644 --- a/crates/rudb-parquet/src/lib.rs +++ b/crates/rudb-parquet/src/lib.rs @@ -51,4 +51,4 @@ pub use metadata::{ pub use page::{Body, DataV1, DataV2, Dictionary, Header}; pub use prune::{Bound, Footer, Op, Test, skips}; pub use reader::{Reader, read}; -pub use write::{Writer, storage}; +pub use write::{ColumnStatistics, Writer, storage}; diff --git a/crates/rudb-parquet/src/write.rs b/crates/rudb-parquet/src/write.rs index ef7280d90..c3a9af5cc 100644 --- a/crates/rudb-parquet/src/write.rs +++ b/crates/rudb-parquet/src/write.rs @@ -38,6 +38,7 @@ struct Column { struct Written { values: i64, nulls: i64, + nans: i64, compressed: i64, uncompressed: i64, min: Option, @@ -54,6 +55,27 @@ pub struct Writer { groups: Vec, rows: i64, created_by: String, + totals: Vec, +} + +/// What one column of a file holds over all of its row groups, which is what a `COPY ... TO` +/// with `RETURN_STATS` reports. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct ColumnStatistics { + /// The name of the column. + pub name: String, + /// How many values it holds, the nulls among them. + pub values: i64, + /// How many of them are null. + pub nulls: i64, + /// How many are NaN, for a `FLOAT` or a `DOUBLE` column, and `None` for any other. + pub nans: Option, + /// How many bytes its column chunks take in the file, page headers and all. + pub size: i64, + /// The smallest value that is neither null nor NaN, of the type it was stored as. + pub min: Option, + /// The largest. + pub max: Option, } /// The type a column of `ty` has to be cast to before it is handed to a [`Writer`]. @@ -147,6 +169,14 @@ impl Writer { })?; columns.push(Column { name: field.name.clone(), ty: field.ty.clone(), physical }); } + let totals = columns + .iter() + .map(|column| ColumnStatistics { + name: column.name.clone(), + nans: matches!(column.physical, Physical::Float | Physical::Double).then_some(0), + ..ColumnStatistics::default() + }) + .collect(); let mut writer = Self { out, at: 0, @@ -155,6 +185,7 @@ impl Writer { groups: Vec::new(), rows: 0, created_by: created_by.to_string(), + totals, }; writer.put(b"PAR1")?; Ok(writer) @@ -193,6 +224,17 @@ impl Writer { total += written.uncompressed; compressed += written.compressed; columns.push(self.chunk_meta(column, first, &written)); + let totals = &mut self.totals[column]; + totals.values += written.values; + totals.nulls += written.nulls; + totals.nans = totals.nans.map(|nans| nans + written.nans); + totals.size += written.compressed; + if let Some(min) = &written.min { + bound(&mut totals.min, min, Ordering::Less); + } + if let Some(max) = &written.max { + bound(&mut totals.max, max, Ordering::Greater); + } } let mut group = thrift::Writer::default(); group.list_of_nested(1, columns); @@ -220,6 +262,11 @@ impl Writer { written.nulls += 1; continue; } + if matches!(value, Value::Double(value) if value.is_nan()) + || matches!(value, Value::Float(value) if value.is_nan()) + { + written.nans += 1; + } defined.push(true); if column.physical == Physical::Boolean { bits.push(matches!(value, Value::Boolean(true))); @@ -292,12 +339,34 @@ impl Writer { chunk } + /// What each column holds over the row groups written so far, in the order of the columns. + #[must_use] + pub fn statistics(&self) -> &[ColumnStatistics] { + &self.totals + } + + /// How many row groups have been written so far. + #[must_use] + pub fn row_groups(&self) -> usize { + self.groups.len() + } + /// Writes the footer and hands back what the file was written to. /// /// # Errors /// /// If the write fails. - pub fn finish(mut self) -> Result { + pub fn finish(self) -> Result { + Ok(self.finish_sized()?.0) + } + + /// Writes the footer and hands back what the file was written to and how many bytes the + /// footer's metadata took, which leaves out its length and the closing magic. + /// + /// # Errors + /// + /// If the write fails. + pub fn finish_sized(mut self) -> Result<(W, u64)> { let mut schema = Vec::with_capacity(self.columns.len() + 1); let mut root = thrift::Writer::default(); root.i32(3, 0); @@ -320,7 +389,7 @@ impl Writer { self.out .flush() .map_err(|error| Error::io(format!("could not write a parquet file: {error}")))?; - Ok(self.out) + Ok((self.out, footer.len() as u64)) } } diff --git a/crates/rudb/src/database.rs b/crates/rudb/src/database.rs index 17ec5de05..715e7715c 100644 --- a/crates/rudb/src/database.rs +++ b/crates/rudb/src/database.rs @@ -7152,10 +7152,13 @@ impl Shared { let zone = session.session_time_zone(); let threads = self.inner.pool.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) + match copy.partitioned.returns { + Returns::Count => QueryResult::changed(rows), + Returns::Files => QueryResult::copied_files( + rows, + files.into_iter().map(|file| file.path).collect(), + ), + Returns::Stats => crate::export::statistics(files), } } Bound::Setting(setting) diff --git a/crates/rudb/src/export.rs b/crates/rudb/src/export.rs index 9b38084d5..6b5bb57db 100644 --- a/crates/rudb/src/export.rs +++ b/crates/rudb/src/export.rs @@ -45,14 +45,30 @@ use crate::QueryResult; /// only between two of them. const BATCH: usize = 2048; +/// One file a `COPY ... TO` wrote, with what `RETURN_FILES` and `RETURN_STATS` report of it. +pub(crate) struct Wrote { + pub(crate) path: String, + pub(crate) rows: usize, + /// How many bytes the file holds. + pub(crate) size: u64, + /// How many bytes a Parquet file's footer metadata takes, and 0 for the other formats. + pub(crate) footer: u64, + /// For a Parquet file, what each column holds. + pub(crate) columns: Vec, + /// For a Parquet file, how many row groups it has. + pub(crate) row_groups: usize, + /// The partition columns and their values cast to VARCHAR, for a partitioned `COPY`. + pub(crate) partition: Option>, +} + /// 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 the paths of the files, in the order they were written. +/// and answers how many there were and the files, in the order they were written. pub(crate) fn write( copy: &CopyTo, result: &QueryResult, zone: SessionTimeZone, threads: usize, -) -> Result<(usize, Vec)> { +) -> Result<(usize, Vec)> { let split = ©.partitioned; let mut files = Vec::new(); if split.skip_empty && result.is_empty() { @@ -83,6 +99,148 @@ pub(crate) fn write( Ok((rows, files)) } +/// What a `COPY ... TO` with `RETURN_STATS` answers: a row for each file, with its rows, its size, +/// the size of its footer, what each column holds, the partition it is of and how many row groups +/// it has, which is what the pin reports for a Parquet file. +pub(crate) fn statistics(files: Vec) -> Result { + let text = LogicalType::Varchar; + let stats = LogicalType::map(text.clone(), text.clone()); + let types = vec![ + text.clone(), + LogicalType::UBigInt, + LogicalType::UBigInt, + LogicalType::UBigInt, + LogicalType::map(text.clone(), stats.clone()), + stats.clone(), + LogicalType::map(text.clone(), LogicalType::Variant), + ]; + // The statistics are keyed in order, the way the pin reports them, and the partition columns + // come in the order they were named. + let map = |value: &LogicalType, entries: Vec<(Value, Value)>| Value::Map { + key: Box::new(text.clone()), + value: Box::new(value.clone()), + entries, + }; + let mut rows = vec![Vec::with_capacity(files.len()); types.len()]; + for file in files { + let mut columns = Vec::with_capacity(file.columns.len()); + for column in &file.columns { + let mut entries = vec![ + ("column_size_bytes", column.size.to_string()), + ("null_count", column.nulls.to_string()), + ("num_values", column.values.to_string()), + ]; + // A float column counts its NaNs over every row group, as the corpus expects of a + // DuckDB newer than the pin, which only kept the last group's has_nan and had no + // nan_count. A column with no values has neither, as in the pin. + if let Some(nans) = column.nans.filter(|_| column.values > 0) { + entries.push(("has_nan", (nans > 0).to_string())); + entries.push(("nan_count", nans.to_string())); + } + if let (Some(min), Some(max)) = (&column.min, &column.max) { + let (min, min_exact) = bound_text(min, false)?; + let (max, max_exact) = bound_text(max, true)?; + entries.push(("min", min)); + entries.push(("min_is_exact", min_exact.to_string())); + entries.push(("max", max)); + entries.push(("max_is_exact", max_exact.to_string())); + } + entries.sort_unstable_by_key(|&(key, _)| key); + let entries = entries + .into_iter() + .map(|(key, value)| (Value::Varchar(key.into()), Value::Varchar(value))) + .collect(); + let name = format!("\"{}\"", column.name.replace('"', "\"\"")); + columns.push((Value::Varchar(name), map(&text, entries))); + } + columns.sort_by_key(|(name, _)| name.to_string()); + let partition = file.partition.map(|keys| { + let entries = keys.into_iter().map(|(key, value)| (Value::Varchar(key), value)); + map(&text, entries.collect()) + }); + let groups = u64::try_from(file.row_groups).unwrap_or(u64::MAX); + let groups = cast_value(&Value::UBigInt(groups), &LogicalType::Variant, false)?; + let row = [ + Value::Varchar(file.path), + Value::UBigInt(u64::try_from(file.rows).unwrap_or(u64::MAX)), + Value::UBigInt(file.size), + Value::UBigInt(file.footer), + map(&stats, columns), + partition.unwrap_or(Value::Null), + map(&LogicalType::Variant, vec![(Value::Varchar("row_group_count".into()), groups)]), + ]; + for (column, value) in rows.iter_mut().zip(row) { + column.push(value); + } + } + let vectors = types + .iter() + .zip(&rows) + .map(|(ty, values)| Vector::from_values(ty.clone(), values)) + .collect::>>()?; + let names = [ + "filename", + "count", + "file_size_bytes", + "footer_size_bytes", + "column_statistics", + "partition_keys", + "extra_info", + ]; + Ok(QueryResult::new( + names.map(str::to_owned).to_vec(), + types, + vec![Chunk::new(vectors)?], + Memory::unlimited().reservation(), + )) +} + +/// How many bytes of a string or a blob a bound keeps, which is the pin's limit. +const BOUND_BYTES: usize = 256; + +/// A bound of a column as the pin reports it, and whether it is exact. A bool is 1 or 0 and a blob +/// is its bytes in hex. A string or a blob longer than [`BOUND_BYTES`] is cut short, a string on a +/// character, and an upper bound then has its last character or byte raised by one so that it is +/// still above every value, and neither is exact any more. +fn bound_text(value: &Value, upper: bool) -> Result<(String, bool)> { + Ok(match value { + Value::Boolean(on) => (u8::from(*on).to_string(), true), + Value::Blob(bytes) if bytes.len() > BOUND_BYTES => { + let mut kept = bytes[..BOUND_BYTES].to_vec(); + if upper { + while kept.last() == Some(&u8::MAX) { + kept.pop(); + } + if let Some(last) = kept.last_mut() { + *last += 1; + } + } + (hex(&kept), false) + } + Value::Blob(bytes) => (hex(bytes), true), + Value::Varchar(text) if text.len() > BOUND_BYTES => { + let mut end = BOUND_BYTES; + while !text.is_char_boundary(end) { + end -= 1; + } + let mut kept = text[..end].to_string(); + if upper && let Some(last) = kept.pop() { + kept.extend(char::from_u32(u32::from(last) + 1)); + } + (kept, false) + } + other => match cast_value(other, &LogicalType::Varchar, false)? { + Value::Varchar(text) => (text, true), + cast => (cast.to_string(), true), + }, + }) +} + +/// Bytes in upper case hex, two digits a byte. +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|byte| format!("{byte:02X}")).collect() +} + /// Makes a directory and the ones above it. fn make_directory(path: &Path) -> Result<()> { std::fs::create_dir_all(path).map_err(|error| { @@ -104,7 +262,7 @@ fn write_files( result: &QueryResult, zone: SessionTimeZone, threads: usize, - files: &mut Vec, + files: &mut Vec, ) -> Result { let split = ©.partitioned; let path = |number: usize| -> Result { @@ -152,7 +310,7 @@ fn write_threads( result: &QueryResult, zone: SessionTimeZone, threads: usize, - files: &mut Vec, + files: &mut Vec, ) -> Result { let batches = batches(result.chunks(), 0)?; let each = batches.len().div_ceil(threads.max(1)).max(1); @@ -204,6 +362,10 @@ fn batches(chunks: &[Chunk], least: u64) -> Result>> { /// What a directory is called for a NULL partition value, which is the name Hive gave it. const NULL_PARTITION: &str = "__HIVE_DEFAULT_PARTITION__"; +/// A partition being written: its directory, the values of its columns, the rows it takes from +/// each chunk, and its values as text, which RETURN_STATS names it by. +type Partition = (String, Vec, Vec>, Vec); + /// Writes the rows of `result` split by the values of the partition columns, a directory a /// column named `column=value`, both percent encoded, with the files in each of the deepest ones. /// @@ -215,13 +377,13 @@ fn write_partitioned( root: &Path, result: &QueryResult, zone: SessionTimeZone, - files: &mut Vec, + files: &mut Vec, ) -> Result { let split = ©.partitioned; let names = result.names(); let chunks = result.chunks(); let mut found: HashMap = HashMap::new(); - let mut partitions: Vec<(String, Vec, Vec>)> = Vec::new(); + let mut partitions: Vec = Vec::new(); for (at, chunk) in chunks.iter().enumerate() { let mut keys = Vec::with_capacity(split.columns.len()); for &column in &split.columns { @@ -230,17 +392,20 @@ fn write_partitioned( } for row in 0..chunk.len() { let mut directory = String::new(); + let mut texts = Vec::with_capacity(keys.len()); 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)? { + let text = key.try_value_at(row)?; + match &text { Value::Null => directory.push_str(NULL_PARTITION), - Value::Varchar(text) => encode_value(&mut directory, &text), + Value::Varchar(text) => encode_value(&mut directory, text), other => encode_value(&mut directory, &other.to_string()), } + texts.push(text); } let partition = match found.get(&directory) { Some(&partition) => partition, @@ -251,7 +416,8 @@ fn write_partitioned( .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()])); + let picked = vec![Vec::new(); chunks.len()]; + partitions.push((directory, values, picked, texts)); partitions.len() - 1 } }; @@ -264,7 +430,7 @@ fn write_partitioned( 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, _)| { + partitions.sort_by(|(_, left, ..), (_, right, ..)| { left.iter() .zip(right) .map(|(left, right)| { @@ -275,7 +441,7 @@ fn write_partitioned( }); } let (mut rows, mut next) = (0, 0); - for (directory, _, picked) in &partitions { + for (directory, _, picked, texts) in &partitions { let mut parts = Vec::new(); for (chunk, rows) in chunks.iter().zip(picked) { if rows.is_empty() { @@ -299,6 +465,11 @@ fn write_partitioned( let before = files.len(); rows += write_files(copy, &place, first, &part, zone, 1, files)?; next += files.len() - before; + let keys = split.columns.iter().map(|&column| names[column].clone()); + let keys = keys.zip(texts.iter().cloned()).collect::>(); + for wrote in &mut files[before..] { + wrote.partition = Some(keys.clone()); + } } Ok(rows) } @@ -599,10 +770,11 @@ impl<'a> Sink<'a> { Ok(()) } - /// 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 { + /// Writes what comes after the last row and closes the file, adds it to `files` and answers + /// how many rows it holds. + fn finish(self, files: &mut Vec) -> Result { let written = failed(&self.path); + let (mut footer, mut columns, mut row_groups) = (0, Vec::new(), 0); match self.body { Body::Csv { mut out, .. } => out.flush().map_err(written)?, Body::Json { mut out, .. } => { @@ -614,10 +786,20 @@ impl<'a> Sink<'a> { } Body::Parquet { mut writer, group, .. } => { writer.write_group(&group)?; - writer.finish()?; + columns = writer.statistics().to_vec(); + row_groups = writer.row_groups(); + footer = writer.finish_sized()?.1; } } - files.push(self.path); + files.push(Wrote { + path: self.path, + rows: self.rows, + size: self.size.get(), + footer, + columns, + row_groups, + partition: None, + }); Ok(self.rows) } } diff --git a/crates/rudb/tests/copy_to.rs b/crates/rudb/tests/copy_to.rs index 332535e67..3f2a0e62c 100644 --- a/crates/rudb/tests/copy_to.rs +++ b/crates/rudb/tests/copy_to.rs @@ -863,3 +863,99 @@ fn return_files_answers_the_count_and_the_files_written_as_the_pin_does() { } let _ = std::fs::remove_dir_all(&root); } + +#[test] +fn return_stats_answers_what_each_parquet_file_holds_as_the_pin_does() { + use rudb_common::Value; + let root = std::env::temp_dir().join(format!("rudb-copy-to-stats-{}", 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 s AS SELECT range i, range % 2 = 0 b, NULLIF(range, 3)::VARCHAR \"a\"\"b\", \ + CASE WHEN range = 1 THEN 'nan'::DOUBLE ELSE range / 2 END f FROM range(5)", + ) + .expect("creates"); + // The sizes are the writer's own and differ from the pin's, so they are masked. + let mask = |text: String| { + let mut out = String::new(); + let mut rest = text.as_str(); + while let Some(at) = rest.find("column_size_bytes=") { + let start = at + "column_size_bytes=".len(); + out.push_str(&rest[..start]); + rest = &rest[start..]; + let digits = rest.len() - rest.trim_start_matches(|c: char| c.is_ascii_digit()).len(); + out.push_str(if &rest[..digits] == "0" { "0" } else { "N" }); + rest = &rest[digits..]; + } + out.push_str(rest); + out + }; + let stats = |sql: &str| { + let sql = sql.replace("DIR", &dir); + let result = db.execute(&sql).unwrap_or_else(|error| panic!("{sql}: {error}")); + let names = [ + "filename", + "count", + "file_size_bytes", + "footer_size_bytes", + "column_statistics", + "partition_keys", + "extra_info", + ]; + assert_eq!(result.names(), names, "{sql}"); + let mut rows = (0..result.len()) + .map(|row| { + let (Value::UBigInt(size), Value::UBigInt(footer)) = + (result.value_at(row, 2), result.value_at(row, 3)) + else { + panic!("{sql}: sizes") + }; + assert!(size > footer && footer > 0, "{sql}: {size} {footer}"); + [0, 1, 4, 5, 6].map(|column| { + mask(result.value_at(row, column).to_string().replace(&dir, "DIR")) + }) + }) + .collect::>(); + rows.sort(); + rows + }; + + assert_eq!( + stats("COPY s TO 'DIR/s.parquet' (RETURN_STATS)"), + [[ + "DIR/s.parquet".into(), + "5".into(), + "{'\"a\"\"b\"'={column_size_bytes=N, max=4, max_is_exact=true, min=0, \ + min_is_exact=true, null_count=1, num_values=5}, '\"b\"'={column_size_bytes=N, max=1, \ + max_is_exact=true, min=0, min_is_exact=true, null_count=0, num_values=5}, \ + '\"f\"'={column_size_bytes=N, has_nan=true, max=2.0, max_is_exact=true, min=0.0, \ + min_is_exact=true, nan_count=1, null_count=0, num_values=5}, \ + '\"i\"'={column_size_bytes=N, max=4, max_is_exact=true, min=0, min_is_exact=true, \ + null_count=0, num_values=5}}" + .to_string(), + "NULL".into(), + "{row_group_count=1}".into(), + ]] + ); + // A partitioned copy names the partition of each file and leaves its column out. + let parts = stats( + "COPY (SELECT i, b FROM s) TO 'DIR/p' (FORMAT parquet, PARTITION_BY b, RETURN_STATS)", + ); + let keys = parts.iter().map(|row| row[3].clone()).collect::>(); + assert_eq!(keys, ["{b=false}", "{b=true}"]); + assert!(parts[0][2].starts_with("{'\"i\"'=") && !parts[0][2].contains("\"b\""), "{parts:?}"); + // A column with no values has no bounds and no NaN count. + assert_eq!( + stats("COPY (SELECT f FROM s LIMIT 0) TO 'DIR/e.parquet' (RETURN_STATS)"), + [[ + "DIR/e.parquet".to_string(), + "0".to_string(), + "{'\"f\"'={column_size_bytes=0, null_count=0, num_values=0}}".to_string(), + "NULL".to_string(), + "{row_group_count=0}".to_string(), + ]] + ); + let _ = std::fs::remove_dir_all(&root); +}