From 077cfe7d90a64e5bc6d49c00410309b37997618e Mon Sep 17 00:00:00 2001 From: osipovartem Date: Thu, 8 Oct 2026 16:14:51 +0300 Subject: [PATCH 1/2] Preserve ordered Snowflake CSV inference across chunks --- Cargo.lock | 34 +-- Cargo.toml | 18 +- datafusion/datasource-csv/src/file_format.rs | 220 +++++++++++++++++-- 3 files changed, 227 insertions(+), 45 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0474e5cafed34..a896e1eccd9fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -165,7 +165,7 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-arith", "arrow-array", @@ -187,7 +187,7 @@ dependencies = [ [[package]] name = "arrow-arith" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-buffer", @@ -200,7 +200,7 @@ dependencies = [ [[package]] name = "arrow-array" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "ahash", "arrow-buffer", @@ -219,7 +219,7 @@ dependencies = [ [[package]] name = "arrow-avro" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-buffer", @@ -242,7 +242,7 @@ dependencies = [ [[package]] name = "arrow-buffer" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "bytes", "half", @@ -253,7 +253,7 @@ dependencies = [ [[package]] name = "arrow-cast" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-buffer", @@ -274,7 +274,7 @@ dependencies = [ [[package]] name = "arrow-csv" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-cast", @@ -288,7 +288,7 @@ dependencies = [ [[package]] name = "arrow-data" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-buffer", "arrow-schema", @@ -300,7 +300,7 @@ dependencies = [ [[package]] name = "arrow-flight" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-arith", "arrow-array", @@ -326,7 +326,7 @@ dependencies = [ [[package]] name = "arrow-ipc" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-buffer", @@ -341,7 +341,7 @@ dependencies = [ [[package]] name = "arrow-json" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-buffer", @@ -365,7 +365,7 @@ dependencies = [ [[package]] name = "arrow-ord" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-buffer", @@ -377,7 +377,7 @@ dependencies = [ [[package]] name = "arrow-row" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-buffer", @@ -389,7 +389,7 @@ dependencies = [ [[package]] name = "arrow-schema" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "bitflags", "serde", @@ -400,7 +400,7 @@ dependencies = [ [[package]] name = "arrow-select" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "ahash", "arrow-array", @@ -413,7 +413,7 @@ dependencies = [ [[package]] name = "arrow-string" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "arrow-array", "arrow-buffer", @@ -4484,7 +4484,7 @@ dependencies = [ [[package]] name = "parquet" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=431221f0c1e824998204bff0ce17a780dae235c2#431221f0c1e824998204bff0ce17a780dae235c2" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" dependencies = [ "ahash", "arrow-array", diff --git a/Cargo.toml b/Cargo.toml index 8677d0fbf282c..064dc9bda945a 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 = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", features = [ +arrow = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", features = [ "prettyprint", "chrono-tz", ] } -arrow-avro = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", default-features = false, features = [ +arrow-avro = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false, features = [ "deflate", "snappy", "zstd", "bzip2", "xz", ] } -arrow-buffer = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", default-features = false } -arrow-data = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", default-features = false } -arrow-flight = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", features = [ +arrow-buffer = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false } +arrow-data = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false } +arrow-flight = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", 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 = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", default-features = false, features = [ +arrow-ipc = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false, features = [ "lz4", "zstd", ] } -arrow-ord = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", default-features = false } -arrow-schema = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", default-features = false } +arrow-ord = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false } +arrow-schema = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", 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 = { git = "https://github.com/Embucket/arrow-rs.git", rev = "431221f0c1e824998204bff0ce17a780dae235c2", default-features = false, features = [ +parquet = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", 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 5c8e769f58b40..a9900dfd9bb0d 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -25,7 +25,7 @@ use crate::source::CsvSource; use arrow::array::RecordBatch; use arrow::csv::{CsvRecordErrorHandler, WriterBuilder}; -use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef, TimeUnit}; use arrow::error::ArrowError; use datafusion_common::config::{ConfigField, ConfigFileType, CsvOptions}; use datafusion_common::file_options::csv_writer::CsvWriterOptions; @@ -400,7 +400,8 @@ impl FileFormat for CsvFormat { store: &Arc, objects: &[ObjectMeta], ) -> Result { - self.infer_schema_impl::(state, store, objects).await + self.infer_schema_impl::(state, store, objects) + .await } async fn infer_stats( @@ -507,11 +508,23 @@ impl CsvFormat { store: &Arc, object: &ObjectMeta, ) -> Result { - self.infer_schema_impl::(state, store, std::slice::from_ref(object)) + self.infer_schema_impl::(state, store, std::slice::from_ref(object)) + .await + } + + /// Infer Snowflake-compatible, order-sensitive types from one CSV object. + /// Ordinary CSV scans and order-independent decimal inference are unchanged. + pub async fn infer_schema_with_snowflake_types( + &self, + state: &dyn Session, + store: &Arc, + object: &ObjectMeta, + ) -> Result { + self.infer_schema_impl::(state, store, std::slice::from_ref(object)) .await } - async fn infer_schema_impl( + async fn infer_schema_impl( &self, state: &dyn Session, store: &Arc, @@ -526,7 +539,7 @@ impl CsvFormat { for object in objects { let stream = self.read_to_delimited_chunks(store, object).await; let (schema, records_read) = self - .infer_schema_from_stream_impl::<_, DECIMAL>( + .infer_schema_from_stream_impl::<_, DECIMAL, SNOWFLAKE>( state, records_to_read, stream, @@ -571,13 +584,18 @@ impl CsvFormat { records_to_read: usize, stream: impl Stream>, ) -> Result<(Schema, usize)> { - self.infer_schema_from_stream_impl::<_, false>(state, records_to_read, stream) - .await + self.infer_schema_from_stream_impl::<_, false, false>( + state, + records_to_read, + stream, + ) + .await } async fn infer_schema_from_stream_impl< S: Stream>, const DECIMAL: bool, + const SNOWFLAKE: bool, >( &self, state: &dyn Session, @@ -587,7 +605,8 @@ impl CsvFormat { let mut total_records_read = 0; let mut column_names = vec![]; let mut column_type_possibilities = vec![]; - let mut decimal_integral_digits = DECIMAL.then(Vec::new); + let mut decimal_integral_digits = (DECIMAL && !SNOWFLAKE).then(Vec::new); + let mut ordered_types = SNOWFLAKE.then(Vec::::new); let mut record_number = -1; let initial_records_to_read = records_to_read; @@ -629,16 +648,56 @@ impl CsvFormat { format = format.with_comment(comment); } - let (Schema { fields, .. }, records_read) = if DECIMAL { + let (Schema { fields, .. }, records_read) = if SNOWFLAKE { + format.infer_schema_with_snowflake_types( + chunk.reader(), + Some(records_to_read), + )? + } else if DECIMAL { format.infer_schema_with_decimal(chunk.reader(), Some(records_to_read))? } else { format.infer_schema(chunk.reader(), Some(records_to_read))? }; + if let Some(ordered) = ordered_types.as_mut() { + for (current, field) in ordered.iter_mut().zip(&fields) { + current.merge(field); + } + ordered.extend( + fields + .iter() + .skip(ordered.len()) + .map(|field| OrderedCsvType::from_field(field)), + ); + } + records_to_read -= records_read; total_records_read += records_read; - if first_chunk { + if SNOWFLAKE { + if first_chunk { + column_names = + fields.iter().map(|field| field.name().clone()).collect(); + } else { + if fields.len() != column_names.len() + && !self.options.truncated_rows.unwrap_or(false) + { + return exec_err!( + "Encountered unequal lengths between records on CSV file whilst inferring schema. \ + Expected {} fields, found {} fields at record {}", + column_names.len(), + fields.len(), + record_number + 1 + ); + } + column_names.extend( + fields + .iter() + .skip(column_names.len()) + .map(|field| field.name().clone()), + ); + } + } else if first_chunk { // set up initial structures for recording inferred schema across chunks (column_names, column_type_possibilities) = fields .into_iter() @@ -703,16 +762,91 @@ impl CsvFormat { } } - let schema = build_schema_helper::( - column_names, - column_type_possibilities, - decimal_integral_digits.as_deref(), - initial_records_to_read == 0, - ); + let schema = if let Some(ordered) = ordered_types { + Schema::new( + column_names + .into_iter() + .zip(ordered) + .map(|(name, inferred)| { + let data_type = if initial_records_to_read == 0 { + DataType::Utf8 + } else { + inferred.data_type + }; + Field::new(name, data_type, true) + }) + .collect::(), + ) + } else { + build_schema_helper::( + column_names, + column_type_possibilities, + decimal_integral_digits.as_deref(), + initial_records_to_read == 0, + ) + }; Ok((schema, total_records_read)) } } +struct OrderedCsvType { + data_type: DataType, + integral_digits: u16, +} + +impl OrderedCsvType { + fn from_field(field: &Field) -> Self { + Self { + data_type: field.data_type().clone(), + integral_digits: csv_decimal_integral_digits(field).unwrap_or(0), + } + } + + fn merge(&mut self, field: &Field) { + let next = Self::from_field(field); + self.data_type = match (&self.data_type, &next.data_type) { + (DataType::Null, other) => other.clone(), + (other, DataType::Null) => other.clone(), + ( + DataType::Decimal128(_, left_scale), + DataType::Decimal128(_, right_scale), + ) if *left_scale >= 0 && *right_scale >= 0 => { + let scale = (*left_scale).max(*right_scale); + let integral_digits = self.integral_digits.max(next.integral_digits); + let precision = + integral_digits.max(1) + u16::try_from(scale).unwrap_or(0); + match u8::try_from(precision) { + Ok(precision @ 1..=38) => DataType::Decimal128(precision, scale), + _ => DataType::Utf8, + } + } + (DataType::Decimal128(_, _), DataType::Float64) => DataType::Utf8, + (DataType::Float64, DataType::Decimal128(_, _)) => DataType::Float64, + (DataType::Date32, DataType::Timestamp(_, _)) => DataType::Date32, + (DataType::Timestamp(_, _), DataType::Date32) => DataType::Utf8, + ( + DataType::Timestamp(left, left_tz), + DataType::Timestamp(right, right_tz), + ) if left_tz == right_tz => DataType::Timestamp( + max_csv_timestamp_unit(*left, *right), + left_tz.clone(), + ), + (left, right) if left == right => left.clone(), + _ => DataType::Utf8, + }; + self.integral_digits = self.integral_digits.max(next.integral_digits); + } +} + +fn max_csv_timestamp_unit(left: TimeUnit, right: TimeUnit) -> TimeUnit { + match (left, right) { + (TimeUnit::Nanosecond, _) | (_, TimeUnit::Nanosecond) => TimeUnit::Nanosecond, + (TimeUnit::Microsecond, _) | (_, TimeUnit::Microsecond) => TimeUnit::Microsecond, + (TimeUnit::Millisecond, _) | (_, TimeUnit::Millisecond) => TimeUnit::Millisecond, + _ => TimeUnit::Second, + } +} + fn csv_decimal_integral_digits(field: &Field) -> Option { let DataType::Decimal128(precision, scale) = field.data_type() else { return None; @@ -1367,7 +1501,11 @@ mod tests { ]); let (schema, records) = CsvFormat::default() .with_has_header(true) - .infer_schema_from_stream_impl::<_, true>(&InferenceSession::new(), 3, chunks) + .infer_schema_from_stream_impl::<_, true, false>( + &InferenceSession::new(), + 3, + chunks, + ) .await .unwrap(); assert_eq!(records, 3); @@ -1381,7 +1519,11 @@ mod tests { ]); let (schema, records) = CsvFormat::default() .with_has_header(true) - .infer_schema_from_stream_impl::<_, true>(&InferenceSession::new(), 2, zeros) + .infer_schema_from_stream_impl::<_, true, false>( + &InferenceSession::new(), + 2, + zeros, + ) .await .unwrap(); assert_eq!(records, 2); @@ -1393,7 +1535,7 @@ mod tests { ]); let (schema, _) = CsvFormat::default() .with_has_header(true) - .infer_schema_from_stream_impl::<_, true>( + .infer_schema_from_stream_impl::<_, true, false>( &InferenceSession::new(), 2, integer_and_tiny, @@ -1402,4 +1544,44 @@ mod tests { .unwrap(); assert_eq!(schema.field(0).data_type(), &DataType::Float64); } + + #[tokio::test] + async fn test_snowflake_ordered_inference_across_csv_chunks() { + let tiny = format!("0.{}1", "0".repeat(37)); + let chunks = stream::iter(vec![ + Ok(Bytes::from(format!( + "INT_EXP,EXP_INT,ZERO_TINY,TINY_ZERO,DATE_TS,TS_DATE,EXACT\n1,1e2,0,{tiny},2024-01-02,2024-01-02 01:02:03,1.20\n" + ))), + Ok(Bytes::from(format!( + "1e2,1,{tiny},0,2024-01-02 01:02:03,2024-01-02,123.400\n" + ))), + ]); + let (schema, records) = CsvFormat::default() + .with_has_header(true) + .infer_schema_from_stream_impl::<_, true, true>( + &InferenceSession::new(), + 2, + chunks, + ) + .await + .unwrap(); + assert_eq!(records, 2); + let actual = schema + .fields() + .iter() + .map(|field| field.data_type().clone()) + .collect::>(); + assert_eq!( + actual, + [ + DataType::Utf8, + DataType::Float64, + DataType::Utf8, + DataType::Float64, + DataType::Date32, + DataType::Utf8, + DataType::Decimal128(6, 3), + ] + ); + } } From 04ce310964bfbdb3e1130ba9750eda7761c38356 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Thu, 8 Oct 2026 16:24:32 +0300 Subject: [PATCH 2/2] Merge typed Arrow CSV inference state across chunks --- Cargo.lock | 34 ++--- Cargo.toml | 18 +-- datafusion/datasource-csv/src/file_format.rs | 130 ++++++++----------- 3 files changed, 80 insertions(+), 102 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a896e1eccd9fc..d39c08372a31f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -165,7 +165,7 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-arith", "arrow-array", @@ -187,7 +187,7 @@ dependencies = [ [[package]] name = "arrow-arith" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-buffer", @@ -200,7 +200,7 @@ dependencies = [ [[package]] name = "arrow-array" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "ahash", "arrow-buffer", @@ -219,7 +219,7 @@ dependencies = [ [[package]] name = "arrow-avro" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-buffer", @@ -242,7 +242,7 @@ dependencies = [ [[package]] name = "arrow-buffer" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "bytes", "half", @@ -253,7 +253,7 @@ dependencies = [ [[package]] name = "arrow-cast" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-buffer", @@ -274,7 +274,7 @@ dependencies = [ [[package]] name = "arrow-csv" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-cast", @@ -288,7 +288,7 @@ dependencies = [ [[package]] name = "arrow-data" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-buffer", "arrow-schema", @@ -300,7 +300,7 @@ dependencies = [ [[package]] name = "arrow-flight" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-arith", "arrow-array", @@ -326,7 +326,7 @@ dependencies = [ [[package]] name = "arrow-ipc" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-buffer", @@ -341,7 +341,7 @@ dependencies = [ [[package]] name = "arrow-json" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-buffer", @@ -365,7 +365,7 @@ dependencies = [ [[package]] name = "arrow-ord" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-buffer", @@ -377,7 +377,7 @@ dependencies = [ [[package]] name = "arrow-row" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-buffer", @@ -389,7 +389,7 @@ dependencies = [ [[package]] name = "arrow-schema" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "bitflags", "serde", @@ -400,7 +400,7 @@ dependencies = [ [[package]] name = "arrow-select" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "ahash", "arrow-array", @@ -413,7 +413,7 @@ dependencies = [ [[package]] name = "arrow-string" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "arrow-array", "arrow-buffer", @@ -4484,7 +4484,7 @@ dependencies = [ [[package]] name = "parquet" version = "59.2.0" -source = "git+https://github.com/Embucket/arrow-rs.git?rev=ec817ac3372ddf79f0232966a392be4b4746d3ac#ec817ac3372ddf79f0232966a392be4b4746d3ac" +source = "git+https://github.com/Embucket/arrow-rs.git?rev=4e743e5dfb8b2b7aba69dade4d887418d01520b5#4e743e5dfb8b2b7aba69dade4d887418d01520b5" dependencies = [ "ahash", "arrow-array", diff --git a/Cargo.toml b/Cargo.toml index 064dc9bda945a..e53ddf895426e 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 = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", features = [ +arrow = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", features = [ "prettyprint", "chrono-tz", ] } -arrow-avro = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false, features = [ +arrow-avro = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", default-features = false, features = [ "deflate", "snappy", "zstd", "bzip2", "xz", ] } -arrow-buffer = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false } -arrow-data = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false } -arrow-flight = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", features = [ +arrow-buffer = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", default-features = false } +arrow-data = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", default-features = false } +arrow-flight = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", 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 = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false, features = [ +arrow-ipc = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", default-features = false, features = [ "lz4", "zstd", ] } -arrow-ord = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false } -arrow-schema = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false } +arrow-ord = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", default-features = false } +arrow-schema = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", 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 = { git = "https://github.com/Embucket/arrow-rs.git", rev = "ec817ac3372ddf79f0232966a392be4b4746d3ac", default-features = false, features = [ +parquet = { git = "https://github.com/Embucket/arrow-rs.git", rev = "4e743e5dfb8b2b7aba69dade4d887418d01520b5", 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 a9900dfd9bb0d..89e2c5ea9f9b7 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -24,8 +24,9 @@ use std::sync::Arc; use crate::source::CsvSource; use arrow::array::RecordBatch; +use arrow::csv::reader::SnowflakeCsvTypeState; use arrow::csv::{CsvRecordErrorHandler, WriterBuilder}; -use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef, TimeUnit}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; use arrow::error::ArrowError; use datafusion_common::config::{ConfigField, ConfigFileType, CsvOptions}; use datafusion_common::file_options::csv_writer::CsvWriterOptions; @@ -606,7 +607,7 @@ impl CsvFormat { let mut column_names = vec![]; let mut column_type_possibilities = vec![]; let mut decimal_integral_digits = (DECIMAL && !SNOWFLAKE).then(Vec::new); - let mut ordered_types = SNOWFLAKE.then(Vec::::new); + let mut ordered_types = SNOWFLAKE.then(Vec::::new); let mut record_number = -1; let initial_records_to_read = records_to_read; @@ -648,27 +649,28 @@ impl CsvFormat { format = format.with_comment(comment); } - let (Schema { fields, .. }, records_read) = if SNOWFLAKE { - format.infer_schema_with_snowflake_types( - chunk.reader(), - Some(records_to_read), - )? + let (Schema { fields, .. }, records_read, chunk_states) = if SNOWFLAKE { + let (schema, states, records) = format + .infer_schema_with_snowflake_types( + chunk.reader(), + Some(records_to_read), + )?; + (schema, records, states) } else if DECIMAL { - format.infer_schema_with_decimal(chunk.reader(), Some(records_to_read))? + let (schema, records) = format + .infer_schema_with_decimal(chunk.reader(), Some(records_to_read))?; + (schema, records, Vec::new()) } else { - format.infer_schema(chunk.reader(), Some(records_to_read))? + let (schema, records) = + format.infer_schema(chunk.reader(), Some(records_to_read))?; + (schema, records, Vec::new()) }; if let Some(ordered) = ordered_types.as_mut() { - for (current, field) in ordered.iter_mut().zip(&fields) { - current.merge(field); + for (current, next) in ordered.iter_mut().zip(&chunk_states) { + current.merge(next); } - ordered.extend( - fields - .iter() - .skip(ordered.len()) - .map(|field| OrderedCsvType::from_field(field)), - ); + ordered.extend(chunk_states.into_iter().skip(ordered.len())); } records_to_read -= records_read; @@ -771,7 +773,7 @@ impl CsvFormat { let data_type = if initial_records_to_read == 0 { DataType::Utf8 } else { - inferred.data_type + inferred.data_type() }; Field::new(name, data_type, true) }) @@ -789,64 +791,6 @@ impl CsvFormat { } } -struct OrderedCsvType { - data_type: DataType, - integral_digits: u16, -} - -impl OrderedCsvType { - fn from_field(field: &Field) -> Self { - Self { - data_type: field.data_type().clone(), - integral_digits: csv_decimal_integral_digits(field).unwrap_or(0), - } - } - - fn merge(&mut self, field: &Field) { - let next = Self::from_field(field); - self.data_type = match (&self.data_type, &next.data_type) { - (DataType::Null, other) => other.clone(), - (other, DataType::Null) => other.clone(), - ( - DataType::Decimal128(_, left_scale), - DataType::Decimal128(_, right_scale), - ) if *left_scale >= 0 && *right_scale >= 0 => { - let scale = (*left_scale).max(*right_scale); - let integral_digits = self.integral_digits.max(next.integral_digits); - let precision = - integral_digits.max(1) + u16::try_from(scale).unwrap_or(0); - match u8::try_from(precision) { - Ok(precision @ 1..=38) => DataType::Decimal128(precision, scale), - _ => DataType::Utf8, - } - } - (DataType::Decimal128(_, _), DataType::Float64) => DataType::Utf8, - (DataType::Float64, DataType::Decimal128(_, _)) => DataType::Float64, - (DataType::Date32, DataType::Timestamp(_, _)) => DataType::Date32, - (DataType::Timestamp(_, _), DataType::Date32) => DataType::Utf8, - ( - DataType::Timestamp(left, left_tz), - DataType::Timestamp(right, right_tz), - ) if left_tz == right_tz => DataType::Timestamp( - max_csv_timestamp_unit(*left, *right), - left_tz.clone(), - ), - (left, right) if left == right => left.clone(), - _ => DataType::Utf8, - }; - self.integral_digits = self.integral_digits.max(next.integral_digits); - } -} - -fn max_csv_timestamp_unit(left: TimeUnit, right: TimeUnit) -> TimeUnit { - match (left, right) { - (TimeUnit::Nanosecond, _) | (_, TimeUnit::Nanosecond) => TimeUnit::Nanosecond, - (TimeUnit::Microsecond, _) | (_, TimeUnit::Microsecond) => TimeUnit::Microsecond, - (TimeUnit::Millisecond, _) | (_, TimeUnit::Millisecond) => TimeUnit::Millisecond, - _ => TimeUnit::Second, - } -} - fn csv_decimal_integral_digits(field: &Field) -> Option { let DataType::Decimal128(precision, scale) = field.data_type() else { return None; @@ -1583,5 +1527,39 @@ mod tests { DataType::Decimal128(6, 3), ] ); + + for (first, later, expected) in [ + ( + "2024-01-02", + "2024-01-02 01:02:03\n2024-01-02", + DataType::Date32, + ), + ("1e2", "1\n1e2", DataType::Float64), + ( + "1234567890123456789012345678901234567", + "0.01", + DataType::Decimal128(38, 1), + ), + ( + "12345678901234567890123456789012345678", + "0.1", + DataType::Decimal128(38, 0), + ), + ] { + let chunks = stream::iter(vec![ + Ok(Bytes::from(format!("VALUE\n{first}\n"))), + Ok(Bytes::from(format!("{later}\n"))), + ]); + let (schema, _) = CsvFormat::default() + .with_has_header(true) + .infer_schema_from_stream_impl::<_, true, true>( + &InferenceSession::new(), + 3, + chunks, + ) + .await + .unwrap(); + assert_eq!(schema.field(0).data_type(), &expected, "{first}; {later}"); + } } }