This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11628-dad78c3243db92fefd57b8b4f3cf65901f874263 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit cdab920c847e62433e24162c897057c74c93d6a9 Author: Spicy Shark <[email protected]> AuthorDate: Sat Oct 17 07:26:21 2026 +0000 fix(csv): return an error instead of panicking on dictionary key overflow (#11628) # Which issue does this PR close? - Closes #11623 # Rationale for this change A CSV column read as `Dictionary(<key>, Utf8)` panics with "dictionary key overflow" when a batch holds more distinct values than the key type can index (for example 129 for `Int8`). A value that does not parse as its column type already comes back as an `Err` from the reader, and so should this. # What changes are included in this PR? The eight per-key-type arms in `parse()` collected into `DictionaryArray<K>`, which goes through the infallible `append_value` of the dictionary builder. They now call one helper, `build_dictionary_array::<K>()`, which uses `StringDictionaryBuilder` and its fallible `append`, so the overflow comes back as `ArrowError::DictionaryKeyOverflowError`. Null handling is unchanged. # Are these changes tested? Yes. `test_csv_dictionary_key_overflow` reads a column with `Int8` keys: 128 distinct values read fine, and 129 return `DictionaryKeyOverflowError`. Without the change the test panics in `GenericByteDictionaryBuilder`. The helper is shared by all eight key types, which `test_csv_with_nullable_dictionary` already reads. # Are there any significant user-facing changes? No. --- arrow-csv/src/reader/mod.rs | 119 ++++++++++++++++++++------------------------ 1 file changed, 54 insertions(+), 65 deletions(-) diff --git a/arrow-csv/src/reader/mod.rs b/arrow-csv/src/reader/mod.rs index cefc9ec01e..6347b3a3b4 100644 --- a/arrow-csv/src/reader/mod.rs +++ b/arrow-csv/src/reader/mod.rs @@ -165,7 +165,7 @@ mod records; -use arrow_array::builder::{NullBuilder, PrimitiveBuilder}; +use arrow_array::builder::{NullBuilder, PrimitiveBuilder, StringDictionaryBuilder}; use arrow_array::types::*; use arrow_array::*; use arrow_cast::parse::{Parser, parse_decimal, string_to_datetime}; @@ -1033,70 +1033,20 @@ fn parse( if value_type.as_ref() == &DataType::Utf8 => { match key_type.as_ref() { - DataType::Int8 => Ok(Arc::new( - rows.iter() - .map(|row| { - let s = row.get(i); - (!null_regex.is_null(s)).then_some(s) - }) - .collect::<DictionaryArray<Int8Type>>(), - ) as ArrayRef), - DataType::Int16 => Ok(Arc::new( - rows.iter() - .map(|row| { - let s = row.get(i); - (!null_regex.is_null(s)).then_some(s) - }) - .collect::<DictionaryArray<Int16Type>>(), - ) as ArrayRef), - DataType::Int32 => Ok(Arc::new( - rows.iter() - .map(|row| { - let s = row.get(i); - (!null_regex.is_null(s)).then_some(s) - }) - .collect::<DictionaryArray<Int32Type>>(), - ) as ArrayRef), - DataType::Int64 => Ok(Arc::new( - rows.iter() - .map(|row| { - let s = row.get(i); - (!null_regex.is_null(s)).then_some(s) - }) - .collect::<DictionaryArray<Int64Type>>(), - ) as ArrayRef), - DataType::UInt8 => Ok(Arc::new( - rows.iter() - .map(|row| { - let s = row.get(i); - (!null_regex.is_null(s)).then_some(s) - }) - .collect::<DictionaryArray<UInt8Type>>(), - ) as ArrayRef), - DataType::UInt16 => Ok(Arc::new( - rows.iter() - .map(|row| { - let s = row.get(i); - (!null_regex.is_null(s)).then_some(s) - }) - .collect::<DictionaryArray<UInt16Type>>(), - ) as ArrayRef), - DataType::UInt32 => Ok(Arc::new( - rows.iter() - .map(|row| { - let s = row.get(i); - (!null_regex.is_null(s)).then_some(s) - }) - .collect::<DictionaryArray<UInt32Type>>(), - ) as ArrayRef), - DataType::UInt64 => Ok(Arc::new( - rows.iter() - .map(|row| { - let s = row.get(i); - (!null_regex.is_null(s)).then_some(s) - }) - .collect::<DictionaryArray<UInt64Type>>(), - ) as ArrayRef), + DataType::Int8 => build_dictionary_array::<Int8Type>(rows, i, null_regex), + DataType::Int16 => build_dictionary_array::<Int16Type>(rows, i, null_regex), + DataType::Int32 => build_dictionary_array::<Int32Type>(rows, i, null_regex), + DataType::Int64 => build_dictionary_array::<Int64Type>(rows, i, null_regex), + DataType::UInt8 => build_dictionary_array::<UInt8Type>(rows, i, null_regex), + DataType::UInt16 => { + build_dictionary_array::<UInt16Type>(rows, i, null_regex) + } + DataType::UInt32 => { + build_dictionary_array::<UInt32Type>(rows, i, null_regex) + } + DataType::UInt64 => { + build_dictionary_array::<UInt64Type>(rows, i, null_regex) + } _ => Err(ArrowError::ParseError(format!( "Unsupported dictionary key type {key_type}" ))), @@ -1282,6 +1232,24 @@ fn build_boolean_array( .map(|e| Arc::new(e) as ArrayRef) } +// parses a specific column (col_idx) into an Arrow dictionary array of strings. +fn build_dictionary_array<K: ArrowDictionaryKeyType>( + rows: &StringRecords<'_>, + col_idx: usize, + null_regex: &NullRegex, +) -> Result<ArrayRef, ArrowError> { + let mut builder = StringDictionaryBuilder::<K>::with_capacity(rows.len(), 256, 1024); + for row in rows.iter() { + let s = row.get(col_idx); + if null_regex.is_null(s) { + builder.append_null(); + } else { + builder.append(s)?; + } + } + Ok(Arc::new(builder.finish())) +} + /// Builder for CSV [`Reader`]s #[derive(Debug)] pub struct ReaderBuilder { @@ -2068,6 +2036,27 @@ mod tests { assert_eq!(strings.value(29), "Uckfield, East Sussex, UK"); } + #[test] + fn test_csv_dictionary_key_overflow() { + let schema = Arc::new(Schema::new(vec![Field::new_dictionary( + "c", + DataType::Int8, + DataType::Utf8, + false, + )])); + let read = |distinct: usize| { + let rows = (0..distinct).map(|i| format!("v{i}")).collect::<Vec<_>>(); + let mut reader = ReaderBuilder::new(schema.clone()) + .build(Cursor::new(rows.join("\n"))) + .unwrap(); + reader.next().unwrap() + }; + + assert_eq!(read(128).unwrap().num_rows(), 128); + let err = read(129).unwrap_err(); + assert!(matches!(err, ArrowError::DictionaryKeyOverflowError)); + } + #[test] fn test_csv_with_nullable_dictionary() { let offset_type = vec