CTTY commented on code in PR #1547: URL: https://github.com/apache/iceberg-rust/pull/1547#discussion_r2229405980
########## crates/iceberg/src/writer/file_writer/rolling_writer.rs: ########## @@ -0,0 +1,331 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::RecordBatch; +use futures::future::try_join_all; + +use crate::runtime::{JoinHandle, spawn}; +use crate::spec::DataFileBuilder; +use crate::writer::CurrentFileStatus; +use crate::writer::file_writer::{FileWriter, FileWriterBuilder}; +use crate::{Error, ErrorKind, Result}; + +/// Builder for creating a `RollingFileWriter` that rolls over to a new file +/// when the data size exceeds a target threshold. +#[derive(Clone)] +pub struct RollingFileWriterBuilder<B: FileWriterBuilder> { + inner_builder: B, + target_size: usize, +} + +impl<B: FileWriterBuilder> RollingFileWriterBuilder<B> { + /// Creates a new `RollingFileWriterBuilder` with the specified inner builder and target size. + /// + /// # Arguments + /// + /// * `inner_builder` - The builder for the underlying file writer + /// * `target_size` - The target size in bytes before rolling over to a new file + pub fn new(inner_builder: B, target_size: usize) -> Self { + Self { + inner_builder, + target_size, + } + } +} + +impl<B: FileWriterBuilder> FileWriterBuilder for RollingFileWriterBuilder<B> { + type R = RollingFileWriter<B>; + + async fn build(self) -> Result<Self::R> { + Ok(RollingFileWriter { + inner: None, + inner_builder: self.inner_builder, + target_size: self.target_size, + close_handles: vec![], + }) + } +} + +/// A writer that automatically rolls over to a new file when the data size +/// exceeds a target threshold. +/// +/// This writer wraps another file writer that tracks the amount of data written. +/// When the data size exceeds the target size, it closes the current file and +/// starts writing to a new one. +pub struct RollingFileWriter<B: FileWriterBuilder> { + inner: Option<B::R>, + inner_builder: B, + target_size: usize, + close_handles: Vec<JoinHandle<Result<Vec<DataFileBuilder>>>>, +} + +impl<B: FileWriterBuilder> RollingFileWriter<B> { + /// Determines if the writer should roll over to a new file. + /// + /// # Arguments + /// + /// * `input_size` - The size in bytes of the incoming data + /// + /// # Returns + /// + /// `true` if a new file should be started, `false` otherwise + pub fn should_roll(&self, input_size: usize) -> bool { + self.current_written_size() + input_size > self.target_size + } +} + +impl<B: FileWriterBuilder> FileWriter for RollingFileWriter<B> { + async fn write(&mut self, input: &RecordBatch) -> Result<()> { + let input_size = input.get_array_memory_size(); + + if self.inner.is_none() { + // initialize inner writer + self.inner = Some(self.inner_builder.clone().build().await?); + } + + if self.should_roll(input_size) { + if let Some(inner) = self.inner.take() { + // close the current writer, roll to a new file + let handle = spawn(async move { inner.close().await }); Review Comment: Good point! I have some rough idea to further improving this: we can use a config to control the maximum parallelism here: ``` struct RollingWriter { ... buffer: Vec<DataFile> } ... while close_handles.len() >= self.max_parallelism() { // wait until some closers complete, and store the data files in a buffer self.buffer.push(future::select(self.close_handles)) } self.close_handles.push(new_handle); ``` I have not thought very clearly on how to prevent `buffer` from eating up the memory as of now, or do we even need it? Either way I agree this can be completed as a follow-up -- 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]
