diff --git a/datafusion_iceberg/src/table/mod.rs b/datafusion_iceberg/src/table/mod.rs index bbcbe0d6..b9e80e63 100644 --- a/datafusion_iceberg/src/table/mod.rs +++ b/datafusion_iceberg/src/table/mod.rs @@ -2522,7 +2522,7 @@ async fn write_parquet_files( Ok(datafiles) } -type PartitionValuesByPath = Arc>>>; +type PartitionValuesByPath = Arc>>>>; type DemuxerTask = ( SpawnedTask>, DemuxedStreamReceiver, @@ -2585,7 +2585,8 @@ async fn partitions_demuxer( table_location: &str, partition_values_by_path: &PartitionValuesByPath, ) -> Result<(), DataFusionError> { - let mut senders: LruCache, mpsc::Sender> = LruCache::unbounded(); + let mut senders: LruCache>, mpsc::Sender> = + LruCache::unbounded(); // Get partition column indices @@ -2718,7 +2719,7 @@ mod tests { prelude::SessionContext, scalar::ScalarValue, }; - use futures::stream; + use futures::{stream, TryStreamExt}; use iceberg_rust::{ catalog::tabular::Tabular, object_store::ObjectStoreBuilder, @@ -2742,7 +2743,7 @@ mod tests { partition::PartitionSpec, view_metadata::{Version, ViewRepresentation}, }, - table::Table, + table::{manifest_list::snapshot_partition_bounds_with_nulls, Table}, view::View, }; use iceberg_sql_catalog::SqlCatalog; @@ -3990,6 +3991,189 @@ mod tests { assert_eq!(table.manifests(None, None).await.unwrap().len(), 2); } + #[tokio::test] + async fn test_nullable_partition_write_and_reload() { + let object_store = ObjectStoreBuilder::memory(); + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", "test", object_store) + .await + .unwrap(), + ); + let schema = Schema::builder() + .with_struct_field(StructField { + id: 1, + name: "id".to_string(), + required: false, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + initial_default: None, + write_default: None, + }) + .with_struct_field(StructField { + id: 2, + name: "payload".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + initial_default: None, + write_default: None, + }) + .build() + .unwrap(); + let partition_spec = PartitionSpec::builder() + .with_partition_field(PartitionField::new( + 1, + 1000, + "id_partition", + Transform::Identity, + )) + .build() + .unwrap(); + let table = Table::builder() + .with_name("nullable_partition") + .with_location("/test/nullable_partition") + .with_schema(schema) + .with_partition_spec(partition_spec) + .build(&["test".to_owned()], Arc::clone(&catalog)) + .await + .unwrap(); + let identifier = table.identifier().clone(); + let ctx = SessionContext::new(); + ctx.register_table("nullable_partition", Arc::new(DataFusionTable::from(table))) + .unwrap(); + + ctx.sql("INSERT INTO nullable_partition VALUES (NULL, 10), (NULL, 11)") + .await + .unwrap() + .collect() + .await + .unwrap(); + let first_commit = match Arc::clone(&catalog) + .load_tabular(&identifier) + .await + .unwrap() + { + Tabular::Table(table) => table, + _ => panic!("expected table"), + }; + let first_manifests = first_commit.manifests(None, None).await.unwrap(); + assert_eq!(first_manifests.len(), 1); + let first_summary = &first_manifests[0].partitions.as_ref().unwrap()[0]; + assert!(first_summary.contains_null); + assert_eq!(first_summary.lower_bound, None); + assert_eq!(first_summary.upper_bound, None); + let first_bounds = snapshot_partition_bounds_with_nulls( + first_commit + .metadata() + .current_snapshot(None) + .unwrap() + .unwrap(), + first_commit.metadata(), + first_commit.object_store().clone(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(first_bounds.contains_null, [true]); + assert!(first_bounds.bounds.min.is_empty()); + + for sql in [ + "INSERT INTO nullable_partition VALUES (1, 20)", + "INSERT INTO nullable_partition VALUES (NULL, 12)", + "INSERT INTO nullable_partition VALUES (NULL, 13), (2, 21)", + ] { + ctx.sql(sql).await.unwrap().collect().await.unwrap(); + } + + let reloaded = match catalog.load_tabular(&identifier).await.unwrap() { + Tabular::Table(table) => table, + _ => panic!("expected table"), + }; + let manifests = reloaded.manifests(None, None).await.unwrap(); + let bounds = snapshot_partition_bounds_with_nulls( + reloaded.metadata().current_snapshot(None).unwrap().unwrap(), + reloaded.metadata(), + reloaded.object_store().clone(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(bounds.contains_null, [true]); + assert_eq!(bounds.bounds.min.as_slice(), [Value::LongInt(1)]); + assert_eq!(bounds.bounds.max.as_slice(), [Value::LongInt(2)]); + assert!(manifests.iter().any(|manifest| { + manifest + .partitions + .as_ref() + .is_some_and(|parts| parts[0].contains_null) + })); + let files = reloaded + .datafiles(&manifests, None, (None, None)) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(files.len(), 5); + assert_eq!( + files + .iter() + .filter( + |(_, entry)| entry.data_file().partition().get("id_partition") == Some(&None) + ) + .count(), + 3 + ); + + let fresh_ctx = SessionContext::new(); + fresh_ctx + .register_table("reloaded", Arc::new(DataFusionTable::from(reloaded))) + .unwrap(); + let batches = fresh_ctx + .sql("SELECT payload FROM reloaded WHERE id IS NULL ORDER BY payload") + .await + .unwrap() + .collect() + .await + .unwrap(); + let values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, [10, 11, 12, 13]); + + let batches = fresh_ctx + .sql("SELECT payload FROM reloaded WHERE id = 2") + .await + .unwrap() + .collect() + .await + .unwrap(); + let nonnull: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(nonnull, [21]); + } + #[tokio::test] pub async fn test_datafusion_table_insert_timestamptz_partitioned() { let object_store = ObjectStoreBuilder::memory(); diff --git a/iceberg-rust/src/arrow/partition.rs b/iceberg-rust/src/arrow/partition.rs index 7f7ffcf8..7e0f0e99 100644 --- a/iceberg-rust/src/arrow/partition.rs +++ b/iceberg-rust/src/arrow/partition.rs @@ -12,11 +12,11 @@ use std::{collections::HashSet, hash::Hash}; use arrow::{ array::{ - as_primitive_array, as_string_array, ArrayRef, BooleanArray, BooleanBufferBuilder, + as_primitive_array, as_string_array, Array, ArrayRef, BooleanArray, BooleanBufferBuilder, PrimitiveArray, Scalar, StringArray, }, compute::{ - and, filter, filter_record_batch, + and, filter, filter_record_batch, is_not_null, is_null, kernels::cmp::{distinct, eq}, }, datatypes::{ @@ -31,6 +31,9 @@ use iceberg_rust_spec::{partition::BoundPartitionField, spec::values::Value}; use super::transform::transform_arrow; +/// A record batch and the partition values shared by all its rows. +pub type PartitionedRecordBatch = (Vec>, RecordBatch); + /// Partitions a record batch according to the given partition fields. /// /// This function takes a record batch and partition field specifications, then splits the batch into @@ -53,7 +56,7 @@ use super::transform::transform_arrow; pub fn partition_record_batch<'a>( record_batch: &'a RecordBatch, partition_fields: &[BoundPartitionField<'_>], -) -> Result, RecordBatch), ArrowError>> + 'a, ArrowError> { +) -> Result> + 'a, ArrowError> { let partition_columns: Vec = partition_fields .iter() .map(|field| { @@ -72,48 +75,56 @@ pub fn partition_record_batch<'a>( let predicates = distinct_values .into_iter() .zip(partition_columns.iter()) - .map(|(distinct, value)| match distinct { - DistinctValues::Int(set) => set - .into_iter() - .map(|x| { - Ok(( - Value::Int(x), - eq(&PrimitiveArray::::new_scalar(x), value)?, - )) - }) - .collect::, ArrowError>>(), - DistinctValues::Long(set) => set - .into_iter() - .map(|x| { - Ok(( - Value::LongInt(x), - eq(&PrimitiveArray::::new_scalar(x), value)?, - )) - }) - .collect::, ArrowError>>(), - DistinctValues::Timestamp(set, timezone) => set - .into_iter() - .map(|x| { - let scalar = Scalar::new( - PrimitiveArray::::from(vec![x]) - .with_timezone_opt(timezone.clone()), - ); - let partition_value = if timezone.is_some() { - Value::TimestampTZ(x) - } else { - Value::Timestamp(x) - }; - Ok((partition_value, eq(&scalar, value)?)) - }) - .collect::, ArrowError>>(), - DistinctValues::String(set) => set - .into_iter() - .map(|x| { - let res = eq(&StringArray::new_scalar(&x), value)?; - Ok((Value::String(x), res)) - }) - .collect::, ArrowError>>(), - }) + .map( + |(distinct, value)| -> Result, BooleanArray)>, ArrowError> { + let mut predicates = match distinct { + DistinctValues::Int(set) => set + .into_iter() + .map(|x| { + Ok(( + Some(Value::Int(x)), + eq(&PrimitiveArray::::new_scalar(x), value)?, + )) + }) + .collect::, ArrowError>>(), + DistinctValues::Long(set) => set + .into_iter() + .map(|x| { + Ok(( + Some(Value::LongInt(x)), + eq(&PrimitiveArray::::new_scalar(x), value)?, + )) + }) + .collect::, ArrowError>>(), + DistinctValues::Timestamp(set, timezone) => set + .into_iter() + .map(|x| { + let scalar = Scalar::new( + PrimitiveArray::::from(vec![x]) + .with_timezone_opt(timezone.clone()), + ); + let partition_value = if timezone.is_some() { + Value::TimestampTZ(x) + } else { + Value::Timestamp(x) + }; + Ok((Some(partition_value), eq(&scalar, value)?)) + }) + .collect::, ArrowError>>(), + DistinctValues::String(set) => set + .into_iter() + .map(|x| { + let res = eq(&StringArray::new_scalar(&x), value)?; + Ok((Some(Value::String(x)), res)) + }) + .collect::, ArrowError>>(), + }?; + if value.null_count() != 0 { + predicates.push((None, is_null(value.as_ref())?)); + } + Ok(predicates) + }, + ) .try_fold( vec![(vec![], BooleanArray::new(true_buffer.finish(), None))], |acc, predicates| { @@ -123,7 +134,7 @@ pub fn partition_record_batch<'a>( Ok((values, and(&x, y)?)) }) .filter_ok(|x| x.1.true_count() != 0) - .collect::, _)>, ArrowError>>() + .collect::>, _)>, ArrowError>>() }, )?; Ok(predicates.into_iter().map(move |(values, predicate)| { @@ -145,11 +156,11 @@ pub fn partition_record_batch<'a>( /// * Int64 - Converted to DistinctValues::Long /// * Utf8 - Converted to DistinctValues::String fn distinct_values(array: ArrayRef) -> Result { - if array.null_count() != 0 { - return Err(ArrowError::ComputeError( - "Null partition values are not supported".to_owned(), - )); - } + let array = if array.null_count() == 0 { + array + } else { + filter(array.as_ref(), &is_not_null(array.as_ref())?)? + }; match array.data_type() { DataType::Int32 => Ok(DistinctValues::Int(distinct_values_primitive::< i32, @@ -186,6 +197,10 @@ fn distinct_values_primitive>( ) -> Result, ArrowError> { let array = as_primitive_array::

(&array); + if array.is_empty() { + return Ok(HashSet::new()); + } + let first = array.value(0); let slice_len = array.len() - 1; @@ -223,10 +238,14 @@ fn distinct_values_primitive>( /// # Returns /// A HashSet containing all unique string values from the array fn distinct_values_string(array: ArrayRef) -> Result, ArrowError> { - let slice_len = array.len() - 1; - let array = as_string_array(&array); + if array.is_empty() { + return Ok(HashSet::new()); + } + + let slice_len = array.len() - 1; + let first = array.value(0).to_owned(); if slice_len == 0 { @@ -305,9 +324,9 @@ mod tests { .collect::, ArrowError>>()?; values.sort(); let expected = if timezone.is_some() { - vec![Value::TimestampTZ(10), Value::TimestampTZ(20)] + vec![Some(Value::TimestampTZ(10)), Some(Value::TimestampTZ(20))] } else { - vec![Value::Timestamp(10), Value::Timestamp(20)] + vec![Some(Value::Timestamp(10)), Some(Value::Timestamp(20))] }; assert_eq!(values, expected); } @@ -315,7 +334,7 @@ mod tests { } #[test] - fn null_partition_values_fail_instead_of_dropping_rows() -> Result<(), ArrowError> { + fn nullable_partition_values_keep_every_row() -> Result<(), ArrowError> { for (array, source_type) in [ ( Arc::new(arrow::array::Int64Array::from(vec![Some(10), None])) as ArrayRef, @@ -331,10 +350,35 @@ mod tests { let source = StructField::new(1, "ts", false, Type::Primitive(source_type), None); let partition = PartitionField::new(1, 1000, "ts", Transform::Identity); let bound = BoundPartitionField::new(&partition, &source); - let result = partition_record_batch(&batch, &[bound]); - assert!( - matches!(result, Err(ArrowError::ComputeError(message)) if message.contains("Null partition values")) - ); + let mut groups = partition_record_batch(&batch, &[bound])? + .map(|group| group.map(|(values, rows)| (values[0].clone(), rows.num_rows()))) + .collect::, ArrowError>>()?; + groups.sort(); + assert_eq!(groups.len(), 2); + assert_eq!(groups[0], (None, 1)); + assert_eq!(groups[1].1, 1); + } + Ok(()) + } + + #[test] + fn all_null_and_empty_partition_batches() -> Result<(), ArrowError> { + for values in [vec![None, None], vec![]] { + let array: ArrayRef = Arc::new(arrow::array::Int64Array::from(values.clone())); + let batch = RecordBatch::try_from_iter(vec![("id", array)])?; + let source = + StructField::new(1, "id", false, Type::Primitive(PrimitiveType::Long), None); + let partition = PartitionField::new(1, 1000, "id", Transform::Identity); + let bound = BoundPartitionField::new(&partition, &source); + let groups = partition_record_batch(&batch, &[bound])? + .collect::, ArrowError>>()?; + if values.is_empty() { + assert!(groups.is_empty()); + } else { + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].0, vec![None]); + assert_eq!(groups[0].1.num_rows(), values.len()); + } } Ok(()) } diff --git a/iceberg-rust/src/arrow/write.rs b/iceberg-rust/src/arrow/write.rs index 64fef9a5..97250cdb 100644 --- a/iceberg-rust/src/arrow/write.rs +++ b/iceberg-rust/src/arrow/write.rs @@ -323,7 +323,7 @@ async fn store_parquet_partitioned( Ok(files) } else { let table_properties = Arc::new(metadata.properties.clone()); - let mut senders: LruCache, Sender>> = + let mut senders: LruCache>, Sender>> = LruCache::unbounded(); let mut set = JoinSet::new(); @@ -444,7 +444,7 @@ async fn write_parquet_files( arrow_schema: &ArrowSchema, partition_fields: &[BoundPartitionField<'_>], partition_path: Option, - partition_values: Option>, + partition_values: Option>>, batches: impl Stream> + Send, object_store: Arc, equality_ids: Option<&[i32]>, @@ -601,14 +601,18 @@ async fn write_parquet_files( #[inline] pub fn generate_partition_path( partition_fields: &[BoundPartitionField<'_>], - partition_values: &[Value], + partition_values: &[Option], ) -> Result { partition_fields .iter() .zip(partition_values.iter()) .map(|(field, value)| { let name = field.name().to_owned(); - Ok(name + "=" + &value.to_string() + "/") + let text = value.as_ref().map_or_else( + || "__HIVE_DEFAULT_PARTITION__".to_owned(), + ToString::to_string, + ); + Ok(name + "=" + &text + "/") }) .collect::>() } @@ -1432,11 +1436,14 @@ mod tests { }; let partfield = PartitionField::new(1, 1001, "month", Transform::Month); let partition_fields = vec![BoundPartitionField::new(&partfield, &field)]; - let partition_values = vec![Value::Int(10)]; + let partition_values = vec![Some(Value::Int(10))]; let result = super::generate_partition_path(&partition_fields, &partition_values); assert!(result.is_ok()); assert_eq!(result.unwrap(), "month=10/"); + + let null_path = super::generate_partition_path(&partition_fields, &[None]); + assert_eq!(null_path.unwrap(), "month=__HIVE_DEFAULT_PARTITION__/"); } } diff --git a/iceberg-rust/src/file_format/parquet.rs b/iceberg-rust/src/file_format/parquet.rs index 58fd6e99..e5794015 100644 --- a/iceberg-rust/src/file_format/parquet.rs +++ b/iceberg-rust/src/file_format/parquet.rs @@ -100,7 +100,7 @@ pub fn parquet_to_datafile_with_partition( partition_fields: &[BoundPartitionField<'_>], equality_ids: Option<&[i32]>, table_properties: &HashMap, - partition_values: &[Value], + partition_values: &[Option], ) -> Result { parquet_to_datafile_impl( location, @@ -123,11 +123,12 @@ fn parquet_to_datafile_impl( partition_fields: &[BoundPartitionField<'_>], equality_ids: Option<&[i32]>, table_properties: &HashMap, - partition_values: Option<&[Value]>, + partition_values: Option<&[Option]>, ) -> Result { let write_distinct_counts = table_properties .get(WRITE_METADATA_METRICS_DISTINCT_COUNTS_ENABLED) .is_some_and(|x| x == "true"); + let infer_partition_from_stats = partition_values.is_none(); let mut partition = if let Some(values) = partition_values { if values.len() != partition_fields.len() { return Err(Error::InvalidFormat("Partition value count".to_owned())); @@ -135,7 +136,7 @@ fn parquet_to_datafile_impl( partition_fields .iter() .zip(values) - .map(|(field, value)| (field.name().to_owned(), Some(value.clone()))) + .map(|(field, value)| (field.name().to_owned(), value.clone())) .collect::() } else { partition_fields @@ -451,33 +452,35 @@ fn parquet_to_datafile_impl( } } - if let Some(partition_field) = partition_fields.get(column_name) { - if let Some(partition_value) = partition.get_mut(partition_field.name()) { - if partition_value.is_none() { - let partition_field = partition_fields - .get(column_name) - .ok_or_else(|| Error::InvalidFormat("transform".to_string()))?; - if let (Some(min_bytes), Some(max_bytes)) = - (statistics.min_bytes_opt(), statistics.max_bytes_opt()) - { - let min = Value::try_from_bytes_with_hint( - min_bytes, - data_type, - physical_type_hint, - )? - .transform(partition_field.transform())?; - let max = Value::try_from_bytes_with_hint( - max_bytes, - data_type, - physical_type_hint, - )? - .transform(partition_field.transform())?; - if min == max { - *partition_value = Some(min) - } else { - return Err(Error::InvalidFormat( - "Partition value of data file".to_owned(), - )); + if infer_partition_from_stats { + if let Some(partition_field) = partition_fields.get(column_name) { + if let Some(partition_value) = partition.get_mut(partition_field.name()) { + if partition_value.is_none() { + let partition_field = partition_fields + .get(column_name) + .ok_or_else(|| Error::InvalidFormat("transform".to_string()))?; + if let (Some(min_bytes), Some(max_bytes)) = + (statistics.min_bytes_opt(), statistics.max_bytes_opt()) + { + let min = Value::try_from_bytes_with_hint( + min_bytes, + data_type, + physical_type_hint, + )? + .transform(partition_field.transform())?; + let max = Value::try_from_bytes_with_hint( + max_bytes, + data_type, + physical_type_hint, + )? + .transform(partition_field.transform())?; + if min == max { + *partition_value = Some(min) + } else { + return Err(Error::InvalidFormat( + "Partition value of data file".to_owned(), + )); + } } } } diff --git a/iceberg-rust/src/table/manifest.rs b/iceberg-rust/src/table/manifest.rs index a3d5cf40..a138acb4 100644 --- a/iceberg-rust/src/table/manifest.rs +++ b/iceberg-rust/src/table/manifest.rs @@ -840,7 +840,12 @@ fn update_partitions( partition_columns: &[PartitionField], ) -> Result<(), Error> { for (field, summary) in partition_columns.iter().zip(partitions.iter_mut()) { - let value = partition_values.get(field.name()).and_then(|x| x.as_ref()); + let value = partition_values + .get(field.name()) + .ok_or_else(|| Error::InvalidFormat("Partition field is missing".to_owned()))?; + if value.is_none() { + summary.contains_null = true; + } if let Some(value) = value { if summary.lower_bound.is_none() { summary.lower_bound = Some(value.clone()); @@ -931,8 +936,32 @@ fn update_partitions( #[cfg(test)] mod tests { + use super::{update_partitions, FieldSummary, PartitionField, Struct, Value}; + use iceberg_rust_spec::partition::Transform; use rstest::rstest; + #[test] + fn manifest_partition_summary_tracks_null_and_non_null_values() { + let fields = [PartitionField::new(1, 1000, "id", Transform::Identity)]; + let mut summaries = [FieldSummary { + contains_null: false, + contains_nan: None, + lower_bound: None, + upper_bound: None, + }]; + let null_partition = Struct::from_iter([("id".to_owned(), None)]); + update_partitions(&mut summaries, &null_partition, &fields).unwrap(); + assert!(summaries[0].contains_null); + assert_eq!(summaries[0].lower_bound, None); + assert_eq!(summaries[0].upper_bound, None); + + let value_partition = Struct::from_iter([("id".to_owned(), Some(Value::Int(7)))]); + update_partitions(&mut summaries, &value_partition, &fields).unwrap(); + assert!(summaries[0].contains_null); + assert_eq!(summaries[0].lower_bound, Some(Value::Int(7))); + assert_eq!(summaries[0].upper_bound, Some(Value::Int(7))); + } + // -- TestManifestReader (15) -- #[rstest] #[case(1)] diff --git a/iceberg-rust/src/table/manifest_list.rs b/iceberg-rust/src/table/manifest_list.rs index 3abddeeb..da98c046 100644 --- a/iceberg-rust/src/table/manifest_list.rs +++ b/iceberg-rust/src/table/manifest_list.rs @@ -338,6 +338,61 @@ pub async fn snapshot_partition_bounds( }) } +/// Partition bounds and NULL presence for each field in the current partition spec. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PartitionBoundsWithNulls { + /// Bounds across the snapshot's manifest summaries. + pub bounds: Rectangle, + /// Whether each partition field contains NULL in any manifest. + pub contains_null: Vec, +} + +/// Returns partition bounds together with NULL presence in one manifest-list read. +/// +/// Returns `None` when a manifest has no summaries or uses a different partition +/// spec; bounds in those cases cannot safely restrict a scan of this snapshot. +pub async fn snapshot_partition_bounds_with_nulls( + snapshot: &Snapshot, + table_metadata: &TableMetadata, + object_store: Arc, +) -> Result, Error> { + let bytes = object_store + .get(&strip_prefix(snapshot.manifest_list()).into()) + .await? + .bytes() + .await?; + let reader = ManifestListReader::new(Cursor::new(bytes), table_metadata)?; + let mut result: Option = None; + for entry in reader { + let entry = entry?; + if entry.partition_spec_id != table_metadata.default_spec_id { + return Ok(None); + } + let Some(partitions) = entry.partitions else { + return Ok(None); + }; + let bounds = summary_to_rectangle(&partitions)?; + if let Some(result) = result.as_mut() { + if result.contains_null.len() != partitions.len() { + return Ok(None); + } + result.bounds.expand(&bounds); + for (contains_null, summary) in result.contains_null.iter_mut().zip(&partitions) { + *contains_null |= summary.contains_null; + } + } else { + result = Some(PartitionBoundsWithNulls { + bounds, + contains_null: partitions + .iter() + .map(|summary| summary.contains_null) + .collect(), + }); + } + } + Ok(result) +} + /// Computes the column bounds (minimum and maximum values) for all primitive fields /// across all data files in a snapshot. /// diff --git a/iceberg-rust/src/table/transaction/append.rs b/iceberg-rust/src/table/transaction/append.rs index aaf476ce..8f213fea 100644 --- a/iceberg-rust/src/table/transaction/append.rs +++ b/iceberg-rust/src/table/transaction/append.rs @@ -36,7 +36,13 @@ fn split_datafiles_once( for manifest_entry in files { let manifest_entry = manifest_entry?; - let position = partition_struct_to_vec(manifest_entry.data_file().partition(), names)?; + let Some(position) = + partition_struct_to_vec(manifest_entry.data_file().partition(), names)? + else { + smaller.push(manifest_entry); + smaller_rect = Some(Rectangle::unbounded()); + continue; + }; // Compare distance to upper and lower bound. Since you can't compute a "norm" for a multidimensional vector where the dimensions have different datatypes, // the dimensions are compared individually and the norm is computed by weighing the earlier columns more than the later. if let Ordering::Greater = cmp_with_priority( diff --git a/iceberg-rust/src/table/transaction/operation.rs b/iceberg-rust/src/table/transaction/operation.rs index 6e9ee203..b8d702bb 100644 --- a/iceberg-rust/src/table/transaction/operation.rs +++ b/iceberg-rust/src/table/transaction/operation.rs @@ -1058,10 +1058,15 @@ pub fn bounding_partition_values<'a>( partition_column_names: &SmallVec<[&str; 4]>, ) -> Result { iter.try_fold(None, |acc, x| { - let node = partition_struct_to_vec(x.partition(), partition_column_names)?; + let Some(node) = partition_struct_to_vec(x.partition(), partition_column_names)? else { + return Ok::<_, Error>(Some(Rectangle::unbounded())); + }; let Some(mut acc) = acc else { return Ok::<_, Error>(Some(Rectangle::new(node.clone(), node))); }; + if acc.min.is_empty() { + return Ok(Some(acc)); + } acc.expand_with_node(node); Ok(Some(acc)) })? diff --git a/iceberg-rust/src/util/mod.rs b/iceberg-rust/src/util/mod.rs index 9d09489a..0f49a28a 100644 --- a/iceberg-rust/src/util/mod.rs +++ b/iceberg-rust/src/util/mod.rs @@ -25,11 +25,20 @@ impl Rectangle { Self { min, max } } + pub(crate) fn unbounded() -> Self { + Self::new(Vec4::new(), Vec4::new()) + } + /// Expands the rectangle to include the given rectangle. /// /// This method updates the minimum and maximum values of the rectangle to include /// the values in the given `rect` rectangle. pub(crate) fn expand(&mut self, rect: &Rectangle) { + if self.min.is_empty() || rect.min.is_empty() { + self.min.clear(); + self.max.clear(); + return; + } for i in 0..self.min.len() { if rect.min[i] < self.min[i] { self.min[i] = rect.min[i].clone(); @@ -59,6 +68,11 @@ impl Rectangle { /// ///Values the earlier columns more than the later. pub(crate) fn cmp_with_priority(&self, other: &Rectangle) -> Result { + match (self.min.is_empty(), other.min.is_empty()) { + (true, false) => return Ok(Ordering::Greater), + (false, true) => return Ok(Ordering::Less), + _ => {} + } let self_iter = self .max .iter() @@ -85,12 +99,17 @@ impl Rectangle { pub(crate) fn partition_struct_to_vec( partition_struct: &Struct, names: &[&str], -) -> Result, Error> { - names +) -> Result>, Error> { + let fields = names .iter() - .map(|x| partition_struct.get(x).and_then(Clone::clone)) - .collect::>>() - .ok_or(Error::InvalidFormat("Partition struct".to_owned())) + .map(|x| { + partition_struct + .get(x) + .cloned() + .ok_or_else(|| Error::InvalidFormat("Partition struct".to_owned())) + }) + .collect::>, Error>>()?; + Ok(fields.into_iter().collect()) } pub(crate) fn summary_to_rectangle(summaries: &[FieldSummary]) -> Result { @@ -98,18 +117,11 @@ pub(crate) fn summary_to_rectangle(summaries: &[FieldSummary]) -> Result Result, Er #[cfg(test)] mod tests { - use iceberg_rust_spec::values::Value; + use iceberg_rust_spec::values::{Struct, Value}; use smallvec::smallvec; use super::*; + #[test] + fn null_partition_bounds_are_conservatively_unbounded() { + let null_tuple = Struct::from_iter([("id".to_owned(), None)]); + assert_eq!(partition_struct_to_vec(&null_tuple, &["id"]).unwrap(), None); + assert!(partition_struct_to_vec(&null_tuple, &["missing"]).is_err()); + + let summary = FieldSummary { + contains_null: true, + contains_nan: None, + lower_bound: None, + upper_bound: None, + }; + let mut bounds = Rectangle::new(smallvec![Value::Int(1)], smallvec![Value::Int(2)]); + bounds.expand(&summary_to_rectangle(&[summary]).unwrap()); + assert!(bounds.min.is_empty()); + assert!(bounds.max.is_empty()); + let bounded = Rectangle::new(smallvec![Value::Int(1)], smallvec![Value::Int(2)]); + assert_eq!( + bounds.cmp_with_priority(&bounded).unwrap(), + Ordering::Greater + ); + } + #[test] fn test_sub_valid() { let left = vec![Value::Int(5), Value::Int(10), Value::Int(15)];