| 83 | } |
| 84 | |
| 85 | Status TupleFileWriter::Write(RuntimeState* state, RowBatch* row_batch) { |
| 86 | DCHECK_EQ(state_, State::InProgress); |
| 87 | SCOPED_TIMER(write_timer_); |
| 88 | // serialize and write row batch |
| 89 | { |
| 90 | SCOPED_TIMER(serialize_timer_); |
| 91 | // Passing in nullptr for 'compression_scratch' disables compression. |
| 92 | RETURN_IF_ERROR( |
| 93 | row_batch->Serialize(out_batch_.get(), /* compression_scratch */ nullptr)); |
| 94 | } |
| 95 | |
| 96 | if (out_batch_->header()->num_rows() == 0) { |
| 97 | DCHECK_EQ(out_batch_->header()->uncompressed_size(), 0); |
| 98 | return Status::OK(); |
| 99 | } |
| 100 | |
| 101 | // Collect all of the pieces that we would want to write, then determine if writing |
| 102 | // them would exceed the max file size. |
| 103 | std::string header_buf; |
| 104 | if (!out_batch_->header()->SerializeToString(&header_buf)) { |
| 105 | return Status(TErrorCode::INTERNAL_ERROR, |
| 106 | "Could not serialize RowBatchHeaderPB to string"); |
| 107 | } |
| 108 | size_t header_len = header_buf.size(); |
| 109 | DCHECK_GT(header_len, 0); |
| 110 | kudu::Slice tuple_data = out_batch_->TupleDataAsSlice(); |
| 111 | kudu::Slice tuple_offsets = out_batch_->TupleOffsetsAsSlice(); |
| 112 | // tuple_data_len is possible to be 0, see IMPALA-13411. |
| 113 | size_t tuple_data_len = tuple_data.size(); |
| 114 | size_t tuple_offsets_len = tuple_offsets.size(); |
| 115 | DCHECK_GT(tuple_offsets_len, 0); |
| 116 | |
| 117 | // We write things in this order (sizes first, then the variable-sized data): |
| 118 | // 1. The size of the header |
| 119 | // 2. The size of the tuple data |
| 120 | // 3. The size of the tuple offsets |
| 121 | // 4. The serialized header |
| 122 | // 5. The tuple data |
| 123 | // 6. The tuple offsets |
| 124 | std::vector<kudu::Slice> slices = { |
| 125 | kudu::Slice(reinterpret_cast<const char*>(&header_len), sizeof(header_len)), |
| 126 | kudu::Slice(reinterpret_cast<const char*>(&tuple_data_len), |
| 127 | sizeof(tuple_data_len)), |
| 128 | kudu::Slice(reinterpret_cast<const char*>(&tuple_offsets_len), |
| 129 | sizeof(tuple_offsets_len)), |
| 130 | kudu::Slice(header_buf), |
| 131 | kudu::Slice(tuple_data), |
| 132 | kudu::Slice(tuple_offsets)}; |
| 133 | |
| 134 | // Enforce the max file size |
| 135 | size_t num_bytes_to_write = 0; |
| 136 | for (auto slice : slices) { |
| 137 | num_bytes_to_write += slice.size(); |
| 138 | } |
| 139 | |
| 140 | if (request_write_size_cb_ != nullptr) { |
| 141 | // If the request_write_size_cb is set, call it before writing to get permission. |
| 142 | // It will return a not-OK status if this writer should stop. |