unikdahal commented on code in PR #6773:
URL: https://github.com/apache/datafusion-comet/pull/6773#discussion_r4243573935
##########
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:
This checks that the metric is wired up and non-zero, but not that the
number is right. The Rust tests pass `WriteMetrics::default()`, so an
off-by-one in `closed + 1` would still pass. One check of the actual count, for
example against the number of extra files, would cover it. No need to pin the
exact file count.
##########
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:
I think this catches more than refusals. With the Spark-backed pools,
`try_grow()` goes through `SparkMemory::try_acquire()` to
`TaskMemoryManager.acquireExecutionMemory`. That call can throw, for example a
`SparkOutOfMemoryError` when spilling another consumer fails, or a
RuntimeException when the task is interrupted. The JNI bridge keeps those as
`DataFusionError::External`, while a normal "no room" is `ResourcesExhausted`.
Right now an `External` error lands in `refused` and we start closing
partitions. Once we've closed enough that the target drops below the current
reservation, the next `try_resize()` is just a shrink: it never calls Spark, so
it succeeds and we return `Ok`. The exception is gone. Spark may also have
handed out part of the request before the spill threw, and the memory doc
(`memory_management.md`, around line 205) says a caller that swallows the
exception has to reconcile that grant.
Before this PR the `?` failed the task. Could we only retry on
`ResourcesExhausted` and return everything else, here and at line 1383?
```rust
Err(e) if matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)) =>
e,
Err(e) => return Err(e),
```
A small test with a pool that returns `External` would keep this from coming
back.
##########
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,
+ };
+ let InnerWriter::Fanout(fanout, _) = self else {
+ return Err(refused);
+ };
+ for (closed, i) in fanout.by_memory().into_iter().enumerate() {
+ pending -= fanout.partitions[i].feed.pending_bytes();
+ fanout.close_early(i, properties).await?;
Review Comment:
Not a blocker, more of a question. After the pool says no, `close_early()`
can still allocate. `finish()` writes out the pending rows first, and a
partition that only has held rows opens a new Parquet writer to encode them. We
don't resize again until that writer has closed, so the pool never sees that
peak.
I know the writers only reserve what they hold between batches, so some
memory used during a call is outside the reservation on purpose. It just feels
worth a quick check since this runs right after a refusal. Do you have any
allocator numbers for this case? `PeakMemoryPool` only sees what we reserve.
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -1057,7 +1111,218 @@ async fn run_write_task(
}
}
-/// Enum-based dispatch over the three iceberg-rust partitioning writers, each
paired with the
+/// A fanout write: like iceberg-rust's `FanoutWriter`, a data file writer
open for every
+/// partition the task has written to, except that a partition can be closed
before the task ends,
+/// to give back the memory it holds. Its next rows then open a new file, with
the same properties.
+/// `FanoutWriter` keeps its writers private and closes them only all at once,
so the fanout path
+/// keeps its own.
+///
+/// iceberg-java's fanout writer keeps every file open until the task ends,
its buffers growing on
+/// the JVM heap. The native writer's buffers count against the task's memory
pool instead, so when
+/// the pool refuses them, closing partitions early lets the write finish with
more, smaller files
+/// rather than fail (see [`InnerWriter::reserve`]).
+struct FanoutPartitions {
+ builder: PartitionWriterBuilder,
+ /// Every partition the task has seen, in the order it first saw them.
+ partitions: Vec<FanoutPartition>,
+ /// Where each partition value sits in `partitions`.
+ index: HashMap<IcebergStruct, usize>,
+ /// Bytes the feeds hold back for their dictionary choice, across all
partitions. Every
Review Comment:
"every partition stays open" reads like it's left over from before.
Partitions can close early now, so maybe say every partition stays in
`FanoutPartitions` instead.
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -1057,7 +1111,218 @@ async fn run_write_task(
}
}
-/// Enum-based dispatch over the three iceberg-rust partitioning writers, each
paired with the
+/// A fanout write: like iceberg-rust's `FanoutWriter`, a data file writer
open for every
+/// partition the task has written to, except that a partition can be closed
before the task ends,
+/// to give back the memory it holds. Its next rows then open a new file, with
the same properties.
+/// `FanoutWriter` keeps its writers private and closes them only all at once,
so the fanout path
+/// keeps its own.
+///
+/// iceberg-java's fanout writer keeps every file open until the task ends,
its buffers growing on
+/// the JVM heap. The native writer's buffers count against the task's memory
pool instead, so when
+/// the pool refuses them, closing partitions early lets the write finish with
more, smaller files
+/// rather than fail (see [`InnerWriter::reserve`]).
+struct FanoutPartitions {
+ builder: PartitionWriterBuilder,
+ /// Every partition the task has seen, in the order it first saw them.
+ partitions: Vec<FanoutPartition>,
+ /// Where each partition value sits in `partitions`.
+ index: HashMap<IcebergStruct, usize>,
+ /// Bytes the feeds hold back for their dictionary choice, across all
partitions. Every
+ /// partition stays open, so the rows they hold back share one limit: once
they reach it, every
+ /// partition still holding makes its choice from the rows it has. A task
fanning out to many
+ /// partitions that each get less than a page would otherwise hold all of
its rows back,
+ /// uncompressed, until it closed.
+ ///
+ /// Always the sum of the feeds' [`PartitionFeed::held_bytes`], kept in
step wherever a feed
+ /// changes what it holds back: a push, the release once the limit is
reached, and an early
+ /// close. `write` checks this in debug builds.
+ held_bytes: usize,
+ /// The data files of the partitions closed early.
+ closed: Vec<DataFile>,
+}
+
+/// One partition of a fanout write.
+struct FanoutPartition {
+ /// What each of the partition's files is opened with: its first, and any
it opens after
+ /// closing one early.
+ key: PartitionKey,
+ feed: PartitionFeed,
+ /// The writer of the partition's open file, if it has one open.
+ writer: Option<PartitionDataFileWriter>,
+ /// What the open file holds, as part of the task's [`OpenFileMemory`].
+ memory: OpenFileMemory,
+ /// The properties the partition's files are written with, taken when its
first file opens and
+ /// kept for the files it opens after closing one early.
+ properties: Option<WriterProperties>,
+}
+
+impl FanoutPartition {
+ /// What the partition holds: its open file and the rows its feed has not
handed over.
+ fn bytes(&self) -> usize {
+ self.memory.bytes() + self.feed.pending_bytes()
+ }
+
+ /// Writes `unit` to the partition's open file, opening one first if it
has none.
+ async fn write(
+ &mut self,
+ builder: &PartitionWriterBuilder,
+ unit: RecordBatch,
+ ) -> iceberg::Result<()> {
+ let writer = match self.writer.as_mut() {
+ Some(writer) => writer,
+ None => {
+ let properties = self
+ .properties
+ .get_or_insert_with(||
builder.properties.take(Some(self.key.data())))
+ .clone();
+ self.writer.insert(
+ builder
+ .build_with(Some(self.key.clone()), properties,
self.memory.clone())
+ .await?,
+ )
+ }
+ };
+ writer.write(unit).await
+ }
+
+ /// Writes out every row the partition still holds back and closes its
file, returning that
+ /// file's data files. A file closed at a partial unit ends off
iceberg-java's row grid, as the
+ /// last file of a task does.
+ async fn close(
+ &mut self,
+ builder: &PartitionWriterBuilder,
+ properties: &PartitionProperties,
+ ) -> DFResult<Vec<DataFile>> {
+ for unit in self.feed.finish(properties)? {
+ self.write(builder, unit).await.map_err(iceberg_err)?;
+ }
+ match self.writer.take() {
+ Some(mut writer) => writer.close().await.map_err(iceberg_err),
+ None => Ok(Vec::new()),
+ }
+ }
+}
+
+impl FanoutPartitions {
+ fn new(builder: PartitionWriterBuilder) -> Self {
+ Self {
+ builder,
+ partitions: Vec::new(),
+ index: HashMap::new(),
+ held_bytes: 0,
+ closed: Vec::new(),
+ }
+ }
+
+ /// Memory held by the rows waiting in the feeds, for every partition the
task has seen.
+ fn pending_bytes(&self) -> usize {
+ self.partitions
+ .iter()
+ .map(|partition| partition.feed.pending_bytes())
+ .sum()
+ }
+
+ /// Writes `batch`, feeding each of its partitions separately.
+ async fn write(
+ &mut self,
+ splitter: &PartitionSplitter,
+ batch: &RecordBatch,
+ properties: &PartitionProperties,
+ ) -> DFResult<()> {
+ for (key, part) in splitter.split_groups(batch)? {
+ let next = self.partitions.len();
+ let i = *self.index.entry(key.data().clone()).or_insert(next);
+ if i == next {
+ self.partitions.push(FanoutPartition {
+ feed: PartitionFeed::new(Some(key.data().clone()),
splitter.slicer),
+ key,
+ writer: None,
+ memory: self.builder.open_files.child(),
+ properties: None,
+ });
+ }
+ let partition = &mut self.partitions[i];
+ let held_before = partition.feed.held_bytes();
+ let units = partition.feed.push(part, properties)?;
+ self.held_bytes = self.held_bytes - held_before +
partition.feed.held_bytes();
+ for unit in units {
+ partition
+ .write(&self.builder, unit)
+ .await
+ .map_err(iceberg_err)?;
+ }
+ }
+ if !properties.chooser.may_hold(self.held_bytes) {
+ for partition in &mut self.partitions {
+ for unit in partition.feed.release(properties)? {
+ partition
+ .write(&self.builder, unit)
+ .await
+ .map_err(iceberg_err)?;
+ }
+ }
+ self.held_bytes = 0;
+ }
+ debug_assert_eq!(
+ self.held_bytes,
+ self.partitions
+ .iter()
+ .map(|partition| partition.feed.held_bytes())
+ .sum::<usize>(),
+ "the held-back total drifted from what the feeds hold back"
+ );
+ Ok(())
+ }
+
+ /// The partitions holding memory, as indexes into `partitions`, the most
first. Of two holding
+ /// as much, the one seen first comes first, so which files a write closes
early does not
+ /// depend on `HashMap` order. Indexes rather than references, because
`reserve` closes each
+ /// partition it picks while it walks the ranking, and a close borrows the
whole
+ /// `FanoutPartitions` mutably.
+ fn by_memory(&self) -> Vec<usize> {
+ let mut holding: Vec<(usize, usize)> = self
+ .partitions
+ .iter()
+ .enumerate()
+ .map(|(i, partition)| (partition.bytes(), i))
+ .filter(|&(bytes, _)| bytes > 0)
+ .collect();
+ holding.sort_unstable_by_key(|&(bytes, i)| (Reverse(bytes), i));
+ holding.into_iter().map(|(_, i)| i).collect()
+ }
+
+ /// Closes the partition at `i` before the task ends, keeping its data
files for the output.
+ async fn close_early(&mut self, i: usize, properties:
&PartitionProperties) -> DFResult<()> {
+ let partition = &mut self.partitions[i];
+ self.held_bytes -= partition.feed.held_bytes();
+ let data_files = partition.close(&self.builder, properties).await?;
+ self.closed.extend(data_files);
Review Comment:
These `DataFile`s stay in memory until the task ends, and their per-column
metrics and bounds aren't counted in the reservation. That was already true for
rolled files, but early close can multiply the file count exactly when memory
is tight. It's probably small, but a quick measurement with many partitions and
wide columns would settle it. Fine as a follow-up.
##########
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:
"at most a row group" isn't quite right on S3 or GCS. Per the memory doc,
the file can also hold its last flushed part there. The point about having
nothing to close still stands.
--
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]
