Skip to content

Add opt-in ordered Snowflake CSV schema inference #15

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
osipovartem merged 2 commits into from
Oct 8, 2026
Merged

Add opt-in ordered Snowflake CSV schema inference #15

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ec817ac
Add opt-in ordered Snowflake CSV type inference
osipovartem Oct 8, 2026
4e743e5
Keep ordered CSV inference state across chunks
osipovartem Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
267 changes: 241 additions & 26 deletions arrow-csv/src/reader/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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::<true>(string);
if self.first_snowflake_type == 0 {
self.first_snowflake_type = value_type;
}
}

fn update_decimal<const SNOWFLAKE: bool>(&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::<i64>().is_err() {
Expand All @@ -371,7 +427,9 @@ impl InferredDataType {
1 << 2 // Float64
} else {
1 << 8 // Utf8
}
};
self.packed |= value_type;
value_type
}
}

Expand Down Expand Up @@ -512,7 +570,9 @@ impl Format {
reader: R,
max_records: Option<usize>,
) -> Result<(Schema, usize), ArrowError> {
self.infer_schema_impl::<R, false>(reader, max_records)
let (schema, _, records) =
self.infer_schema_impl::<R, false, false>(reader, max_records)?;
Ok((schema, records))
}

/// Infer fixed-point CSV values as `Decimal128(precision, scale)` instead
Expand All @@ -523,16 +583,29 @@ impl Format {
reader: R,
max_records: Option<usize>,
) -> Result<(Schema, usize), ArrowError> {
self.infer_schema_impl::<R, true>(reader, max_records)
let (schema, _, records) = self.infer_schema_impl::<R, true, false>(reader, max_records)?;
Ok((schema, records))
}

fn infer_schema_impl<R: Read, const DECIMAL: bool>(
/// 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<R: Read>(
&self,
reader: R,
max_records: Option<usize>,
) -> Result<(Schema, usize), ArrowError> {
) -> Result<(Schema, Vec<SnowflakeCsvTypeState>, usize), ArrowError> {
self.infer_schema_impl::<R, true, true>(reader, max_records)
}

fn infer_schema_impl<R: Read, const DECIMAL: bool, const SNOWFLAKE: bool>(
&self,
reader: R,
max_records: Option<usize>,
) -> Result<(Schema, Vec<SnowflakeCsvTypeState>, usize), ArrowError> {
if self.preserve_quoted_empty && self.null_regex.is_null("") {
return self.infer_schema_with_quoted_empty::<R, DECIMAL>(reader, max_records);
return self
.infer_schema_with_quoted_empty::<R, DECIMAL, SNOWFLAKE>(reader, max_records);
}
let mut csv_reader = self.build_reader(reader);

Expand Down Expand Up @@ -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::<false>(string);
} else {
column_type.update(string);
}
Expand All @@ -581,17 +656,23 @@ impl Format {
let fields: Fields = column_types
.iter()
.zip(&headers)
.map(|(inferred, field_name)| inferred_field::<DECIMAL>(field_name, inferred))
.map(|(inferred, field_name)| {
inferred_field::<DECIMAL, SNOWFLAKE>(field_name, inferred)
})
.collect();

Ok((Schema::new(fields), records_count))
Ok((
Schema::new(fields),
snowflake_states::<SNOWFLAKE>(&column_types),
records_count,
))
}

fn infer_schema_with_quoted_empty<R: Read, const DECIMAL: bool>(
fn infer_schema_with_quoted_empty<R: Read, const DECIMAL: bool, const SNOWFLAKE: bool>(
&self,
reader: R,
max_records: Option<usize>,
) -> Result<(Schema, usize), ArrowError> {
) -> Result<(Schema, Vec<SnowflakeCsvTypeState>, usize), ArrowError> {
let recording = RecordingReader {
inner: reader,
prefix: Vec::new(),
Expand All @@ -613,7 +694,11 @@ impl Format {
let mut column_types: Vec<InferredDataType> = vec![Default::default(); headers.len()];
let max_records = max_records.unwrap_or(usize::MAX);
if max_records == 0 {
return Ok((inferred_schema::<DECIMAL>(&headers, &column_types), 0));
return Ok((
inferred_schema::<DECIMAL, SNOWFLAKE>(&headers, &column_types),
snowflake_states::<SNOWFLAKE>(&column_types),
0,
));
}

let schema = Arc::new(Schema::new(
Expand Down Expand Up @@ -648,16 +733,19 @@ 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::<false>(value);
} else {
column_type.update(value);
}
}
}
}
Ok((
inferred_schema::<DECIMAL>(&headers, &column_types),
inferred_schema::<DECIMAL, SNOWFLAKE>(&headers, &column_types),
snowflake_states::<SNOWFLAKE>(&column_types),
records_count,
))
}
Expand Down Expand Up @@ -703,21 +791,40 @@ impl Format {
}
}

fn inferred_schema<const DECIMAL: bool>(
fn snowflake_states<const SNOWFLAKE: bool>(
column_types: &[InferredDataType],
) -> Vec<SnowflakeCsvTypeState> {
if SNOWFLAKE {
column_types
.iter()
.copied()
.map(SnowflakeCsvTypeState)
.collect()
} else {
Vec::new()
}
}

fn inferred_schema<const DECIMAL: bool, const SNOWFLAKE: bool>(
headers: &[String],
column_types: &[InferredDataType],
) -> Schema {
Schema::new(
column_types
.iter()
.zip(headers)
.map(|(inferred, name)| inferred_field::<DECIMAL>(name, inferred))
.map(|(inferred, name)| inferred_field::<DECIMAL, SNOWFLAKE>(name, inferred))
.collect::<Fields>(),
)
}

fn inferred_field<const DECIMAL: bool>(name: &str, inferred: &InferredDataType) -> Field {
let data_type = if DECIMAL {
fn inferred_field<const DECIMAL: bool, const SNOWFLAKE: bool>(
name: &str,
inferred: &InferredDataType,
) -> Field {
let data_type = if SNOWFLAKE {
inferred.get_snowflake()
} else if DECIMAL {
inferred.get_decimal()
} else {
inferred.get()
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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:?}")
}
Expand Down
Loading