Read header and column names from buffer, create column builders Returns the # of bytes consumed
| 580 | // Read header and column names from buffer, create column builders |
| 581 | // Returns the # of bytes consumed |
| 582 | Result<int64_t> ProcessHeader(const std::shared_ptr<Buffer>& buf, |
| 583 | std::shared_ptr<Buffer>* rest) { |
| 584 | const uint8_t* data = buf->data(); |
| 585 | const auto data_end = data + buf->size(); |
| 586 | DCHECK_GT(data_end - data, 0); |
| 587 | int64_t num_rows_seen = 1; |
| 588 | |
| 589 | if (read_options_.skip_rows) { |
| 590 | // Skip initial rows (potentially invalid CSV data) |
| 591 | auto num_skipped_rows = SkipRows(data, static_cast<uint32_t>(data_end - data), |
| 592 | read_options_.skip_rows, &data); |
| 593 | if (num_skipped_rows < read_options_.skip_rows) { |
| 594 | return Status::Invalid( |
| 595 | "Could not skip initial ", read_options_.skip_rows, |
| 596 | " rows from CSV file, " |
| 597 | "either file is too short or header is larger than block size"); |
| 598 | } |
| 599 | if (count_rows_) { |
| 600 | num_rows_seen += num_skipped_rows; |
| 601 | } |
| 602 | } |
| 603 | |
| 604 | if (read_options_.column_names.empty()) { |
| 605 | // Parse one row (either to read column names or to know the number of columns) |
| 606 | BlockParser parser(io_context_.pool(), parse_options_, /*num_cols=*/-1, |
| 607 | /*first_row=*/num_rows_seen, /*max_num_rows=*/1); |
| 608 | uint32_t parsed_size = 0; |
| 609 | RETURN_NOT_OK(parser.Parse( |
| 610 | std::string_view(reinterpret_cast<const char*>(data), data_end - data), |
| 611 | &parsed_size)); |
| 612 | if (parser.num_rows() != 1) { |
| 613 | return Status::Invalid( |
| 614 | "Could not read first row from CSV file, either " |
| 615 | "file is too short or header is larger than block size"); |
| 616 | } |
| 617 | if (parser.num_cols() == 0) { |
| 618 | return Status::Invalid("No columns in CSV file"); |
| 619 | } |
| 620 | |
| 621 | if (read_options_.autogenerate_column_names) { |
| 622 | column_names_ = GenerateColumnNames(parser.num_cols()); |
| 623 | } else { |
| 624 | // Read column names from header row |
| 625 | auto visit = [&](const uint8_t* data, uint32_t size, bool quoted) -> Status { |
| 626 | column_names_.emplace_back(reinterpret_cast<const char*>(data), size); |
| 627 | return Status::OK(); |
| 628 | }; |
| 629 | RETURN_NOT_OK(parser.VisitLastRow(visit)); |
| 630 | DCHECK_EQ(static_cast<size_t>(parser.num_cols()), column_names_.size()); |
| 631 | // Skip parsed header row |
| 632 | data += parsed_size; |
| 633 | if (count_rows_) { |
| 634 | ++num_rows_seen; |
| 635 | } |
| 636 | } |
| 637 | } else { |
| 638 | column_names_ = read_options_.column_names; |
| 639 | } |