diff --git a/Cargo.lock b/Cargo.lock index 5f47369e0d50a..4d6706897298d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -165,8 +165,7 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61d285d16bce7d0be61912f7928342b673067b6b7d7ef6cc179258ba7de1fecf" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-arith", "arrow-array", @@ -188,8 +187,7 @@ dependencies = [ [[package]] name = "arrow-arith" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "757ef1836251e88222542a7da2623bc1c9cb9e20afefa6db2c41e79991cd91d4" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-buffer", @@ -202,8 +200,7 @@ dependencies = [ [[package]] name = "arrow-array" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc9a4a4b2b5ecd0e04df03471661cb61f28bed3c7fd50994715129b01b2edb97" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "ahash", "arrow-buffer", @@ -222,8 +219,7 @@ dependencies = [ [[package]] name = "arrow-avro" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fb45cd6bd2b25c0965793b83200eaca82214273a8030fbbc2d783e4c7c65a61" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-buffer", @@ -246,8 +242,7 @@ dependencies = [ [[package]] name = "arrow-buffer" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c12b576ef18c1deb80925a248b25ad84f419198d791b8e293fc6aaa60441fe90" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "bytes", "half", @@ -258,8 +253,7 @@ dependencies = [ [[package]] name = "arrow-cast" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68338a9096a5dc9bc11927c58c43a8526d96bf6abd2012ef6c0c9f505991cc79" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-buffer", @@ -280,8 +274,7 @@ dependencies = [ [[package]] name = "arrow-csv" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25011b52b346407d497ef0030e12b45e4f2d0cc279efc09c4f3d09106db30e36" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-cast", @@ -295,8 +288,7 @@ dependencies = [ [[package]] name = "arrow-data" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "723fe4aeed7604e00b9883a465af4ff0a0e6c44c03e41a68c3d1cbc403e0e44d" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-buffer", "arrow-schema", @@ -308,8 +300,7 @@ dependencies = [ [[package]] name = "arrow-flight" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2bebfacc9d71f0728f6774164e4d4254b5e504d2b46812d0512d8290ec119a64" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-arith", "arrow-array", @@ -335,8 +326,7 @@ dependencies = [ [[package]] name = "arrow-ipc" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "149437b14371f5b9ec60f5ddc751483ae99d7a7072653c0075e5e469156eea7b" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-buffer", @@ -351,8 +341,7 @@ dependencies = [ [[package]] name = "arrow-json" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f18b9123ccfec418a663f821c9a034af339711678c11ffe00d3ec07da5ff9f7e" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-buffer", @@ -376,8 +365,7 @@ dependencies = [ [[package]] name = "arrow-ord" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c08dff0686cf23ca4f562803f191ccbeb726dbae6309cd4b4aaf65e0f2c979" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-buffer", @@ -389,8 +377,7 @@ dependencies = [ [[package]] name = "arrow-row" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbec439386df71ad570e6758a946111322b9e9dc8db83b5527321f0b4c9119c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-buffer", @@ -402,8 +389,7 @@ dependencies = [ [[package]] name = "arrow-schema" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6fed2ca0d1eade57e811cbe73b98ad50cc08a1183e13b2d2aa43a7df593f40e" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "bitflags", "serde", @@ -414,8 +400,7 @@ dependencies = [ [[package]] name = "arrow-select" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "466b19cf75130b891dc1b23a84b343c714c62c64c9c62e365c76aa0ff90a53fb" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "ahash", "arrow-array", @@ -428,8 +413,7 @@ dependencies = [ [[package]] name = "arrow-string" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c838a25bb3691e919e0f617616ac51a4ff8517a952e29ca133cf0c22b2ce65b1" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "arrow-array", "arrow-buffer", @@ -4499,8 +4483,7 @@ dependencies = [ [[package]] name = "parquet" version = "59.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7065842956a20c2a536924ce8e4d9955f7422451511b9eb7500d7bfe5077e59c" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4aee61ac7892a32a832d0cef2d21886e2a0adcc0#4aee61ac7892a32a832d0cef2d21886e2a0adcc0" dependencies = [ "ahash", "arrow-array", diff --git a/Cargo.toml b/Cargo.toml index d03909a042800..8febefb88ad81 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -89,30 +89,30 @@ version = "55.1.0" # # See for more details: https://github.com/rust-lang/cargo/issues/11329 apache-avro = { version = "0.21", default-features = false } -arrow = { version = "59.2.0", features = [ +arrow = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", features = [ "prettyprint", "chrono-tz", ] } -arrow-avro = { version = "59.2.0", default-features = false, features = [ +arrow-avro = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", default-features = false, features = [ "deflate", "snappy", "zstd", "bzip2", "xz", ] } -arrow-buffer = { version = "59.2.0", default-features = false } -arrow-data = { version = "59.2.0", default-features = false } -arrow-flight = { version = "59.2.0", features = [ +arrow-buffer = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", default-features = false } +arrow-data = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", default-features = false } +arrow-flight = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", features = [ "flight-sql-experimental", ] } # Both codecs are required here to make sure that code paths like # file-spilling have access to all compression codecs. -arrow-ipc = { version = "59.2.0", default-features = false, features = [ +arrow-ipc = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", default-features = false, features = [ "lz4", "zstd", ] } -arrow-ord = { version = "59.2.0", default-features = false } -arrow-schema = { version = "59.2.0", default-features = false } +arrow-ord = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", default-features = false } +arrow-schema = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", default-features = false } async-trait = "0.1.89" bigdecimal = "0.4.8" bytes = "1.11" @@ -178,7 +178,7 @@ memchr = "2.8.1" num-traits = { version = "0.2" } object_store = { version = "0.13.2", default-features = false } parking_lot = "0.12" -parquet = { version = "59.2.0", default-features = false, features = [ +parquet = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4aee61ac7892a32a832d0cef2d21886e2a0adcc0", default-features = false, features = [ "arrow", "async", "object_store", diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index c0d22b80f08d0..2e2310a250129 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use crate::source::CsvSource; use arrow::array::RecordBatch; -use arrow::csv::WriterBuilder; +use arrow::csv::{CsvRecordErrorHandler, WriterBuilder}; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; use arrow::error::ArrowError; use datafusion_common::config::{ConfigField, ConfigFileType, CsvOptions}; @@ -134,6 +134,7 @@ impl GetExt for CsvFormatFactory { #[derive(Debug, Default)] pub struct CsvFormat { options: CsvOptions, + record_error_handler: Option>, } impl CsvFormat { @@ -200,6 +201,18 @@ impl CsvFormat { self } + /// Skip malformed field-count records and report them to `handler`. + /// + /// This recovery path retains source record bytes and disables file range + /// repartitioning. The default strict path is unchanged. + pub fn with_record_error_handler( + mut self, + handler: Arc, + ) -> Self { + self.record_error_handler = Some(handler); + self + } + /// Retrieve the csv options pub fn options(&self) -> &CsvOptions { &self.options @@ -494,7 +507,11 @@ impl FileFormat for CsvFormat { if csv_options.has_header.is_none() { csv_options.has_header = Some(true); } - Arc::new(CsvSource::new(table_schema).with_csv_options(csv_options)) + let mut source = CsvSource::new(table_schema).with_csv_options(csv_options); + if let Some(handler) = &self.record_error_handler { + source = source.with_record_error_handler(Arc::clone(handler)); + } + Arc::new(source) } } diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 08e4607498e62..5f83735f5a230 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -86,6 +86,7 @@ use tokio::io::AsyncWriteExt; #[derive(Debug, Clone)] pub struct CsvSource { options: CsvOptions, + record_error_handler: Option>, batch_size: Option, table_schema: TableSchema, projection: SplitProjection, @@ -98,6 +99,7 @@ impl CsvSource { let table_schema = table_schema.into(); Self { options: CsvOptions::default(), + record_error_handler: None, projection: SplitProjection::unprojected(&table_schema), table_schema, batch_size: None, @@ -111,6 +113,15 @@ impl CsvSource { self } + /// Skip malformed field-count records and report them to `handler`. + pub fn with_record_error_handler( + mut self, + handler: Arc, + ) -> Self { + self.record_error_handler = Some(handler); + self + } + /// true if the first line of each file is a header pub fn has_header(&self) -> bool { self.options.has_header.unwrap_or(true) @@ -205,6 +216,9 @@ impl CsvSource { if let Some(comment) = self.comment() { builder = builder.with_comment(comment); } + if let Some(handler) = &self.record_error_handler { + builder = builder.with_record_error_handler(Arc::clone(handler)); + } builder } @@ -298,7 +312,8 @@ impl FileSource for CsvSource { fn supports_repartitioning(&self) -> bool { // Cannot repartition if values may contain newlines, as record // boundaries cannot be determined by byte offset alone - !self.options.newlines_in_values.unwrap_or(false) + self.record_error_handler.is_none() + && !self.options.newlines_in_values.unwrap_or(false) } fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { @@ -329,6 +344,12 @@ impl FileSource for CsvSource { use datafusion_proto_models::protobuf; use protobuf::physical_plan_node::PhysicalPlanType; + if self.record_error_handler.is_some() { + return datafusion_common::not_impl_err!( + "CSV record error handlers cannot be serialized" + ); + } + let node = protobuf::CsvScanExecNode { base_conf: Some(base.try_to_proto(ctx)?), has_header: self.has_header(), @@ -648,3 +669,73 @@ impl CsvSource { Ok(DataSourceExec::from_data_source(conf)) } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int64Array, StringArray}; + use arrow::csv::{CsvRecordError, CsvRecordErrorHandler}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::error::ArrowError; + use std::io::Cursor; + use std::sync::Mutex; + + #[derive(Debug, Default)] + struct CollectRecordErrors(Mutex>); + + impl CsvRecordErrorHandler for CollectRecordErrors { + fn handle(&self, error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + self.0.lock().unwrap().push(( + error.line_number, + error.byte_offset, + error.actual_fields, + )); + Ok(()) + } + } + + #[test] + fn record_error_handler_skips_malformed_rows_and_disables_repartitioning() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("value", DataType::Utf8, false), + ])); + let handler = Arc::new(CollectRecordErrors::default()); + let options = CsvOptions { + has_header: Some(false), + ..Default::default() + }; + let mut source = CsvSource::new(schema) + .with_csv_options(options) + .with_record_error_handler(handler.clone()); + source.batch_size = Some(1024); + + assert!(!source.supports_repartitioning()); + let batches = source + .open(Cursor::new(b"1,ok\n2,extra,value\n3\n4,after\n")) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[1, 4] + ); + assert_eq!( + batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>(), + [Some("ok"), Some("after")] + ); + assert_eq!(*handler.0.lock().unwrap(), [(2, 5, 3), (3, 19, 1)]); + } +}