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