Return the inferred schema reading up to records_to_read from a stream of delimited chunks returning the inferred schema and the number of lines that were read. This method can handle CSV files with different numbers of columns. The inferred schema will be the union of all columns found across all files. Files with fewer columns will have missing columns filled with null values. # Example If yo
(
&self,
state: &dyn Session,
mut records_to_read: usize,
stream: impl Stream<Item = Result<Bytes>>,
)
| 517 | /// The inferred schema will contain all 5 columns, with files that don't |
| 518 | /// have columns 4 and 5 having null values for those columns. |
| 519 | pub async fn infer_schema_from_stream( |
| 520 | &self, |
| 521 | state: &dyn Session, |
| 522 | mut records_to_read: usize, |
| 523 | stream: impl Stream<Item = Result<Bytes>>, |
| 524 | ) -> Result<(Schema, usize)> { |
| 525 | let mut total_records_read = 0; |
| 526 | let mut column_names = vec![]; |
| 527 | let mut column_type_possibilities = vec![]; |
| 528 | let mut record_number = -1; |
| 529 | let initial_records_to_read = records_to_read; |
| 530 | |
| 531 | pin_mut!(stream); |
| 532 | |
| 533 | while let Some(chunk) = stream.next().await.transpose()? { |
| 534 | record_number += 1; |
| 535 | let first_chunk = record_number == 0; |
| 536 | let mut format = arrow::csv::reader::Format::default() |
| 537 | .with_header( |
| 538 | first_chunk |
| 539 | && self |
| 540 | .options |
| 541 | .has_header |
| 542 | .unwrap_or_else(|| state.config_options().catalog.has_header), |
| 543 | ) |
| 544 | .with_delimiter(self.options.delimiter) |
| 545 | .with_quote(self.options.quote) |
| 546 | .with_truncated_rows(self.options.truncated_rows.unwrap_or(false)); |
| 547 | |
| 548 | if let Some(null_regex) = &self.options.null_regex { |
| 549 | let regex = Regex::new(null_regex.as_str()) |
| 550 | .expect("Unable to parse CSV null regex."); |
| 551 | format = format.with_null_regex(regex); |
| 552 | } |
| 553 | |
| 554 | if let Some(escape) = self.options.escape { |
| 555 | format = format.with_escape(escape); |
| 556 | } |
| 557 | |
| 558 | if let Some(comment) = self.options.comment { |
| 559 | format = format.with_comment(comment); |
| 560 | } |
| 561 | |
| 562 | let (Schema { fields, .. }, records_read) = |
| 563 | format.infer_schema(chunk.reader(), Some(records_to_read))?; |
| 564 | |
| 565 | records_to_read -= records_read; |
| 566 | total_records_read += records_read; |
| 567 | |
| 568 | if first_chunk { |
| 569 | // set up initial structures for recording inferred schema across chunks |
| 570 | (column_names, column_type_possibilities) = fields |
| 571 | .into_iter() |
| 572 | .map(|field| { |
| 573 | let mut possibilities = HashSet::new(); |
| 574 | if records_read > 0 { |
| 575 | // at least 1 data row read, record the inferred datatype |
| 576 | possibilities.insert(field.data_type().clone()); |
no test coverage detected