Skip to content

Commit 84841ea

Browse files
committed
Refactor: Use WriteAndClearBuffer helper to prevent stale content
Addresses review feedback from HuaHuaY: Instead of clearing the buffer in TranslateMinimalBatch for empty batches, use a WriteAndClearBuffer() helper that writes and clears the buffer in all write paths. This is cleaner because: - Every write follows the same pattern (write -> clear) - Easier to reason about for future write stages - The invariant is explicit: buffer is always clean after flush
1 parent 664e11d commit 84841ea

1 file changed

Lines changed: 10 additions & 5 deletions

File tree

cpp/src/arrow/csv/writer.cc

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -541,7 +541,7 @@ class CSVWriterImpl : public ipc::RecordBatchWriter {
541541
for (auto maybe_slice : iterator) {
542542
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<RecordBatch> slice, maybe_slice);
543543
RETURN_NOT_OK(TranslateMinimalBatch(*slice));
544-
RETURN_NOT_OK(sink_->Write(data_buffer_));
544+
RETURN_NOT_OK(WriteAndClearBuffer());
545545
stats_.num_record_batches++;
546546
}
547547
return Status::OK();
@@ -554,7 +554,7 @@ class CSVWriterImpl : public ipc::RecordBatchWriter {
554554
RETURN_NOT_OK(reader.ReadNext(&batch));
555555
while (batch != nullptr) {
556556
RETURN_NOT_OK(TranslateMinimalBatch(*batch));
557-
RETURN_NOT_OK(sink_->Write(data_buffer_));
557+
RETURN_NOT_OK(WriteAndClearBuffer());
558558
RETURN_NOT_OK(reader.ReadNext(&batch));
559559
stats_.num_record_batches++;
560560
}
@@ -590,6 +590,13 @@ class CSVWriterImpl : public ipc::RecordBatchWriter {
590590
return Status::OK();
591591
}
592592

593+
// GH-36889: Write buffer to sink and clear it to avoid stale content
594+
// being written again if the next batch is empty.
595+
Status WriteAndClearBuffer() {
596+
RETURN_NOT_OK(sink_->Write(data_buffer_));
597+
return data_buffer_->Resize(0, /*shrink_to_fit=*/false);
598+
}
599+
593600
int64_t CalculateHeaderSize(QuotingStyle quoting_style) const {
594601
int64_t header_length = 0;
595602
for (int col = 0; col < schema_->num_fields(); col++) {
@@ -654,13 +661,11 @@ class CSVWriterImpl : public ipc::RecordBatchWriter {
654661
next += options_.eol.size();
655662
DCHECK_EQ(reinterpret_cast<uint8_t*>(next),
656663
data_buffer_->data() + data_buffer_->size());
657-
return sink_->Write(data_buffer_);
664+
return WriteAndClearBuffer();
658665
}
659666

660667
Status TranslateMinimalBatch(const RecordBatch& batch) {
661668
if (batch.num_rows() == 0) {
662-
// GH-36889: Clear buffer to avoid writing stale content (e.g., header)
663-
RETURN_NOT_OK(data_buffer_->Resize(0, /*shrink_to_fit=*/false));
664669
return Status::OK();
665670
}
666671
offsets_.resize(batch.num_rows());

0 commit comments

Comments
 (0)