andygrove commented on code in PR #6773:
URL: https://github.com/apache/datafusion-comet/pull/6773#discussion_r4248040468
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -1087,17 +1348,47 @@ impl InnerWriter {
fn pending_bytes(&self) -> usize {
match self {
InnerWriter::Unpartitioned(_, feed) => feed.pending_bytes(),
- InnerWriter::Fanout(_, _, fanout) => fanout
- .feeds
- .values()
- .map(|(_, feed)| feed.pending_bytes())
- .sum(),
+ InnerWriter::Fanout(fanout, _) => fanout.pending_bytes(),
InnerWriter::Clustered(_, _, live) => {
live.as_ref().map_or(0, |(_, feed)| feed.pending_bytes())
}
}
}
+ /// Resizes `reservation` to what the task's writers hold between batches:
the open files'
+ /// shares (`open_files`, see [`MeteredParquetWriter`]) and the rows
waiting in their
+ /// [`PartitionFeed`]s. When the pool refuses, a fanout write writes out
and closes its
+ /// partitions, the one holding the most first, until what is left fits,
and this returns how
+ /// many it closed. The unpartitioned and clustered writers hold one file
open, at most a row
Review Comment:
Right, the size was wrong for S3 and GCS, and the comment doesn't need it.
In 9270e1745a it just says those writers hold one file open and cannot close it
early.
##########
docs/source/contributor-guide/iceberg-writes.md:
##########
@@ -281,9 +282,19 @@ hold in memory. It hands `ParquetWriter` each file's
`OutputFile` behind a `Coun
writer counts the bytes that leave memory on their way to storage, and a file
reports what it has
written less those. When they leave depends on the storage (`StorageWrites`).
After every batch
`run_write_task` resizes the task's reservation to what the open files report
plus the rows each
-`PartitionFeed` holds, for the dictionary choice or for pacing, and a resize
the pool refuses fails
-the task. What that figure covers, and what it misses, is described under
-[Native writers](memory_management.md#native-writers).
+`PartitionFeed` holds, for the dictionary choice or for pacing
(`InnerWriter::reserve`). When the
+pool refuses a fanout write, `reserve` writes out and closes the partitions
holding the most, in
+their open file and their feed, until the resize succeeds, and
`run_write_task` counts them in the
+`files_closed_early` metric. Taking the largest first frees the most memory
per file closed, so
+each refusal closes as few files as it can. Each partition's files report to a
child of the task's
+`OpenFileMemory`, which is how `reserve` finds them. That is why the fanout
path uses
+`FanoutPartitions` rather than iceberg-rust's `FanoutWriter`, which cannot
close one partition's
+writer. A closed partition's next rows open a new file with the properties its
first file used,
+which the partition keeps. The pool cannot fail a fanout write: everything the
write reserves
Review Comment:
Done in 9270e1745a. It now says a refusal from the pool cannot fail a fanout
write, and that any other error from a resize fails the task, with the failed
spill and the killed task as examples.
##########
docs/source/contributor-guide/memory_management.md:
##########
@@ -378,8 +378,10 @@ An operator that never calls `try_grow` is invisible to
the pool no matter how m
### Native writers
Both native writers reserve what they hold between batches through a single
consumer per task,
-`ParquetWriterExec[N]` or `IcebergWriteExec[N]`, resized after every batch.
Neither can spill, so
-when the pool refuses a resize the task fails with a `CometNativeException`
whose message starts
+`ParquetWriterExec[N]` or `IcebergWriteExec[N]`, resized after every batch.
Neither can spill. A
+fanout Iceberg write gives memory back instead, by closing partitions early
when the pool refuses a
+resize, as described below, so the pool never fails it. When the pool refuses
any other write's
Review Comment:
Done in 9270e1745a, with the same split: a refusal never fails it, and an
exception Spark throws during a resize still fails the task.
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -4894,32 +5149,165 @@ mod tests {
assert_eq!(pool.reserved(), 0, "the write kept a reservation");
}
- /// A write whose open files outgrow what the pool grants fails with
the pool's error,
- /// rather than holding memory nothing accounts for, and deletes the
files it had opened.
+ /// Rows per partition in each data file of `data_files`, keyed by
partition value.
+ fn rows_per_partition(data_files: &[DataFile]) ->
HashMap<IcebergStruct, u64> {
+ let mut rows = HashMap::new();
+ for file in data_files {
+ *rows.entry(file.partition().clone()).or_default() +=
file.record_count();
+ }
+ rows
+ }
+
+ /// Writes `batches` into a fanout table of `schema` twice: with a
pool to spare, then with
+ /// a pool of the most that first write reserved divided by `divisor`,
which it can only
+ /// fit by closing partitions early. Returns both writes' data files
and the second write's
+ /// directory, after checking that the second fit and gave its
reservation back.
+ async fn fanout_write_in_part_of_the_pool(
+ divisor: usize,
+ schema: Schema,
+ batches: impl Fn() -> Vec<RecordBatch>,
+ properties: impl Fn() -> WriterProperties,
+ ) -> (Vec<DataFile>, Vec<DataFile>, TempDir) {
+ let roomy = Arc::new(PeakMemoryPool::new(usize::MAX));
+ let (_dir, written) = write_reserving_from(
+ &roomy,
+ schema.clone(),
+ ProtoIcebergWriterMode::IcebergWriterFanout,
+ batches(),
+ properties(),
+ TARGET_FILE_SIZE,
+ )
+ .await;
+ let roomy_files = written.unwrap();
+ let tight = Arc::new(PeakMemoryPool::new(roomy.peak() / divisor));
+ let (dir, written) = write_reserving_from(
+ &tight,
+ schema,
+ ProtoIcebergWriterMode::IcebergWriterFanout,
+ batches(),
+ properties(),
+ TARGET_FILE_SIZE,
+ )
+ .await;
+ let tight_files = written.expect("closing partitions early lets
the write fit");
+ assert_eq!(tight.reserved(), 0, "the write kept a reservation");
+ (roomy_files, tight_files, dir)
+ }
+
+ /// A fanout write whose partitions outgrow what the pool grants
writes out and closes the
+ /// ones holding the most until what is left fits, and carries on: a
closed partition's
+ /// next rows open a new file. Every row still lands, in more files
than partitions.
+ #[tokio::test]
+ async fn
a_fanout_write_the_pool_cannot_hold_closes_partitions_until_it_fits() {
+ // Four batches, each a unit for every one of 16 partitions, so a
partition closed
+ // early gets more rows after it. One unit is one page, so every
unit goes straight to
+ // its partition's file.
+ let (roomy, tight, dir) = fanout_write_in_part_of_the_pool(
+ 2,
+ iceberg_user_schema(),
+ || {
+ (0..4)
+ .map(|batch| {
+ round_robin_batch_from(batch * 16 * ROWS_DIVISOR,
16 * ROWS_DIVISOR, 16)
+ })
+ .collect()
+ },
+ || {
+ WriterProperties::builder()
+ .set_data_page_row_count_limit(ROWS_DIVISOR)
+ .build()
+ },
+ )
+ .await;
+ assert_eq!(roomy.len(), 16);
+ assert!(tight.len() > 16, "{} files for 16 partitions",
tight.len());
+ assert_eq!(rows_per_partition(&tight), rows_per_partition(&roomy));
+ assert_eq!(parquet_files_under(dir.path()).len(), tight.len());
+ }
+
+ /// A partition's memory counts the rows its feed still holds back,
and closing it writes
+ /// them out first. Here every partition's rows wait for their
dictionary choice, so no
+ /// file is open when the pool refuses: the write only fits by writing
partitions out.
+ #[tokio::test]
+ async fn closing_a_partition_early_writes_out_the_rows_it_holds_back()
{
+ let (roomy, tight, _dir) = fanout_write_in_part_of_the_pool(
+ 2,
+ iceberg_user_schema(),
+ || vec![round_robin_batch(16 * ROWS_DIVISOR, 16)],
+ || WriterProperties::builder().build(),
+ )
+ .await;
+ assert_eq!(record_counts(&tight), record_counts(&roomy));
+ }
+
+ /// Closing the partitions holding the most frees the most memory per
file closed, so a
+ /// write short of memory closes few files. Here 16 or 64 partitions
take turns through 32
+ /// batches, so every partition keeps getting rows, and the pool
grants an eighth of what
+ /// the write needs. The write ends with between four and five files
per partition,
+ /// however many partitions it has. Closing every partition at once
would leave seven, and
+ /// closing them in the order they were first seen, or smallest first,
seventeen to twenty.
+ #[tokio::test]
+ async fn
a_fanout_write_short_of_memory_closes_the_partitions_holding_the_most() {
+ for partitions in [16, 64] {
+ let (roomy, tight, _dir) = fanout_write_in_part_of_the_pool(
+ 8,
+ iceberg_payload_schema(),
+ || {
+ (0..32)
+ .map(|batch| {
+ round_robin_payload_batch_from(batch * 1024,
1024, partitions, 100)
+ })
+ .collect()
+ },
+ || WriterProperties::builder().build(),
+ )
+ .await;
+ assert_eq!(roomy.len(), partitions);
+ assert_eq!(rows_per_partition(&tight),
rows_per_partition(&roomy));
+ assert!(
+ tight.len() < 6 * partitions,
+ "{} files for {partitions} partitions",
+ tight.len()
+ );
+ }
+ }
+
+ /// A write whose open file outgrows what the pool grants, with no
partition it can close
+ /// early, fails with the pool's error rather than holding memory
nothing accounts for,
+ /// and deletes the files it had opened. An unpartitioned write keeps
its one file open.
#[tokio::test]
async fn a_write_the_pool_cannot_hold_fails_and_deletes_its_files() {
Review Comment:
Added in 9270e1745a as
`a_fanout_write_fails_on_a_resize_error_that_is_not_a_refusal`, much as you
describe. Its pool refuses the first two resizes, so two partitions close early
and their files are complete on disk, and then fails the third with an error
that is not a refusal. The test asserts that the write returns that error and
that `parquet_files_under(dir)` is empty, and checks that both refusals were
used, so the early closes did happen before the error.
--
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]
