1mod cache;
21use cache::*;
22mod context;
23use context::*;
24mod task;
25
26use std::sync::Arc;
27
28use arrow_array::RecordBatch;
29use futures::channel::mpsc::{Sender, channel};
30use futures::stream::BoxStream;
31use futures::{SinkExt, StreamExt, TryStreamExt};
32pub use task::*;
33
34use crate::arrow::ArrowReaderBuilder;
35pub use crate::arrow::{ScanMetrics, ScanResult};
36use crate::delete_file_index::DeleteFileIndex;
37use crate::expr::visitors::inclusive_metrics_evaluator::InclusiveMetricsEvaluator;
38use crate::expr::{Bind, BoundPredicate, Predicate};
39use crate::io::FileIO;
40use crate::metadata_columns::{
41 RESERVED_FIELD_ID_PARTITION, get_metadata_field_id, is_metadata_column_name,
42};
43use crate::partitioning::compute_unified_partition_type;
44use crate::runtime::Runtime;
45use crate::spec::{DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, NameMapping, SnapshotRef};
46use crate::table::Table;
47use crate::util::available_parallelism;
48use crate::{Error, ErrorKind, Result};
49
50pub type ArrowRecordBatchStream = BoxStream<'static, Result<RecordBatch>>;
52
53pub struct TableScanBuilder<'a> {
55 table: &'a Table,
56 column_names: Option<Vec<String>>,
58 snapshot_id: Option<i64>,
59 batch_size: Option<usize>,
60 case_sensitive: bool,
61 filter: Option<Predicate>,
62 concurrency_limit_data_files: usize,
63 concurrency_limit_manifest_entries: usize,
64 concurrency_limit_manifest_files: usize,
65 row_group_filtering_enabled: bool,
66 row_selection_enabled: bool,
67}
68
69impl<'a> TableScanBuilder<'a> {
70 pub(crate) fn new(table: &'a Table) -> Self {
71 let num_cpus = available_parallelism().get();
72
73 Self {
74 table,
75 column_names: None,
76 snapshot_id: None,
77 batch_size: None,
78 case_sensitive: true,
79 filter: None,
80 concurrency_limit_data_files: num_cpus,
81 concurrency_limit_manifest_entries: num_cpus,
82 concurrency_limit_manifest_files: num_cpus,
83 row_group_filtering_enabled: true,
84 row_selection_enabled: false,
85 }
86 }
87
88 pub fn with_batch_size(mut self, batch_size: Option<usize>) -> Self {
91 self.batch_size = batch_size;
92 self
93 }
94
95 pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self {
97 self.case_sensitive = case_sensitive;
98 self
99 }
100
101 pub fn with_filter(mut self, predicate: Predicate) -> Self {
103 self.filter = Some(predicate.rewrite_not());
106 self
107 }
108
109 pub fn select_all(mut self) -> Self {
111 self.column_names = None;
112 self
113 }
114
115 pub fn select_empty(mut self) -> Self {
117 self.column_names = Some(vec![]);
118 self
119 }
120
121 pub fn select(mut self, column_names: impl IntoIterator<Item = impl ToString>) -> Self {
123 self.column_names = Some(
124 column_names
125 .into_iter()
126 .map(|item| item.to_string())
127 .collect(),
128 );
129 self
130 }
131
132 pub fn snapshot_id(mut self, snapshot_id: i64) -> Self {
134 self.snapshot_id = Some(snapshot_id);
135 self
136 }
137
138 pub fn with_concurrency_limit(mut self, limit: usize) -> Self {
141 self.concurrency_limit_manifest_files = limit;
142 self.concurrency_limit_manifest_entries = limit;
143 self.concurrency_limit_data_files = limit;
144 self
145 }
146
147 pub fn with_data_file_concurrency_limit(mut self, limit: usize) -> Self {
149 self.concurrency_limit_data_files = limit;
150 self
151 }
152
153 pub fn with_manifest_entry_concurrency_limit(mut self, limit: usize) -> Self {
155 self.concurrency_limit_manifest_entries = limit;
156 self
157 }
158
159 pub fn with_row_group_filtering_enabled(mut self, row_group_filtering_enabled: bool) -> Self {
168 self.row_group_filtering_enabled = row_group_filtering_enabled;
169 self
170 }
171
172 pub fn with_row_selection_enabled(mut self, row_selection_enabled: bool) -> Self {
187 self.row_selection_enabled = row_selection_enabled;
188 self
189 }
190
191 pub fn build(self) -> Result<TableScan> {
193 let snapshot = match self.snapshot_id {
194 Some(snapshot_id) => self
195 .table
196 .metadata()
197 .snapshot_by_id(snapshot_id)
198 .ok_or_else(|| {
199 Error::new(
200 ErrorKind::DataInvalid,
201 format!("Snapshot with id {snapshot_id} not found"),
202 )
203 })?
204 .clone(),
205 None => {
206 let Some(current_snapshot_id) = self.table.metadata().current_snapshot() else {
207 return Ok(TableScan {
208 batch_size: self.batch_size,
209 column_names: self.column_names,
210 file_io: self.table.file_io().clone(),
211 plan_context: None,
212 concurrency_limit_data_files: self.concurrency_limit_data_files,
213 concurrency_limit_manifest_entries: self.concurrency_limit_manifest_entries,
214 concurrency_limit_manifest_files: self.concurrency_limit_manifest_files,
215 row_group_filtering_enabled: self.row_group_filtering_enabled,
216 row_selection_enabled: self.row_selection_enabled,
217 runtime: self.table.runtime().clone(),
218 });
219 };
220 current_snapshot_id.clone()
221 }
222 };
223
224 let schema = snapshot.schema(self.table.metadata())?;
225
226 if let Some(column_names) = self.column_names.as_ref() {
228 for column_name in column_names {
229 if is_metadata_column_name(column_name) {
231 continue;
232 }
233 if schema.field_by_name(column_name).is_none() {
234 return Err(Error::new(
235 ErrorKind::DataInvalid,
236 format!("Column {column_name} not found in table. Schema: {schema}"),
237 ));
238 }
239 }
240 }
241
242 let mut field_ids = vec![];
243 let column_names = self.column_names.clone().unwrap_or_else(|| {
244 schema
245 .as_struct()
246 .fields()
247 .iter()
248 .map(|f| f.name.clone())
249 .collect()
250 });
251
252 for column_name in column_names.iter() {
253 if is_metadata_column_name(column_name) {
255 field_ids.push(get_metadata_field_id(column_name)?);
256 continue;
257 }
258
259 let field_id = schema.field_id_by_name(column_name).ok_or_else(|| {
260 Error::new(
261 ErrorKind::DataInvalid,
262 format!("Column {column_name} not found in table. Schema: {schema}"),
263 )
264 })?;
265
266 schema
267 .as_struct()
268 .field_by_id(field_id)
269 .ok_or_else(|| {
270 Error::new(
271 ErrorKind::FeatureUnsupported,
272 format!(
273 "Column {column_name} is not a direct child of schema but a nested field, which is not supported now. Schema: {schema}"
274 ),
275 )
276 })?;
277
278 field_ids.push(field_id);
279 }
280
281 let snapshot_bound_predicate = if let Some(ref predicates) = self.filter {
282 Some(predicates.bind(schema.clone(), true)?)
283 } else {
284 None
285 };
286
287 let name_mapping = self
288 .table
289 .metadata()
290 .properties()
291 .get(DEFAULT_SCHEMA_NAME_MAPPING)
292 .map(|raw| {
293 serde_json::from_str::<NameMapping>(raw).map_err(|e| {
294 Error::new(
295 ErrorKind::DataInvalid,
296 format!(
297 "Failed to parse table property {DEFAULT_SCHEMA_NAME_MAPPING} as a NameMapping"
298 ),
299 )
300 .with_source(e)
301 })
302 })
303 .transpose()?
304 .map(Arc::new);
305
306 let unified_partition_type = if field_ids.contains(&RESERVED_FIELD_ID_PARTITION) {
308 let partition_type = compute_unified_partition_type(
309 self.table
310 .metadata()
311 .partition_specs_iter()
312 .map(|s| s.as_ref()),
313 &schema,
314 )?;
315 Some(Arc::new(partition_type))
316 } else {
317 None
318 };
319
320 let plan_context = PlanContext {
321 snapshot,
322 table_metadata: self.table.metadata_ref(),
323 snapshot_schema: schema,
324 case_sensitive: self.case_sensitive,
325 predicate: self.filter.map(Arc::new),
326 snapshot_bound_predicate: snapshot_bound_predicate.map(Arc::new),
327 object_cache: self.table.object_cache(),
328 field_ids: Arc::new(field_ids),
329 name_mapping,
330 partition_filter_cache: Arc::new(PartitionFilterCache::new()),
331 manifest_evaluator_cache: Arc::new(ManifestEvaluatorCache::new()),
332 expression_evaluator_cache: Arc::new(ExpressionEvaluatorCache::new()),
333 unified_partition_type,
334 };
335
336 Ok(TableScan {
337 batch_size: self.batch_size,
338 column_names: self.column_names,
339 file_io: self.table.file_io().clone(),
340 plan_context: Some(plan_context),
341 concurrency_limit_data_files: self.concurrency_limit_data_files,
342 concurrency_limit_manifest_entries: self.concurrency_limit_manifest_entries,
343 concurrency_limit_manifest_files: self.concurrency_limit_manifest_files,
344 row_group_filtering_enabled: self.row_group_filtering_enabled,
345 row_selection_enabled: self.row_selection_enabled,
346 runtime: self.table.runtime().clone(),
347 })
348 }
349}
350
351#[derive(Debug)]
353pub struct TableScan {
354 plan_context: Option<PlanContext>,
358 batch_size: Option<usize>,
359 file_io: FileIO,
360 column_names: Option<Vec<String>>,
361 concurrency_limit_manifest_files: usize,
364
365 concurrency_limit_manifest_entries: usize,
368
369 concurrency_limit_data_files: usize,
372
373 row_group_filtering_enabled: bool,
374 row_selection_enabled: bool,
375
376 runtime: Runtime,
377}
378
379impl TableScan {
380 pub async fn plan_files(&self) -> Result<FileScanTaskStream> {
382 let Some(plan_context) = self.plan_context.as_ref() else {
383 return Ok(Box::pin(futures::stream::empty()));
384 };
385
386 let concurrency_limit_manifest_files = self.concurrency_limit_manifest_files;
387 let concurrency_limit_manifest_entries = self.concurrency_limit_manifest_entries;
388
389 let (manifest_entry_data_ctx_tx, manifest_entry_data_ctx_rx) =
391 channel(concurrency_limit_manifest_files);
392 let (manifest_entry_delete_ctx_tx, manifest_entry_delete_ctx_rx) =
393 channel(concurrency_limit_manifest_files);
394
395 let (file_scan_task_tx, file_scan_task_rx) = channel(concurrency_limit_manifest_entries);
397
398 let (delete_file_idx, delete_file_tx) = DeleteFileIndex::new(self.runtime.clone());
399
400 let manifest_list = plan_context.get_manifest_list().await?;
401
402 let manifest_file_contexts = plan_context.build_manifest_file_contexts(
406 manifest_list,
407 manifest_entry_data_ctx_tx,
408 delete_file_idx.clone(),
409 manifest_entry_delete_ctx_tx,
410 )?;
411
412 let mut channel_for_manifest_error = file_scan_task_tx.clone();
413 let mut channel_for_data_manifest_entry_error = file_scan_task_tx.clone();
414 let mut channel_for_delete_manifest_entry_error = file_scan_task_tx.clone();
415
416 let rt = self.runtime.clone();
417
418 rt.io().spawn(async move {
420 let result = futures::stream::iter(manifest_file_contexts)
421 .try_for_each_concurrent(concurrency_limit_manifest_files, |ctx| async move {
422 ctx.fetch_manifest_and_stream_manifest_entries().await
423 })
424 .await;
425
426 if let Err(error) = result {
427 let _ = channel_for_manifest_error.send(Err(error)).await;
428 }
429 });
430
431 {
433 let rt = rt.clone();
434 let rt_inner = rt.clone();
435 rt.cpu().spawn(async move {
436 let result = manifest_entry_delete_ctx_rx
437 .map(|me_ctx| Ok((me_ctx, delete_file_tx.clone())))
438 .try_for_each_concurrent(
439 concurrency_limit_manifest_entries,
440 |(manifest_entry_context, tx)| {
441 let rt_inner = rt_inner.clone();
442 async move {
443 rt_inner
444 .cpu()
445 .spawn(async move {
446 Self::process_delete_manifest_entry(
447 manifest_entry_context,
448 tx,
449 )
450 .await
451 })
452 .await?
453 }
454 },
455 )
456 .await;
457
458 if let Err(error) = result {
459 let _ = channel_for_delete_manifest_entry_error
460 .send(Err(error))
461 .await;
462 }
463 });
464 }
465
466 {
468 let rt_inner = rt.clone();
469 rt.cpu().spawn(async move {
470 let result = manifest_entry_data_ctx_rx
471 .map(|me_ctx| Ok((me_ctx, file_scan_task_tx.clone())))
472 .try_for_each_concurrent(
473 concurrency_limit_manifest_entries,
474 |(manifest_entry_context, tx)| {
475 let rt_inner = rt_inner.clone();
476 async move {
477 rt_inner
478 .cpu()
479 .spawn(async move {
480 Self::process_data_manifest_entry(
481 manifest_entry_context,
482 tx,
483 )
484 .await
485 })
486 .await?
487 }
488 },
489 )
490 .await;
491
492 if let Err(error) = result {
493 let _ = channel_for_data_manifest_entry_error.send(Err(error)).await;
494 }
495 });
496 }
497
498 Ok(file_scan_task_rx.boxed())
499 }
500
501 pub async fn to_arrow(&self) -> Result<ArrowRecordBatchStream> {
503 let mut arrow_reader_builder =
504 ArrowReaderBuilder::new(self.file_io.clone(), self.runtime.clone())
505 .with_data_file_concurrency_limit(self.concurrency_limit_data_files)
506 .with_row_group_filtering_enabled(self.row_group_filtering_enabled)
507 .with_row_selection_enabled(self.row_selection_enabled);
508
509 if let Some(batch_size) = self.batch_size {
510 arrow_reader_builder = arrow_reader_builder.with_batch_size(batch_size);
511 }
512
513 arrow_reader_builder
514 .build()
515 .read(self.plan_files().await?)
516 .map(|result| result.stream())
517 }
518
519 pub fn column_names(&self) -> Option<&[String]> {
521 self.column_names.as_deref()
522 }
523
524 pub fn snapshot(&self) -> Option<&SnapshotRef> {
526 self.plan_context.as_ref().map(|x| &x.snapshot)
527 }
528
529 async fn process_data_manifest_entry(
530 manifest_entry_context: ManifestEntryContext,
531 mut file_scan_task_tx: Sender<Result<FileScanTask>>,
532 ) -> Result<()> {
533 if !manifest_entry_context.manifest_entry.is_alive() {
535 return Ok(());
536 }
537
538 if manifest_entry_context.manifest_entry.content_type() != DataContentType::Data {
540 return Err(Error::new(
541 ErrorKind::FeatureUnsupported,
542 "Encountered an entry for a delete file in a data file manifest",
543 ));
544 }
545
546 if let Some(ref bound_predicates) = manifest_entry_context.bound_predicates {
547 let BoundPredicates {
548 snapshot_bound_predicate,
549 partition_bound_predicate,
550 } = bound_predicates.as_ref();
551
552 let expression_evaluator_cache =
553 manifest_entry_context.expression_evaluator_cache.as_ref();
554
555 let expression_evaluator = expression_evaluator_cache.get(
556 manifest_entry_context.partition_spec_id,
557 partition_bound_predicate,
558 )?;
559
560 if !expression_evaluator.eval(manifest_entry_context.manifest_entry.data_file())? {
563 return Ok(());
564 }
565
566 if !InclusiveMetricsEvaluator::eval(
568 snapshot_bound_predicate,
569 manifest_entry_context.manifest_entry.data_file(),
570 false,
571 )? {
572 return Ok(());
573 }
574 }
575
576 file_scan_task_tx
580 .send(Ok(manifest_entry_context.into_file_scan_task().await?))
581 .await?;
582
583 Ok(())
584 }
585
586 async fn process_delete_manifest_entry(
587 manifest_entry_context: ManifestEntryContext,
588 mut delete_file_ctx_tx: Sender<DeleteFileContext>,
589 ) -> Result<()> {
590 if !manifest_entry_context.manifest_entry.is_alive() {
592 return Ok(());
593 }
594
595 if manifest_entry_context.manifest_entry.content_type() == DataContentType::Data {
597 return Err(Error::new(
598 ErrorKind::FeatureUnsupported,
599 "Encountered an entry for a data file in a delete manifest",
600 ));
601 }
602
603 if let Some(ref bound_predicates) = manifest_entry_context.bound_predicates {
604 let expression_evaluator_cache =
605 manifest_entry_context.expression_evaluator_cache.as_ref();
606
607 let expression_evaluator = expression_evaluator_cache.get(
608 manifest_entry_context.partition_spec_id,
609 &bound_predicates.partition_bound_predicate,
610 )?;
611
612 if !expression_evaluator.eval(manifest_entry_context.manifest_entry.data_file())? {
615 return Ok(());
616 }
617 }
618
619 delete_file_ctx_tx
620 .send(DeleteFileContext {
621 manifest_entry: manifest_entry_context.manifest_entry.clone(),
622 partition_spec_id: manifest_entry_context.partition_spec_id,
623 })
624 .await?;
625
626 Ok(())
627 }
628}
629
630pub(crate) struct BoundPredicates {
631 partition_bound_predicate: BoundPredicate,
632 snapshot_bound_predicate: BoundPredicate,
633}
634
635#[cfg(test)]
636pub mod tests {
637 #![allow(missing_docs)]
639
640 use std::collections::HashMap;
641 use std::fs;
642 use std::fs::File;
643 use std::sync::Arc;
644
645 use arrow_array::cast::AsArray;
646 use arrow_array::types::Int32Type;
647 use arrow_array::{
648 Array, ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, RunArray,
649 StringArray,
650 };
651 use futures::{TryStreamExt, stream};
652 use minijinja::value::Value;
653 use minijinja::{AutoEscape, Environment, context};
654 use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
655 use parquet::basic::Compression;
656 use parquet::file::properties::WriterProperties;
657 use tempfile::TempDir;
658 use uuid::Uuid;
659
660 use crate::arrow::ArrowReaderBuilder;
661 use crate::expr::{BoundPredicate, Reference};
662 use crate::io::{FileIO, OutputFile};
663 use crate::metadata_columns::{
664 RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS,
665 RESERVED_COL_NAME_FILE, RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
666 RESERVED_COL_NAME_POS, RESERVED_COL_NAME_SPEC_ID, RESERVED_FIELD_ID_DELETE_FILE_PATH,
667 RESERVED_FIELD_ID_DELETE_FILE_POS, RESERVED_FIELD_ID_POS,
668 };
669 use crate::scan::FileScanTask;
670 use crate::spec::{
671 DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, DataFileBuilder, DataFileFormat, Datum,
672 FormatVersion, Literal, MAIN_BRANCH, ManifestEntry, ManifestListWriter, ManifestStatus,
673 ManifestWriterBuilder, NestedField, Operation, PartitionSpec, PrimitiveType, Schema,
674 Snapshot, Struct, StructType, Summary, TableMetadata, TableMetadataBuilder, Type,
675 UnboundPartitionSpec,
676 };
677 use crate::table::Table;
678 use crate::test_utils::test_runtime;
679 use crate::{ErrorKind, TableIdent};
680
681 fn render_template(template: &str, ctx: Value) -> String {
682 let mut env = Environment::new();
683 env.set_auto_escape_callback(|_| AutoEscape::None);
684 env.render_str(template, ctx).unwrap()
685 }
686
687 fn assert_last_updated_seq_all(batches: &[RecordBatch], expected: Option<i64>) {
691 use arrow_cast::cast;
692 use arrow_schema::DataType;
693 for batch in batches {
694 let col = batch
695 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
696 .expect("_last_updated_sequence_number column should be present");
697 let logical = cast(col, &DataType::Int64).unwrap();
698 let values = logical.as_primitive::<arrow_array::types::Int64Type>();
699 for i in 0..values.len() {
700 let actual = (!values.is_null(i)).then(|| values.value(i));
701 assert_eq!(actual, expected, "row {i}");
702 }
703 }
704 }
705
706 pub struct TableTestFixture {
707 pub table_location: String,
708 pub table: Table,
709 }
710
711 impl TableTestFixture {
712 #[allow(clippy::new_without_default)]
713 pub fn new() -> Self {
714 let tmp_dir = TempDir::new().unwrap();
715 let table_location = tmp_dir.path().join("table1");
716 let manifest_list1_location = table_location.join("metadata/manifests_list_1.avro");
717 let manifest_list2_location = table_location.join("metadata/manifests_list_2.avro");
718 let table_metadata1_location = table_location.join("metadata/v1.json");
719
720 let file_io = FileIO::new_with_fs();
721
722 let table_metadata = {
723 let template_json_str = fs::read_to_string(format!(
724 "{}/testdata/example_table_metadata_v2.json",
725 env!("CARGO_MANIFEST_DIR")
726 ))
727 .unwrap();
728 let metadata_json = render_template(&template_json_str, context! {
729 table_location => &table_location,
730 manifest_list_1_location => &manifest_list1_location,
731 manifest_list_2_location => &manifest_list2_location,
732 table_metadata_1_location => &table_metadata1_location,
733 });
734 serde_json::from_str::<TableMetadata>(&metadata_json).unwrap()
735 };
736
737 let table = Table::builder()
738 .metadata(table_metadata)
739 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
740 .file_io(file_io.clone())
741 .metadata_location(table_metadata1_location.as_os_str().to_str().unwrap())
742 .runtime(test_runtime())
743 .build()
744 .unwrap();
745
746 Self {
747 table_location: table_location.to_str().unwrap().to_string(),
748 table,
749 }
750 }
751
752 #[allow(clippy::new_without_default)]
753 pub fn new_empty() -> Self {
754 let tmp_dir = TempDir::new().unwrap();
755 let table_location = tmp_dir.path().join("table1");
756 let table_metadata1_location = table_location.join("metadata/v1.json");
757
758 let file_io = FileIO::new_with_fs();
759
760 let table_metadata = {
761 let template_json_str = fs::read_to_string(format!(
762 "{}/testdata/example_empty_table_metadata_v2.json",
763 env!("CARGO_MANIFEST_DIR")
764 ))
765 .unwrap();
766 let metadata_json = render_template(&template_json_str, context! {
767 table_location => &table_location,
768 table_metadata_1_location => &table_metadata1_location,
769 });
770 serde_json::from_str::<TableMetadata>(&metadata_json).unwrap()
771 };
772
773 let table = Table::builder()
774 .metadata(table_metadata)
775 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
776 .file_io(file_io.clone())
777 .metadata_location(table_metadata1_location.as_os_str().to_str().unwrap())
778 .runtime(test_runtime())
779 .build()
780 .unwrap();
781
782 Self {
783 table_location: table_location.to_str().unwrap().to_string(),
784 table,
785 }
786 }
787
788 pub fn new_with_deep_history() -> Self {
792 let tmp_dir = TempDir::new().unwrap();
793 let table_location = tmp_dir.path().join("table1");
794 let table_metadata1_location = table_location.join("metadata/v1.json");
795
796 let file_io = FileIO::new_with_fs();
797
798 let table_metadata = {
799 let json_str = fs::read_to_string(format!(
800 "{}/testdata/example_table_metadata_v2_deep_history.json",
801 env!("CARGO_MANIFEST_DIR")
802 ))
803 .unwrap();
804 serde_json::from_str::<TableMetadata>(&json_str).unwrap()
805 };
806
807 let table = Table::builder()
808 .metadata(table_metadata)
809 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
810 .file_io(file_io.clone())
811 .metadata_location(table_metadata1_location.as_os_str().to_str().unwrap())
812 .runtime(test_runtime())
813 .build()
814 .unwrap();
815
816 Self {
817 table_location: table_location.to_str().unwrap().to_string(),
818 table,
819 }
820 }
821
822 pub fn new_unpartitioned() -> Self {
823 let tmp_dir = TempDir::new().unwrap();
824 let table_location = tmp_dir.path().join("table1");
825 let manifest_list1_location = table_location.join("metadata/manifests_list_1.avro");
826 let manifest_list2_location = table_location.join("metadata/manifests_list_2.avro");
827 let table_metadata1_location = table_location.join("metadata/v1.json");
828
829 let file_io = FileIO::new_with_fs();
830
831 let mut table_metadata = {
832 let template_json_str = fs::read_to_string(format!(
833 "{}/testdata/example_table_metadata_v2.json",
834 env!("CARGO_MANIFEST_DIR")
835 ))
836 .unwrap();
837 let metadata_json = render_template(&template_json_str, context! {
838 table_location => &table_location,
839 manifest_list_1_location => &manifest_list1_location,
840 manifest_list_2_location => &manifest_list2_location,
841 table_metadata_1_location => &table_metadata1_location,
842 });
843 serde_json::from_str::<TableMetadata>(&metadata_json).unwrap()
844 };
845
846 table_metadata.default_spec = Arc::new(PartitionSpec::unpartition_spec());
847 table_metadata.partition_specs.clear();
848 table_metadata.default_partition_type = StructType::new(vec![]);
849 table_metadata
850 .partition_specs
851 .insert(0, table_metadata.default_spec.clone());
852
853 let table = Table::builder()
854 .metadata(table_metadata)
855 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
856 .file_io(file_io.clone())
857 .metadata_location(table_metadata1_location.to_str().unwrap())
858 .runtime(test_runtime())
859 .build()
860 .unwrap();
861
862 Self {
863 table_location: table_location.to_str().unwrap().to_string(),
864 table,
865 }
866 }
867
868 pub fn new_with_partition_evolution() -> Self {
869 let table = Self::new().table;
870 let table_location = table.metadata().location.clone();
871
872 let manifest_list1_location =
873 format!("{}/metadata/manifests_list_1.avro", table_location);
874 let manifest_list2_location =
875 format!("{}/metadata/manifests_list_2.avro", table_location);
876 let manifest_list3_location =
877 format!("{}/metadata/manifests_list_3.avro", table_location);
878 let table_metadata1_location = format!("{}/metadata/v1.json", table_location);
879
880 let new_table_metadata = {
881 let template_json_str = fs::read_to_string(format!(
882 "{}/testdata/example_table_metadata_v2_partition_evolution.json",
883 env!("CARGO_MANIFEST_DIR")
884 ))
885 .unwrap();
886 let metadata_json = render_template(&template_json_str, context! {
887 table_location => &table_location,
888 manifest_list_1_location => &manifest_list1_location,
889 manifest_list_2_location => &manifest_list2_location,
890 manifest_list_3_location => &manifest_list3_location,
891 table_metadata_1_location => &table_metadata1_location,
892 });
893 Arc::new(serde_json::from_str::<TableMetadata>(&metadata_json).unwrap())
894 };
895
896 Self {
897 table_location,
898 table: table.with_metadata(new_table_metadata),
899 }
900 }
901
902 fn next_manifest_file(&self) -> OutputFile {
903 self.table
904 .file_io()
905 .new_output(format!(
906 "{}/metadata/manifest_{}.avro",
907 self.table_location,
908 Uuid::new_v4()
909 ))
910 .unwrap()
911 }
912
913 pub async fn setup_manifest_files(&mut self) {
914 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
915 let parent_snapshot = current_snapshot
916 .parent_snapshot(self.table.metadata())
917 .unwrap();
918 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
919 let current_partition_spec = self.table.metadata().default_partition_spec();
920
921 let parquet_file_size = self.write_parquet_data_files();
923
924 let mut writer = ManifestWriterBuilder::new(
925 self.next_manifest_file(),
926 Some(current_snapshot.snapshot_id()),
927 current_schema.clone(),
928 current_partition_spec.as_ref().clone(),
929 )
930 .build_v2_data();
931 writer
932 .add_entry(
933 ManifestEntry::builder()
934 .status(ManifestStatus::Added)
935 .data_file(
936 DataFileBuilder::default()
937 .partition_spec_id(0)
938 .content(DataContentType::Data)
939 .file_path(format!("{}/1.parquet", &self.table_location))
940 .file_format(DataFileFormat::Parquet)
941 .file_size_in_bytes(parquet_file_size)
942 .record_count(1)
943 .partition(Struct::from_iter([Some(Literal::long(100))]))
944 .key_metadata(None)
945 .build()
946 .unwrap(),
947 )
948 .build(),
949 )
950 .unwrap();
951 writer
952 .add_delete_entry(
953 ManifestEntry::builder()
954 .status(ManifestStatus::Deleted)
955 .snapshot_id(parent_snapshot.snapshot_id())
956 .sequence_number(parent_snapshot.sequence_number())
957 .file_sequence_number(parent_snapshot.sequence_number())
958 .data_file(
959 DataFileBuilder::default()
960 .partition_spec_id(0)
961 .content(DataContentType::Data)
962 .file_path(format!("{}/2.parquet", &self.table_location))
963 .file_format(DataFileFormat::Parquet)
964 .file_size_in_bytes(parquet_file_size)
965 .record_count(1)
966 .partition(Struct::from_iter([Some(Literal::long(200))]))
967 .build()
968 .unwrap(),
969 )
970 .build(),
971 )
972 .unwrap();
973 writer
974 .add_existing_entry(
975 ManifestEntry::builder()
976 .status(ManifestStatus::Existing)
977 .snapshot_id(parent_snapshot.snapshot_id())
978 .sequence_number(parent_snapshot.sequence_number())
979 .file_sequence_number(parent_snapshot.sequence_number())
980 .data_file(
981 DataFileBuilder::default()
982 .partition_spec_id(0)
983 .content(DataContentType::Data)
984 .file_path(format!("{}/3.parquet", &self.table_location))
985 .file_format(DataFileFormat::Parquet)
986 .file_size_in_bytes(parquet_file_size)
987 .record_count(1)
988 .partition(Struct::from_iter([Some(Literal::long(300))]))
989 .build()
990 .unwrap(),
991 )
992 .build(),
993 )
994 .unwrap();
995 let data_file_manifest = writer.write_manifest_file().await.unwrap();
996
997 let manifest_list_writer = self
999 .table
1000 .file_io()
1001 .new_output(current_snapshot.manifest_list())
1002 .unwrap()
1003 .writer()
1004 .await
1005 .unwrap();
1006 let mut manifest_list_write = ManifestListWriter::v2(
1007 manifest_list_writer,
1008 current_snapshot.snapshot_id(),
1009 current_snapshot.parent_snapshot_id(),
1010 current_snapshot.sequence_number(),
1011 );
1012 manifest_list_write
1013 .add_manifests(vec![data_file_manifest].into_iter())
1014 .unwrap();
1015 manifest_list_write.close().await.unwrap();
1016 }
1017
1018 pub async fn setup_v3_manifest_files(&mut self) {
1022 let metadata = TableMetadataBuilder::new_from_metadata(
1023 self.table.metadata().clone(),
1024 self.table.metadata_location().map(str::to_string),
1025 )
1026 .upgrade_format_version(FormatVersion::V3)
1027 .unwrap()
1028 .build()
1029 .unwrap()
1030 .metadata;
1031 self.table = Table::builder()
1032 .metadata(metadata)
1033 .identifier(self.table.identifier().clone())
1034 .file_io(self.table.file_io().clone())
1035 .metadata_location(self.table.metadata_location().unwrap().to_string())
1036 .runtime(test_runtime())
1037 .build()
1038 .unwrap();
1039
1040 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
1041 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
1042 let current_partition_spec = self.table.metadata().default_partition_spec();
1043
1044 let parquet_file_size = self.write_parquet_data_files();
1045
1046 let mut writer = ManifestWriterBuilder::new(
1047 self.next_manifest_file(),
1048 Some(current_snapshot.snapshot_id()),
1049 current_schema.clone(),
1050 current_partition_spec.as_ref().clone(),
1051 )
1052 .build_v3_data();
1053 writer
1054 .add_entry(
1055 ManifestEntry::builder()
1056 .status(ManifestStatus::Added)
1057 .data_file(
1058 DataFileBuilder::default()
1059 .partition_spec_id(0)
1060 .content(DataContentType::Data)
1061 .file_path(format!("{}/1.parquet", &self.table_location))
1062 .file_format(DataFileFormat::Parquet)
1063 .file_size_in_bytes(parquet_file_size)
1064 .record_count(1)
1065 .partition(Struct::from_iter([Some(Literal::long(100))]))
1066 .key_metadata(None)
1067 .build()
1068 .unwrap(),
1069 )
1070 .build(),
1071 )
1072 .unwrap();
1073 let data_file_manifest = writer.write_manifest_file().await.unwrap();
1074
1075 let manifest_list_writer = self
1076 .table
1077 .file_io()
1078 .new_output(current_snapshot.manifest_list())
1079 .unwrap()
1080 .writer()
1081 .await
1082 .unwrap();
1083 let mut manifest_list_write = ManifestListWriter::v3(
1084 manifest_list_writer,
1085 current_snapshot.snapshot_id(),
1086 current_snapshot.parent_snapshot_id(),
1087 current_snapshot.sequence_number(),
1088 Some(42),
1089 );
1090 manifest_list_write
1091 .add_manifests(vec![data_file_manifest].into_iter())
1092 .unwrap();
1093 manifest_list_write.close().await.unwrap();
1094 }
1095
1096 pub async fn setup_manifest_files_with_partition_evolution(&mut self) {
1097 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
1098 let parent_snapshot = current_snapshot
1099 .parent_snapshot(self.table.metadata())
1100 .unwrap();
1101 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
1102 let current_partition_spec = self.table.metadata().default_partition_spec();
1103
1104 let parquet_file_size = self.write_parquet_data_files();
1106
1107 let mut writer = ManifestWriterBuilder::new(
1108 self.next_manifest_file(),
1109 Some(current_snapshot.snapshot_id()),
1110 current_schema.clone(),
1111 current_partition_spec.as_ref().clone(),
1112 )
1113 .build_v2_data();
1114 writer
1115 .add_entry(
1116 ManifestEntry::builder()
1117 .status(ManifestStatus::Added)
1118 .data_file(
1119 DataFileBuilder::default()
1120 .partition_spec_id(1)
1121 .content(DataContentType::Data)
1122 .file_path(format!("{}/1.parquet", &self.table_location))
1123 .file_format(DataFileFormat::Parquet)
1124 .file_size_in_bytes(parquet_file_size)
1125 .record_count(1)
1126 .partition(Struct::from_iter([
1127 Some(Literal::long(100)),
1128 Some(Literal::string("apa")),
1129 Some(Literal::int(27)),
1130 ]))
1131 .key_metadata(None)
1132 .build()
1133 .unwrap(),
1134 )
1135 .build(),
1136 )
1137 .unwrap();
1138 writer
1139 .add_delete_entry(
1140 ManifestEntry::builder()
1141 .status(ManifestStatus::Deleted)
1142 .snapshot_id(parent_snapshot.snapshot_id())
1143 .sequence_number(parent_snapshot.sequence_number())
1144 .file_sequence_number(parent_snapshot.sequence_number())
1145 .data_file(
1146 DataFileBuilder::default()
1147 .partition_spec_id(1)
1148 .content(DataContentType::Data)
1149 .file_path(format!("{}/2.parquet", &self.table_location))
1150 .file_format(DataFileFormat::Parquet)
1151 .file_size_in_bytes(parquet_file_size)
1152 .record_count(1)
1153 .partition(Struct::from_iter([
1154 Some(Literal::long(200)),
1155 Some(Literal::string("ice")),
1156 Some(Literal::int(5)),
1157 ]))
1158 .build()
1159 .unwrap(),
1160 )
1161 .build(),
1162 )
1163 .unwrap();
1164 writer
1165 .add_existing_entry(
1166 ManifestEntry::builder()
1167 .status(ManifestStatus::Existing)
1168 .snapshot_id(parent_snapshot.snapshot_id())
1169 .sequence_number(parent_snapshot.sequence_number())
1170 .file_sequence_number(parent_snapshot.sequence_number())
1171 .data_file(
1172 DataFileBuilder::default()
1173 .partition_spec_id(1)
1174 .content(DataContentType::Data)
1175 .file_path(format!("{}/3.parquet", &self.table_location))
1176 .file_format(DataFileFormat::Parquet)
1177 .file_size_in_bytes(parquet_file_size)
1178 .record_count(1)
1179 .partition(Struct::from_iter([
1180 Some(Literal::long(300)),
1181 Some(Literal::string("apa")),
1182 Some(Literal::int(19)),
1183 ]))
1184 .build()
1185 .unwrap(),
1186 )
1187 .build(),
1188 )
1189 .unwrap();
1190 let data_file_manifest = writer.write_manifest_file().await.unwrap();
1191
1192 let manifest_list_writer = self
1194 .table
1195 .file_io()
1196 .new_output(current_snapshot.manifest_list())
1197 .unwrap()
1198 .writer()
1199 .await
1200 .unwrap();
1201 let mut manifest_list_write = ManifestListWriter::v2(
1202 manifest_list_writer,
1203 current_snapshot.snapshot_id(),
1204 current_snapshot.parent_snapshot_id(),
1205 current_snapshot.sequence_number(),
1206 );
1207 manifest_list_write
1208 .add_manifests(vec![data_file_manifest].into_iter())
1209 .unwrap();
1210 manifest_list_write.close().await.unwrap();
1211 }
1212
1213 fn write_parquet_data_files(&self) -> u64 {
1216 fs::create_dir_all(&self.table_location).unwrap();
1217
1218 let schema = {
1219 let fields = vec![
1220 arrow_schema::Field::new("x", arrow_schema::DataType::Int64, false)
1221 .with_metadata(HashMap::from([(
1222 PARQUET_FIELD_ID_META_KEY.to_string(),
1223 "1".to_string(),
1224 )])),
1225 arrow_schema::Field::new("y", arrow_schema::DataType::Int64, false)
1226 .with_metadata(HashMap::from([(
1227 PARQUET_FIELD_ID_META_KEY.to_string(),
1228 "2".to_string(),
1229 )])),
1230 arrow_schema::Field::new("z", arrow_schema::DataType::Int64, false)
1231 .with_metadata(HashMap::from([(
1232 PARQUET_FIELD_ID_META_KEY.to_string(),
1233 "3".to_string(),
1234 )])),
1235 arrow_schema::Field::new("a", arrow_schema::DataType::Utf8, false)
1236 .with_metadata(HashMap::from([(
1237 PARQUET_FIELD_ID_META_KEY.to_string(),
1238 "4".to_string(),
1239 )])),
1240 arrow_schema::Field::new("dbl", arrow_schema::DataType::Float64, false)
1241 .with_metadata(HashMap::from([(
1242 PARQUET_FIELD_ID_META_KEY.to_string(),
1243 "5".to_string(),
1244 )])),
1245 arrow_schema::Field::new("i32", arrow_schema::DataType::Int32, false)
1246 .with_metadata(HashMap::from([(
1247 PARQUET_FIELD_ID_META_KEY.to_string(),
1248 "6".to_string(),
1249 )])),
1250 arrow_schema::Field::new("i64", arrow_schema::DataType::Int64, false)
1251 .with_metadata(HashMap::from([(
1252 PARQUET_FIELD_ID_META_KEY.to_string(),
1253 "7".to_string(),
1254 )])),
1255 arrow_schema::Field::new("bool", arrow_schema::DataType::Boolean, false)
1256 .with_metadata(HashMap::from([(
1257 PARQUET_FIELD_ID_META_KEY.to_string(),
1258 "8".to_string(),
1259 )])),
1260 ];
1261 Arc::new(arrow_schema::Schema::new(fields))
1262 };
1263 let col1 = Arc::new(Int64Array::from_iter_values(vec![1; 1024])) as ArrayRef;
1265
1266 let mut values = vec![2; 512];
1267 values.append(vec![3; 200].as_mut());
1268 values.append(vec![4; 300].as_mut());
1269 values.append(vec![5; 12].as_mut());
1270
1271 let col2 = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
1273
1274 let mut values = vec![3; 512];
1275 values.append(vec![4; 512].as_mut());
1276
1277 let col3 = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
1279
1280 let mut values = vec!["Apache"; 512];
1282 values.append(vec!["Iceberg"; 512].as_mut());
1283 let col4 = Arc::new(StringArray::from_iter_values(values)) as ArrayRef;
1284
1285 let mut values = vec![100.0f64; 512];
1287 values.append(vec![150.0f64; 12].as_mut());
1288 values.append(vec![200.0f64; 500].as_mut());
1289 let col5 = Arc::new(Float64Array::from_iter_values(values)) as ArrayRef;
1290
1291 let mut values = vec![100i32; 512];
1293 values.append(vec![150i32; 12].as_mut());
1294 values.append(vec![200i32; 500].as_mut());
1295 let col6 = Arc::new(Int32Array::from_iter_values(values)) as ArrayRef;
1296
1297 let mut values = vec![100i64; 512];
1299 values.append(vec![150i64; 12].as_mut());
1300 values.append(vec![200i64; 500].as_mut());
1301 let col7 = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
1302
1303 let mut values = vec![false; 512];
1305 values.append(vec![true; 512].as_mut());
1306 let values: BooleanArray = values.into();
1307 let col8 = Arc::new(values) as ArrayRef;
1308
1309 let to_write = RecordBatch::try_new(schema.clone(), vec![
1310 col1, col2, col3, col4, col5, col6, col7, col8,
1311 ])
1312 .unwrap();
1313
1314 let props = WriterProperties::builder()
1316 .set_compression(Compression::SNAPPY)
1317 .build();
1318
1319 for n in 1..=3 {
1320 let file = File::create(format!("{}/{}.parquet", &self.table_location, n)).unwrap();
1321 let mut writer =
1322 ArrowWriter::try_new(file, to_write.schema(), Some(props.clone())).unwrap();
1323
1324 writer.write(&to_write).expect("Writing batch");
1325
1326 writer.close().unwrap();
1328 }
1329
1330 fs::metadata(format!("{}/1.parquet", &self.table_location))
1331 .unwrap()
1332 .len()
1333 }
1334
1335 pub async fn setup_unpartitioned_manifest_files(&mut self) {
1336 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
1337 let parent_snapshot = current_snapshot
1338 .parent_snapshot(self.table.metadata())
1339 .unwrap();
1340 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
1341 let current_partition_spec = Arc::new(PartitionSpec::unpartition_spec());
1342
1343 let parquet_file_size = self.write_parquet_data_files();
1345
1346 let mut writer = ManifestWriterBuilder::new(
1348 self.next_manifest_file(),
1349 Some(current_snapshot.snapshot_id()),
1350 current_schema.clone(),
1351 current_partition_spec.as_ref().clone(),
1352 )
1353 .build_v2_data();
1354
1355 let empty_partition = Struct::empty();
1357
1358 writer
1359 .add_entry(
1360 ManifestEntry::builder()
1361 .status(ManifestStatus::Added)
1362 .data_file(
1363 DataFileBuilder::default()
1364 .partition_spec_id(0)
1365 .content(DataContentType::Data)
1366 .file_path(format!("{}/1.parquet", &self.table_location))
1367 .file_format(DataFileFormat::Parquet)
1368 .file_size_in_bytes(parquet_file_size)
1369 .record_count(1)
1370 .partition(empty_partition.clone())
1371 .key_metadata(None)
1372 .build()
1373 .unwrap(),
1374 )
1375 .build(),
1376 )
1377 .unwrap();
1378
1379 writer
1380 .add_delete_entry(
1381 ManifestEntry::builder()
1382 .status(ManifestStatus::Deleted)
1383 .snapshot_id(parent_snapshot.snapshot_id())
1384 .sequence_number(parent_snapshot.sequence_number())
1385 .file_sequence_number(parent_snapshot.sequence_number())
1386 .data_file(
1387 DataFileBuilder::default()
1388 .partition_spec_id(0)
1389 .content(DataContentType::Data)
1390 .file_path(format!("{}/2.parquet", &self.table_location))
1391 .file_format(DataFileFormat::Parquet)
1392 .file_size_in_bytes(parquet_file_size)
1393 .record_count(1)
1394 .partition(empty_partition.clone())
1395 .build()
1396 .unwrap(),
1397 )
1398 .build(),
1399 )
1400 .unwrap();
1401
1402 writer
1403 .add_existing_entry(
1404 ManifestEntry::builder()
1405 .status(ManifestStatus::Existing)
1406 .snapshot_id(parent_snapshot.snapshot_id())
1407 .sequence_number(parent_snapshot.sequence_number())
1408 .file_sequence_number(parent_snapshot.sequence_number())
1409 .data_file(
1410 DataFileBuilder::default()
1411 .partition_spec_id(0)
1412 .content(DataContentType::Data)
1413 .file_path(format!("{}/3.parquet", &self.table_location))
1414 .file_format(DataFileFormat::Parquet)
1415 .file_size_in_bytes(parquet_file_size)
1416 .record_count(1)
1417 .partition(empty_partition.clone())
1418 .build()
1419 .unwrap(),
1420 )
1421 .build(),
1422 )
1423 .unwrap();
1424
1425 let data_file_manifest = writer.write_manifest_file().await.unwrap();
1426
1427 let manifest_list_writer = self
1429 .table
1430 .file_io()
1431 .new_output(current_snapshot.manifest_list())
1432 .unwrap()
1433 .writer()
1434 .await
1435 .unwrap();
1436 let mut manifest_list_write = ManifestListWriter::v2(
1437 manifest_list_writer,
1438 current_snapshot.snapshot_id(),
1439 current_snapshot.parent_snapshot_id(),
1440 current_snapshot.sequence_number(),
1441 );
1442 manifest_list_write
1443 .add_manifests(vec![data_file_manifest].into_iter())
1444 .unwrap();
1445 manifest_list_write.close().await.unwrap();
1446 }
1447
1448 pub async fn setup_deadlock_manifests(&mut self) {
1449 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
1450 let _parent_snapshot = current_snapshot
1451 .parent_snapshot(self.table.metadata())
1452 .unwrap();
1453 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
1454 let current_partition_spec = self.table.metadata().default_partition_spec();
1455
1456 let mut writer = ManifestWriterBuilder::new(
1458 self.next_manifest_file(),
1459 Some(current_snapshot.snapshot_id()),
1460 current_schema.clone(),
1461 current_partition_spec.as_ref().clone(),
1462 )
1463 .build_v2_data();
1464
1465 for i in 0..10 {
1467 writer
1468 .add_entry(
1469 ManifestEntry::builder()
1470 .status(ManifestStatus::Added)
1471 .data_file(
1472 DataFileBuilder::default()
1473 .partition_spec_id(0)
1474 .content(DataContentType::Data)
1475 .file_path(format!("{}/{}.parquet", &self.table_location, i))
1476 .file_format(DataFileFormat::Parquet)
1477 .file_size_in_bytes(100)
1478 .record_count(1)
1479 .partition(Struct::from_iter([Some(Literal::long(100))]))
1480 .key_metadata(None)
1481 .build()
1482 .unwrap(),
1483 )
1484 .build(),
1485 )
1486 .unwrap();
1487 }
1488 let data_manifest = writer.write_manifest_file().await.unwrap();
1489
1490 let mut writer = ManifestWriterBuilder::new(
1492 self.next_manifest_file(),
1493 Some(current_snapshot.snapshot_id()),
1494 current_schema.clone(),
1495 current_partition_spec.as_ref().clone(),
1496 )
1497 .build_v2_deletes();
1498
1499 writer
1500 .add_entry(
1501 ManifestEntry::builder()
1502 .status(ManifestStatus::Added)
1503 .data_file(
1504 DataFileBuilder::default()
1505 .partition_spec_id(0)
1506 .content(DataContentType::PositionDeletes)
1507 .file_path(format!("{}/del.parquet", &self.table_location))
1508 .file_format(DataFileFormat::Parquet)
1509 .file_size_in_bytes(100)
1510 .record_count(1)
1511 .partition(Struct::from_iter([Some(Literal::long(100))]))
1512 .build()
1513 .unwrap(),
1514 )
1515 .build(),
1516 )
1517 .unwrap();
1518 let delete_manifest = writer.write_manifest_file().await.unwrap();
1519
1520 let manifest_list_writer = self
1523 .table
1524 .file_io()
1525 .new_output(current_snapshot.manifest_list())
1526 .unwrap()
1527 .writer()
1528 .await
1529 .unwrap();
1530 let mut manifest_list_write = ManifestListWriter::v2(
1531 manifest_list_writer,
1532 current_snapshot.snapshot_id(),
1533 current_snapshot.parent_snapshot_id(),
1534 current_snapshot.sequence_number(),
1535 );
1536 manifest_list_write
1537 .add_manifests(vec![data_manifest, delete_manifest].into_iter())
1538 .unwrap();
1539 manifest_list_write.close().await.unwrap();
1540 }
1541
1542 pub async fn setup_multi_row_group_manifest(&mut self, delete_positions: &[i64]) {
1551 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
1552 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
1553 let current_partition_spec = self.table.metadata().default_partition_spec();
1554
1555 let partition = Struct::from_iter([Some(Literal::long(1000))]);
1560
1561 let (data_file_path, data_file_size) = self.write_multi_row_group_data_file();
1562
1563 let mut data_writer = ManifestWriterBuilder::new(
1564 self.next_manifest_file(),
1565 Some(current_snapshot.snapshot_id()),
1566 current_schema.clone(),
1567 current_partition_spec.as_ref().clone(),
1568 )
1569 .build_v2_data();
1570 data_writer
1571 .add_entry(
1572 ManifestEntry::builder()
1573 .status(ManifestStatus::Added)
1574 .data_file(
1575 DataFileBuilder::default()
1576 .partition_spec_id(0)
1577 .content(DataContentType::Data)
1578 .file_path(data_file_path.clone())
1579 .file_format(DataFileFormat::Parquet)
1580 .file_size_in_bytes(data_file_size)
1581 .record_count(300)
1582 .partition(partition.clone())
1583 .key_metadata(None)
1584 .build()
1585 .unwrap(),
1586 )
1587 .build(),
1588 )
1589 .unwrap();
1590 let data_manifest = data_writer.write_manifest_file().await.unwrap();
1591
1592 let mut manifests = vec![data_manifest];
1593
1594 if !delete_positions.is_empty() {
1595 let (del_path, del_size) =
1596 self.write_positional_delete_file(&data_file_path, delete_positions);
1597
1598 let mut delete_writer = ManifestWriterBuilder::new(
1599 self.next_manifest_file(),
1600 Some(current_snapshot.snapshot_id()),
1601 current_schema.clone(),
1602 current_partition_spec.as_ref().clone(),
1603 )
1604 .build_v2_deletes();
1605 delete_writer
1606 .add_entry(
1607 ManifestEntry::builder()
1608 .status(ManifestStatus::Added)
1609 .data_file(
1610 DataFileBuilder::default()
1611 .partition_spec_id(0)
1612 .content(DataContentType::PositionDeletes)
1613 .file_path(del_path)
1614 .file_format(DataFileFormat::Parquet)
1615 .file_size_in_bytes(del_size)
1616 .record_count(delete_positions.len() as u64)
1617 .partition(partition.clone())
1618 .build()
1619 .unwrap(),
1620 )
1621 .build(),
1622 )
1623 .unwrap();
1624 manifests.push(delete_writer.write_manifest_file().await.unwrap());
1625 }
1626
1627 let manifest_list_writer = self
1628 .table
1629 .file_io()
1630 .new_output(current_snapshot.manifest_list())
1631 .unwrap()
1632 .writer()
1633 .await
1634 .unwrap();
1635 let mut manifest_list_write = ManifestListWriter::v2(
1636 manifest_list_writer,
1637 current_snapshot.snapshot_id(),
1638 current_snapshot.parent_snapshot_id(),
1639 current_snapshot.sequence_number(),
1640 );
1641 manifest_list_write
1642 .add_manifests(manifests.into_iter())
1643 .unwrap();
1644 manifest_list_write.close().await.unwrap();
1645 }
1646
1647 fn write_multi_row_group_data_file(&self) -> (String, u64) {
1651 fs::create_dir_all(&self.table_location).unwrap();
1652
1653 let arrow_schema = Arc::new(arrow_schema::Schema::new(vec![
1654 arrow_schema::Field::new("x", arrow_schema::DataType::Int64, false).with_metadata(
1655 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
1656 ),
1657 arrow_schema::Field::new("y", arrow_schema::DataType::Int64, false).with_metadata(
1658 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())]),
1659 ),
1660 ]));
1661
1662 let path = format!("{}/mrg.parquet", &self.table_location);
1663 let max_row_group_row_count = 100;
1664 let props = WriterProperties::builder()
1665 .set_compression(Compression::SNAPPY)
1666 .set_max_row_group_row_count(Some(max_row_group_row_count))
1667 .build();
1668
1669 let file = File::create(&path).unwrap();
1670 let mut writer = ArrowWriter::try_new(file, arrow_schema.clone(), Some(props)).unwrap();
1671 for group in 0..3i64 {
1672 let base = 1000 + group * max_row_group_row_count as i64;
1673 let col = Arc::new(Int64Array::from_iter_values(
1674 base..base + max_row_group_row_count as i64,
1675 )) as ArrayRef;
1676 let batch =
1677 RecordBatch::try_new(arrow_schema.clone(), vec![col.clone(), col]).unwrap();
1678 writer.write(&batch).unwrap();
1679 }
1680 writer.close().unwrap();
1681
1682 let size = fs::metadata(&path).unwrap().len();
1683 (path, size)
1684 }
1685
1686 fn write_positional_delete_file(
1689 &self,
1690 data_path: &str,
1691 positions: &[i64],
1692 ) -> (String, u64) {
1693 let del_schema = Arc::new(arrow_schema::Schema::new(vec![
1694 arrow_schema::Field::new(
1695 RESERVED_COL_NAME_DELETE_FILE_PATH,
1696 arrow_schema::DataType::Utf8,
1697 false,
1698 )
1699 .with_metadata(HashMap::from([(
1700 PARQUET_FIELD_ID_META_KEY.to_string(),
1701 RESERVED_FIELD_ID_DELETE_FILE_PATH.to_string(), )])),
1703 arrow_schema::Field::new(
1704 RESERVED_COL_NAME_DELETE_FILE_POS,
1705 arrow_schema::DataType::Int64,
1706 false,
1707 )
1708 .with_metadata(HashMap::from([(
1709 PARQUET_FIELD_ID_META_KEY.to_string(),
1710 RESERVED_FIELD_ID_DELETE_FILE_POS.to_string(), )])),
1712 ]));
1713
1714 let batch = RecordBatch::try_new(del_schema.clone(), vec![
1715 Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
1716 data_path.to_string(),
1717 positions.len(),
1718 ))) as ArrayRef,
1719 Arc::new(Int64Array::from_iter_values(positions.iter().copied())) as ArrayRef,
1720 ])
1721 .unwrap();
1722
1723 let path = format!("{}/pos-del.parquet", &self.table_location);
1724 let props = WriterProperties::builder()
1725 .set_compression(Compression::SNAPPY)
1726 .build();
1727 let file = File::create(&path).unwrap();
1728 let mut writer = ArrowWriter::try_new(file, del_schema, Some(props)).unwrap();
1729 writer.write(&batch).unwrap();
1730 writer.close().unwrap();
1731
1732 let size = fs::metadata(&path).unwrap().len();
1733 (path, size)
1734 }
1735 }
1736
1737 #[tokio::test]
1738 async fn test_table_scan_columns() {
1739 let table = TableTestFixture::new().table;
1740
1741 let table_scan = table.scan().select(["x", "y"]).build().unwrap();
1742 assert_eq!(
1743 Some(vec!["x".to_string(), "y".to_string()]),
1744 table_scan.column_names
1745 );
1746
1747 let table_scan = table
1748 .scan()
1749 .select(["x", "y"])
1750 .select(["z"])
1751 .build()
1752 .unwrap();
1753 assert_eq!(Some(vec!["z".to_string()]), table_scan.column_names);
1754 }
1755
1756 #[tokio::test]
1757 async fn test_select_all() {
1758 let table = TableTestFixture::new().table;
1759
1760 let table_scan = table.scan().select_all().build().unwrap();
1761 assert!(table_scan.column_names.is_none());
1762 }
1763
1764 #[test]
1765 fn test_select_no_exist_column() {
1766 let table = TableTestFixture::new().table;
1767
1768 let table_scan = table.scan().select(["x", "y", "z", "a", "b"]).build();
1769 assert!(table_scan.is_err());
1770 }
1771
1772 #[tokio::test]
1773 async fn test_table_scan_default_snapshot_id() {
1774 let table = TableTestFixture::new().table;
1775
1776 let table_scan = table.scan().build().unwrap();
1777 assert_eq!(
1778 table.metadata().current_snapshot().unwrap().snapshot_id(),
1779 table_scan.snapshot().unwrap().snapshot_id()
1780 );
1781 }
1782
1783 #[test]
1784 fn test_table_scan_non_exist_snapshot_id() {
1785 let table = TableTestFixture::new().table;
1786
1787 let table_scan = table.scan().snapshot_id(1024).build();
1788 assert!(table_scan.is_err());
1789 }
1790
1791 #[tokio::test]
1792 async fn test_table_scan_with_snapshot_id() {
1793 let table = TableTestFixture::new().table;
1794
1795 let table_scan = table
1796 .scan()
1797 .snapshot_id(3051729675574597004)
1798 .with_row_selection_enabled(true)
1799 .build()
1800 .unwrap();
1801 assert_eq!(
1802 table_scan.snapshot().unwrap().snapshot_id(),
1803 3051729675574597004
1804 );
1805 }
1806
1807 fn table_with_property(key: &str, value: &str) -> Table {
1808 let fixture = TableTestFixture::new();
1809 let mut metadata = fixture.table.metadata().clone();
1810 metadata
1811 .properties
1812 .insert(key.to_string(), value.to_string());
1813 Table::builder()
1814 .metadata(metadata)
1815 .identifier(fixture.table.identifier().clone())
1816 .file_io(fixture.table.file_io().clone())
1817 .metadata_location(fixture.table.metadata_location().unwrap().to_string())
1818 .runtime(test_runtime())
1819 .build()
1820 .unwrap()
1821 }
1822
1823 #[test]
1824 fn test_table_scan_without_name_mapping_property() {
1825 let table = TableTestFixture::new().table;
1826
1827 let table_scan = table.scan().build().unwrap();
1828 assert!(
1829 table_scan
1830 .plan_context
1831 .as_ref()
1832 .unwrap()
1833 .name_mapping
1834 .is_none()
1835 );
1836 }
1837
1838 #[test]
1839 fn test_table_scan_with_name_mapping_property() {
1840 let mapping_json = r#"[{"field-id":1,"names":["id","record_id"]}]"#;
1841 let table = table_with_property(DEFAULT_SCHEMA_NAME_MAPPING, mapping_json);
1842
1843 let table_scan = table.scan().build().unwrap();
1844 let mapping = table_scan
1845 .plan_context
1846 .as_ref()
1847 .unwrap()
1848 .name_mapping
1849 .as_ref()
1850 .expect("name_mapping should be parsed from the table property");
1851 let fields = mapping.fields();
1852 assert_eq!(fields.len(), 1);
1853 assert_eq!(fields[0].field_id(), Some(1));
1854 assert_eq!(fields[0].names(), &[
1855 "id".to_string(),
1856 "record_id".to_string()
1857 ]);
1858 }
1859
1860 #[test]
1861 fn test_table_scan_with_malformed_name_mapping_property() {
1862 let table = table_with_property(DEFAULT_SCHEMA_NAME_MAPPING, "{ not valid json");
1863
1864 let err = table
1865 .scan()
1866 .build()
1867 .expect_err("malformed name mapping should fail to parse");
1868 assert_eq!(err.kind(), ErrorKind::DataInvalid);
1869 }
1870
1871 #[tokio::test]
1872 async fn test_plan_files_carries_name_mapping_into_file_scan_task() {
1873 let mut fixture = TableTestFixture::new();
1874 fixture.setup_manifest_files().await;
1875
1876 let mapping_json = r#"[{"field-id":1,"names":["id","record_id"]}]"#;
1877 let mut metadata = fixture.table.metadata().clone();
1878 metadata.properties.insert(
1879 DEFAULT_SCHEMA_NAME_MAPPING.to_string(),
1880 mapping_json.to_string(),
1881 );
1882 let table = Table::builder()
1883 .metadata(metadata)
1884 .identifier(fixture.table.identifier().clone())
1885 .file_io(fixture.table.file_io().clone())
1886 .metadata_location(fixture.table.metadata_location().unwrap().to_string())
1887 .runtime(test_runtime())
1888 .build()
1889 .unwrap();
1890
1891 let tasks: Vec<_> = table
1892 .scan()
1893 .build()
1894 .unwrap()
1895 .plan_files()
1896 .await
1897 .unwrap()
1898 .try_collect()
1899 .await
1900 .unwrap();
1901
1902 assert!(!tasks.is_empty(), "expected at least one FileScanTask");
1903 for task in &tasks {
1904 let mapping = task
1905 .name_mapping
1906 .as_ref()
1907 .expect("name_mapping should reach the FileScanTask");
1908 assert_eq!(mapping.fields().len(), 1);
1909 assert_eq!(mapping.fields()[0].field_id(), Some(1));
1910 }
1911 }
1912
1913 #[tokio::test]
1914 async fn test_plan_files_on_table_without_any_snapshots() {
1915 let table = TableTestFixture::new_empty().table;
1916 let batch_stream = table.scan().build().unwrap().to_arrow().await.unwrap();
1917 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1918 assert!(batches.is_empty());
1919 }
1920
1921 #[tokio::test]
1922 async fn test_plan_files_no_deletions() {
1923 let mut fixture = TableTestFixture::new();
1924 fixture.setup_manifest_files().await;
1925
1926 let table_scan = fixture
1928 .table
1929 .scan()
1930 .with_row_selection_enabled(true)
1931 .build()
1932 .unwrap();
1933
1934 let mut tasks = table_scan
1935 .plan_files()
1936 .await
1937 .unwrap()
1938 .try_fold(vec![], |mut acc, task| async move {
1939 acc.push(task);
1940 Ok(acc)
1941 })
1942 .await
1943 .unwrap();
1944
1945 assert_eq!(tasks.len(), 2);
1946
1947 tasks.sort_by_key(|t| t.data_file_path.to_string());
1948
1949 assert_eq!(
1951 tasks[0].data_file_path,
1952 format!("{}/1.parquet", &fixture.table_location)
1953 );
1954
1955 assert_eq!(
1957 tasks[1].data_file_path,
1958 format!("{}/3.parquet", &fixture.table_location)
1959 );
1960 }
1961
1962 #[tokio::test]
1963 async fn test_plan_files_carries_row_lineage_into_file_scan_task() {
1964 let mut fixture = TableTestFixture::new();
1965 fixture.setup_manifest_files().await;
1966
1967 let mut tasks: Vec<_> = fixture
1968 .table
1969 .scan()
1970 .build()
1971 .unwrap()
1972 .plan_files()
1973 .await
1974 .unwrap()
1975 .try_collect()
1976 .await
1977 .unwrap();
1978
1979 assert_eq!(tasks.len(), 2);
1980 tasks.sort_by_key(|task| task.data_file_path.to_string());
1981
1982 assert_eq!(
1985 tasks[0].data_file_path,
1986 format!("{}/1.parquet", &fixture.table_location)
1987 );
1988 assert_eq!(tasks[0].data_sequence_number, Some(1));
1989 assert_eq!(
1990 tasks[1].data_file_path,
1991 format!("{}/3.parquet", &fixture.table_location)
1992 );
1993 assert_eq!(tasks[1].data_sequence_number, Some(0));
1994
1995 assert!(tasks.iter().all(|task| task.first_row_id.is_none()));
1997 }
1998
1999 #[tokio::test]
2000 async fn test_plan_files_carries_row_lineage_from_v3_manifest() {
2001 let mut fixture = TableTestFixture::new();
2002 fixture.setup_v3_manifest_files().await;
2003
2004 let task = fixture
2005 .table
2006 .scan()
2007 .build()
2008 .unwrap()
2009 .plan_files()
2010 .await
2011 .unwrap()
2012 .try_collect::<Vec<_>>()
2013 .await
2014 .unwrap()
2015 .into_iter()
2016 .next()
2017 .expect("expected one FileScanTask");
2018
2019 assert_eq!(task.first_row_id, Some(42));
2022 assert_eq!(task.data_sequence_number, Some(1));
2024 }
2025
2026 #[tokio::test]
2027 async fn test_filtered_scan_with_dropped_partition_source_column() {
2028 let mut fixture = TableTestFixture::new();
2029 fixture.setup_manifest_files().await;
2030
2031 let baseline = scan_y_gte_5(&fixture.table).await;
2033 assert!(!baseline.is_empty());
2034 assert!(baseline.iter().all(|y| *y >= 5));
2035
2036 let current_schema = fixture.table.metadata().current_schema();
2040 let evolved_schema = Schema::builder()
2041 .with_fields(
2042 current_schema
2043 .as_struct()
2044 .fields()
2045 .iter()
2046 .filter(|field| field.id != 1)
2047 .cloned(),
2048 )
2049 .with_identifier_field_ids(vec![2])
2050 .build()
2051 .unwrap();
2052 let evolved =
2053 TableMetadataBuilder::new_from_metadata(fixture.table.metadata().clone(), None)
2054 .add_default_partition_spec(UnboundPartitionSpec::builder().build())
2055 .unwrap()
2056 .add_current_schema(evolved_schema)
2057 .unwrap()
2058 .build()
2059 .unwrap()
2060 .metadata;
2061
2062 let parent = evolved.current_snapshot().unwrap().clone();
2065 let snapshot = Snapshot::builder()
2066 .with_snapshot_id(parent.snapshot_id() + 1)
2067 .with_parent_snapshot_id(Some(parent.snapshot_id()))
2068 .with_sequence_number(evolved.last_sequence_number() + 1)
2069 .with_timestamp_ms(evolved.last_updated_ms + 1)
2070 .with_schema_id(evolved.current_schema_id())
2071 .with_manifest_list(parent.manifest_list())
2072 .with_summary(Summary {
2073 operation: Operation::Append,
2074 additional_properties: HashMap::new(),
2075 })
2076 .build();
2077 let metadata = TableMetadataBuilder::new_from_metadata(evolved, None)
2078 .set_branch_snapshot(snapshot, MAIN_BRANCH)
2079 .unwrap()
2080 .build()
2081 .unwrap()
2082 .metadata;
2083 let table = fixture.table.clone().with_metadata(Arc::new(metadata));
2084
2085 let evolved = scan_y_gte_5(&table).await;
2088 assert_eq!(evolved, baseline);
2089 }
2090
2091 async fn scan_y_gte_5(table: &Table) -> Vec<i64> {
2092 let table_scan = table
2093 .scan()
2094 .select(["y"])
2095 .with_filter(Reference::new("y").greater_than_or_equal_to(Datum::long(5)))
2096 .build()
2097 .unwrap();
2098 let batches: Vec<_> = table_scan
2099 .to_arrow()
2100 .await
2101 .unwrap()
2102 .try_collect()
2103 .await
2104 .unwrap();
2105
2106 let mut values: Vec<i64> = batches
2107 .iter()
2108 .flat_map(|batch| {
2109 let col = batch.column_by_name("y").unwrap();
2110 let arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
2111 (0..arr.len()).map(|i| arr.value(i)).collect::<Vec<_>>()
2112 })
2113 .collect();
2114 values.sort_unstable();
2115 values
2116 }
2117
2118 #[tokio::test]
2119 async fn test_open_parquet_no_deletions() {
2120 let mut fixture = TableTestFixture::new();
2121 fixture.setup_manifest_files().await;
2122
2123 let table_scan = fixture
2125 .table
2126 .scan()
2127 .with_row_selection_enabled(true)
2128 .build()
2129 .unwrap();
2130
2131 let batch_stream = table_scan.to_arrow().await.unwrap();
2132
2133 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2134
2135 let col = batches[0].column_by_name("x").unwrap();
2136
2137 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
2138 assert_eq!(int64_arr.value(0), 1);
2139 }
2140
2141 #[tokio::test]
2142 async fn test_open_parquet_no_deletions_by_separate_reader() {
2143 let mut fixture = TableTestFixture::new();
2144 fixture.setup_manifest_files().await;
2145
2146 let table_scan = fixture
2148 .table
2149 .scan()
2150 .with_row_selection_enabled(true)
2151 .build()
2152 .unwrap();
2153
2154 let mut plan_task: Vec<_> = table_scan
2155 .plan_files()
2156 .await
2157 .unwrap()
2158 .try_collect()
2159 .await
2160 .unwrap();
2161 assert_eq!(plan_task.len(), 2);
2162
2163 let reader = ArrowReaderBuilder::new(
2164 fixture.table.file_io().clone(),
2165 fixture.table.runtime().clone(),
2166 )
2167 .build();
2168 let batch_stream = reader
2169 .clone()
2170 .read(Box::pin(stream::iter(vec![Ok(plan_task.remove(0))])))
2171 .unwrap()
2172 .stream();
2173 let batch_1: Vec<_> = batch_stream.try_collect().await.unwrap();
2174
2175 let reader = ArrowReaderBuilder::new(
2176 fixture.table.file_io().clone(),
2177 fixture.table.runtime().clone(),
2178 )
2179 .build();
2180 let batch_stream = reader
2181 .read(Box::pin(stream::iter(vec![Ok(plan_task.remove(0))])))
2182 .unwrap()
2183 .stream();
2184 let batch_2: Vec<_> = batch_stream.try_collect().await.unwrap();
2185
2186 assert_eq!(batch_1, batch_2);
2187 }
2188
2189 #[tokio::test]
2190 async fn test_open_parquet_with_projection() {
2191 let mut fixture = TableTestFixture::new();
2192 fixture.setup_manifest_files().await;
2193
2194 let table_scan = fixture
2196 .table
2197 .scan()
2198 .select(["x", "z"])
2199 .with_row_selection_enabled(true)
2200 .build()
2201 .unwrap();
2202
2203 let batch_stream = table_scan.to_arrow().await.unwrap();
2204
2205 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2206
2207 assert_eq!(batches[0].num_columns(), 2);
2208
2209 let col1 = batches[0].column_by_name("x").unwrap();
2210 let int64_arr = col1.as_any().downcast_ref::<Int64Array>().unwrap();
2211 assert_eq!(int64_arr.value(0), 1);
2212
2213 let col2 = batches[0].column_by_name("z").unwrap();
2214 let int64_arr = col2.as_any().downcast_ref::<Int64Array>().unwrap();
2215 assert_eq!(int64_arr.value(0), 3);
2216
2217 let table_scan = fixture.table.scan().select_empty().build().unwrap();
2219 let batch_stream = table_scan.to_arrow().await.unwrap();
2220 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2221
2222 assert_eq!(batches[0].num_columns(), 0);
2223 assert_eq!(batches[0].num_rows(), 1024);
2224 }
2225
2226 #[tokio::test]
2227 async fn test_filter_on_arrow_lt() {
2228 let mut fixture = TableTestFixture::new();
2229 fixture.setup_manifest_files().await;
2230
2231 let mut builder = fixture.table.scan();
2233 let predicate = Reference::new("y").less_than(Datum::long(3));
2234 builder = builder
2235 .with_filter(predicate)
2236 .with_row_selection_enabled(true);
2237 let table_scan = builder.build().unwrap();
2238
2239 let batch_stream = table_scan.to_arrow().await.unwrap();
2240
2241 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2242
2243 assert_eq!(batches[0].num_rows(), 512);
2244
2245 let col = batches[0].column_by_name("x").unwrap();
2246 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
2247 assert_eq!(int64_arr.value(0), 1);
2248
2249 let col = batches[0].column_by_name("y").unwrap();
2250 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
2251 assert_eq!(int64_arr.value(0), 2);
2252 }
2253
2254 #[tokio::test]
2255 async fn test_filter_on_arrow_gt_eq() {
2256 let mut fixture = TableTestFixture::new();
2257 fixture.setup_manifest_files().await;
2258
2259 let mut builder = fixture.table.scan();
2261 let predicate = Reference::new("y").greater_than_or_equal_to(Datum::long(5));
2262 builder = builder
2263 .with_filter(predicate)
2264 .with_row_selection_enabled(true);
2265 let table_scan = builder.build().unwrap();
2266
2267 let batch_stream = table_scan.to_arrow().await.unwrap();
2268
2269 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2270
2271 assert_eq!(batches[0].num_rows(), 12);
2272
2273 let col = batches[0].column_by_name("x").unwrap();
2274 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
2275 assert_eq!(int64_arr.value(0), 1);
2276
2277 let col = batches[0].column_by_name("y").unwrap();
2278 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
2279 assert_eq!(int64_arr.value(0), 5);
2280 }
2281
2282 #[tokio::test]
2283 async fn test_filter_double_eq() {
2284 let mut fixture = TableTestFixture::new();
2285 fixture.setup_manifest_files().await;
2286
2287 let mut builder = fixture.table.scan();
2289 let predicate = Reference::new("dbl").equal_to(Datum::double(150.0f64));
2290 builder = builder
2291 .with_filter(predicate)
2292 .with_row_selection_enabled(true);
2293 let table_scan = builder.build().unwrap();
2294
2295 let batch_stream = table_scan.to_arrow().await.unwrap();
2296
2297 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2298
2299 assert_eq!(batches.len(), 2);
2300 assert_eq!(batches[0].num_rows(), 12);
2301
2302 let col = batches[0].column_by_name("dbl").unwrap();
2303 let f64_arr = col.as_any().downcast_ref::<Float64Array>().unwrap();
2304 assert_eq!(f64_arr.value(1), 150.0f64);
2305 }
2306
2307 #[tokio::test]
2308 async fn test_filter_int_eq() {
2309 let mut fixture = TableTestFixture::new();
2310 fixture.setup_manifest_files().await;
2311
2312 let mut builder = fixture.table.scan();
2314 let predicate = Reference::new("i32").equal_to(Datum::int(150i32));
2315 builder = builder
2316 .with_filter(predicate)
2317 .with_row_selection_enabled(true);
2318 let table_scan = builder.build().unwrap();
2319
2320 let batch_stream = table_scan.to_arrow().await.unwrap();
2321
2322 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2323
2324 assert_eq!(batches.len(), 2);
2325 assert_eq!(batches[0].num_rows(), 12);
2326
2327 let col = batches[0].column_by_name("i32").unwrap();
2328 let i32_arr = col.as_any().downcast_ref::<Int32Array>().unwrap();
2329 assert_eq!(i32_arr.value(1), 150i32);
2330 }
2331
2332 #[tokio::test]
2333 async fn test_filter_long_eq() {
2334 let mut fixture = TableTestFixture::new();
2335 fixture.setup_manifest_files().await;
2336
2337 let mut builder = fixture.table.scan();
2339 let predicate = Reference::new("i64").equal_to(Datum::long(150i64));
2340 builder = builder
2341 .with_filter(predicate)
2342 .with_row_selection_enabled(true);
2343 let table_scan = builder.build().unwrap();
2344
2345 let batch_stream = table_scan.to_arrow().await.unwrap();
2346
2347 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2348
2349 assert_eq!(batches.len(), 2);
2350 assert_eq!(batches[0].num_rows(), 12);
2351
2352 let col = batches[0].column_by_name("i64").unwrap();
2353 let i64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
2354 assert_eq!(i64_arr.value(1), 150i64);
2355 }
2356
2357 #[tokio::test]
2358 async fn test_filter_bool_eq() {
2359 let mut fixture = TableTestFixture::new();
2360 fixture.setup_manifest_files().await;
2361
2362 let mut builder = fixture.table.scan();
2364 let predicate = Reference::new("bool").equal_to(Datum::bool(true));
2365 builder = builder
2366 .with_filter(predicate)
2367 .with_row_selection_enabled(true);
2368 let table_scan = builder.build().unwrap();
2369
2370 let batch_stream = table_scan.to_arrow().await.unwrap();
2371
2372 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2373
2374 assert_eq!(batches.len(), 2);
2375 assert_eq!(batches[0].num_rows(), 512);
2376
2377 let col = batches[0].column_by_name("bool").unwrap();
2378 let bool_arr = col.as_any().downcast_ref::<BooleanArray>().unwrap();
2379 assert!(bool_arr.value(1));
2380 }
2381
2382 #[tokio::test]
2383 async fn test_filter_on_arrow_is_null() {
2384 let mut fixture = TableTestFixture::new();
2385 fixture.setup_manifest_files().await;
2386
2387 let mut builder = fixture.table.scan();
2389 let predicate = Reference::new("y").is_null();
2390 builder = builder
2391 .with_filter(predicate)
2392 .with_row_selection_enabled(true);
2393 let table_scan = builder.build().unwrap();
2394
2395 let batch_stream = table_scan.to_arrow().await.unwrap();
2396
2397 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2398 assert_eq!(batches.len(), 0);
2399 }
2400
2401 #[tokio::test]
2402 async fn test_filter_on_arrow_is_not_null() {
2403 let mut fixture = TableTestFixture::new();
2404 fixture.setup_manifest_files().await;
2405
2406 let mut builder = fixture.table.scan();
2408 let predicate = Reference::new("y").is_not_null();
2409 builder = builder
2410 .with_filter(predicate)
2411 .with_row_selection_enabled(true);
2412 let table_scan = builder.build().unwrap();
2413
2414 let batch_stream = table_scan.to_arrow().await.unwrap();
2415
2416 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2417 assert_eq!(batches[0].num_rows(), 1024);
2418 }
2419
2420 #[tokio::test]
2421 async fn test_filter_on_arrow_lt_and_gt() {
2422 let mut fixture = TableTestFixture::new();
2423 fixture.setup_manifest_files().await;
2424
2425 let mut builder = fixture.table.scan();
2427 let predicate = Reference::new("y")
2428 .less_than(Datum::long(5))
2429 .and(Reference::new("z").greater_than_or_equal_to(Datum::long(4)));
2430 builder = builder
2431 .with_filter(predicate)
2432 .with_row_selection_enabled(true);
2433 let table_scan = builder.build().unwrap();
2434
2435 let batch_stream = table_scan.to_arrow().await.unwrap();
2436
2437 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2438 assert_eq!(batches[0].num_rows(), 500);
2439
2440 let col = batches[0].column_by_name("x").unwrap();
2441 let expected_x = Arc::new(Int64Array::from_iter_values(vec![1; 500])) as ArrayRef;
2442 assert_eq!(col, &expected_x);
2443
2444 let col = batches[0].column_by_name("y").unwrap();
2445 let mut values = vec![];
2446 values.append(vec![3; 200].as_mut());
2447 values.append(vec![4; 300].as_mut());
2448 let expected_y = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
2449 assert_eq!(col, &expected_y);
2450
2451 let col = batches[0].column_by_name("z").unwrap();
2452 let expected_z = Arc::new(Int64Array::from_iter_values(vec![4; 500])) as ArrayRef;
2453 assert_eq!(col, &expected_z);
2454 }
2455
2456 #[tokio::test]
2457 async fn test_filter_on_arrow_lt_or_gt() {
2458 let mut fixture = TableTestFixture::new();
2459 fixture.setup_manifest_files().await;
2460
2461 let mut builder = fixture.table.scan();
2463 let predicate = Reference::new("y")
2464 .less_than(Datum::long(5))
2465 .or(Reference::new("z").greater_than_or_equal_to(Datum::long(4)));
2466 builder = builder
2467 .with_filter(predicate)
2468 .with_row_selection_enabled(true);
2469 let table_scan = builder.build().unwrap();
2470
2471 let batch_stream = table_scan.to_arrow().await.unwrap();
2472
2473 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2474 assert_eq!(batches[0].num_rows(), 1024);
2475
2476 let col = batches[0].column_by_name("x").unwrap();
2477 let expected_x = Arc::new(Int64Array::from_iter_values(vec![1; 1024])) as ArrayRef;
2478 assert_eq!(col, &expected_x);
2479
2480 let col = batches[0].column_by_name("y").unwrap();
2481 let mut values = vec![2; 512];
2482 values.append(vec![3; 200].as_mut());
2483 values.append(vec![4; 300].as_mut());
2484 values.append(vec![5; 12].as_mut());
2485 let expected_y = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
2486 assert_eq!(col, &expected_y);
2487
2488 let col = batches[0].column_by_name("z").unwrap();
2489 let mut values = vec![3; 512];
2490 values.append(vec![4; 512].as_mut());
2491 let expected_z = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
2492 assert_eq!(col, &expected_z);
2493 }
2494
2495 #[tokio::test]
2496 async fn test_filter_on_arrow_startswith() {
2497 let mut fixture = TableTestFixture::new();
2498 fixture.setup_manifest_files().await;
2499
2500 let mut builder = fixture.table.scan();
2502 let predicate = Reference::new("a").starts_with(Datum::string("Ice"));
2503 builder = builder
2504 .with_filter(predicate)
2505 .with_row_selection_enabled(true);
2506 let table_scan = builder.build().unwrap();
2507
2508 let batch_stream = table_scan.to_arrow().await.unwrap();
2509
2510 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2511
2512 assert_eq!(batches[0].num_rows(), 512);
2513
2514 let col = batches[0].column_by_name("a").unwrap();
2515 let string_arr = col.as_any().downcast_ref::<StringArray>().unwrap();
2516 assert_eq!(string_arr.value(0), "Iceberg");
2517 }
2518
2519 #[tokio::test]
2520 async fn test_filter_on_arrow_not_startswith() {
2521 let mut fixture = TableTestFixture::new();
2522 fixture.setup_manifest_files().await;
2523
2524 let mut builder = fixture.table.scan();
2526 let predicate = Reference::new("a").not_starts_with(Datum::string("Ice"));
2527 builder = builder
2528 .with_filter(predicate)
2529 .with_row_selection_enabled(true);
2530 let table_scan = builder.build().unwrap();
2531
2532 let batch_stream = table_scan.to_arrow().await.unwrap();
2533
2534 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2535
2536 assert_eq!(batches[0].num_rows(), 512);
2537
2538 let col = batches[0].column_by_name("a").unwrap();
2539 let string_arr = col.as_any().downcast_ref::<StringArray>().unwrap();
2540 assert_eq!(string_arr.value(0), "Apache");
2541 }
2542
2543 #[tokio::test]
2544 async fn test_filter_on_arrow_in() {
2545 let mut fixture = TableTestFixture::new();
2546 fixture.setup_manifest_files().await;
2547
2548 let mut builder = fixture.table.scan();
2550 let predicate =
2551 Reference::new("a").is_in([Datum::string("Sioux"), Datum::string("Iceberg")]);
2552 builder = builder
2553 .with_filter(predicate)
2554 .with_row_selection_enabled(true);
2555 let table_scan = builder.build().unwrap();
2556
2557 let batch_stream = table_scan.to_arrow().await.unwrap();
2558
2559 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2560
2561 assert_eq!(batches[0].num_rows(), 512);
2562
2563 let col = batches[0].column_by_name("a").unwrap();
2564 let string_arr = col.as_any().downcast_ref::<StringArray>().unwrap();
2565 assert_eq!(string_arr.value(0), "Iceberg");
2566 }
2567
2568 #[tokio::test]
2569 async fn test_filter_on_arrow_not_in() {
2570 let mut fixture = TableTestFixture::new();
2571 fixture.setup_manifest_files().await;
2572
2573 let mut builder = fixture.table.scan();
2575 let predicate =
2576 Reference::new("a").is_not_in([Datum::string("Sioux"), Datum::string("Iceberg")]);
2577 builder = builder
2578 .with_filter(predicate)
2579 .with_row_selection_enabled(true);
2580 let table_scan = builder.build().unwrap();
2581
2582 let batch_stream = table_scan.to_arrow().await.unwrap();
2583
2584 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2585
2586 assert_eq!(batches[0].num_rows(), 512);
2587
2588 let col = batches[0].column_by_name("a").unwrap();
2589 let string_arr = col.as_any().downcast_ref::<StringArray>().unwrap();
2590 assert_eq!(string_arr.value(0), "Apache");
2591 }
2592
2593 #[test]
2594 fn test_file_scan_task_serialize_deserialize() {
2595 let test_fn = |task: FileScanTask| {
2596 let serialized = serde_json::to_string(&task).unwrap();
2597 let deserialized: FileScanTask = serde_json::from_str(&serialized).unwrap();
2598
2599 assert_eq!(task.data_file_path, deserialized.data_file_path);
2600 assert_eq!(task.start, deserialized.start);
2601 assert_eq!(task.length, deserialized.length);
2602 assert_eq!(task.project_field_ids, deserialized.project_field_ids);
2603 assert_eq!(task.predicate, deserialized.predicate);
2604 assert_eq!(task.schema, deserialized.schema);
2605 assert_eq!(task.first_row_id, deserialized.first_row_id);
2606 assert_eq!(task.data_sequence_number, deserialized.data_sequence_number);
2607 };
2608
2609 let schema = Arc::new(
2611 Schema::builder()
2612 .with_fields(vec![Arc::new(NestedField::required(
2613 1,
2614 "x",
2615 Type::Primitive(PrimitiveType::Binary),
2616 ))])
2617 .build()
2618 .unwrap(),
2619 );
2620 let task = FileScanTask::builder()
2621 .with_data_file_path("data_file_path".to_string())
2622 .with_file_size_in_bytes(0)
2623 .with_start(0)
2624 .with_length(100)
2625 .with_project_field_ids(vec![1, 2, 3])
2626 .with_schema(schema.clone())
2627 .with_record_count(Some(100))
2628 .with_first_row_id(Some(1000))
2629 .with_data_sequence_number(Some(5))
2630 .with_data_file_format(DataFileFormat::Parquet)
2631 .with_case_sensitive(false)
2632 .build();
2633 test_fn(task);
2634
2635 let task = FileScanTask::builder()
2637 .with_data_file_path("data_file_path".to_string())
2638 .with_file_size_in_bytes(0)
2639 .with_start(0)
2640 .with_length(100)
2641 .with_project_field_ids(vec![1, 2, 3])
2642 .with_predicate(Some(BoundPredicate::AlwaysTrue))
2643 .with_schema(schema)
2644 .with_data_file_format(DataFileFormat::Avro)
2645 .with_case_sensitive(false)
2646 .build();
2647 test_fn(task);
2648 }
2649
2650 #[tokio::test]
2651 async fn test_select_with_file_column() {
2652 let mut fixture = TableTestFixture::new();
2653 fixture.setup_manifest_files().await;
2654
2655 let table_scan = fixture
2657 .table
2658 .scan()
2659 .select(["x", RESERVED_COL_NAME_FILE])
2660 .with_row_selection_enabled(true)
2661 .build()
2662 .unwrap();
2663
2664 let batch_stream = table_scan.to_arrow().await.unwrap();
2665 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2666
2667 assert_eq!(batches[0].num_columns(), 2);
2669
2670 let x_col = batches[0].column_by_name("x").unwrap();
2672 let x_arr = x_col.as_primitive::<arrow_array::types::Int64Type>();
2673 assert_eq!(x_arr.value(0), 1);
2674
2675 let file_col = batches[0].column_by_name(RESERVED_COL_NAME_FILE);
2677 assert!(
2678 file_col.is_some(),
2679 "_file column should be present in the batch"
2680 );
2681
2682 let file_col = file_col.unwrap();
2684 assert!(
2685 matches!(
2686 file_col.data_type(),
2687 arrow_schema::DataType::RunEndEncoded(_, _)
2688 ),
2689 "_file column should use RunEndEncoded type"
2690 );
2691
2692 let run_array = file_col
2694 .as_any()
2695 .downcast_ref::<RunArray<Int32Type>>()
2696 .expect("_file column should be a RunArray");
2697
2698 let values = run_array.values();
2699 let string_values = values.as_string::<i32>();
2700 assert_eq!(string_values.len(), 1, "Should have a single file path");
2701
2702 let file_path = string_values.value(0);
2703 assert!(
2704 file_path.ends_with(".parquet"),
2705 "File path should end with .parquet, got: {file_path}"
2706 );
2707 }
2708
2709 #[tokio::test]
2710 async fn test_select_file_column_position() {
2711 let mut fixture = TableTestFixture::new();
2712 fixture.setup_manifest_files().await;
2713
2714 let table_scan = fixture
2716 .table
2717 .scan()
2718 .select(["x", RESERVED_COL_NAME_FILE, "z"])
2719 .with_row_selection_enabled(true)
2720 .build()
2721 .unwrap();
2722
2723 let batch_stream = table_scan.to_arrow().await.unwrap();
2724 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2725
2726 assert_eq!(batches[0].num_columns(), 3);
2727
2728 let schema = batches[0].schema();
2730 assert_eq!(schema.field(0).name(), "x");
2731 assert_eq!(schema.field(1).name(), RESERVED_COL_NAME_FILE);
2732 assert_eq!(schema.field(2).name(), "z");
2733
2734 assert!(batches[0].column_by_name("x").is_some());
2736 assert!(batches[0].column_by_name(RESERVED_COL_NAME_FILE).is_some());
2737 assert!(batches[0].column_by_name("z").is_some());
2738 }
2739
2740 #[tokio::test]
2741 async fn test_select_file_column_only() {
2742 let mut fixture = TableTestFixture::new();
2743 fixture.setup_manifest_files().await;
2744
2745 let table_scan = fixture
2747 .table
2748 .scan()
2749 .select([RESERVED_COL_NAME_FILE])
2750 .with_row_selection_enabled(true)
2751 .build()
2752 .unwrap();
2753
2754 let batch_stream = table_scan.to_arrow().await.unwrap();
2755 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2756
2757 assert_eq!(batches[0].num_columns(), 1);
2759
2760 let schema = batches[0].schema();
2762 assert_eq!(schema.field(0).name(), RESERVED_COL_NAME_FILE);
2763
2764 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
2768 assert_eq!(total_rows, 2048);
2769 }
2770
2771 #[tokio::test]
2772 async fn test_file_column_with_multiple_files() {
2773 use std::collections::HashSet;
2774
2775 let mut fixture = TableTestFixture::new();
2776 fixture.setup_manifest_files().await;
2777
2778 let table_scan = fixture
2780 .table
2781 .scan()
2782 .select(["x", RESERVED_COL_NAME_FILE])
2783 .with_row_selection_enabled(true)
2784 .build()
2785 .unwrap();
2786
2787 let batch_stream = table_scan.to_arrow().await.unwrap();
2788 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2789
2790 let mut file_paths = HashSet::new();
2792 for batch in &batches {
2793 let file_col = batch.column_by_name(RESERVED_COL_NAME_FILE).unwrap();
2794 let run_array = file_col
2795 .as_any()
2796 .downcast_ref::<RunArray<Int32Type>>()
2797 .expect("_file column should be a RunArray");
2798
2799 let values = run_array.values();
2800 let string_values = values.as_string::<i32>();
2801 for i in 0..string_values.len() {
2802 file_paths.insert(string_values.value(i).to_string());
2803 }
2804 }
2805
2806 assert!(!file_paths.is_empty(), "Should have at least one file path");
2808
2809 for path in &file_paths {
2811 assert!(
2812 path.ends_with(".parquet"),
2813 "All file paths should end with .parquet, got: {path}"
2814 );
2815 }
2816 }
2817
2818 #[tokio::test]
2819 async fn test_file_column_at_start() {
2820 let mut fixture = TableTestFixture::new();
2821 fixture.setup_manifest_files().await;
2822
2823 let table_scan = fixture
2825 .table
2826 .scan()
2827 .select([RESERVED_COL_NAME_FILE, "x", "y"])
2828 .with_row_selection_enabled(true)
2829 .build()
2830 .unwrap();
2831
2832 let batch_stream = table_scan.to_arrow().await.unwrap();
2833 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2834
2835 assert_eq!(batches[0].num_columns(), 3);
2836
2837 let schema = batches[0].schema();
2839 assert_eq!(schema.field(0).name(), RESERVED_COL_NAME_FILE);
2840 assert_eq!(schema.field(1).name(), "x");
2841 assert_eq!(schema.field(2).name(), "y");
2842 }
2843
2844 #[tokio::test]
2845 async fn test_file_column_at_end() {
2846 let mut fixture = TableTestFixture::new();
2847 fixture.setup_manifest_files().await;
2848
2849 let table_scan = fixture
2851 .table
2852 .scan()
2853 .select(["x", "y", RESERVED_COL_NAME_FILE])
2854 .with_row_selection_enabled(true)
2855 .build()
2856 .unwrap();
2857
2858 let batch_stream = table_scan.to_arrow().await.unwrap();
2859 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2860
2861 assert_eq!(batches[0].num_columns(), 3);
2862
2863 let schema = batches[0].schema();
2865 assert_eq!(schema.field(0).name(), "x");
2866 assert_eq!(schema.field(1).name(), "y");
2867 assert_eq!(schema.field(2).name(), RESERVED_COL_NAME_FILE);
2868 }
2869
2870 #[tokio::test]
2871 async fn test_select_with_repeated_column_names() {
2872 let mut fixture = TableTestFixture::new();
2873 fixture.setup_manifest_files().await;
2874
2875 let table_scan = fixture
2878 .table
2879 .scan()
2880 .select([
2881 "x",
2882 RESERVED_COL_NAME_FILE,
2883 "x", "y",
2885 RESERVED_COL_NAME_FILE, "y", ])
2888 .with_row_selection_enabled(true)
2889 .build()
2890 .unwrap();
2891
2892 let batch_stream = table_scan.to_arrow().await.unwrap();
2893 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2894
2895 assert_eq!(
2897 batches[0].num_columns(),
2898 6,
2899 "Should have exactly 6 columns with duplicates"
2900 );
2901
2902 let schema = batches[0].schema();
2903
2904 assert_eq!(schema.field(0).name(), "x", "Column 0 should be x");
2906 assert_eq!(
2907 schema.field(1).name(),
2908 RESERVED_COL_NAME_FILE,
2909 "Column 1 should be _file"
2910 );
2911 assert_eq!(
2912 schema.field(2).name(),
2913 "x",
2914 "Column 2 should be x (duplicate)"
2915 );
2916 assert_eq!(schema.field(3).name(), "y", "Column 3 should be y");
2917 assert_eq!(
2918 schema.field(4).name(),
2919 RESERVED_COL_NAME_FILE,
2920 "Column 4 should be _file (duplicate)"
2921 );
2922 assert_eq!(
2923 schema.field(5).name(),
2924 "y",
2925 "Column 5 should be y (duplicate)"
2926 );
2927
2928 assert!(
2930 matches!(schema.field(0).data_type(), arrow_schema::DataType::Int64),
2931 "Column x should be Int64"
2932 );
2933 assert!(
2934 matches!(schema.field(2).data_type(), arrow_schema::DataType::Int64),
2935 "Column x (duplicate) should be Int64"
2936 );
2937 assert!(
2938 matches!(schema.field(3).data_type(), arrow_schema::DataType::Int64),
2939 "Column y should be Int64"
2940 );
2941 assert!(
2942 matches!(schema.field(5).data_type(), arrow_schema::DataType::Int64),
2943 "Column y (duplicate) should be Int64"
2944 );
2945 assert!(
2946 matches!(
2947 schema.field(1).data_type(),
2948 arrow_schema::DataType::RunEndEncoded(_, _)
2949 ),
2950 "_file column should use RunEndEncoded type"
2951 );
2952 assert!(
2953 matches!(
2954 schema.field(4).data_type(),
2955 arrow_schema::DataType::RunEndEncoded(_, _)
2956 ),
2957 "_file column (duplicate) should use RunEndEncoded type"
2958 );
2959 }
2960
2961 fn table_with_data_column(column_name: &str) -> Table {
2967 use crate::spec::{
2968 FormatVersion, MAIN_BRANCH, Operation, Snapshot, SnapshotReference, SnapshotRetention,
2969 SortOrder, Summary, TableMetadataBuilder, UnboundPartitionSpec,
2970 };
2971
2972 let schema = Schema::builder()
2973 .with_schema_id(0)
2974 .with_fields(vec![
2975 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2976 NestedField::required(2, column_name, Type::Primitive(PrimitiveType::Int)).into(),
2977 ])
2978 .build()
2979 .unwrap();
2980
2981 let snapshot = Snapshot::builder()
2982 .with_snapshot_id(1)
2983 .with_timestamp_ms(1)
2984 .with_sequence_number(0)
2985 .with_schema_id(0)
2986 .with_manifest_list("/snap-1.avro")
2987 .with_summary(Summary {
2988 operation: Operation::Append,
2989 additional_properties: HashMap::new(),
2990 })
2991 .build();
2992
2993 let metadata = TableMetadataBuilder::new(
2994 schema,
2995 UnboundPartitionSpec::builder().with_spec_id(0).build(),
2996 SortOrder::unsorted_order(),
2997 "s3://bucket/table".to_string(),
2998 FormatVersion::V2,
2999 HashMap::new(),
3000 )
3001 .unwrap()
3002 .add_snapshot(snapshot)
3003 .unwrap()
3004 .set_ref(MAIN_BRANCH, SnapshotReference {
3005 snapshot_id: 1,
3006 retention: SnapshotRetention::Branch {
3007 min_snapshots_to_keep: None,
3008 max_snapshot_age_ms: None,
3009 max_ref_age_ms: None,
3010 },
3011 })
3012 .unwrap()
3013 .build()
3014 .unwrap()
3015 .metadata;
3016
3017 Table::builder()
3018 .metadata(metadata)
3019 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
3020 .file_io(FileIO::new_with_fs())
3021 .runtime(test_runtime())
3022 .build()
3023 .unwrap()
3024 }
3025
3026 #[test]
3030 fn test_scan_projects_data_column_named_like_delete_file_column() {
3031 for column_name in ["pos", "file_path"] {
3032 let table = table_with_data_column(column_name);
3033
3034 let table_scan = table
3037 .scan()
3038 .select([column_name])
3039 .build()
3040 .unwrap_or_else(|e| panic!("scan of data column `{column_name}` failed: {e}"));
3041
3042 assert_eq!(
3043 table_scan.plan_context.as_ref().unwrap().field_ids.as_ref(),
3044 &[2]
3045 );
3046
3047 let default_scan = table.scan().build().unwrap();
3050 assert_eq!(
3051 default_scan
3052 .plan_context
3053 .as_ref()
3054 .unwrap()
3055 .field_ids
3056 .as_ref(),
3057 &[1, 2]
3058 );
3059 }
3060 }
3061
3062 #[test]
3065 fn test_scan_rejects_unknown_column_named_like_delete_file_column() {
3066 let table = table_with_data_column("file_path");
3068
3069 let err = table
3070 .scan()
3071 .select(["pos"])
3072 .build()
3073 .expect_err("projecting an absent column should fail");
3074 assert_eq!(err.kind(), ErrorKind::DataInvalid);
3075 assert!(err.to_string().contains("not found"));
3076 }
3077
3078 #[tokio::test]
3079 async fn test_scan_deadlock() {
3080 let mut fixture = TableTestFixture::new();
3081 fixture.setup_deadlock_manifests().await;
3082
3083 let table_scan = fixture
3090 .table
3091 .scan()
3092 .with_concurrency_limit(1)
3093 .build()
3094 .unwrap();
3095
3096 let result = tokio::time::timeout(std::time::Duration::from_secs(5), async {
3099 table_scan
3100 .plan_files()
3101 .await
3102 .unwrap()
3103 .try_collect::<Vec<_>>()
3104 .await
3105 })
3106 .await;
3107
3108 assert!(result.is_ok(), "Scan timed out - deadlock detected");
3110 }
3111
3112 #[tokio::test]
3113 async fn test_select_with_spec_id_column() {
3114 let mut fixture = TableTestFixture::new();
3115 fixture.setup_manifest_files().await;
3116
3117 let table_scan = fixture
3119 .table
3120 .scan()
3121 .select(["x", RESERVED_COL_NAME_SPEC_ID, "z"])
3122 .with_row_selection_enabled(true)
3123 .build()
3124 .unwrap();
3125
3126 let batch_stream = table_scan.to_arrow().await.unwrap();
3127 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
3128
3129 assert_eq!(batches[0].num_columns(), 3);
3131
3132 let col1 = batches[0].column_by_name("x").unwrap();
3134 let int64_arr = col1.as_any().downcast_ref::<Int64Array>().unwrap();
3135 assert_eq!(int64_arr.value(0), 1);
3136
3137 let spec_id_col = batches[0].column_by_name(RESERVED_COL_NAME_SPEC_ID);
3139 assert!(
3140 spec_id_col.is_some(),
3141 "_spec_id column should be present in the batch"
3142 );
3143
3144 let spec_id_col = spec_id_col.unwrap();
3146 assert!(
3147 matches!(
3148 spec_id_col.data_type(),
3149 arrow_schema::DataType::RunEndEncoded(_, _)
3150 ),
3151 "_spec_id column should use RunEndEncoded type"
3152 );
3153
3154 let run_array = spec_id_col
3156 .as_any()
3157 .downcast_ref::<RunArray<Int32Type>>()
3158 .expect("_spec_id column should be a RunArray");
3159
3160 let values = run_array.values();
3161 let int_values = values.as_primitive::<Int32Type>();
3162 assert_eq!(int_values.len(), 1, "Should have a single _spec_id");
3163
3164 let spec_id = int_values.value(0);
3165 assert_eq!(spec_id, 0, "_spec_id should be 0, got: {spec_id}");
3166
3167 assert!(batches[0].column_by_name("z").is_some());
3169 }
3170
3171 #[tokio::test]
3172 async fn test_select_with_last_updated_sequence_number_column() {
3173 let mut fixture = TableTestFixture::new();
3178 fixture.setup_manifest_files().await;
3179
3180 let table_scan = fixture
3181 .table
3182 .scan()
3183 .select(["x", RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER])
3184 .with_row_selection_enabled(true)
3185 .build()
3186 .unwrap();
3187
3188 let batches: Vec<_> = table_scan
3189 .to_arrow()
3190 .await
3191 .unwrap()
3192 .try_collect()
3193 .await
3194 .unwrap();
3195
3196 assert_last_updated_seq_all(&batches, None);
3198 }
3199
3200 #[tokio::test]
3201 async fn test_select_with_last_updated_sequence_number_column_v3() {
3202 let mut fixture = TableTestFixture::new();
3208 fixture.setup_v3_manifest_files().await;
3209
3210 let expected_seq = fixture
3213 .table
3214 .metadata()
3215 .current_snapshot()
3216 .unwrap()
3217 .sequence_number();
3218
3219 let table_scan = fixture
3220 .table
3221 .scan()
3222 .select(["x", RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER])
3223 .with_row_selection_enabled(true)
3224 .build()
3225 .unwrap();
3226
3227 let batches: Vec<_> = table_scan
3228 .to_arrow()
3229 .await
3230 .unwrap()
3231 .try_collect()
3232 .await
3233 .unwrap();
3234
3235 assert_last_updated_seq_all(&batches, Some(expected_seq));
3236 }
3237
3238 #[tokio::test]
3239 async fn test_select_with_spec_id_column_from_unpartitioned_table() {
3240 let mut fixture = TableTestFixture::new_unpartitioned();
3241 fixture.setup_unpartitioned_manifest_files().await;
3242
3243 let table_scan = fixture
3245 .table
3246 .scan()
3247 .select(["x", RESERVED_COL_NAME_SPEC_ID])
3248 .with_row_selection_enabled(true)
3249 .build()
3250 .unwrap();
3251
3252 let batch_stream = table_scan.to_arrow().await.unwrap();
3253 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
3254
3255 assert_eq!(batches[0].num_columns(), 2);
3257
3258 let spec_id_col = batches[0].column_by_name(RESERVED_COL_NAME_SPEC_ID);
3260 assert!(
3261 spec_id_col.is_some(),
3262 "_spec_id column should be present in the batch"
3263 );
3264
3265 let spec_id_col = spec_id_col.unwrap();
3267 assert!(
3268 matches!(
3269 spec_id_col.data_type(),
3270 arrow_schema::DataType::RunEndEncoded(_, _)
3271 ),
3272 "_spec_id column should use RunEndEncoded type"
3273 );
3274
3275 let run_array = spec_id_col
3277 .as_any()
3278 .downcast_ref::<RunArray<Int32Type>>()
3279 .expect("_spec_id column should be a RunArray");
3280
3281 let values = run_array.values();
3282 let int_values = values.as_primitive::<Int32Type>();
3283 assert_eq!(int_values.len(), 1, "Should have a single _spec_id");
3284
3285 let spec_id = int_values.value(0);
3286 assert_eq!(spec_id, 0, "_spec_id should be 0, got: {spec_id}");
3287 }
3288
3289 #[tokio::test]
3290 async fn test_select_with_spec_id_column_with_partition_evolution() {
3291 let mut fixture = TableTestFixture::new_with_partition_evolution();
3292 fixture
3293 .setup_manifest_files_with_partition_evolution()
3294 .await;
3295
3296 let table_scan = fixture
3298 .table
3299 .scan()
3300 .select(["x", RESERVED_COL_NAME_SPEC_ID, "z"])
3301 .with_row_selection_enabled(true)
3302 .build()
3303 .unwrap();
3304
3305 let batch_stream = table_scan.to_arrow().await.unwrap();
3306 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
3307
3308 let col1 = batches[0].column_by_name("x").unwrap();
3310 let int64_arr = col1.as_any().downcast_ref::<Int64Array>().unwrap();
3311 assert_eq!(int64_arr.value(0), 1);
3312
3313 let spec_id_col = batches[0].column_by_name(RESERVED_COL_NAME_SPEC_ID);
3315 assert!(
3316 spec_id_col.is_some(),
3317 "_spec_id column should be present in the batch"
3318 );
3319
3320 let spec_id_col = spec_id_col.unwrap();
3322 assert!(
3323 matches!(
3324 spec_id_col.data_type(),
3325 arrow_schema::DataType::RunEndEncoded(_, _)
3326 ),
3327 "_spec_id column should use RunEndEncoded type"
3328 );
3329
3330 let run_array = spec_id_col
3332 .as_any()
3333 .downcast_ref::<RunArray<Int32Type>>()
3334 .expect("_spec_id column should be a RunArray");
3335
3336 let values = run_array.values();
3337 let int_values = values.as_primitive::<Int32Type>();
3338 assert_eq!(int_values.len(), 1, "Should have a single _spec_id");
3339
3340 let spec_id = int_values.value(0);
3341 assert_eq!(spec_id, 2, "_spec_id should be 2, got: {spec_id}");
3342 }
3343
3344 #[tokio::test]
3345 async fn test_select_with_pos_and_file_columns() {
3346 use arrow_array::cast::AsArray;
3347
3348 let mut fixture = TableTestFixture::new();
3349 fixture.setup_manifest_files().await;
3350
3351 let table_scan = fixture
3353 .table
3354 .scan()
3355 .select(["x", RESERVED_COL_NAME_POS, RESERVED_COL_NAME_FILE])
3356 .with_row_selection_enabled(true)
3357 .build()
3358 .unwrap();
3359
3360 let batch_stream = table_scan.to_arrow().await.unwrap();
3361 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
3362 assert_eq!(batches.len(), 2);
3363
3364 for batch in batches.iter() {
3366 assert_eq!(batch.num_columns(), 3);
3368
3369 let x_col = batch.column_by_name("x").unwrap();
3371 let x_arr = x_col.as_primitive::<arrow_array::types::Int64Type>();
3372 assert_eq!(x_arr.value(0), 1);
3373
3374 let pos_col = batch.column(1);
3376 let pos_array: &Int64Array = pos_col
3377 .as_any()
3378 .downcast_ref::<Int64Array>()
3379 .expect("_pos column should be a Int64Array");
3380 assert_eq!(*pos_array, Int64Array::from_iter_values(0i64..1024));
3381
3382 let file_col = batch.column_by_name(RESERVED_COL_NAME_FILE);
3384 assert!(
3385 file_col.is_some(),
3386 "_file column should be present in the batch"
3387 );
3388 }
3389 }
3390
3391 #[tokio::test]
3392 async fn test_pos_column_at_start_with_filters() {
3393 let mut fixture = TableTestFixture::new();
3394 fixture.setup_manifest_files().await;
3395
3396 let predicate = Reference::new("y")
3398 .greater_than(Datum::long(4i64))
3399 .and(Reference::new("y").less_than_or_equal_to(Datum::long(5i64)));
3400 let table_scan = fixture
3402 .table
3403 .scan()
3404 .select([RESERVED_COL_NAME_POS, "x", "y"])
3405 .with_filter(predicate)
3406 .with_row_selection_enabled(true)
3407 .build()
3408 .unwrap();
3409
3410 let batch_stream = table_scan.to_arrow().await.unwrap();
3411 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
3412 assert_eq!(batches.len(), 2);
3413
3414 for batch in batches.iter() {
3416 assert_eq!(batch.num_columns(), 3);
3417 assert_eq!(batch.num_rows(), 12);
3418
3419 let schema = batch.schema();
3421 assert_eq!(schema.field(0).name(), RESERVED_COL_NAME_POS);
3422 assert_eq!(schema.field(1).name(), "x");
3423 assert_eq!(schema.field(2).name(), "y");
3424
3425 let pos_col = batch.column(0);
3426 let pos_array: &Int64Array = pos_col
3427 .as_any()
3428 .downcast_ref::<Int64Array>()
3429 .expect("_pos column should be a Int64Array");
3430 assert_eq!(*pos_array, Int64Array::from_iter_values(1012i64..1024));
3431 }
3432 }
3433
3434 #[tokio::test]
3435 async fn test_repeated_pos_column_with_filter() {
3436 let mut fixture = TableTestFixture::new();
3437 fixture.setup_manifest_files().await;
3438
3439 let predicate = Reference::new("a").not_starts_with(Datum::string("Apa"));
3441 let table_scan = fixture
3443 .table
3444 .scan()
3445 .select([RESERVED_COL_NAME_POS, "a", RESERVED_COL_NAME_POS, "x"])
3446 .with_row_selection_enabled(true)
3447 .with_filter(predicate)
3448 .build()
3449 .unwrap();
3450
3451 let batch_stream = table_scan.to_arrow().await.unwrap();
3452
3453 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
3454 assert_eq!(batches.len(), 2);
3455
3456 for batch in batches.iter() {
3458 assert_eq!(batch.num_rows(), 512);
3459
3460 let pos_col = batch
3462 .column_by_name("_pos")
3463 .expect("_pos column should be present in the batch");
3464 let pos_array: &Int64Array = pos_col
3465 .as_any()
3466 .downcast_ref::<Int64Array>()
3467 .expect("_pos column should be a Int64Array");
3468 assert_eq!(*pos_array, Int64Array::from_iter_values(512i64..1024));
3469
3470 let pos_col = batch.column(2);
3472 let pos_array: &Int64Array = pos_col
3473 .as_any()
3474 .downcast_ref::<Int64Array>()
3475 .expect("_pos column should be a Int64Array");
3476 assert_eq!(*pos_array, Int64Array::from_iter_values(512i64..1024));
3477 }
3478 }
3479
3480 #[tokio::test]
3484 async fn test_pos_across_row_groups_via_table_scan() {
3485 let mut fixture = TableTestFixture::new();
3486 fixture.setup_multi_row_group_manifest(&[]).await;
3487
3488 let tasks: Vec<_> = fixture
3491 .table
3492 .scan()
3493 .select(["x", RESERVED_COL_NAME_POS])
3494 .build()
3495 .unwrap()
3496 .plan_files()
3497 .await
3498 .unwrap()
3499 .try_collect()
3500 .await
3501 .unwrap();
3502 assert_eq!(tasks.len(), 1, "expected a single FileScanTask");
3503 let task = &tasks[0];
3504 assert!(
3505 task.project_field_ids.contains(&RESERVED_FIELD_ID_POS),
3506 "_pos field id must be projected into the FileScanTask"
3507 );
3508 assert_eq!(task.start, 0, "TableScan should plan whole-file tasks");
3509 assert_eq!(task.length, task.file_size_in_bytes);
3510 assert!(task.deletes.is_empty());
3511
3512 let batches: Vec<_> = fixture
3514 .table
3515 .scan()
3516 .select(["x", RESERVED_COL_NAME_POS])
3517 .build()
3518 .unwrap()
3519 .to_arrow()
3520 .await
3521 .unwrap()
3522 .try_collect()
3523 .await
3524 .unwrap();
3525
3526 let pos: Vec<i64> = batches
3527 .iter()
3528 .flat_map(|b| {
3529 b.column_by_name(RESERVED_COL_NAME_POS)
3530 .expect("_pos column should be present")
3531 .as_any()
3532 .downcast_ref::<Int64Array>()
3533 .expect("_pos column should be a Int64Array")
3534 .values()
3535 .to_vec()
3536 })
3537 .collect();
3538 assert_eq!(pos, (0..300).collect::<Vec<i64>>());
3539
3540 let x: Vec<i64> = batches
3542 .iter()
3543 .flat_map(|b| {
3544 b.column_by_name("x")
3545 .unwrap()
3546 .as_primitive::<arrow_array::types::Int64Type>()
3547 .values()
3548 .to_vec()
3549 })
3550 .collect();
3551 assert_eq!(x, (1000..1300).collect::<Vec<i64>>());
3552 }
3553
3554 #[tokio::test]
3557 async fn test_pos_with_positional_deletes_via_table_scan() {
3558 let mut fixture = TableTestFixture::new();
3559 fixture.setup_multi_row_group_manifest(&[150, 299]).await;
3561
3562 let tasks: Vec<_> = fixture
3564 .table
3565 .scan()
3566 .select(["x", RESERVED_COL_NAME_POS])
3567 .build()
3568 .unwrap()
3569 .plan_files()
3570 .await
3571 .unwrap()
3572 .try_collect()
3573 .await
3574 .unwrap();
3575 assert_eq!(tasks.len(), 1);
3576 assert_eq!(
3577 tasks[0].deletes.len(),
3578 1,
3579 "positional delete file should be planned into the task"
3580 );
3581 assert_eq!(
3582 tasks[0].deletes[0].file_type,
3583 DataContentType::PositionDeletes
3584 );
3585
3586 let batches: Vec<_> = fixture
3588 .table
3589 .scan()
3590 .select(["x", RESERVED_COL_NAME_POS])
3591 .build()
3592 .unwrap()
3593 .to_arrow()
3594 .await
3595 .unwrap()
3596 .try_collect()
3597 .await
3598 .unwrap();
3599
3600 let pos: Vec<i64> = batches
3601 .iter()
3602 .flat_map(|b| {
3603 b.column_by_name(RESERVED_COL_NAME_POS)
3604 .expect("_pos column should be present")
3605 .as_any()
3606 .downcast_ref::<Int64Array>()
3607 .expect("_pos column should be a Int64Array")
3608 .values()
3609 .to_vec()
3610 })
3611 .collect();
3612
3613 let total: usize = batches.iter().map(|b| b.num_rows()).sum();
3614 assert_eq!(
3615 total, 298,
3616 "two rows should be removed by positional deletes"
3617 );
3618 assert!(!pos.contains(&150) && !pos.contains(&299), "got {pos:?}");
3619 let expected: Vec<i64> = (0..150).chain(151..299).collect();
3620 assert_eq!(pos, expected);
3621 }
3622
3623 #[tokio::test]
3631 async fn test_pos_reads_only_middle_row_group_via_filter() {
3632 let mut fixture = TableTestFixture::new();
3633 fixture.setup_multi_row_group_manifest(&[]).await;
3634
3635 let predicate = Reference::new("y")
3637 .greater_than_or_equal_to(Datum::long(1100))
3638 .and(Reference::new("y").less_than(Datum::long(1200)));
3639
3640 let batches: Vec<_> = fixture
3641 .table
3642 .scan()
3643 .select(["y", RESERVED_COL_NAME_POS])
3644 .with_filter(predicate)
3645 .with_row_group_filtering_enabled(true)
3646 .build()
3647 .unwrap()
3648 .to_arrow()
3649 .await
3650 .unwrap()
3651 .try_collect()
3652 .await
3653 .unwrap();
3654
3655 let total: usize = batches.iter().map(|b| b.num_rows()).sum();
3656 assert_eq!(total, 100, "only the middle row group should be read");
3657
3658 let pos: Vec<i64> = batches
3659 .iter()
3660 .flat_map(|b| {
3661 b.column_by_name(RESERVED_COL_NAME_POS)
3662 .expect("_pos column should be present")
3663 .as_any()
3664 .downcast_ref::<Int64Array>()
3665 .expect("_pos column should be a Int64Array")
3666 .values()
3667 .to_vec()
3668 })
3669 .collect();
3670 assert_eq!(
3671 pos,
3672 (100..200).collect::<Vec<i64>>(),
3673 "_pos must be file-absolute for the middle row group"
3674 );
3675
3676 let y: Vec<i64> = batches
3678 .iter()
3679 .flat_map(|b| {
3680 b.column_by_name("y")
3681 .unwrap()
3682 .as_primitive::<arrow_array::types::Int64Type>()
3683 .values()
3684 .to_vec()
3685 })
3686 .collect();
3687 assert_eq!(y, (1100..1200).collect::<Vec<i64>>());
3688 }
3689}