| 404 | std::shared_ptr<::arrow::Schema> schema() const override { return schema_; } |
| 405 | |
| 406 | Status WriteTable(const Table& table, int64_t chunk_size) override { |
| 407 | RETURN_NOT_OK(CheckClosed()); |
| 408 | RETURN_NOT_OK(table.Validate()); |
| 409 | |
| 410 | if (chunk_size <= 0 && table.num_rows() > 0) { |
| 411 | return Status::Invalid("chunk size per row_group must be greater than 0"); |
| 412 | } else if (!table.schema()->Equals(*schema_, false)) { |
| 413 | return Status::Invalid("table schema does not match this writer's. table:'", |
| 414 | table.schema()->ToString(), "' this:'", schema_->ToString(), |
| 415 | "'"); |
| 416 | } else if (chunk_size > this->properties().max_row_group_length()) { |
| 417 | chunk_size = this->properties().max_row_group_length(); |
| 418 | } |
| 419 | |
| 420 | auto WriteRowGroup = [&](int64_t offset, int64_t size) { |
| 421 | RETURN_NOT_OK(NewRowGroup()); |
| 422 | for (int i = 0; i < table.num_columns(); i++) { |
| 423 | RETURN_NOT_OK(WriteColumnChunk(table.column(i), offset, size)); |
| 424 | } |
| 425 | return Status::OK(); |
| 426 | }; |
| 427 | |
| 428 | if (table.num_rows() == 0) { |
| 429 | // Append a row group with 0 rows |
| 430 | RETURN_NOT_OK( |
| 431 | WriteRowGroup(0, 0).OrElse([&](auto&&) { PARQUET_IGNORE_NOT_OK(Close()); })); |
| 432 | } |
| 433 | |
| 434 | for (int chunk = 0; chunk * chunk_size < table.num_rows(); chunk++) { |
| 435 | int64_t offset = chunk * chunk_size; |
| 436 | RETURN_NOT_OK(WriteRowGroup(offset, std::min(chunk_size, table.num_rows() - offset)) |
| 437 | .OrElse([&](auto&&) { PARQUET_IGNORE_NOT_OK(Close()); })); |
| 438 | } |
| 439 | return Status::OK(); |
| 440 | } |
| 441 | |
| 442 | Status NewBufferedRowGroup() override { |
| 443 | RETURN_NOT_OK(CheckClosed()); |