liurenjie1024 commented on code in PR #1547: URL: https://github.com/apache/iceberg-rust/pull/1547#discussion_r2228119893
########## 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(); Review Comment: This is incorrect, usually parquet files written size is much smaller than arrow's in memory array size since parquet will do a lot of compression. The `target_size` is not for exact control, so it's fine to write file a little larger that this size. ########## 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, Review Comment: ```suggestion target_file_size: usize, ``` ########## 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: This is an interesting optimization, but I would suggest not do it for now. A writer usually consumes some resources like memory, connections etc. Closing them in an async approach make things difficult to understand in production, for example, we may have a lot of unclosed writers which consumes a lot memory and leading to oom. -- 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]
