andygrove commented on code in PR #6773:
URL: https://github.com/apache/datafusion-comet/pull/6773#discussion_r4248038335
##########
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:
Thanks, both of you. Fixed in 9270e1745a. Both `try_resize` calls now retry
only when the error's root is `ResourcesExhausted`, through a small
`is_refusal` helper, and pass anything else straight back, so an exception from
Spark fails the task as it did before this PR. I checked that every pool we
build reports a plain refusal that way: `greedy_unified` and `fair_unified`
both return `ResourcesExhausted`, the `TrackConsumersPool` and
`SpillReplayPool` wrappers keep the variant, and `unbounded` never refuses. So
early close still happens under each of them, and the
`CometIcebergWriteActionSuite` fanout test, which runs against the real pool,
still closes partitions early.
The new `a_fanout_write_fails_on_a_resize_error_that_is_not_a_refusal` uses
a pool that refuses twice, so two partitions close early, and then fails the
third resize. It runs once with an `External` error and once with an
`Execution` one, so it pins "retry only on refusals" rather than "skip
`External`". Treating every error as a refusal, as before, fails both cases,
and skipping only `External` fails the second.
##########
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:
I measured it by adding exact peak tracking to Comet's accounting allocator
in a scratch build, and recording the peak allocation during each `close_early`
above what was allocated just before it, over every partition closed in six
tight-pool writes. A partition holding only held-back rows, tens to hundreds of
KB, peaked at 1.4 to 2.3 times what it was charged, with a floor of about 50 KB
for the Parquet writer itself, so an 11 KB partition peaked at 51 KB. A
partition holding an open file, or MBs of held-back rows, peaked at 0.5 to 1.05
times what it was charged. The largest peak was 3 MB, in a write whose pool was
19 MB.
Each close then gave back about what the partition was charged, or a little
more, and the closes run one partition at a time, so the extra never stacks.
It's the same work the task's final close already does for every partition,
just earlier. So I'd leave it outside the reservation, like the rest of what a
write allocates within a batch.
##########
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:
I measured it with the same scratch allocator by cloning each `DataFile`.
It's about 1 KB for a 3-column table and 25 KB for a 101-column one, mostly the
per-column sizes, counts and bounds. Next to what a close frees that's small:
25 KB kept for a median 330 KB freed with 101 columns, and for 2.3 MB freed
once the partitions have files open. It does add up over a task, though. In a
101-column write of 64 partitions given a quarter of the memory it needed, the
89 files closed early kept 2.2 MB until the task ended, about a fifth of its 10
MB pool, none of it reserved. Rolled files already had the same gap, so I filed
#6885 to count the kept `DataFile`s in the reservation.
##########
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:
Went with viirya's version in e9d8f76cfc. `fanout_write_in_part_of_the_pool`
now gives its tight write a real `WriteMetrics` and checks `closed_early <=
files <= closed_early + partitions`, so the three early-close tests all check
the count. Counting `closed` instead of `closed + 1` fails
`a_fanout_write_short_of_memory_closes_the_partitions_holding_the_most` with 28
counted for 70 files over 16 partitions, and counting one too many fails it
with 84. The Scala test still only checks that the metric reaches the JVM.
##########
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:
Fixed in d6d4be6587. It now says every partition stays in `partitions` until
the task ends, which is what makes the held-back rows share one limit.
--
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]
