Skip to content

feat: adding elapsed_compute to the writer #5143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
comphead wants to merge 4 commits into
base: main
Choose a base branch
from
Open

feat: adding elapsed_compute to the writer #5143

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
52b590d
feat: adding `elapsed_compute` to the writer
comphead Jul 30, 2026
46fff0d
feat: adding `elapsed_compute` to the writer
comphead Jul 30, 2026
1a889bd
feat: adding `elapsed_compute` to the writer
comphead Jul 30, 2026
8937de0
fmt
comphead Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions native/core/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ lazy_static = "1.4"
assertables = "10"
hex = "0.4.3"
datafusion-functions-nested = { version = "54.1.0" }
opendal = { version = "0.57.0", features = ["services-memory"] }

[features]
backtrace = ["datafusion/backtrace"]
Expand Down
233 changes: 155 additions & 78 deletions native/core/src/execution/operators/parquet_writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,12 +91,12 @@ enum ParquetWriter {
/// an Arrow writer writes to in-memory buffer the data converted to Parquet format
/// The opendal::Writer is created lazily on first write
#[cfg(feature = "hdfs-opendal")]
Remote(
ArrowWriter<Cursor<Vec<u8>>>,
Option<opendal::Writer>,
Operator,
String,
),
Remote {
arrow_writer: ArrowWriter<Cursor<Vec<u8>>>,
hdfs_writer: Option<opendal::Writer>,
op: Operator,
output_path: String,
},
}

impl ParquetWriter {
Expand All @@ -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()?;

@andygrove andygrove Aug 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

arrow_writer.sync()?;
let cursor = arrow_writer.inner_mut();
let current_data = std::mem::take(cursor.get_mut());
cursor.set_position(0);

// Create HDFS writer lazily on first write
if hdfs_writer_opt.is_none() {
if hdfs_writer.is_none() {
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);
*hdfs_writer = Some(writer);
}

// Write the accumulated data to HDFS
if let Some(hdfs_writer) = hdfs_writer_opt {
hdfs_writer.write(current_data).await.map_err(|e| {
if let Some(w) = hdfs_writer {
w.write(current_data).await.map_err(|e| {
parquet::errors::ParquetError::External(
format!(
"Failed to write batch to HDFS file '{}': {}",
Expand All @@ -146,68 +151,67 @@ impl ParquetWriter {
})?;
}

// Clear the buffer after upload
cursor.get_mut().clear();
cursor.set_position(0);
Comment on lines -149 to -151

@andygrove andygrove Jul 30, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it intentional that this is removed? The PR description doesn't talk about this change


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() {

@andygrove andygrove Aug 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

if let Some(mut hdfs_writer) = hdfs_writer_opt {
hdfs_writer.write(final_data).await.map_err(|e| {
let mut w = match hdfs_writer.take() {
Some(w) => w,
None => op.writer(output_path.as_str()).await.map_err(|e| {
parquet::errors::ParquetError::External(
format!(
"Failed to write final data to HDFS file '{}': {}",
"Failed to create HDFS writer for '{}': {}",
output_path, e
)
.into(),
)
})?;

// Close the HDFS writer
hdfs_writer.close().await.map_err(|e| {
parquet::errors::ParquetError::External(
format!("Failed to close HDFS writer for '{}': {}", output_path, e)
.into(),
})?,
};
w.write(final_data).await.map_err(|e| {
parquet::errors::ParquetError::External(
format!(
"Failed to write final data to HDFS file '{}': {}",
output_path, e
)
})?;
}
.into(),
)
})?;
w.close().await.map_err(|e| {
parquet::errors::ParquetError::External(
format!("Failed to close HDFS writer for '{}': {}", output_path, e)
.into(),
)
})?;
}

Ok(())
Ok(total_bytes)
}
}
}
Expand Down Expand Up @@ -337,12 +341,12 @@ impl ParquetWriterExec {

// HDFS writer will be created lazily on first write
// Use the path from prepare_object_store_with_configs
Ok(ParquetWriter::Remote(
arrow_parquet_buffer_writer,
None,
Ok(ParquetWriter::Remote {
arrow_writer: arrow_parquet_buffer_writer,
hdfs_writer: None,
op,
object_store_path.to_string(),
))
output_path: object_store_path.to_string(),
})
}
#[cfg(not(feature = "hdfs-opendal"))]
{
Expand Down Expand Up @@ -472,6 +476,7 @@ impl ExecutionPlan for ParquetWriterExec {
let files_written = MetricBuilder::new(&self.metrics).counter("files_written", partition);
let bytes_written = MetricBuilder::new(&self.metrics).counter("bytes_written", partition);
let rows_written = MetricBuilder::new(&self.metrics).counter("rows_written", partition);
let elapsed_compute = MetricBuilder::new(&self.metrics).elapsed_compute(partition);

let runtime_env = context.runtime_env();
let input = self.input.execute(partition, context)?;
Expand Down Expand Up @@ -528,6 +533,8 @@ impl ExecutionPlan for ParquetWriterExec {
while let Some(batch_result) = stream.try_next().await.transpose() {
let batch = batch_result?;

let mut timer = elapsed_compute.timer();

// Track row count
total_rows += batch.num_rows() as i64;

Expand All @@ -547,20 +554,16 @@ impl ExecutionPlan for ParquetWriterExec {
writer.write(&renamed_batch).await.map_err(|e| {
DataFusionError::Execution(format!("Failed to write batch: {}", e))
})?;
}

writer.close().await.map_err(|e| {
DataFusionError::Execution(format!("Failed to close writer: {}", e))
})?;
timer.stop();
}

// Get file size - strip file:// prefix if present for local filesystem access
let local_path = part_file
.strip_prefix("file://")
.or_else(|| part_file.strip_prefix("file:"))
.unwrap_or(&part_file);
let file_size = std::fs::metadata(local_path)
.map(|m| m.len() as i64)
.unwrap_or(0);
let mut timer = elapsed_compute.timer();
let file_size =
writer.close().await.map_err(|e| {
DataFusionError::Execution(format!("Failed to close writer: {}", e))
})? as i64;
timer.stop();

// Update metrics with write statistics
files_written.add(1);
Expand Down Expand Up @@ -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")]

@andygrove andygrove Aug 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

async fn test_parquet_writer_remote_memory_backend() -> Result<()> {
use opendal::services::Memory;
use opendal::Operator;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

let op = Operator::new(Memory::default())
.map_err(|e| {
DataFusionError::Execution(format!("Failed to create memory operator: {}", e))
})?
.finish();
let output_path = "test/data.parquet".to_string();

let schema = create_test_record_batch(1)?.schema();
let props = WriterProperties::builder()
.set_compression(Compression::UNCOMPRESSED)
.build();
let buffer = Vec::new();
let cursor = Cursor::new(buffer);
let arrow_writer = ArrowWriter::try_new(cursor, Arc::clone(&schema), Some(props))
.map_err(|e| DataFusionError::Execution(format!("try_new failed: {}", e)))?;

let mut writer = ParquetWriter::Remote {
arrow_writer,
hdfs_writer: None,
op: op.clone(),
output_path: output_path.clone(),
};

let mut expected_rows: i64 = 0;
for i in 1..=3 {
let batch = create_test_record_batch(i)?;
expected_rows += batch.num_rows() as i64;
writer.write(&batch).await.map_err(|e| {
DataFusionError::Execution(format!("Failed to write batch {}: {}", i, e))
})?;
}

let reported_bytes = writer
.close()
.await
.map_err(|e| DataFusionError::Execution(format!("Failed to close writer: {}", e)))?;

let uploaded = op.read(&output_path).await.map_err(|e| {
DataFusionError::Execution(format!("Failed to read uploaded object: {}", e))
})?;
let uploaded_bytes = uploaded.to_vec();
assert_eq!(
reported_bytes as usize,
uploaded_bytes.len(),
"bytes_written must match uploaded object length"
);

let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(uploaded_bytes))
.map_err(|e| DataFusionError::Execution(format!("Reader builder failed: {}", e)))?
.build()
.map_err(|e| DataFusionError::Execution(format!("Reader build failed: {}", e)))?;
let mut actual_rows: i64 = 0;
for batch in reader {
let batch =
batch.map_err(|e| DataFusionError::Execution(format!("Read error: {}", e)))?;
actual_rows += batch.num_rows() as i64;
}
assert_eq!(actual_rows, expected_rows);

Ok(())
}

/// Helper function to create a test RecordBatch with 1000 rows of (int, string) data
/// Example batch_id 1 -> 0..1000, 2 -> 1001..2000
#[allow(dead_code)]
Expand Down
5 changes: 4 additions & 1 deletion spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,10 @@ case class CometNativeWriteExec(
override lazy val metrics: Map[String, SQLMetric] = Map(
"files_written" -> SQLMetrics.createMetric(sparkContext, "number of written data files"),
"bytes_written" -> SQLMetrics.createSizeMetric(sparkContext, "written data"),
"rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows"))
"rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows"),
"elapsed_compute" -> SQLMetrics.createNanoTimingMetric(
sparkContext,
"total time (in ms) spent in this operator"))

override def doExecute(): RDD[InternalRow] = {
// Setup job if committer is present
Expand Down
Loading