Merge branch 'JanKaul:main' into main

Author: gruuya

Cargo.lock

@@ -35,8 +35,8 @@ derive_builder = "0.20"
 futures = "0.3.31"
 getrandom = { version = "0.3.1", features = ["std"] }
 itertools = "0.14.0"
+lru = "0.16.0"
 object_store = { version = "0.12", features = ["aws", "gcp"] }
-once_map = "0.4"
 parquet = { version = "55", features = ["async", "object_store"] }
 pin-project-lite = "0.2"
 serde = "^1.0"

Cargo.toml

@@ -21,8 +21,8 @@ futures = { workspace = true }
 getrandom = { workspace = true }
 iceberg-rust-spec = { path = "../iceberg-rust-spec", version = "0.8.0" }
 itertools = { workspace = true }
+lru = { workspace = true }
 object_store = { workspace = true }
-once_map = { workspace = true }
 parquet = { workspace = true }
 pin-project-lite = { workspace = true }
 serde = { workspace = true }
@@ -37,6 +37,7 @@ tokio = { version = "1.43", features = ["sync"] }
 url = { workspace = true }
 uuid = { workspace = true }
 
+
 [dev-dependencies]
 chrono = { workspace = true }
 iceberg-sql-catalog = { path = "../catalogs/iceberg-sql-catalog" }

iceberg-rust/Cargo.toml

@@ -13,15 +13,18 @@ use std::{
     hash::Hash,
     pin::Pin,
     task::{Context, Poll},
+    thread::available_parallelism,
 };
 
 use arrow::{
     array::{
-        as_primitive_array, as_string_array, ArrayIter, ArrayRef, BooleanArray,
-        BooleanBufferBuilder, PrimitiveArray, StringArray,
+        as_primitive_array, as_string_array, ArrayRef, BooleanArray, BooleanBufferBuilder,
+        PrimitiveArray, StringArray,
+    },
+    compute::{
+        and, filter, filter_record_batch,
+        kernels::cmp::{distinct, eq},
     },
-    compute::kernels::cmp::eq,
-    compute::{and, filter_record_batch},
     datatypes::{ArrowPrimitiveType, DataType, Int32Type, Int64Type},
     error::ArrowError,
     record_batch::RecordBatch,
@@ -33,7 +36,7 @@ use futures::{
 use itertools::{iproduct, Itertools};
 
 use iceberg_rust_spec::{partition::BoundPartitionField, spec::values::Value};
-use once_map::OnceMap;
+use lru::LruCache;
 use pin_project_lite::pin_project;
 
 use crate::error::Error;
@@ -58,7 +61,7 @@ pin_project! {
         #[pin]
         record_batches: Pin<Box<dyn Stream<Item = Result<RecordBatch, ArrowError>> + Send>>,
         partition_fields: &'a [BoundPartitionField<'a>],
-        partition_streams: OnceMap<Vec<Value>, RecordBatchSender>,
+        partition_streams: LruCache<Vec<Value>, RecordBatchSender>,
         queue: VecDeque<Result<(Vec<Value>, RecordBatchReceiver), Error>>,
         sends: Vec<(RecordBatchSender, RecordBatch)>,
     }
@@ -72,7 +75,7 @@ impl<'a> PartitionStream<'a> {
         Self {
             record_batches,
             partition_fields,
-            partition_streams: OnceMap::new(),
+            partition_streams: LruCache::unbounded(),
             queue: VecDeque::new(),
             sends: Vec::new(),
         }
@@ -109,13 +112,22 @@ impl Stream for PartitionStream<'_> {
                 }
             }
 
+            // Limit the number of open partition_streams by available parallelism
+            if this.partition_streams.len() > available_parallelism().unwrap().get() {
+                if let Some((_, mut sender)) = this.partition_streams.pop_lru() {
+                    sender.close_channel();
+                }
+            }
+
             match this.record_batches.as_mut().poll_next(cx) {
                 Poll::Pending => {
                     break Poll::Pending;
                 }
                 Poll::Ready(None) => {
-                    for sender in this.partition_streams.read_only_view().values() {
-                        sender.clone().close_channel();
+                    while let Some((_, sender)) = this.partition_streams.pop_lru() {
+                        if !sender.is_closed() {
+                            sender.clone().close_channel();
+                        }
                     }
                     break Poll::Ready(None);
                 }
@@ -127,16 +139,16 @@ impl Stream for PartitionStream<'_> {
                         let (partition_values, batch) = result?;
 
                         let sender = if let Some(sender) =
-                            this.partition_streams.get_cloned(&partition_values)
+                            this.partition_streams.get(&partition_values).cloned()
                         {
                             sender
                         } else {
+                            let (sender, reciever) = channel(1);
+                            this.queue
+                                .push_back(Ok((partition_values.clone(), reciever)));
                             this.partition_streams
-                                .insert_cloned(partition_values, |key| {
-                                    let (sender, reciever) = channel(1);
-                                    this.queue.push_back(Ok((key.clone(), reciever)));
-                                    sender
-                                })
+                                .push(partition_values, sender.clone());
+                            sender
                         };
 
                         this.sends.push((sender, batch));
@@ -250,12 +262,12 @@ fn distinct_values(array: ArrayRef) -> Result<DistinctValues, ArrowError> {
         DataType::Int32 => Ok(DistinctValues::Int(distinct_values_primitive::<
             i32,
             Int32Type,
-        >(array))),
+        >(array)?)),
         DataType::Int64 => Ok(DistinctValues::Long(distinct_values_primitive::<
             i64,
             Int64Type,
-        >(array))),
-        DataType::Utf8 => Ok(DistinctValues::String(distinct_values_string(array))),
+        >(array)?)),
+        DataType::Utf8 => Ok(DistinctValues::String(distinct_values_string(array)?)),
         _ => Err(ArrowError::ComputeError(
             "Datatype not supported for transform.".to_string(),
         )),
@@ -275,15 +287,36 @@ fn distinct_values(array: ArrayRef) -> Result<DistinctValues, ArrowError> {
 /// A HashSet containing all unique values from the array
 fn distinct_values_primitive<T: Eq + Hash, P: ArrowPrimitiveType<Native = T>>(
     array: ArrayRef,
-) -> HashSet<P::Native> {
-    let mut set = HashSet::new();
+) -> Result<HashSet<P::Native>, ArrowError> {
     let array = as_primitive_array::<P>(&array);
-    for value in ArrayIter::new(array).flatten() {
-        if !set.contains(&value) {
-            set.insert(value);
-        }
+
+    let first = array.value(0);
+
+    let slice_len = array.len() - 1;
+
+    if slice_len == 0 {
+        return Ok(HashSet::from_iter([first]));
     }
-    set
+
+    let v1 = array.slice(0, slice_len);
+    let v2 = array.slice(1, slice_len);
+
+    // Which consecutive entries are different
+    let mask = distinct(&v1, &v2)?;
+
+    let unique = filter(&v2, &mask)?;
+
+    let unique = as_primitive_array::<P>(&unique);
+
+    let set = unique
+        .iter()
+        .fold(HashSet::from_iter([first]), |mut acc, x| {
+            if let Some(x) = x {
+                acc.insert(x);
+            }
+            acc
+        });
+    Ok(set)
 }
 
 /// Extracts distinct string values from an Arrow array into a HashSet
@@ -293,15 +326,36 @@ fn distinct_values_primitive<T: Eq + Hash, P: ArrowPrimitiveType<Native = T>>(
 ///
 /// # Returns
 /// A HashSet containing all unique string values from the array
-fn distinct_values_string(array: ArrayRef) -> HashSet<String> {
-    let mut set = HashSet::new();
+fn distinct_values_string(array: ArrayRef) -> Result<HashSet<String>, ArrowError> {
+    let slice_len = array.len() - 1;
+
     let array = as_string_array(&array);
-    for value in ArrayIter::new(array).flatten() {
-        if !set.contains(value) {
-            set.insert(value.to_owned());
-        }
+
+    let first = array.value(0).to_owned();
+
+    if slice_len == 0 {
+        return Ok(HashSet::from_iter([first]));
     }
-    set
+
+    let v1 = array.slice(0, slice_len);
+    let v2 = array.slice(1, slice_len);
+
+    // Which consecutive entries are different
+    let mask = distinct(&v1, &v2)?;
+
+    let unique = filter(&v2, &mask)?;
+
+    let unique = as_string_array(&unique);
+
+    let set = unique
+        .iter()
+        .fold(HashSet::from_iter([first]), |mut acc, x| {
+            if let Some(x) = x {
+                acc.insert(x.to_owned());
+            }
+            acc
+        });
+    Ok(set)
 }
 
 /// Represents distinct values found in Arrow arrays during partitioning

iceberg-rust/src/arrow/partition.rs

@@ -65,6 +65,7 @@ use crate::{
 use super::partition::PartitionStream;
 
 const MAX_PARQUET_SIZE: usize = 512_000_000;
+const COMPRESSION_FACTOR: usize = 200;
 
 #[instrument(skip(table, batches), fields(table_name = %table.identifier().name()))]
 /// Writes Arrow record batches as partitioned Parquet files.
@@ -341,7 +342,7 @@ async fn write_parquet_files(
                     let batch_size = record_batch_size(&batch);
                     let new_size = state.bytes_written + batch_size;
 
-                    if new_size > MAX_PARQUET_SIZE {
+                    if new_size > COMPRESSION_FACTOR * MAX_PARQUET_SIZE {
                         // Send current writer to channel
                         let finished_writer = state.writer;
                         let file = finished_writer.1.close().await?;
@@ -362,9 +363,6 @@ async fn write_parquet_files(
                         state.bytes_written = batch_size;
                     } else {
                         state.bytes_written = new_size;
-                        if new_size % 64_000_000 >= 32_000_000 {
-                            state.writer.1.flush().await?;
-                        }
                     }
 
                     state.writer.1.write(&batch).await?;

iceberg-rust/src/arrow/write.rs

@@ -419,7 +419,10 @@ pub fn apply_table_updates(
 ) -> Result<(), Error> {
     for update in updates {
         match update {
-            TableUpdate::UpgradeFormatVersion { format_version: _ } => {
+            TableUpdate::UpgradeFormatVersion { format_version } => {
+                if u8::from(metadata.format_version) == format_version as u8 {
+                    return Ok(());
+                }
                 unimplemented!();
             }
             TableUpdate::AssignUuid { uuid } => {