viirya commented on code in PR #6773:
URL: https://github.com/apache/datafusion-comet/pull/6773#discussion_r4243795347
##########
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:
Same as in `iceberg-writes.md`: "so the pool never fails it" should be
limited to refusals. An error from Spark during the resize should still fail
the task once `reserve` passes it through.
##########
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:
Once `reserve` passes through errors that are not refusals, "The pool cannot
fail a fanout write" is too strong. What holds is that a refusal cannot fail
it. An exception from Spark, such as a failed spill or a killed task, still
fails the task. Could this say "A refusal from the pool cannot fail a fanout
write"?
##########
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:
Now that this test uses an unpartitioned write, nothing checks that files a
fanout write closed early are deleted when the task fails later. The mechanism
should already handle it, since `TrackingLocationGenerator` records each
location when it is generated. Could we add a test that pins it, together with
the error handling in `reserve`? A test pool that refuses the first few resizes
and then returns a non-`ResourcesExhausted` error would make some partitions
close early and then hit the error. The test could assert that the write
returns that error and that `parquet_files_under(dir)` is empty afterwards.
##########
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
+ /// group, and cannot close it early, so the pool's refusal stands.
+ async fn reserve(
+ &mut self,
+ reservation: &MemoryReservation,
+ open_files: &OpenFileMemory,
+ properties: &PartitionProperties,
+ ) -> DFResult<usize> {
+ let mut pending = self.pending_bytes();
+ let mut refused = match reservation.try_resize(open_files.bytes() +
pending) {
+ Ok(()) => return Ok(0),
+ Err(refused) => refused,
Review Comment:
+1, and I think this needs to be fixed before merge. Tracing it through
`SparkMemory::try_acquire`, a Java throwable from `acquireExecutionMemory` is
cleared by `check_exception` and comes back as `DataFusionError::External`. Any
other `CometError`, such as `Internal`, becomes `Execution`. So the filter
should be "retry only on `ResourcesExhausted`" rather than "skip `External`",
as in your snippet.
Besides a failed spill, this also swallows the `InterruptedException` a
killed task gets while it waits in `ExecutionMemoryPool.acquireMemory`. A
speculative loser or a cancelled task would then carry on writing files until
something else notices the kill.
The same check is needed at both `try_resize` calls, the first one and the
one inside the loop.
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -3109,10 +3109,69 @@ class CometIcebergWriteActionSuite
}
withNativeEnabled {
+ // withSQLConf returns Unit before Spark 4.0, so the assertions run
inside it.
withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key ->
"0.002") {
- val (failedPlans, error) = captureFailedPlans(spark) {
+ val snapshot = captureWrite("fanout_oom") {
spark.sql(s"INSERT INTO $catalog.$ns.fanout_oom SELECT * FROM
fanout_oom_src")
}
+ assert(
+ snapshot.snapshotDelta == 1L,
+ s"expected 1 commit, got ${snapshot.snapshotDelta}")
+ val writers = snapshot.plans.flatMap { plan =>
+ collectWithSubqueries(plan) { case w: CometIcebergWriteExec => w }
+ }
+ assert(
+ writers.nonEmpty,
+ s"the write did not run
natively:\n${snapshot.plans.mkString("\n--\n")}")
+ assert(
+ writers.map(_.metrics("files_closed_early").value).sum > 0,
Review Comment:
One way to check the count without pinning a file count is in the Rust
tests. Pass a real `WriteMetrics` into `fanout_write_in_part_of_the_pool` and
assert `closed_early <= tight.len() && tight.len() <= closed_early +
partitions`. At `TARGET_FILE_SIZE` no file rolls, so every file is either
closed early or closed at the end, at most one per partition. Counting `closed`
instead of `closed + 1` fails that whenever a refusal closes a single partition.
--
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]
