andygrove commented on code in PR #5143:
URL: https://github.com/apache/datafusion-comet/pull/5143#discussion_r3713900512
##########
native/core/src/execution/operators/parquet_writer.rs:
##########
@@ -146,68 +151,67 @@ impl ParquetWriter {
})?;
}
- // Clear the buffer after upload
- cursor.get_mut().clear();
- cursor.set_position(0);
-
Ok(())
}
}
}
- /// Close the writer and finalize the file
- async fn close(self) -> std::result::Result<(),
parquet::errors::ParquetError> {
+ /// Close the writer and finalize the file, returning the total bytes
written.
+ async fn close(self) -> std::result::Result<u64,
parquet::errors::ParquetError> {
match self {
- ParquetWriter::LocalFile(writer) => {
- writer.close()?;
- Ok(())
+ ParquetWriter::LocalFile(mut writer) => {
+ writer.finish()?;
+ Ok(writer.bytes_written() as u64)
}
#[cfg(feature = "hdfs-opendal")]
- ParquetWriter::Remote(
- arrow_parquet_buffer_writer,
- mut hdfs_writer_opt,
+ ParquetWriter::Remote {
+ mut arrow_writer,
+ mut hdfs_writer,
op,
output_path,
- ) => {
- // Close the arrow writer to finalize parquet format
- let cursor = arrow_parquet_buffer_writer.into_inner()?;
- let final_data = cursor.into_inner();
-
- // Create HDFS writer if not already created
- if hdfs_writer_opt.is_none() && !final_data.is_empty() {
- let writer =
op.writer(output_path.as_str()).await.map_err(|e| {
- parquet::errors::ParquetError::External(
- format!("Failed to create HDFS writer for '{}':
{}", output_path, e)
- .into(),
- )
- })?;
- hdfs_writer_opt = Some(writer);
- }
+ } => {
+ // Finalize the Parquet footer into the in-memory cursor.
`bytes_written()`
+ // reports the authoritative file size once `finish()` has
flushed the footer.
+ // We cannot call `into_inner()` after `finish()`: `finish()`
marks the
+ // underlying `SerializedFileWriter` as finished, and
`into_inner()` then fails
+ // with `SerializedFileWriter already finished`. Pull the
bytes out through
+ // `inner_mut()` instead - `finish()` has already flushed the
buffered writer
+ // into the cursor.
+ arrow_writer.finish()?;
+ let total_bytes = arrow_writer.bytes_written() as u64;
+ let final_data =
std::mem::take(arrow_writer.inner_mut().get_mut());
- // Write any remaining data
if !final_data.is_empty() {
Review Comment:
This guard can never be false. I checked the empty-partition case, where no
batches are written at all, and `finish()` still emits a 355 byte footer, so
`close()` returns `Ok(355)` and a valid empty Parquet file lands in the store.
That case works correctly.
The reason I would rather see the guard gone than left alone is the branch
that never runs. If `final_data` were ever empty we would drop `hdfs_writer`
without calling `close()` on it, and for a multipart upload that means the
object is silently never committed. Silent data loss behind a currently
unreachable condition is worth deleting while these lines are already being
rewritten. Could we just always take or create the writer and close it
unconditionally?
##########
native/core/src/execution/operators/parquet_writer.rs:
##########
@@ -108,34 +108,39 @@ impl ParquetWriter {
match self {
ParquetWriter::LocalFile(writer) => writer.write(batch),
#[cfg(feature = "hdfs-opendal")]
- ParquetWriter::Remote(
- arrow_parquet_buffer_writer,
- hdfs_writer_opt,
+ ParquetWriter::Remote {
+ arrow_writer,
+ hdfs_writer,
op,
output_path,
- ) => {
+ } => {
// Write batch to in-memory buffer
- arrow_parquet_buffer_writer.write(batch)?;
-
- // Flush and get the current buffer content
- arrow_parquet_buffer_writer.flush()?;
- let cursor = arrow_parquet_buffer_writer.inner_mut();
- let current_data = cursor.get_ref().clone();
+ arrow_writer.write(batch)?;
+
+ // `flush()` closes the in-progress row group but leaves bytes
in the internal
+ // `BufWriter`. `sync()` pushes those bytes down into our
cursor so the upload
+ // is genuinely incremental. Then take ownership of the
cursor's buffer and
+ // reset it to empty for the next batch (no clone, no explicit
clear).
+ arrow_writer.flush()?;
Review Comment:
The `sync()` addition is correct. I confirmed `ArrowWriter::sync()` is
`self.writer.flush()`, which reaches `SerializedFileWriter::flush()` and then
`TrackedWrite::flush()`, so the upload is genuinely incremental now and your
comment is accurate.
Separately, I noticed `flush()` here creates a new row group on every batch.
I confirmed this with your new test, where three batches produce exactly three
row groups. The local path does not do this, since it just calls
`writer.write(batch)` and lets `ArrowWriter` manage row group boundaries at its
default of 1,048,576 rows. So a file written to HDFS ends up with a row group
every 8192 rows, roughly 128 times more row groups than the same data written
locally, which hurts compression and bloats the footer.
This predates your PR so I am not asking you to fix it here. Given that you
are measuring write performance, could you file a tracking issue for it and
link it from this thread?
##########
native/core/src/execution/operators/parquet_writer.rs:
##########
@@ -618,6 +621,80 @@ mod tests {
);
}
+ /// Exercise the `ParquetWriter::Remote` write/close path against an
in-memory
+ /// opendal `Operator`, so the remote path has real automated coverage
without
+ /// requiring an HDFS cluster. Writes a handful of batches, reads the
uploaded
+ /// bytes back with `ParquetRecordBatchReaderBuilder`, and asserts the
returned
+ /// row count and the reported `bytes_written` both match the upload.
+ #[tokio::test]
+ #[cfg(feature = "hdfs-opendal")]
Review Comment:
Thanks for adding this. I ran it and it passes, and I confirmed it has
teeth, since it fails without the `into_inner()` fix and again without the
`set_position(0)` reset. It also runs in CI by default, because `hdfs-opendal`
is in the default feature set.
Since this is the only automated coverage the remote writer has, would you
mind also comparing the read-back values against what was written rather than
just the row count? `create_test_record_batch` already generates distinct data
per batch, so it is only a few extra lines, and buffer and cursor-position
juggling is exactly the kind of change that can corrupt values.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
