Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion crates/iceberg/public-api.txt
Original file line number Diff line number Diff line change
Expand Up @@ -1395,7 +1395,7 @@ pub fn iceberg::scan::FileScanTaskDeleteFile::content_size_in_bytes(&self) -> co
pub fn iceberg::scan::FileScanTaskDeleteFile::equality_ids(&self) -> core::option::Option<&[i32]>
pub fn iceberg::scan::FileScanTaskDeleteFile::file_format(&self) -> iceberg::spec::DataFileFormat
pub fn iceberg::scan::FileScanTaskDeleteFile::file_path(&self) -> &str
pub fn iceberg::scan::FileScanTaskDeleteFile::file_size_in_bytes(&self) -> u64
pub fn iceberg::scan::FileScanTaskDeleteFile::file_size_in_bytes(&self) -> core::option::Option<u64>
pub fn iceberg::scan::FileScanTaskDeleteFile::file_type(&self) -> iceberg::spec::DataContentType
pub fn iceberg::scan::FileScanTaskDeleteFile::key_metadata(&self) -> core::option::Option<&[u8]>
pub fn iceberg::scan::FileScanTaskDeleteFile::partition_spec_id(&self) -> i32
Expand Down
310 changes: 307 additions & 3 deletions crates/iceberg/src/arrow/caching_delete_file_loader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -823,23 +823,48 @@ mod tests {
use std::collections::HashMap;
use std::fs::File;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use arrow_array::cast::AsArray;
use arrow_array::{
ArrayRef, BinaryArray, Int32Array, Int64Array, RecordBatch, StringArray, StructArray,
};
use arrow_schema::{DataType, Field, Fields};
use futures::stream::BoxStream;
use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
use parquet::basic::Compression;
use parquet::file::properties::WriterProperties;
use tempfile::TempDir;

use super::*;
use crate::arrow::delete_filter::tests::setup;
use crate::io::{
FileIOBuilder, FileMetadata, FileRead, FileWrite, InputFile, LocalFsStorage, OutputFile,
Storage, StorageConfig, StorageFactory,
};
use crate::scan::FileScanTaskDeleteFile;
use crate::spec::{DataContentType, DataFileFormat, Schema};
use crate::test_utils::encode_dv_blob;

#[test]
fn test_delete_file_without_size_round_trips_through_serde() {
let task = FileScanTaskDeleteFile::builder()
.with_file_path("pos-del.parquet".to_string())
.with_file_type(DataContentType::PositionDeletes)
.with_file_format(DataFileFormat::Parquet)
.with_partition_spec_id(0)
.build()
.unwrap();
assert_eq!(task.file_size_in_bytes(), None);

// The unknown size is left out when serializing, and a task without the
// field deserializes to `None`.
let json = serde_json::to_value(&task).unwrap();
assert!(json.get("file_size_in_bytes").is_none(), "{json}");
let round_tripped: FileScanTaskDeleteFile = serde_json::from_value(json).unwrap();
assert_eq!(round_tripped, task);
}

#[test]
fn test_deserialized_deletion_vector_rejects_missing_fields() {
// A deletion vector that omits its required coordinates fails validation
Expand Down Expand Up @@ -873,7 +898,7 @@ mod tests {
let record_batch_stream = basic_delete_file_loader
.parquet_to_batch_stream(
&eq_delete_file_path,
std::fs::metadata(&eq_delete_file_path).unwrap().len(),
Some(std::fs::metadata(&eq_delete_file_path).unwrap().len()),
None,
)
.await
Expand Down Expand Up @@ -1129,6 +1154,285 @@ mod tests {
assert!(result.is_none()); // no pos dels for file 3
}

/// Local filesystem storage that counts metadata requests.
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
struct MetadataCountingStorage {
#[serde(skip)]
inner: LocalFsStorage,
#[serde(skip)]
calls: Arc<AtomicUsize>,
/// Fail every metadata request with a retryable error.
#[serde(skip)]
fail_metadata: bool,
}

#[async_trait::async_trait]
#[typetag::serde(name = "test_metadata_counting_storage")]
impl Storage for MetadataCountingStorage {
async fn metadata(&self, path: &str) -> Result<FileMetadata> {
self.calls.fetch_add(1, Ordering::SeqCst);
if self.fail_metadata {
return Err(
Error::new(ErrorKind::Unexpected, "metadata temporarily unavailable")
.with_retryable(true),
);
}
self.inner.metadata(path).await
}

fn new_input(&self, path: &str) -> Result<InputFile> {
Ok(InputFile::new(Arc::new(self.clone()), path.to_string()))
}

fn new_output(&self, path: &str) -> Result<OutputFile> {
Ok(OutputFile::new(Arc::new(self.clone()), path.to_string()))
}

async fn exists(&self, path: &str) -> Result<bool> {
self.inner.exists(path).await
}

async fn read(&self, path: &str) -> Result<Bytes> {
self.inner.read(path).await
}

async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> {
self.inner.reader(path).await
}

async fn write(&self, path: &str, bs: Bytes) -> Result<()> {
self.inner.write(path, bs).await
}

async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> {
self.inner.writer(path).await
}

async fn delete(&self, path: &str) -> Result<()> {
self.inner.delete(path).await
}

async fn delete_prefix(&self, path: &str) -> Result<()> {
self.inner.delete_prefix(path).await
}

async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> {
self.inner.delete_stream(paths).await
}
}

#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
struct MetadataCountingStorageFactory {
#[serde(skip)]
calls: Arc<AtomicUsize>,
#[serde(skip)]
fail_metadata: bool,
}

#[typetag::serde(name = "test_metadata_counting_storage_factory")]
impl StorageFactory for MetadataCountingStorageFactory {
fn build(&self, _config: &StorageConfig) -> Result<Arc<dyn Storage>> {
Ok(Arc::new(MetadataCountingStorage {
inner: LocalFsStorage,
calls: Arc::clone(&self.calls),
fail_metadata: self.fail_metadata,
}))
}
}
Comment on lines +1157 to +1241

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could the test storage import LocalFsStorage, AtomicUsize, Ordering and the other crate::io types at the top of the test module, and hold one LocalFsStorage in a field? Each method currently builds a new crate::io::LocalFsStorage::new() behind a fully qualified path, which makes the 90 lines harder to scan for the two methods that differ (metadata and the constructors).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in b2981e2: the crate::io types are imported at the top and the storage holds a LocalFsStorage.


/// Rebuilds `delete` at `path` without a file size.
fn without_size(delete: &FileScanTaskDeleteFile, path: &str) -> FileScanTaskDeleteFile {
FileScanTaskDeleteFile::builder()
.with_file_path(path.to_string())
.with_file_type(delete.file_type())
.with_file_format(delete.file_format())
.with_partition_spec_id(delete.partition_spec_id())
.with_equality_ids(delete.equality_ids().map(<[i32]>::to_vec))
.with_referenced_data_file(delete.referenced_data_file().map(str::to_string))
.with_content_offset(delete.content_offset())
.with_content_size_in_bytes(delete.content_size_in_bytes())
.with_record_count(delete.record_count())
.with_key_metadata(delete.key_metadata().map(Box::from))
.build()
.unwrap()
}

#[tokio::test]
async fn test_unknown_delete_sizes_are_resolved_once_per_file() {
let tmp_dir = TempDir::new().unwrap();
let file_scan_tasks = setup(tmp_dir.path());
let factory = Arc::new(MetadataCountingStorageFactory::default());
let calls = Arc::clone(&factory.calls);
let file_io = FileIOBuilder::new(factory).build();
let task = &file_scan_tasks[0];
let unknown: Vec<_> = task
.deletes()
.iter()
.map(|delete| without_size(delete, delete.file_path()))
.collect();
let positions = |filter: &DeleteFilter| -> Vec<u64> {
filter
.get_delete_vector(task)
.unwrap()
.lock()
.unwrap()
.iter()
.collect()
};

// Known sizes never request metadata.
let known = CachingDeleteFileLoader::new(file_io.clone(), 10, Runtime::current())
.load_deletes(task.deletes(), task.schema_ref())
.await
.unwrap()
.unwrap();
assert_eq!(calls.load(Ordering::SeqCst), 0);

// Unknown sizes: one request per delete file, and none again when
// another task references an already loaded delete file.
let loader = CachingDeleteFileLoader::new(file_io, 10, Runtime::current());
let lazy = loader
.load_deletes(&unknown, task.schema_ref())
.await
.unwrap()
.unwrap();
assert_eq!(calls.load(Ordering::SeqCst), unknown.len());
loader
.load_deletes(&unknown[1..], task.schema_ref())

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this also run two load_deletes on the same unknown-size files concurrently, say with futures::join!, and check calls stays at unknown.len()? Here the second load starts after the first finishes, so it takes the AlreadyLoaded path. The only WaitFor test is state-level and doesn't go through load_deletes. A succeeding stat is enough, since a failing one hangs today (#3434).

.await
.unwrap()
.unwrap();
assert_eq!(calls.load(Ordering::SeqCst), unknown.len());

// The same positions are deleted either way.
let expected = positions(&known);
assert_eq!(expected.len(), 12);
assert_eq!(positions(&lazy), expected);
}

#[tokio::test]
async fn test_unknown_size_equality_deletes_are_resolved_once_and_cached() {
use crate::scan::FileScanTask;
use crate::spec::DataFileFormat;

let tmp_dir = TempDir::new().unwrap();
let table_location = tmp_dir.path().to_str().unwrap();
let eq_delete_path = setup_write_equality_delete_file_1(table_location);
let schema = Arc::new(
Schema::builder()
.with_fields(vec![
NestedField::optional(2, "y", Type::Primitive(PrimitiveType::Long)).into(),
NestedField::optional(3, "z", Type::Primitive(PrimitiveType::Long)).into(),
])
.build()
.unwrap(),
);
let equality_delete = |file_size_in_bytes| {
FileScanTaskDeleteFile::builder()
.with_file_path(eq_delete_path.clone())
.with_file_size_in_bytes_opt(file_size_in_bytes)
.with_file_type(DataContentType::EqualityDeletes)
.with_file_format(DataFileFormat::Parquet)
.with_partition_spec_id(0)
.with_equality_ids(Some(vec![2, 3]))
.build()
.unwrap()
};
let data_task = |name: &str, delete| {
FileScanTask::builder()
.with_file_size_in_bytes(0)
.with_start(0)
.with_length(0)
.with_data_file_path(format!("{table_location}/{name}"))
.with_data_file_format(DataFileFormat::Parquet)
.with_schema(Arc::clone(&schema))
.with_project_field_ids(vec![2, 3])
.with_deletes(vec![delete])
.with_case_sensitive(false)
.build()
.unwrap()
};
let known_size = std::fs::metadata(&eq_delete_path).unwrap().len();

let factory = Arc::new(MetadataCountingStorageFactory::default());
let calls = Arc::clone(&factory.calls);
let file_io = FileIOBuilder::new(factory).build();

// Known size: no metadata request.
let known_task = data_task("known.parquet", equality_delete(Some(known_size)));
let known = CachingDeleteFileLoader::new(file_io.clone(), 10, Runtime::current())
.load_deletes(known_task.deletes(), known_task.schema_ref())
.await
.unwrap()
.unwrap()
.build_equality_delete_predicate(&known_task)
.await
.unwrap();
assert!(known.is_some());
assert_eq!(calls.load(Ordering::SeqCst), 0);

// Unknown size: one metadata request, and none when a second data
// file referencing the same equality delete file is loaded.
let loader = CachingDeleteFileLoader::new(file_io, 10, Runtime::current());
for (name, expected_calls) in [("first.parquet", 1), ("second.parquet", 1)] {
let task = data_task(name, equality_delete(None));
let predicate = loader
.load_deletes(task.deletes(), task.schema_ref())
.await
.unwrap()
.unwrap()
.build_equality_delete_predicate(&task)
.await
.unwrap();
assert_eq!(predicate, known, "{name}");
assert_eq!(calls.load(Ordering::SeqCst), expected_calls, "{name}");
}
}

#[tokio::test]
async fn test_load_deletes_reports_unknown_size_failures() {
let tmp_dir = TempDir::new().unwrap();
let file_scan_tasks = setup(tmp_dir.path());
let task = &file_scan_tasks[0];
let unknown_size = |path: &str| without_size(&task.deletes()[0], path);

// A missing delete file fails the load instead of dropping its deletes.
let missing = unknown_size(&format!(
"{}/missing-pos-del.parquet",
tmp_dir.path().to_str().unwrap()
));
let error = CachingDeleteFileLoader::new(FileIO::new_with_fs(), 10, Runtime::current())
.load_deletes(&[missing], task.schema_ref())
.await
.unwrap()
.err()
.unwrap();
assert!(
error.to_string().contains("Failed to stat delete file"),
"{error}"
);

// A retryable storage failure stays retryable through the loader.
let factory = Arc::new(MetadataCountingStorageFactory {
fail_metadata: true,
..Default::default()
});
let file_io = FileIOBuilder::new(factory).build();
let present = unknown_size(task.deletes()[0].file_path());
let error = CachingDeleteFileLoader::new(file_io, 10, Runtime::current())
.load_deletes(&[present], task.schema_ref())
.await
.unwrap()
.err()
.unwrap();
assert!(
error.to_string().contains("Failed to stat delete file"),
"{error}"
);
assert_eq!(error.kind(), ErrorKind::Unexpected);
assert!(error.retryable());
}

#[tokio::test]
async fn test_parse_positional_deletes_rejects_negative_positions() {
let schema = crate::arrow::delete_filter::tests::create_pos_del_schema();
Expand Down Expand Up @@ -1312,7 +1616,7 @@ mod tests {
let batch_stream = basic_delete_file_loader
.parquet_to_batch_stream(
&delete_file_path,
std::fs::metadata(&delete_file_path).unwrap().len(),
Some(std::fs::metadata(&delete_file_path).unwrap().len()),
None,
)
.await
Expand Down Expand Up @@ -1490,7 +1794,7 @@ mod tests {
let basic_delete_file_loader =
BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new());
let record_batch_stream = basic_delete_file_loader
.parquet_to_batch_stream(&path, std::fs::metadata(&path).unwrap().len(), None)
.parquet_to_batch_stream(&path, Some(std::fs::metadata(&path).unwrap().len()), None)
.await
.expect("could not get batch stream");

Expand Down
Loading
Loading