Xuanwo commented on code in PR #176:
URL: https://github.com/apache/iceberg-rust/pull/176#discussion_r1483994701
##########
crates/iceberg/src/io.rs:
##########
@@ -240,9 +241,9 @@ impl InputFile {
}
/// Trait for writing file.
-pub trait FileWrite: AsyncWrite {}
+pub trait FileWrite: AsyncWrite + TokioAsyncWrite + Send + Unpin {}
Review Comment:
It's quite weird that `FileWrite` requires to have both `AsyncWrite` and
`TokioAsyncWrite`. It increases the burden for the downstream to implement.
Did I missed something here?
##########
crates/iceberg/src/writer/file_writer/location_generator.rs:
##########
@@ -0,0 +1,99 @@
+// 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.
+
+//! This module contains the location generator and file name generator for
generating path of data file.
+
+use std::sync::{atomic::AtomicU64, Arc};
+
+use crate::spec::DataFileFormat;
+
+/// `LocationGenerator` used to generate the location of data file.
+pub trait LocationGenerator: Clone + Send + 'static {
Review Comment:
Can we merge `LocationGenerator` and `FileNameGenerator`? Both of them only
have one easy function to implement. By merging them, we can reduce an extra
generic param for users to use.
Or maybe we can go deeper by generating the final location in one shot.
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -0,0 +1,385 @@
+// 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.
+
+//! The module contains the file writer for parquet file format.
+
+use std::{
+ cmp::max,
+ collections::HashMap,
+ sync::{atomic::AtomicI64, Arc},
+};
+
+use crate::{io::FileIO, Result};
+use crate::{
+ io::OutputFile,
+ spec::{DataFileBuilder, DataFileFormat},
+ writer::CurrentFileStatus,
+ Error,
+};
+use arrow_schema::SchemaRef as ArrowSchemaRef;
+use parquet::{arrow::AsyncArrowWriter, format::FileMetaData};
+use parquet::{arrow::PARQUET_FIELD_ID_META_KEY,
file::properties::WriterProperties};
+
+use super::{
+ location_generator::{FileNameGenerator, LocationGenerator},
+ track_writer::TrackWriter,
+ FileWriter, FileWriterBuilder,
+};
+
+/// ParquetWriterBuilder is used to builder a [`ParquetWriter`]
+#[derive(Clone)]
+pub struct ParquetWriterBuilder<T: LocationGenerator, F: FileNameGenerator> {
+ /// `buffer_size` determines the initial size of the intermediate buffer.
+ /// The intermediate buffer will automatically be resized if necessary
+ init_buffer_size: usize,
+ props: WriterProperties,
+ schema: ArrowSchemaRef,
+
+ file_io: FileIO,
+ location_generator: T,
+ file_name_generator: F,
+}
+
+impl<T: LocationGenerator, F: FileNameGenerator> ParquetWriterBuilder<T, F> {
+ /// To avoid EntiryTooSmall error, we set the minimum buffer size to 8MB
if the given buffer size is smaller than it.
+ const MIN_BUFFER_SIZE: usize = 8 * 1024 * 1024;
+
+ /// Create a new `ParquetWriterBuilder`
+ /// To construct the write result, the schema should contain the
`PARQUET_FIELD_ID_META_KEY` metadata for each field.
+ pub fn new(
+ init_buffer_size: usize,
+ props: WriterProperties,
+ schema: ArrowSchemaRef,
+ file_io: FileIO,
+ location_generator: T,
+ file_name_generator: F,
+ ) -> Self {
+ Self {
+ init_buffer_size,
+ props,
+ schema,
+ file_io,
+ location_generator,
+ file_name_generator,
+ }
+ }
+}
+
+impl<T: LocationGenerator, F: FileNameGenerator> FileWriterBuilder for
ParquetWriterBuilder<T, F> {
+ type R = ParquetWriter;
+
+ async fn build(self) -> crate::Result<Self::R> {
+ // Fetch field id from schema
+ let field_ids = self
+ .schema
+ .fields()
+ .iter()
+ .map(|field| {
+ field
+ .metadata()
+ .get(PARQUET_FIELD_ID_META_KEY)
+ .ok_or_else(|| {
+ Error::new(
+ crate::ErrorKind::Unexpected,
+ "Field id not found in arrow schema metadata.",
+ )
+ })?
+ .parse::<i32>()
+ .map_err(|err| {
+ Error::new(crate::ErrorKind::Unexpected, "Failed to
parse field id.")
+ .with_source(err)
+ })
+ })
+ .collect::<crate::Result<Vec<_>>>()?;
+
+ let written_size = Arc::new(AtomicI64::new(0));
+ let out_file = self.file_io.new_output(
+ self.location_generator
+
.generate_location(&self.file_name_generator.generate_file_name()),
Review Comment:
For example, we seem to always use them together.
##########
crates/iceberg/src/writer/file_writer/track_writer.rs:
##########
@@ -0,0 +1,72 @@
+// 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 std::{
+ pin::Pin,
+ sync::{atomic::AtomicI64, Arc},
+};
+
+use tokio::io::AsyncWrite;
+
+use crate::io::FileWrite;
+
+/// `TrackWriter` is used to track the written size.
+pub(crate) struct TrackWriter {
+ inner: Box<dyn FileWrite>,
+ written_size: Arc<AtomicI64>,
Review Comment:
Writer only has `&mut` API so we can simplely use an `i64` here instead.
--
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: issues-unsubscr...@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org
For additional commands, e-mail: issues-h...@iceberg.apache.org
