diff --git a/arrow-csv/src/reader/mod.rs b/arrow-csv/src/reader/mod.rs index d4a2bab9f5fa..c254fc548285 100644 --- a/arrow-csv/src/reader/mod.rs +++ b/arrow-csv/src/reader/mod.rs @@ -286,6 +286,29 @@ struct InferredDataType { packed: u16, max_integral_digits: u8, max_scale: u8, + first_snowflake_type: u16, +} + +/// Ordered, per-column CSV inference state. Merge states in file order to +/// preserve Snowflake's asymmetric numeric and temporal type transitions. +#[derive(Clone, Copy, Default)] +pub struct SnowflakeCsvTypeState(InferredDataType); + +impl SnowflakeCsvTypeState { + /// Merge a later chunk's inferred state into this state. + pub fn merge(&mut self, later: &Self) { + if self.0.first_snowflake_type == 0 { + self.0.first_snowflake_type = later.0.first_snowflake_type; + } + self.0.packed |= later.0.packed; + self.0.max_integral_digits = self.0.max_integral_digits.max(later.0.max_integral_digits); + self.0.max_scale = self.0.max_scale.max(later.0.max_scale); + } + + /// Return the type inferred from all states merged so far. + pub fn data_type(&self) -> DataType { + self.0.get_snowflake() + } } /// Field metadata set by decimal CSV inference when all observed numeric @@ -339,26 +362,59 @@ impl InferredDataType { self.get() } - fn update_decimal(&mut self, string: &str) { + fn get_snowflake(&self) -> DataType { + const TIMESTAMP_TYPES: u16 = (1 << 4) | (1 << 5) | (1 << 6) | (1 << 7); + if (self.first_snowflake_type == (1 << 1) && self.packed & (1 << 2) != 0) + || (self.first_snowflake_type & TIMESTAMP_TYPES != 0 && self.packed & (1 << 3) != 0) + { + return DataType::Utf8; + } + const TEMPORAL_TYPES: u16 = (1 << 3) | (1 << 4) | (1 << 5) | (1 << 6) | (1 << 7); + if self.first_snowflake_type == (1 << 3) && self.packed & !TEMPORAL_TYPES == 0 { + return DataType::Date32; + } + if self.packed == (1 << 1) { + let integral_digits = self.max_integral_digits.max(1); + let scale = self.max_scale.min(38 - integral_digits); + return DataType::Decimal128(integral_digits + scale, scale as i8); + } + self.get_decimal() + } + + fn update_snowflake(&mut self, string: &str) { + let value_type = self.update_decimal::(string); + if self.first_snowflake_type == 0 { + self.first_snowflake_type = value_type; + } + } + + fn update_decimal(&mut self, string: &str) -> u16 { if let Some((integral_digits, scale)) = decimal_shape(string) { - let precision = integral_digits + scale; + let precision = if SNOWFLAKE { + integral_digits.max(1) + scale + } else { + integral_digits + scale + }; if precision > 38 { self.packed |= 1 << 2; // Float64 for numbers outside Decimal128's range. + 1 << 2 } else { self.packed |= 1 << 1; self.max_integral_digits = self.max_integral_digits.max(integral_digits as u8); self.max_scale = self.max_scale.max(scale as u8); + 1 << 1 } } else if string.is_ascii() { - self.update(string); + self.update(string) } else { self.packed |= 1 << 8; // Unicode digits are not parsed by Decimal128. + 1 << 8 } } /// Updates the [`InferredDataType`] with the given string - fn update(&mut self, string: &str) { - self.packed |= if string.starts_with('"') { + fn update(&mut self, string: &str) -> u16 { + let value_type = if string.starts_with('"') { 1 << 8 // Utf8 } else if let Some(m) = REGEX_SET.matches(string).into_iter().next() { if m == 1 && string.len() >= 19 && string.parse::().is_err() { @@ -371,7 +427,9 @@ impl InferredDataType { 1 << 2 // Float64 } else { 1 << 8 // Utf8 - } + }; + self.packed |= value_type; + value_type } } @@ -512,7 +570,9 @@ impl Format { reader: R, max_records: Option, ) -> Result<(Schema, usize), ArrowError> { - self.infer_schema_impl::(reader, max_records) + let (schema, _, records) = + self.infer_schema_impl::(reader, max_records)?; + Ok((schema, records)) } /// Infer fixed-point CSV values as `Decimal128(precision, scale)` instead @@ -523,16 +583,29 @@ impl Format { reader: R, max_records: Option, ) -> Result<(Schema, usize), ArrowError> { - self.infer_schema_impl::(reader, max_records) + let (schema, _, records) = self.infer_schema_impl::(reader, max_records)?; + Ok((schema, records)) } - fn infer_schema_impl( + /// Infer exact decimals with Snowflake's order-sensitive CSV type promotion. + /// The returned per-column states can be merged in order when a file is + /// scanned in multiple chunks. Ordinary inference methods are unchanged. + pub fn infer_schema_with_snowflake_types( &self, reader: R, max_records: Option, - ) -> Result<(Schema, usize), ArrowError> { + ) -> Result<(Schema, Vec, usize), ArrowError> { + self.infer_schema_impl::(reader, max_records) + } + + fn infer_schema_impl( + &self, + reader: R, + max_records: Option, + ) -> Result<(Schema, Vec, usize), ArrowError> { if self.preserve_quoted_empty && self.null_regex.is_null("") { - return self.infer_schema_with_quoted_empty::(reader, max_records); + return self + .infer_schema_with_quoted_empty::(reader, max_records); } let mut csv_reader = self.build_reader(reader); @@ -567,8 +640,10 @@ impl Format { for (i, column_type) in column_types.iter_mut().enumerate().take(header_length) { if let Some(string) = record.get(i) { if !self.null_regex.is_null(string) { - if DECIMAL { - column_type.update_decimal(string); + if SNOWFLAKE { + column_type.update_snowflake(string); + } else if DECIMAL { + column_type.update_decimal::(string); } else { column_type.update(string); } @@ -581,17 +656,23 @@ impl Format { let fields: Fields = column_types .iter() .zip(&headers) - .map(|(inferred, field_name)| inferred_field::(field_name, inferred)) + .map(|(inferred, field_name)| { + inferred_field::(field_name, inferred) + }) .collect(); - Ok((Schema::new(fields), records_count)) + Ok(( + Schema::new(fields), + snowflake_states::(&column_types), + records_count, + )) } - fn infer_schema_with_quoted_empty( + fn infer_schema_with_quoted_empty( &self, reader: R, max_records: Option, - ) -> Result<(Schema, usize), ArrowError> { + ) -> Result<(Schema, Vec, usize), ArrowError> { let recording = RecordingReader { inner: reader, prefix: Vec::new(), @@ -613,7 +694,11 @@ impl Format { let mut column_types: Vec = vec![Default::default(); headers.len()]; let max_records = max_records.unwrap_or(usize::MAX); if max_records == 0 { - return Ok((inferred_schema::(&headers, &column_types), 0)); + return Ok(( + inferred_schema::(&headers, &column_types), + snowflake_states::(&column_types), + 0, + )); } let schema = Arc::new(Schema::new( @@ -648,8 +733,10 @@ impl Format { ArrowError::CsvError("CSV inference expected a string column".to_owned()) })?; for value in strings.iter().flatten() { - if DECIMAL { - column_type.update_decimal(value); + if SNOWFLAKE { + column_type.update_snowflake(value); + } else if DECIMAL { + column_type.update_decimal::(value); } else { column_type.update(value); } @@ -657,7 +744,8 @@ impl Format { } } Ok(( - inferred_schema::(&headers, &column_types), + inferred_schema::(&headers, &column_types), + snowflake_states::(&column_types), records_count, )) } @@ -703,7 +791,21 @@ impl Format { } } -fn inferred_schema( +fn snowflake_states( + column_types: &[InferredDataType], +) -> Vec { + if SNOWFLAKE { + column_types + .iter() + .copied() + .map(SnowflakeCsvTypeState) + .collect() + } else { + Vec::new() + } +} + +fn inferred_schema( headers: &[String], column_types: &[InferredDataType], ) -> Schema { @@ -711,13 +813,18 @@ fn inferred_schema( column_types .iter() .zip(headers) - .map(|(inferred, name)| inferred_field::(name, inferred)) + .map(|(inferred, name)| inferred_field::(name, inferred)) .collect::(), ) } -fn inferred_field(name: &str, inferred: &InferredDataType) -> Field { - let data_type = if DECIMAL { +fn inferred_field( + name: &str, + inferred: &InferredDataType, +) -> Field { + let data_type = if SNOWFLAKE { + inferred.get_snowflake() + } else if DECIMAL { inferred.get_decimal() } else { inferred.get() @@ -1750,6 +1857,114 @@ mod tests { assert_eq!(schema.field(2).data_type(), &DataType::Decimal128(1, 0)); } + #[test] + fn test_snowflake_ordered_inference_matches_numeric_and_temporal_probes() { + let tiny = format!("0.{}1", "0".repeat(37)); + let large = "123456789012345678901234567890123456789"; + let cases = [ + ("1\n1e2\n".to_owned(), DataType::Utf8), + ("1e2\n1\n".to_owned(), DataType::Float64), + (format!("1\n{large}\n"), DataType::Utf8), + (format!("{large}\n1\n"), DataType::Float64), + (format!("0\n{tiny}\n"), DataType::Utf8), + (format!("{tiny}\n0\n"), DataType::Float64), + ( + "2024-01-02\n2024-01-02 01:02:03\n".to_owned(), + DataType::Date32, + ), + ( + "2024-01-02 01:02:03\n2024-01-02\n".to_owned(), + DataType::Utf8, + ), + ( + format!("0\n0.{}1\n", "0".repeat(36)), + DataType::Decimal128(38, 37), + ), + ("1.20\n123.400\n".to_owned(), DataType::Decimal128(6, 3)), + ]; + for (values, expected) in cases { + let csv = format!("VALUE\n{values}"); + let (schema, _, records) = Format::default() + .with_header(true) + .infer_schema_with_snowflake_types(Cursor::new(csv), None) + .unwrap(); + assert_eq!(records, 2); + assert_eq!(schema.field(0).data_type(), &expected, "{values}"); + } + + let (limited, _, records) = Format::default() + .with_header(true) + .infer_schema_with_snowflake_types(Cursor::new("VALUE\n1\n1e2\n"), Some(1)) + .unwrap(); + assert_eq!(records, 1); + assert_eq!(limited.field(0).data_type(), &DataType::Decimal128(1, 0)); + } + + #[test] + fn test_snowflake_ordered_inference_preserves_quoted_empty() { + let tiny = format!("0.{}1", "0".repeat(37)); + let csv = format!("VALUE\n0\n\"\"\n{tiny}\n"); + let (schema, _, records) = Format::default() + .with_header(true) + .with_preserve_quoted_empty(true) + .infer_schema_with_snowflake_types(Cursor::new(csv), None) + .unwrap(); + assert_eq!(records, 3); + assert_eq!(schema.field(0).data_type(), &DataType::Utf8); + } + + #[test] + fn test_snowflake_states_merge_independent_of_chunk_boundaries() { + let cases = [ + ("2024-01-02", "2024-01-02 01:02:03\n2024-01-02"), + ("2024-01-02 01:02:03", "2024-01-02\n2024-01-02 01:02:03"), + ("1e2", "1\n1e2"), + ("1", "1e2\n1"), + ("1234567890123456789012345678901234567", "0.01"), + ("0.01", "1234567890123456789012345678901234567"), + ("12345678901234567890123456789012345678", "0.1"), + ("0.1", "12345678901234567890123456789012345678"), + ]; + let format = Format::default().with_header(true); + for (first, later) in cases { + let (direct, _, _) = format + .infer_schema_with_snowflake_types( + Cursor::new(format!("VALUE\n{first}\n{later}\n")), + None, + ) + .unwrap(); + let (_, mut first_state, _) = format + .infer_schema_with_snowflake_types(Cursor::new(format!("VALUE\n{first}\n")), None) + .unwrap(); + let (_, later_state, _) = format + .infer_schema_with_snowflake_types(Cursor::new(format!("VALUE\n{later}\n")), None) + .unwrap(); + first_state[0].merge(&later_state[0]); + assert_eq!(first_state[0].data_type(), *direct.field(0).data_type()); + } + + for (first, later, expected) in [ + ( + "1234567890123456789012345678901234567", + "0.01", + DataType::Decimal128(38, 1), + ), + ( + "12345678901234567890123456789012345678", + "0.1", + DataType::Decimal128(38, 0), + ), + ] { + let (schema, _, _) = format + .infer_schema_with_snowflake_types( + Cursor::new(format!("VALUE\n{first}\n{later}\n")), + None, + ) + .unwrap(); + assert_eq!(schema.field(0).data_type(), &expected); + } + } + #[test] fn test_decimal_inference_schema_reads_exact_values() { let tiny = format!("0.{}1", "0".repeat(37)); @@ -3691,7 +3906,7 @@ mod tests { for (values, expected) in cases { let mut t = InferredDataType::default(); for v in *values { - t.update(v) + t.update(v); } assert_eq!(&t.get(), expected, "{values:?}") }