This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new d4a5d3951a Infer CSV headers from sampled column types (#10361)
d4a5d3951a is described below
commit d4a5d3951a507ff1c0e96681885e2a8dd6f2640f
Author: Johan Vaz <[email protected]>
AuthorDate: Thu Sep 24 12:15:56 2026 +0800
Infer CSV headers from sampled column types (#10361)
# Which issue does this PR close?
Closes #6882.
# Rationale for this change
When callers do not know whether a CSV begins with a header, they
currently need to pre-scan it or choose a setting themselves. This adds
an opt-in, conservative format-inference pass. The existing header
default and ordinary schema-inference rules remain unchanged.
# What changes are included in this PR?
- Add `Format::infer_format`, returning the updated format and number of
records consumed, including the initial candidate row.
- Infer a header from text in the first record followed by a consistent
non-text, non-null column type. Empty, single-record, and all-text
samples remain ambiguous and headerless.
- Reuse the configured CSV parser, null handling, and type inference.
The optional limit counts records after the first candidate row.
- Do not use numeric-looking first-row values such as `+1`, `+1.5`, or
an overflowing integer as header evidence. This conservative check is
local to header detection; it does not alter the schema's numeric
regexes or inferred column types.
- Merge current upstream and preserve its Miri test annotation. The
shared leading-plus numeric-inference change has been removed from this
PR and submitted separately in #11105.
The reader is consumed, as with `infer_schema`; callers must rewind or
create another reader before parsing.
# Are these changes tested?
Validated on `8f66150ed` against upstream `c60ac3fc6` with the
repository-pinned Rust 1.98.1 and test fixture submodules:
- `cargo test -p arrow-csv --all-features`: 91 unit tests and 14
doctests passed.
- `cargo test -p arrow --features csv`: 159 unit/integration tests and
11 doctests passed.
- `cargo clippy -p arrow-csv --all-targets --all-features -- -D
warnings`: passed.
- `cargo fmt --all -- --check`: passed.
- `git diff upstream/main --check`: passed.
Regression coverage checks typed headers, headerless/ambiguous inputs,
delimiter settings, bounded sampling, consumed row counts, and
numeric-looking first rows. The latter explicitly compares the schema
with ordinary `Format::default().infer_schema` and verifies no records
are skipped.
# Are there any user-facing changes?
A new opt-in public API. It is a heuristic: all-text data and other
ambiguous samples may need an explicit header setting. No shared
schema-inference behavior changes are included.
# AI assistance disclosure
AI assistance was used for implementation, regression tests, conflict
resolution, and this description. The final diff was inspected against
current upstream; the listed checks were executed on the final code.
Maintainer acceptance is pending.
---
arrow-csv/src/reader/mod.rs | 196 ++++++++++++++++++++++++++++++++++++++++++++
1 file changed, 196 insertions(+)
diff --git a/arrow-csv/src/reader/mod.rs b/arrow-csv/src/reader/mod.rs
index df4c450c14..cefc9ec01e 100644
--- a/arrow-csv/src/reader/mod.rs
+++ b/arrow-csv/src/reader/mod.rs
@@ -351,6 +351,98 @@ impl Format {
self
}
+ /// Infer format settings from the CSV records in `reader`
+ ///
+ /// This currently infers whether the first record is a header. Up to
+ /// `max_records` records after the first record are inspected; if `None`,
all
+ /// records are read. Detection is conservative and returns no header when
the
+ /// sampled records do not provide type evidence. Returns the updated
format
+ /// and the number of records read, including the first header candidate.
+ ///
+ /// # Example
+ ///
+ /// ```
+ /// use arrow_csv::reader::Format;
+ /// use std::io::Cursor;
+ ///
+ /// let csv = "name,count\nalice,1\nbob,2\n";
+ /// let (format, format_records_read) =
+ /// Format::default().infer_format(Cursor::new(csv), Some(10))?;
+ /// let (schema, records_read) = format.infer_schema(Cursor::new(csv),
None)?;
+ ///
+ /// assert_eq!(schema.field(0).name(), "name");
+ /// assert_eq!(format_records_read, 3);
+ /// assert_eq!(records_read, 2);
+ /// # Ok::<_, arrow_schema::ArrowError>(())
+ /// ```
+ pub fn infer_format<R: Read>(
+ mut self,
+ reader: R,
+ max_records: Option<usize>,
+ ) -> Result<(Self, usize), ArrowError> {
+ let (header, records_read) = self.infer_header(reader, max_records)?;
+ self.header = header;
+ Ok((self, records_read))
+ }
+
+ /// Infer whether the first CSV record is a header
+ ///
+ /// Inspects up to `max_records` records after the first record. Returns
`true`
+ /// when a value in the first record is text while the remaining values in
the
+ /// same column have a consistent non-text type.
+ fn infer_header<R: Read>(
+ &self,
+ reader: R,
+ max_records: Option<usize>,
+ ) -> Result<(bool, usize), ArrowError> {
+ let mut format = self.clone();
+ format.header = false;
+ let mut csv_reader = format.build_reader(reader);
+
+ let mut first_record = StringRecord::new();
+ if !csv_reader
+ .read_record(&mut first_record)
+ .map_err(map_csv_error)?
+ {
+ return Ok((false, 0));
+ }
+
+ let mut first_types = vec![InferredDataType::default();
first_record.len()];
+ for (value, inferred) in first_record.iter().zip(&mut first_types) {
+ if !self.null_regex.is_null(value) {
+ inferred.update(value);
+ }
+ }
+
+ let mut column_types = vec![InferredDataType::default();
first_record.len()];
+ let mut record = StringRecord::new();
+ let mut records_count = 0;
+ let max_records = max_records.unwrap_or(usize::MAX);
+ while records_count < max_records
+ && csv_reader.read_record(&mut record).map_err(map_csv_error)?
+ {
+ records_count += 1;
+ for (value, inferred) in record.iter().zip(&mut column_types) {
+ if !self.null_regex.is_null(value) {
+ inferred.update(value);
+ }
+ }
+ }
+
+ let has_header = first_types
+ .iter()
+ .zip(&column_types)
+ .zip(first_record.iter())
+ .any(|((first, rest), value)| {
+ // Numeric-looking values (e.g. +1) are not header evidence,
even
+ // when ordinary schema inference conservatively treats them
as text.
+ first.get() == DataType::Utf8
+ && value.parse::<f64>().is_err()
+ && !matches!(rest.get(), DataType::Utf8 | DataType::Null)
+ });
+ Ok((has_header, records_count + 1))
+ }
+
/// Infer schema of CSV records from the provided `reader`
///
/// If `max_records` is `None`, all records will be read, otherwise up to
`max_records`
@@ -1634,6 +1726,110 @@ mod tests {
assert_eq!(3, batch.num_columns());
}
+ #[test]
+ fn test_infer_format_with_typed_columns() {
+ let csv = "name,count,active\nalice,1,true\nbob,2,false\n";
+
+ let (format, format_records_read) = Format::default()
+ .infer_format(Cursor::new(csv), None)
+ .unwrap();
+ let (schema, records_read) = format.infer_schema(Cursor::new(csv),
None).unwrap();
+
+ assert_eq!(schema.field(0).name(), "name");
+ assert_eq!(schema.field(1).name(), "count");
+ assert_eq!(schema.field(2).name(), "active");
+ assert_eq!(format_records_read, 3);
+ assert_eq!(records_read, 2);
+ }
+
+ #[test]
+ fn test_infer_format_without_header() {
+ let csv = "1,true\n2,false\n";
+
+ let (format, format_records_read) = Format::default()
+ .infer_format(Cursor::new(csv), None)
+ .unwrap();
+ let (schema, records_read) = format.infer_schema(Cursor::new(csv),
None).unwrap();
+
+ assert_eq!(schema.field(0).name(), "column_1");
+ assert_eq!(schema.field(1).name(), "column_2");
+ assert_eq!(format_records_read, 2);
+ assert_eq!(records_read, 2);
+ }
+
+ #[test]
+ fn test_infer_format_returns_no_header_when_ambiguous() {
+ for csv in ["name,count\n", "alice,london\nbob,paris\n"] {
+ let (format, _) = Format::default()
+ .infer_format(Cursor::new(csv), None)
+ .unwrap();
+ let (schema, _) = format.infer_schema(Cursor::new(csv),
None).unwrap();
+ assert_eq!(schema.field(0).name(), "column_1", "CSV: {csv:?}");
+ }
+
+ let (format, format_records_read) = Format::default()
+ .infer_format(Cursor::new(""), None)
+ .unwrap();
+ let (schema, records_read) = format.infer_schema(Cursor::new(""),
None).unwrap();
+ assert!(schema.fields().is_empty());
+ assert_eq!(format_records_read, 0);
+ assert_eq!(records_read, 0);
+ }
+
+ #[test]
+ fn test_infer_format_honors_format_options() {
+ let csv = "name;count\nalice;1\nbob;2\n";
+ let (format, _) = Format::default()
+ .with_delimiter(b';')
+ .infer_format(Cursor::new(csv), None)
+ .unwrap();
+ let (schema, records_read) = format.infer_schema(Cursor::new(csv),
None).unwrap();
+
+ assert_eq!(schema.field(0).name(), "name");
+ assert_eq!(schema.field(1).name(), "count");
+ assert_eq!(records_read, 2);
+ }
+
+ #[test]
+ fn test_infer_format_respects_max_records() {
+ let csv = "name,count\nalice,1\nbob,unknown\n";
+ let infer = |max_records| {
+ let (format, records_read) = Format::default()
+ .infer_format(Cursor::new(csv), max_records)
+ .unwrap();
+ let (schema, _) = format.infer_schema(Cursor::new(csv),
None).unwrap();
+ (schema.field(0).name().clone(), records_read)
+ };
+
+ assert_eq!(infer(Some(1)), ("name".to_string(), 2));
+ assert_eq!(infer(None), ("column_1".to_string(), 3));
+ assert_eq!(infer(Some(0)), ("column_1".to_string(), 1));
+ }
+
+ #[test]
+ fn test_infer_format_numeric_text_is_not_header() {
+ for csv in [
+ "+1\n2\n3\n",
+ "+1.5\n2.5\n3.5\n",
+ "+1e3\n2e3\n3e3\n",
+ "9223372036854775808\n2\n3\n",
+ ] {
+ let (format, records_read) = Format::default()
+ .infer_format(Cursor::new(csv), None)
+ .unwrap();
+ let (schema, schema_records_read) =
+ format.infer_schema(Cursor::new(csv), None).unwrap();
+ let (ordinary_schema, _) = Format::default()
+ .infer_schema(Cursor::new(csv), None)
+ .unwrap();
+
+ assert_eq!(schema.field(0).name(), "column_1", "CSV: {csv:?}");
+ assert_eq!(schema, ordinary_schema, "CSV: {csv:?}");
+ assert_eq!(records_read, 3);
+ assert_eq!(schema_records_read, 3);
+ }
+ }
+
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn test_csv_with_schema_inference() {
