dhruvarya-db commented on code in PR #3260:
URL: https://github.com/apache/iceberg-rust/pull/3260#discussion_r4128560292
##########
crates/iceberg/src/scan/mod.rs:
##########
@@ -1956,6 +1986,329 @@ pub mod tests {
);
}
+ #[test]
+ fn test_table_scan_as_of_time_after_rollback() {
+ for elapsed_ms in [0, 1000] {
+ let table = TableTestFixture::new().table;
+ let mut metadata = table.metadata().clone();
+ let original = metadata.history()[0].snapshot_id;
+ let previous = metadata.history()[1].clone();
+ let rollback_time = previous.timestamp_ms + elapsed_ms;
+ metadata.snapshot_log.push(crate::spec::SnapshotLog {
+ timestamp_ms: rollback_time,
+ snapshot_id: original,
+ });
+ metadata.current_snapshot_id = Some(original);
+ metadata.refs.get_mut(MAIN_BRANCH).unwrap().snapshot_id = original;
+ let table = table.with_metadata(Arc::new(metadata));
+ let scan = table.scan().as_of_time(rollback_time).build().unwrap();
+ // A rollback at the same timestamp keeps the earlier history
entry.
+ let expected = if elapsed_ms == 0 {
+ previous.snapshot_id
+ } else {
+ original
+ };
+ assert_eq!(scan.snapshot().unwrap().snapshot_id(), expected);
+ }
+ }
+
+ #[test]
+ fn test_table_scan_as_of_time_rejects_empty_history() {
+ let table = TableTestFixture::new_empty().table;
+ assert!(table.scan().build().unwrap().snapshot().is_none());
+ let err = table.scan().as_of_time(0).build().unwrap_err();
+ assert_eq!(err.kind(), ErrorKind::DataInvalid);
+ assert!(err.message().contains("No snapshot history"));
+ }
+
+ #[test]
+ fn test_table_scan_as_of_time_rejects_before_history() {
+ let table = TableTestFixture::new().table;
+ let before = table.metadata().history()[0].timestamp_ms - 1;
+ let err = table.scan().as_of_time(before).build().unwrap_err();
+ assert_eq!(err.kind(), ErrorKind::DataInvalid);
+ assert!(err.message().contains(&before.to_string()));
+ }
+
+ #[test]
+ fn test_table_scan_as_of_time_rejects_all_qualifying_snapshots_expired() {
+ let table = TableTestFixture::new_with_deep_history().table;
+ let mut metadata = table.metadata().clone();
+ let entry = metadata.history()[3].clone();
+ let expired_ids: Vec<_> = metadata
+ .history()
+ .iter()
+ .filter(|log| log.timestamp_ms <= entry.timestamp_ms)
+ .map(|log| log.snapshot_id)
+ .collect();
+ assert_eq!(expired_ids.len(), 4);
+ for snapshot_id in expired_ids {
+ metadata.snapshots.remove(&snapshot_id);
+ }
+ let table = table.with_metadata(Arc::new(metadata));
+ // History resolution succeeds even though every qualifying snapshot
expired.
+ assert_eq!(
+ snapshot_id_as_of_time(&table.metadata_ref(),
entry.timestamp_ms).unwrap(),
+ entry.snapshot_id
+ );
+ assert!(table.scan().build().unwrap().snapshot().is_some());
+ let err = table
+ .scan()
+ .as_of_time(entry.timestamp_ms)
+ .build()
+ .unwrap_err();
+ assert_eq!(err.kind(), ErrorKind::DataInvalid);
+ assert_eq!(
+ err.message(),
+ format!("Snapshot with id {} not found", entry.snapshot_id)
+ );
+ }
+
+ #[test]
+ fn test_table_scan_as_of_time_conflicts_with_snapshot_id() {
+ for id_first in [true, false] {
+ let table = TableTestFixture::new().table;
+ let entry = &table.metadata().history()[0];
+ let scan = table.scan();
+ let scan = if id_first {
+ scan.snapshot_id(entry.snapshot_id)
+ .as_of_time(entry.timestamp_ms)
+ } else {
+ scan.as_of_time(entry.timestamp_ms)
+ .snapshot_id(entry.snapshot_id)
+ };
+ let err = scan.build().unwrap_err();
+ assert_eq!(err.kind(), ErrorKind::DataInvalid);
+ assert_eq!(err.message(), "Cannot combine snapshot_id and
as_of_time");
+ }
+ }
+
+ #[test]
+ fn test_table_scan_as_of_time_last_timestamp_wins() {
+ let table = TableTestFixture::new().table;
+ let first = &table.metadata().history()[0];
+ let scan = table
+ .scan()
+ .as_of_time(i64::MAX)
+ .as_of_time(first.timestamp_ms)
+ .build()
+ .unwrap();
+ assert_eq!(scan.snapshot().unwrap().snapshot_id(), first.snapshot_id);
+ // Keep the existing repeated snapshot_id setter behavior as well.
+ let scan = table
+ .scan()
+ .snapshot_id(-1)
+ .snapshot_id(first.snapshot_id)
+ .build()
+ .unwrap();
+ assert_eq!(scan.snapshot().unwrap().snapshot_id(), first.snapshot_id);
+ }
+
+ #[test]
+ fn test_table_scan_as_of_time_uses_snapshot_schema() {
+ for select_current in [false, true] {
+ let table = TableTestFixture::new().table;
+ let mut metadata = table.metadata().clone();
+ let entry =
metadata.history()[usize::from(select_current)].clone();
+ // Simulate a schema-only update after this snapshot: schema 0 has
x only,
+ // while the current table schema also contains y and other
columns.
+
Arc::make_mut(metadata.snapshots.get_mut(&entry.snapshot_id).unwrap()).schema_id
=
+ Some(0);
+ let table = table.with_metadata(Arc::new(metadata));
+ let scan = table
+ .scan()
+ .as_of_time(entry.timestamp_ms)
+ .select(["x"])
+ .with_filter(Reference::new("x").greater_than(Datum::long(0)))
+ .build()
+ .unwrap();
+ let context = scan.plan_context.as_ref().unwrap();
+ assert_eq!(context.snapshot_schema.schema_id(), 0);
+ assert!(context.snapshot_bound_predicate.is_some());
+ assert!(
+ table
+ .scan()
+ .as_of_time(entry.timestamp_ms)
+ .select(["y"])
+ .build()
+ .is_err()
+ );
+ assert!(
+ table
+ .scan()
+ .as_of_time(entry.timestamp_ms)
+
.with_filter(Reference::new("y").greater_than(Datum::long(0)))
+ .build()
+ .is_err()
+ );
+ }
+ }
+
+ #[test]
+ fn test_table_scan_as_of_time_schema_compatibility() {
+ let table = TableTestFixture::new().table;
+ let entry = table.metadata().history()[0].clone();
+ // Older snapshots may omit schema-id and use the current-schema
fallback.
+ assert!(
+ table
+ .metadata()
+ .snapshot_by_id(entry.snapshot_id)
+ .unwrap()
+ .schema_id()
+ .is_none()
+ );
+ let scan =
table.scan().as_of_time(entry.timestamp_ms).build().unwrap();
+ assert_eq!(
+ scan.plan_context.as_ref().unwrap().snapshot_schema,
+ *table.metadata().current_schema()
+ );
+
+ let mut metadata = table.metadata().clone();
+
Arc::make_mut(metadata.snapshots.get_mut(&entry.snapshot_id).unwrap()).schema_id
=
+ Some(1234);
+ let table = table.with_metadata(Arc::new(metadata));
+ let err = table
+ .scan()
+ .as_of_time(entry.timestamp_ms)
+ .build()
+ .unwrap_err();
+ assert_eq!(err.kind(), ErrorKind::DataInvalid);
+ assert!(err.message().contains("Schema with id 1234 not found"));
+ }
+
+ #[tokio::test]
+ async fn test_table_scan_as_of_time_reads_historical_rows() {
+ let mut fixture = TableTestFixture::new();
+ fixture.setup_manifest_files().await;
+ let entry = fixture.table.metadata().history()[0].clone();
+
+ // The older snapshot has only x, while the current snapshot has eight
columns.
+ let mut metadata = fixture.table.metadata().clone();
+
Arc::make_mut(metadata.snapshots.get_mut(&entry.snapshot_id).unwrap()).schema_id
= Some(0);
+ fixture.table = fixture.table.with_metadata(Arc::new(metadata));
+ let snapshot = fixture
+ .table
+ .metadata()
+ .snapshot_by_id(entry.snapshot_id)
+ .unwrap();
+ let schema = snapshot.schema(fixture.table.metadata()).unwrap();
+ let arrow_schema =
Arc::new(crate::arrow::schema_to_arrow_schema(&schema).unwrap());
+
+ // Write distinct historical data: x = 100, versus x = 1 in the
current files.
+ let path = format!("{}/historical.parquet", fixture.table_location);
+ let batch = RecordBatch::try_new(arrow_schema.clone(),
vec![Arc::new(Int64Array::from(
+ vec![100],
+ ))])
+ .unwrap();
+ let mut writer =
+ ArrowWriter::try_new(File::create(&path).unwrap(), arrow_schema,
None).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ let mut manifest_writer = ManifestWriterBuilder::new(
+ fixture.next_manifest_file(),
+ Some(snapshot.snapshot_id()),
+ schema,
+ fixture
+ .table
+ .metadata()
+ .default_partition_spec()
+ .as_ref()
+ .clone(),
+ )
+ .build_v2_data();
+ manifest_writer
+ .add_entry(
+ ManifestEntry::builder()
+ .status(ManifestStatus::Added)
+ .data_file(
+ DataFileBuilder::default()
+ .partition_spec_id(0)
+ .content(DataContentType::Data)
+ .file_format(DataFileFormat::Parquet)
+
.file_size_in_bytes(fs::metadata(&path).unwrap().len())
+ .file_path(path)
+ .record_count(1)
+
.partition(Struct::from_iter([Some(Literal::long(100))]))
+ .build()
+ .unwrap(),
+ )
+ .build(),
+ )
+ .unwrap();
+ let manifest = manifest_writer.write_manifest_file().await.unwrap();
+ let output = fixture
+ .table
+ .file_io()
+ .new_output(snapshot.manifest_list())
+ .unwrap()
+ .writer()
+ .await
+ .unwrap();
+ let mut list_writer = ManifestListWriter::v2(
+ output,
+ snapshot.snapshot_id(),
+ snapshot.parent_snapshot_id(),
+ snapshot.sequence_number(),
+ );
+ list_writer.add_manifests([manifest].into_iter()).unwrap();
+ list_writer.close().await.unwrap();
+
+ let table = fixture.table;
+ let [by_time, by_id] = [
+ table.scan().as_of_time(entry.timestamp_ms),
+ table.scan().snapshot_id(entry.snapshot_id),
+ ]
+ .map(|scan| {
+
scan.with_filter(Reference::new("x").greater_than_or_equal_to(Datum::long(100)))
+ .build()
+ .unwrap()
+ });
+ let current = table.scan().build().unwrap();
+ let mut planned_files = Vec::new();
+ let mut rows = Vec::new();
+ let mut schemas = Vec::new();
+ for scan in [by_time, by_id, current] {
+ let mut tasks: Vec<_> = scan
+ .plan_files()
+ .await
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ tasks.sort_by_key(|task| task.data_file_path().to_string());
+ planned_files.push(tasks);
+ let batches: Vec<_> =
scan.to_arrow().await.unwrap().try_collect().await.unwrap();
+ assert!(!batches.is_empty());
+ schemas.push(batches[0].schema());
+ let mut values: Vec<i64> = batches
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column_by_name("x")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values()
+ .iter()
+ .copied()
+ })
+ .collect();
+ values.sort_unstable();
+ rows.push(values);
+ }
+ assert!(!planned_files[0].is_empty());
+ assert_eq!(planned_files[0], planned_files[1]);
+ assert_eq!(rows[0], vec![100]);
+ assert_eq!(rows[0], rows[1]);
+ // Two live fixture files contain 1,024 rows each.
+ assert_eq!(rows[2], vec![1; 2048]);
Review Comment:
Makes sense, updated to use `ROWS_PER_FIXTURE_FILE`
--
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]
