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::{get_metadata_field_id, is_metadata_column_name};
41use crate::runtime::Runtime;
42use crate::spec::{DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, NameMapping, SnapshotRef};
43use crate::table::Table;
44use crate::util::available_parallelism;
45use crate::{Error, ErrorKind, Result};
46
47pub type ArrowRecordBatchStream = BoxStream<'static, Result<RecordBatch>>;
49
50pub struct TableScanBuilder<'a> {
52 table: &'a Table,
53 column_names: Option<Vec<String>>,
55 snapshot_id: Option<i64>,
56 batch_size: Option<usize>,
57 case_sensitive: bool,
58 filter: Option<Predicate>,
59 concurrency_limit_data_files: usize,
60 concurrency_limit_manifest_entries: usize,
61 concurrency_limit_manifest_files: usize,
62 row_group_filtering_enabled: bool,
63 row_selection_enabled: bool,
64}
65
66impl<'a> TableScanBuilder<'a> {
67 pub(crate) fn new(table: &'a Table) -> Self {
68 let num_cpus = available_parallelism().get();
69
70 Self {
71 table,
72 column_names: None,
73 snapshot_id: None,
74 batch_size: None,
75 case_sensitive: true,
76 filter: None,
77 concurrency_limit_data_files: num_cpus,
78 concurrency_limit_manifest_entries: num_cpus,
79 concurrency_limit_manifest_files: num_cpus,
80 row_group_filtering_enabled: true,
81 row_selection_enabled: false,
82 }
83 }
84
85 pub fn with_batch_size(mut self, batch_size: Option<usize>) -> Self {
88 self.batch_size = batch_size;
89 self
90 }
91
92 pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self {
94 self.case_sensitive = case_sensitive;
95 self
96 }
97
98 pub fn with_filter(mut self, predicate: Predicate) -> Self {
100 self.filter = Some(predicate.rewrite_not());
103 self
104 }
105
106 pub fn select_all(mut self) -> Self {
108 self.column_names = None;
109 self
110 }
111
112 pub fn select_empty(mut self) -> Self {
114 self.column_names = Some(vec![]);
115 self
116 }
117
118 pub fn select(mut self, column_names: impl IntoIterator<Item = impl ToString>) -> Self {
120 self.column_names = Some(
121 column_names
122 .into_iter()
123 .map(|item| item.to_string())
124 .collect(),
125 );
126 self
127 }
128
129 pub fn snapshot_id(mut self, snapshot_id: i64) -> Self {
131 self.snapshot_id = Some(snapshot_id);
132 self
133 }
134
135 pub fn with_concurrency_limit(mut self, limit: usize) -> Self {
138 self.concurrency_limit_manifest_files = limit;
139 self.concurrency_limit_manifest_entries = limit;
140 self.concurrency_limit_data_files = limit;
141 self
142 }
143
144 pub fn with_data_file_concurrency_limit(mut self, limit: usize) -> Self {
146 self.concurrency_limit_data_files = limit;
147 self
148 }
149
150 pub fn with_manifest_entry_concurrency_limit(mut self, limit: usize) -> Self {
152 self.concurrency_limit_manifest_entries = limit;
153 self
154 }
155
156 pub fn with_row_group_filtering_enabled(mut self, row_group_filtering_enabled: bool) -> Self {
165 self.row_group_filtering_enabled = row_group_filtering_enabled;
166 self
167 }
168
169 pub fn with_row_selection_enabled(mut self, row_selection_enabled: bool) -> Self {
184 self.row_selection_enabled = row_selection_enabled;
185 self
186 }
187
188 pub fn build(self) -> Result<TableScan> {
190 let snapshot = match self.snapshot_id {
191 Some(snapshot_id) => self
192 .table
193 .metadata()
194 .snapshot_by_id(snapshot_id)
195 .ok_or_else(|| {
196 Error::new(
197 ErrorKind::DataInvalid,
198 format!("Snapshot with id {snapshot_id} not found"),
199 )
200 })?
201 .clone(),
202 None => {
203 let Some(current_snapshot_id) = self.table.metadata().current_snapshot() else {
204 return Ok(TableScan {
205 batch_size: self.batch_size,
206 column_names: self.column_names,
207 file_io: self.table.file_io().clone(),
208 plan_context: None,
209 concurrency_limit_data_files: self.concurrency_limit_data_files,
210 concurrency_limit_manifest_entries: self.concurrency_limit_manifest_entries,
211 concurrency_limit_manifest_files: self.concurrency_limit_manifest_files,
212 row_group_filtering_enabled: self.row_group_filtering_enabled,
213 row_selection_enabled: self.row_selection_enabled,
214 runtime: self.table.runtime().clone(),
215 });
216 };
217 current_snapshot_id.clone()
218 }
219 };
220
221 let schema = snapshot.schema(self.table.metadata())?;
222
223 if let Some(column_names) = self.column_names.as_ref() {
225 for column_name in column_names {
226 if is_metadata_column_name(column_name) {
228 continue;
229 }
230 if schema.field_by_name(column_name).is_none() {
231 return Err(Error::new(
232 ErrorKind::DataInvalid,
233 format!("Column {column_name} not found in table. Schema: {schema}"),
234 ));
235 }
236 }
237 }
238
239 let mut field_ids = vec![];
240 let column_names = self.column_names.clone().unwrap_or_else(|| {
241 schema
242 .as_struct()
243 .fields()
244 .iter()
245 .map(|f| f.name.clone())
246 .collect()
247 });
248
249 for column_name in column_names.iter() {
250 if is_metadata_column_name(column_name) {
252 field_ids.push(get_metadata_field_id(column_name)?);
253 continue;
254 }
255
256 let field_id = schema.field_id_by_name(column_name).ok_or_else(|| {
257 Error::new(
258 ErrorKind::DataInvalid,
259 format!("Column {column_name} not found in table. Schema: {schema}"),
260 )
261 })?;
262
263 schema
264 .as_struct()
265 .field_by_id(field_id)
266 .ok_or_else(|| {
267 Error::new(
268 ErrorKind::FeatureUnsupported,
269 format!(
270 "Column {column_name} is not a direct child of schema but a nested field, which is not supported now. Schema: {schema}"
271 ),
272 )
273 })?;
274
275 field_ids.push(field_id);
276 }
277
278 let snapshot_bound_predicate = if let Some(ref predicates) = self.filter {
279 Some(predicates.bind(schema.clone(), true)?)
280 } else {
281 None
282 };
283
284 let name_mapping = self
285 .table
286 .metadata()
287 .properties()
288 .get(DEFAULT_SCHEMA_NAME_MAPPING)
289 .map(|raw| {
290 serde_json::from_str::<NameMapping>(raw).map_err(|e| {
291 Error::new(
292 ErrorKind::DataInvalid,
293 format!(
294 "Failed to parse table property {DEFAULT_SCHEMA_NAME_MAPPING} as a NameMapping"
295 ),
296 )
297 .with_source(e)
298 })
299 })
300 .transpose()?
301 .map(Arc::new);
302
303 let plan_context = PlanContext {
304 snapshot,
305 table_metadata: self.table.metadata_ref(),
306 snapshot_schema: schema,
307 case_sensitive: self.case_sensitive,
308 predicate: self.filter.map(Arc::new),
309 snapshot_bound_predicate: snapshot_bound_predicate.map(Arc::new),
310 object_cache: self.table.object_cache(),
311 field_ids: Arc::new(field_ids),
312 name_mapping,
313 partition_filter_cache: Arc::new(PartitionFilterCache::new()),
314 manifest_evaluator_cache: Arc::new(ManifestEvaluatorCache::new()),
315 expression_evaluator_cache: Arc::new(ExpressionEvaluatorCache::new()),
316 };
317
318 Ok(TableScan {
319 batch_size: self.batch_size,
320 column_names: self.column_names,
321 file_io: self.table.file_io().clone(),
322 plan_context: Some(plan_context),
323 concurrency_limit_data_files: self.concurrency_limit_data_files,
324 concurrency_limit_manifest_entries: self.concurrency_limit_manifest_entries,
325 concurrency_limit_manifest_files: self.concurrency_limit_manifest_files,
326 row_group_filtering_enabled: self.row_group_filtering_enabled,
327 row_selection_enabled: self.row_selection_enabled,
328 runtime: self.table.runtime().clone(),
329 })
330 }
331}
332
333#[derive(Debug)]
335pub struct TableScan {
336 plan_context: Option<PlanContext>,
340 batch_size: Option<usize>,
341 file_io: FileIO,
342 column_names: Option<Vec<String>>,
343 concurrency_limit_manifest_files: usize,
346
347 concurrency_limit_manifest_entries: usize,
350
351 concurrency_limit_data_files: usize,
354
355 row_group_filtering_enabled: bool,
356 row_selection_enabled: bool,
357
358 runtime: Runtime,
359}
360
361impl TableScan {
362 pub async fn plan_files(&self) -> Result<FileScanTaskStream> {
364 let Some(plan_context) = self.plan_context.as_ref() else {
365 return Ok(Box::pin(futures::stream::empty()));
366 };
367
368 let concurrency_limit_manifest_files = self.concurrency_limit_manifest_files;
369 let concurrency_limit_manifest_entries = self.concurrency_limit_manifest_entries;
370
371 let (manifest_entry_data_ctx_tx, manifest_entry_data_ctx_rx) =
373 channel(concurrency_limit_manifest_files);
374 let (manifest_entry_delete_ctx_tx, manifest_entry_delete_ctx_rx) =
375 channel(concurrency_limit_manifest_files);
376
377 let (file_scan_task_tx, file_scan_task_rx) = channel(concurrency_limit_manifest_entries);
379
380 let (delete_file_idx, delete_file_tx) = DeleteFileIndex::new(self.runtime.clone());
381
382 let manifest_list = plan_context.get_manifest_list().await?;
383
384 let manifest_file_contexts = plan_context.build_manifest_file_contexts(
388 manifest_list,
389 manifest_entry_data_ctx_tx,
390 delete_file_idx.clone(),
391 manifest_entry_delete_ctx_tx,
392 )?;
393
394 let mut channel_for_manifest_error = file_scan_task_tx.clone();
395 let mut channel_for_data_manifest_entry_error = file_scan_task_tx.clone();
396 let mut channel_for_delete_manifest_entry_error = file_scan_task_tx.clone();
397
398 let rt = self.runtime.clone();
399
400 rt.io().spawn(async move {
402 let result = futures::stream::iter(manifest_file_contexts)
403 .try_for_each_concurrent(concurrency_limit_manifest_files, |ctx| async move {
404 ctx.fetch_manifest_and_stream_manifest_entries().await
405 })
406 .await;
407
408 if let Err(error) = result {
409 let _ = channel_for_manifest_error.send(Err(error)).await;
410 }
411 });
412
413 {
415 let rt = rt.clone();
416 let rt_inner = rt.clone();
417 rt.cpu().spawn(async move {
418 let result = manifest_entry_delete_ctx_rx
419 .map(|me_ctx| Ok((me_ctx, delete_file_tx.clone())))
420 .try_for_each_concurrent(
421 concurrency_limit_manifest_entries,
422 |(manifest_entry_context, tx)| {
423 let rt_inner = rt_inner.clone();
424 async move {
425 rt_inner
426 .cpu()
427 .spawn(async move {
428 Self::process_delete_manifest_entry(
429 manifest_entry_context,
430 tx,
431 )
432 .await
433 })
434 .await?
435 }
436 },
437 )
438 .await;
439
440 if let Err(error) = result {
441 let _ = channel_for_delete_manifest_entry_error
442 .send(Err(error))
443 .await;
444 }
445 });
446 }
447
448 {
450 let rt_inner = rt.clone();
451 rt.cpu().spawn(async move {
452 let result = manifest_entry_data_ctx_rx
453 .map(|me_ctx| Ok((me_ctx, file_scan_task_tx.clone())))
454 .try_for_each_concurrent(
455 concurrency_limit_manifest_entries,
456 |(manifest_entry_context, tx)| {
457 let rt_inner = rt_inner.clone();
458 async move {
459 rt_inner
460 .cpu()
461 .spawn(async move {
462 Self::process_data_manifest_entry(
463 manifest_entry_context,
464 tx,
465 )
466 .await
467 })
468 .await?
469 }
470 },
471 )
472 .await;
473
474 if let Err(error) = result {
475 let _ = channel_for_data_manifest_entry_error.send(Err(error)).await;
476 }
477 });
478 }
479
480 Ok(file_scan_task_rx.boxed())
481 }
482
483 pub async fn to_arrow(&self) -> Result<ArrowRecordBatchStream> {
485 let mut arrow_reader_builder =
486 ArrowReaderBuilder::new(self.file_io.clone(), self.runtime.clone())
487 .with_data_file_concurrency_limit(self.concurrency_limit_data_files)
488 .with_row_group_filtering_enabled(self.row_group_filtering_enabled)
489 .with_row_selection_enabled(self.row_selection_enabled);
490
491 if let Some(batch_size) = self.batch_size {
492 arrow_reader_builder = arrow_reader_builder.with_batch_size(batch_size);
493 }
494
495 arrow_reader_builder
496 .build()
497 .read(self.plan_files().await?)
498 .map(|result| result.stream())
499 }
500
501 pub fn column_names(&self) -> Option<&[String]> {
503 self.column_names.as_deref()
504 }
505
506 pub fn snapshot(&self) -> Option<&SnapshotRef> {
508 self.plan_context.as_ref().map(|x| &x.snapshot)
509 }
510
511 async fn process_data_manifest_entry(
512 manifest_entry_context: ManifestEntryContext,
513 mut file_scan_task_tx: Sender<Result<FileScanTask>>,
514 ) -> Result<()> {
515 if !manifest_entry_context.manifest_entry.is_alive() {
517 return Ok(());
518 }
519
520 if manifest_entry_context.manifest_entry.content_type() != DataContentType::Data {
522 return Err(Error::new(
523 ErrorKind::FeatureUnsupported,
524 "Encountered an entry for a delete file in a data file manifest",
525 ));
526 }
527
528 if let Some(ref bound_predicates) = manifest_entry_context.bound_predicates {
529 let BoundPredicates {
530 snapshot_bound_predicate,
531 partition_bound_predicate,
532 } = bound_predicates.as_ref();
533
534 let expression_evaluator_cache =
535 manifest_entry_context.expression_evaluator_cache.as_ref();
536
537 let expression_evaluator = expression_evaluator_cache.get(
538 manifest_entry_context.partition_spec_id,
539 partition_bound_predicate,
540 )?;
541
542 if !expression_evaluator.eval(manifest_entry_context.manifest_entry.data_file())? {
545 return Ok(());
546 }
547
548 if !InclusiveMetricsEvaluator::eval(
550 snapshot_bound_predicate,
551 manifest_entry_context.manifest_entry.data_file(),
552 false,
553 )? {
554 return Ok(());
555 }
556 }
557
558 file_scan_task_tx
562 .send(Ok(manifest_entry_context.into_file_scan_task().await?))
563 .await?;
564
565 Ok(())
566 }
567
568 async fn process_delete_manifest_entry(
569 manifest_entry_context: ManifestEntryContext,
570 mut delete_file_ctx_tx: Sender<DeleteFileContext>,
571 ) -> Result<()> {
572 if !manifest_entry_context.manifest_entry.is_alive() {
574 return Ok(());
575 }
576
577 if manifest_entry_context.manifest_entry.content_type() == DataContentType::Data {
579 return Err(Error::new(
580 ErrorKind::FeatureUnsupported,
581 "Encountered an entry for a data file in a delete manifest",
582 ));
583 }
584
585 if let Some(ref bound_predicates) = manifest_entry_context.bound_predicates {
586 let expression_evaluator_cache =
587 manifest_entry_context.expression_evaluator_cache.as_ref();
588
589 let expression_evaluator = expression_evaluator_cache.get(
590 manifest_entry_context.partition_spec_id,
591 &bound_predicates.partition_bound_predicate,
592 )?;
593
594 if !expression_evaluator.eval(manifest_entry_context.manifest_entry.data_file())? {
597 return Ok(());
598 }
599 }
600
601 delete_file_ctx_tx
602 .send(DeleteFileContext {
603 manifest_entry: manifest_entry_context.manifest_entry.clone(),
604 partition_spec_id: manifest_entry_context.partition_spec_id,
605 })
606 .await?;
607
608 Ok(())
609 }
610}
611
612pub(crate) struct BoundPredicates {
613 partition_bound_predicate: BoundPredicate,
614 snapshot_bound_predicate: BoundPredicate,
615}
616
617#[cfg(test)]
618pub mod tests {
619 #![allow(missing_docs)]
621
622 use std::collections::HashMap;
623 use std::fs;
624 use std::fs::File;
625 use std::sync::Arc;
626
627 use arrow_array::cast::AsArray;
628 use arrow_array::{
629 Array, ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch,
630 StringArray,
631 };
632 use futures::{TryStreamExt, stream};
633 use minijinja::value::Value;
634 use minijinja::{AutoEscape, Environment, context};
635 use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
636 use parquet::basic::Compression;
637 use parquet::file::properties::WriterProperties;
638 use tempfile::TempDir;
639 use uuid::Uuid;
640
641 use crate::arrow::ArrowReaderBuilder;
642 use crate::expr::{BoundPredicate, Reference};
643 use crate::io::{FileIO, OutputFile};
644 use crate::metadata_columns::RESERVED_COL_NAME_FILE;
645 use crate::scan::FileScanTask;
646 use crate::spec::{
647 DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, DataFileBuilder, DataFileFormat, Datum,
648 Literal, ManifestEntry, ManifestListWriter, ManifestStatus, ManifestWriterBuilder,
649 NestedField, PartitionSpec, PrimitiveType, Schema, Struct, StructType, TableMetadata, Type,
650 };
651 use crate::table::Table;
652 use crate::test_utils::test_runtime;
653 use crate::{ErrorKind, TableIdent};
654
655 fn render_template(template: &str, ctx: Value) -> String {
656 let mut env = Environment::new();
657 env.set_auto_escape_callback(|_| AutoEscape::None);
658 env.render_str(template, ctx).unwrap()
659 }
660
661 pub struct TableTestFixture {
662 pub table_location: String,
663 pub table: Table,
664 }
665
666 impl TableTestFixture {
667 #[allow(clippy::new_without_default)]
668 pub fn new() -> Self {
669 let tmp_dir = TempDir::new().unwrap();
670 let table_location = tmp_dir.path().join("table1");
671 let manifest_list1_location = table_location.join("metadata/manifests_list_1.avro");
672 let manifest_list2_location = table_location.join("metadata/manifests_list_2.avro");
673 let table_metadata1_location = table_location.join("metadata/v1.json");
674
675 let file_io = FileIO::new_with_fs();
676
677 let table_metadata = {
678 let template_json_str = fs::read_to_string(format!(
679 "{}/testdata/example_table_metadata_v2.json",
680 env!("CARGO_MANIFEST_DIR")
681 ))
682 .unwrap();
683 let metadata_json = render_template(&template_json_str, context! {
684 table_location => &table_location,
685 manifest_list_1_location => &manifest_list1_location,
686 manifest_list_2_location => &manifest_list2_location,
687 table_metadata_1_location => &table_metadata1_location,
688 });
689 serde_json::from_str::<TableMetadata>(&metadata_json).unwrap()
690 };
691
692 let table = Table::builder()
693 .metadata(table_metadata)
694 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
695 .file_io(file_io.clone())
696 .metadata_location(table_metadata1_location.as_os_str().to_str().unwrap())
697 .runtime(test_runtime())
698 .build()
699 .unwrap();
700
701 Self {
702 table_location: table_location.to_str().unwrap().to_string(),
703 table,
704 }
705 }
706
707 #[allow(clippy::new_without_default)]
708 pub fn new_empty() -> Self {
709 let tmp_dir = TempDir::new().unwrap();
710 let table_location = tmp_dir.path().join("table1");
711 let table_metadata1_location = table_location.join("metadata/v1.json");
712
713 let file_io = FileIO::new_with_fs();
714
715 let table_metadata = {
716 let template_json_str = fs::read_to_string(format!(
717 "{}/testdata/example_empty_table_metadata_v2.json",
718 env!("CARGO_MANIFEST_DIR")
719 ))
720 .unwrap();
721 let metadata_json = render_template(&template_json_str, context! {
722 table_location => &table_location,
723 table_metadata_1_location => &table_metadata1_location,
724 });
725 serde_json::from_str::<TableMetadata>(&metadata_json).unwrap()
726 };
727
728 let table = Table::builder()
729 .metadata(table_metadata)
730 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
731 .file_io(file_io.clone())
732 .metadata_location(table_metadata1_location.as_os_str().to_str().unwrap())
733 .runtime(test_runtime())
734 .build()
735 .unwrap();
736
737 Self {
738 table_location: table_location.to_str().unwrap().to_string(),
739 table,
740 }
741 }
742
743 pub fn new_with_deep_history() -> Self {
747 let tmp_dir = TempDir::new().unwrap();
748 let table_location = tmp_dir.path().join("table1");
749 let table_metadata1_location = table_location.join("metadata/v1.json");
750
751 let file_io = FileIO::new_with_fs();
752
753 let table_metadata = {
754 let json_str = fs::read_to_string(format!(
755 "{}/testdata/example_table_metadata_v2_deep_history.json",
756 env!("CARGO_MANIFEST_DIR")
757 ))
758 .unwrap();
759 serde_json::from_str::<TableMetadata>(&json_str).unwrap()
760 };
761
762 let table = Table::builder()
763 .metadata(table_metadata)
764 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
765 .file_io(file_io.clone())
766 .metadata_location(table_metadata1_location.as_os_str().to_str().unwrap())
767 .runtime(test_runtime())
768 .build()
769 .unwrap();
770
771 Self {
772 table_location: table_location.to_str().unwrap().to_string(),
773 table,
774 }
775 }
776
777 pub fn new_unpartitioned() -> Self {
778 let tmp_dir = TempDir::new().unwrap();
779 let table_location = tmp_dir.path().join("table1");
780 let manifest_list1_location = table_location.join("metadata/manifests_list_1.avro");
781 let manifest_list2_location = table_location.join("metadata/manifests_list_2.avro");
782 let table_metadata1_location = table_location.join("metadata/v1.json");
783
784 let file_io = FileIO::new_with_fs();
785
786 let mut table_metadata = {
787 let template_json_str = fs::read_to_string(format!(
788 "{}/testdata/example_table_metadata_v2.json",
789 env!("CARGO_MANIFEST_DIR")
790 ))
791 .unwrap();
792 let metadata_json = render_template(&template_json_str, context! {
793 table_location => &table_location,
794 manifest_list_1_location => &manifest_list1_location,
795 manifest_list_2_location => &manifest_list2_location,
796 table_metadata_1_location => &table_metadata1_location,
797 });
798 serde_json::from_str::<TableMetadata>(&metadata_json).unwrap()
799 };
800
801 table_metadata.default_spec = Arc::new(PartitionSpec::unpartition_spec());
802 table_metadata.partition_specs.clear();
803 table_metadata.default_partition_type = StructType::new(vec![]);
804 table_metadata
805 .partition_specs
806 .insert(0, table_metadata.default_spec.clone());
807
808 let table = Table::builder()
809 .metadata(table_metadata)
810 .identifier(TableIdent::from_strs(["db", "table1"]).unwrap())
811 .file_io(file_io.clone())
812 .metadata_location(table_metadata1_location.to_str().unwrap())
813 .runtime(test_runtime())
814 .build()
815 .unwrap();
816
817 Self {
818 table_location: table_location.to_str().unwrap().to_string(),
819 table,
820 }
821 }
822
823 fn next_manifest_file(&self) -> OutputFile {
824 self.table
825 .file_io()
826 .new_output(format!(
827 "{}/metadata/manifest_{}.avro",
828 self.table_location,
829 Uuid::new_v4()
830 ))
831 .unwrap()
832 }
833
834 pub async fn setup_manifest_files(&mut self) {
835 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
836 let parent_snapshot = current_snapshot
837 .parent_snapshot(self.table.metadata())
838 .unwrap();
839 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
840 let current_partition_spec = self.table.metadata().default_partition_spec();
841
842 let parquet_file_size = self.write_parquet_data_files();
844
845 let mut writer = ManifestWriterBuilder::new(
846 self.next_manifest_file(),
847 Some(current_snapshot.snapshot_id()),
848 current_schema.clone(),
849 current_partition_spec.as_ref().clone(),
850 )
851 .build_v2_data();
852 writer
853 .add_entry(
854 ManifestEntry::builder()
855 .status(ManifestStatus::Added)
856 .data_file(
857 DataFileBuilder::default()
858 .partition_spec_id(0)
859 .content(DataContentType::Data)
860 .file_path(format!("{}/1.parquet", &self.table_location))
861 .file_format(DataFileFormat::Parquet)
862 .file_size_in_bytes(parquet_file_size)
863 .record_count(1)
864 .partition(Struct::from_iter([Some(Literal::long(100))]))
865 .key_metadata(None)
866 .build()
867 .unwrap(),
868 )
869 .build(),
870 )
871 .unwrap();
872 writer
873 .add_delete_entry(
874 ManifestEntry::builder()
875 .status(ManifestStatus::Deleted)
876 .snapshot_id(parent_snapshot.snapshot_id())
877 .sequence_number(parent_snapshot.sequence_number())
878 .file_sequence_number(parent_snapshot.sequence_number())
879 .data_file(
880 DataFileBuilder::default()
881 .partition_spec_id(0)
882 .content(DataContentType::Data)
883 .file_path(format!("{}/2.parquet", &self.table_location))
884 .file_format(DataFileFormat::Parquet)
885 .file_size_in_bytes(parquet_file_size)
886 .record_count(1)
887 .partition(Struct::from_iter([Some(Literal::long(200))]))
888 .build()
889 .unwrap(),
890 )
891 .build(),
892 )
893 .unwrap();
894 writer
895 .add_existing_entry(
896 ManifestEntry::builder()
897 .status(ManifestStatus::Existing)
898 .snapshot_id(parent_snapshot.snapshot_id())
899 .sequence_number(parent_snapshot.sequence_number())
900 .file_sequence_number(parent_snapshot.sequence_number())
901 .data_file(
902 DataFileBuilder::default()
903 .partition_spec_id(0)
904 .content(DataContentType::Data)
905 .file_path(format!("{}/3.parquet", &self.table_location))
906 .file_format(DataFileFormat::Parquet)
907 .file_size_in_bytes(parquet_file_size)
908 .record_count(1)
909 .partition(Struct::from_iter([Some(Literal::long(300))]))
910 .build()
911 .unwrap(),
912 )
913 .build(),
914 )
915 .unwrap();
916 let data_file_manifest = writer.write_manifest_file().await.unwrap();
917
918 let manifest_list_writer = self
920 .table
921 .file_io()
922 .new_output(current_snapshot.manifest_list())
923 .unwrap()
924 .writer()
925 .await
926 .unwrap();
927 let mut manifest_list_write = ManifestListWriter::v2(
928 manifest_list_writer,
929 current_snapshot.snapshot_id(),
930 current_snapshot.parent_snapshot_id(),
931 current_snapshot.sequence_number(),
932 );
933 manifest_list_write
934 .add_manifests(vec![data_file_manifest].into_iter())
935 .unwrap();
936 manifest_list_write.close().await.unwrap();
937 }
938
939 fn write_parquet_data_files(&self) -> u64 {
942 std::fs::create_dir_all(&self.table_location).unwrap();
943
944 let schema = {
945 let fields = vec![
946 arrow_schema::Field::new("x", arrow_schema::DataType::Int64, false)
947 .with_metadata(HashMap::from([(
948 PARQUET_FIELD_ID_META_KEY.to_string(),
949 "1".to_string(),
950 )])),
951 arrow_schema::Field::new("y", arrow_schema::DataType::Int64, false)
952 .with_metadata(HashMap::from([(
953 PARQUET_FIELD_ID_META_KEY.to_string(),
954 "2".to_string(),
955 )])),
956 arrow_schema::Field::new("z", arrow_schema::DataType::Int64, false)
957 .with_metadata(HashMap::from([(
958 PARQUET_FIELD_ID_META_KEY.to_string(),
959 "3".to_string(),
960 )])),
961 arrow_schema::Field::new("a", arrow_schema::DataType::Utf8, false)
962 .with_metadata(HashMap::from([(
963 PARQUET_FIELD_ID_META_KEY.to_string(),
964 "4".to_string(),
965 )])),
966 arrow_schema::Field::new("dbl", arrow_schema::DataType::Float64, false)
967 .with_metadata(HashMap::from([(
968 PARQUET_FIELD_ID_META_KEY.to_string(),
969 "5".to_string(),
970 )])),
971 arrow_schema::Field::new("i32", arrow_schema::DataType::Int32, false)
972 .with_metadata(HashMap::from([(
973 PARQUET_FIELD_ID_META_KEY.to_string(),
974 "6".to_string(),
975 )])),
976 arrow_schema::Field::new("i64", arrow_schema::DataType::Int64, false)
977 .with_metadata(HashMap::from([(
978 PARQUET_FIELD_ID_META_KEY.to_string(),
979 "7".to_string(),
980 )])),
981 arrow_schema::Field::new("bool", arrow_schema::DataType::Boolean, false)
982 .with_metadata(HashMap::from([(
983 PARQUET_FIELD_ID_META_KEY.to_string(),
984 "8".to_string(),
985 )])),
986 ];
987 Arc::new(arrow_schema::Schema::new(fields))
988 };
989 let col1 = Arc::new(Int64Array::from_iter_values(vec![1; 1024])) as ArrayRef;
991
992 let mut values = vec![2; 512];
993 values.append(vec![3; 200].as_mut());
994 values.append(vec![4; 300].as_mut());
995 values.append(vec![5; 12].as_mut());
996
997 let col2 = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
999
1000 let mut values = vec![3; 512];
1001 values.append(vec![4; 512].as_mut());
1002
1003 let col3 = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
1005
1006 let mut values = vec!["Apache"; 512];
1008 values.append(vec!["Iceberg"; 512].as_mut());
1009 let col4 = Arc::new(StringArray::from_iter_values(values)) as ArrayRef;
1010
1011 let mut values = vec![100.0f64; 512];
1013 values.append(vec![150.0f64; 12].as_mut());
1014 values.append(vec![200.0f64; 500].as_mut());
1015 let col5 = Arc::new(Float64Array::from_iter_values(values)) as ArrayRef;
1016
1017 let mut values = vec![100i32; 512];
1019 values.append(vec![150i32; 12].as_mut());
1020 values.append(vec![200i32; 500].as_mut());
1021 let col6 = Arc::new(Int32Array::from_iter_values(values)) as ArrayRef;
1022
1023 let mut values = vec![100i64; 512];
1025 values.append(vec![150i64; 12].as_mut());
1026 values.append(vec![200i64; 500].as_mut());
1027 let col7 = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
1028
1029 let mut values = vec![false; 512];
1031 values.append(vec![true; 512].as_mut());
1032 let values: BooleanArray = values.into();
1033 let col8 = Arc::new(values) as ArrayRef;
1034
1035 let to_write = RecordBatch::try_new(schema.clone(), vec![
1036 col1, col2, col3, col4, col5, col6, col7, col8,
1037 ])
1038 .unwrap();
1039
1040 let props = WriterProperties::builder()
1042 .set_compression(Compression::SNAPPY)
1043 .build();
1044
1045 for n in 1..=3 {
1046 let file = File::create(format!("{}/{}.parquet", &self.table_location, n)).unwrap();
1047 let mut writer =
1048 ArrowWriter::try_new(file, to_write.schema(), Some(props.clone())).unwrap();
1049
1050 writer.write(&to_write).expect("Writing batch");
1051
1052 writer.close().unwrap();
1054 }
1055
1056 std::fs::metadata(format!("{}/1.parquet", &self.table_location))
1057 .unwrap()
1058 .len()
1059 }
1060
1061 pub async fn setup_unpartitioned_manifest_files(&mut self) {
1062 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
1063 let parent_snapshot = current_snapshot
1064 .parent_snapshot(self.table.metadata())
1065 .unwrap();
1066 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
1067 let current_partition_spec = Arc::new(PartitionSpec::unpartition_spec());
1068
1069 let parquet_file_size = self.write_parquet_data_files();
1071
1072 let mut writer = ManifestWriterBuilder::new(
1074 self.next_manifest_file(),
1075 Some(current_snapshot.snapshot_id()),
1076 current_schema.clone(),
1077 current_partition_spec.as_ref().clone(),
1078 )
1079 .build_v2_data();
1080
1081 let empty_partition = Struct::empty();
1083
1084 writer
1085 .add_entry(
1086 ManifestEntry::builder()
1087 .status(ManifestStatus::Added)
1088 .data_file(
1089 DataFileBuilder::default()
1090 .partition_spec_id(0)
1091 .content(DataContentType::Data)
1092 .file_path(format!("{}/1.parquet", &self.table_location))
1093 .file_format(DataFileFormat::Parquet)
1094 .file_size_in_bytes(parquet_file_size)
1095 .record_count(1)
1096 .partition(empty_partition.clone())
1097 .key_metadata(None)
1098 .build()
1099 .unwrap(),
1100 )
1101 .build(),
1102 )
1103 .unwrap();
1104
1105 writer
1106 .add_delete_entry(
1107 ManifestEntry::builder()
1108 .status(ManifestStatus::Deleted)
1109 .snapshot_id(parent_snapshot.snapshot_id())
1110 .sequence_number(parent_snapshot.sequence_number())
1111 .file_sequence_number(parent_snapshot.sequence_number())
1112 .data_file(
1113 DataFileBuilder::default()
1114 .partition_spec_id(0)
1115 .content(DataContentType::Data)
1116 .file_path(format!("{}/2.parquet", &self.table_location))
1117 .file_format(DataFileFormat::Parquet)
1118 .file_size_in_bytes(parquet_file_size)
1119 .record_count(1)
1120 .partition(empty_partition.clone())
1121 .build()
1122 .unwrap(),
1123 )
1124 .build(),
1125 )
1126 .unwrap();
1127
1128 writer
1129 .add_existing_entry(
1130 ManifestEntry::builder()
1131 .status(ManifestStatus::Existing)
1132 .snapshot_id(parent_snapshot.snapshot_id())
1133 .sequence_number(parent_snapshot.sequence_number())
1134 .file_sequence_number(parent_snapshot.sequence_number())
1135 .data_file(
1136 DataFileBuilder::default()
1137 .partition_spec_id(0)
1138 .content(DataContentType::Data)
1139 .file_path(format!("{}/3.parquet", &self.table_location))
1140 .file_format(DataFileFormat::Parquet)
1141 .file_size_in_bytes(parquet_file_size)
1142 .record_count(1)
1143 .partition(empty_partition.clone())
1144 .build()
1145 .unwrap(),
1146 )
1147 .build(),
1148 )
1149 .unwrap();
1150
1151 let data_file_manifest = writer.write_manifest_file().await.unwrap();
1152
1153 let manifest_list_writer = self
1155 .table
1156 .file_io()
1157 .new_output(current_snapshot.manifest_list())
1158 .unwrap()
1159 .writer()
1160 .await
1161 .unwrap();
1162 let mut manifest_list_write = ManifestListWriter::v2(
1163 manifest_list_writer,
1164 current_snapshot.snapshot_id(),
1165 current_snapshot.parent_snapshot_id(),
1166 current_snapshot.sequence_number(),
1167 );
1168 manifest_list_write
1169 .add_manifests(vec![data_file_manifest].into_iter())
1170 .unwrap();
1171 manifest_list_write.close().await.unwrap();
1172 }
1173
1174 pub async fn setup_deadlock_manifests(&mut self) {
1175 let current_snapshot = self.table.metadata().current_snapshot().unwrap();
1176 let _parent_snapshot = current_snapshot
1177 .parent_snapshot(self.table.metadata())
1178 .unwrap();
1179 let current_schema = current_snapshot.schema(self.table.metadata()).unwrap();
1180 let current_partition_spec = self.table.metadata().default_partition_spec();
1181
1182 let mut writer = ManifestWriterBuilder::new(
1184 self.next_manifest_file(),
1185 Some(current_snapshot.snapshot_id()),
1186 current_schema.clone(),
1187 current_partition_spec.as_ref().clone(),
1188 )
1189 .build_v2_data();
1190
1191 for i in 0..10 {
1193 writer
1194 .add_entry(
1195 ManifestEntry::builder()
1196 .status(ManifestStatus::Added)
1197 .data_file(
1198 DataFileBuilder::default()
1199 .partition_spec_id(0)
1200 .content(DataContentType::Data)
1201 .file_path(format!("{}/{}.parquet", &self.table_location, i))
1202 .file_format(DataFileFormat::Parquet)
1203 .file_size_in_bytes(100)
1204 .record_count(1)
1205 .partition(Struct::from_iter([Some(Literal::long(100))]))
1206 .key_metadata(None)
1207 .build()
1208 .unwrap(),
1209 )
1210 .build(),
1211 )
1212 .unwrap();
1213 }
1214 let data_manifest = writer.write_manifest_file().await.unwrap();
1215
1216 let mut writer = ManifestWriterBuilder::new(
1218 self.next_manifest_file(),
1219 Some(current_snapshot.snapshot_id()),
1220 current_schema.clone(),
1221 current_partition_spec.as_ref().clone(),
1222 )
1223 .build_v2_deletes();
1224
1225 writer
1226 .add_entry(
1227 ManifestEntry::builder()
1228 .status(ManifestStatus::Added)
1229 .data_file(
1230 DataFileBuilder::default()
1231 .partition_spec_id(0)
1232 .content(DataContentType::PositionDeletes)
1233 .file_path(format!("{}/del.parquet", &self.table_location))
1234 .file_format(DataFileFormat::Parquet)
1235 .file_size_in_bytes(100)
1236 .record_count(1)
1237 .partition(Struct::from_iter([Some(Literal::long(100))]))
1238 .build()
1239 .unwrap(),
1240 )
1241 .build(),
1242 )
1243 .unwrap();
1244 let delete_manifest = writer.write_manifest_file().await.unwrap();
1245
1246 let manifest_list_writer = self
1249 .table
1250 .file_io()
1251 .new_output(current_snapshot.manifest_list())
1252 .unwrap()
1253 .writer()
1254 .await
1255 .unwrap();
1256 let mut manifest_list_write = ManifestListWriter::v2(
1257 manifest_list_writer,
1258 current_snapshot.snapshot_id(),
1259 current_snapshot.parent_snapshot_id(),
1260 current_snapshot.sequence_number(),
1261 );
1262 manifest_list_write
1263 .add_manifests(vec![data_manifest, delete_manifest].into_iter())
1264 .unwrap();
1265 manifest_list_write.close().await.unwrap();
1266 }
1267 }
1268
1269 #[tokio::test]
1270 async fn test_table_scan_columns() {
1271 let table = TableTestFixture::new().table;
1272
1273 let table_scan = table.scan().select(["x", "y"]).build().unwrap();
1274 assert_eq!(
1275 Some(vec!["x".to_string(), "y".to_string()]),
1276 table_scan.column_names
1277 );
1278
1279 let table_scan = table
1280 .scan()
1281 .select(["x", "y"])
1282 .select(["z"])
1283 .build()
1284 .unwrap();
1285 assert_eq!(Some(vec!["z".to_string()]), table_scan.column_names);
1286 }
1287
1288 #[tokio::test]
1289 async fn test_select_all() {
1290 let table = TableTestFixture::new().table;
1291
1292 let table_scan = table.scan().select_all().build().unwrap();
1293 assert!(table_scan.column_names.is_none());
1294 }
1295
1296 #[test]
1297 fn test_select_no_exist_column() {
1298 let table = TableTestFixture::new().table;
1299
1300 let table_scan = table.scan().select(["x", "y", "z", "a", "b"]).build();
1301 assert!(table_scan.is_err());
1302 }
1303
1304 #[tokio::test]
1305 async fn test_table_scan_default_snapshot_id() {
1306 let table = TableTestFixture::new().table;
1307
1308 let table_scan = table.scan().build().unwrap();
1309 assert_eq!(
1310 table.metadata().current_snapshot().unwrap().snapshot_id(),
1311 table_scan.snapshot().unwrap().snapshot_id()
1312 );
1313 }
1314
1315 #[test]
1316 fn test_table_scan_non_exist_snapshot_id() {
1317 let table = TableTestFixture::new().table;
1318
1319 let table_scan = table.scan().snapshot_id(1024).build();
1320 assert!(table_scan.is_err());
1321 }
1322
1323 #[tokio::test]
1324 async fn test_table_scan_with_snapshot_id() {
1325 let table = TableTestFixture::new().table;
1326
1327 let table_scan = table
1328 .scan()
1329 .snapshot_id(3051729675574597004)
1330 .with_row_selection_enabled(true)
1331 .build()
1332 .unwrap();
1333 assert_eq!(
1334 table_scan.snapshot().unwrap().snapshot_id(),
1335 3051729675574597004
1336 );
1337 }
1338
1339 fn table_with_property(key: &str, value: &str) -> Table {
1340 let fixture = TableTestFixture::new();
1341 let mut metadata = fixture.table.metadata().clone();
1342 metadata
1343 .properties
1344 .insert(key.to_string(), value.to_string());
1345 Table::builder()
1346 .metadata(metadata)
1347 .identifier(fixture.table.identifier().clone())
1348 .file_io(fixture.table.file_io().clone())
1349 .metadata_location(fixture.table.metadata_location().unwrap().to_string())
1350 .runtime(test_runtime())
1351 .build()
1352 .unwrap()
1353 }
1354
1355 #[test]
1356 fn test_table_scan_without_name_mapping_property() {
1357 let table = TableTestFixture::new().table;
1358
1359 let table_scan = table.scan().build().unwrap();
1360 assert!(
1361 table_scan
1362 .plan_context
1363 .as_ref()
1364 .unwrap()
1365 .name_mapping
1366 .is_none()
1367 );
1368 }
1369
1370 #[test]
1371 fn test_table_scan_with_name_mapping_property() {
1372 let mapping_json = r#"[{"field-id":1,"names":["id","record_id"]}]"#;
1373 let table = table_with_property(DEFAULT_SCHEMA_NAME_MAPPING, mapping_json);
1374
1375 let table_scan = table.scan().build().unwrap();
1376 let mapping = table_scan
1377 .plan_context
1378 .as_ref()
1379 .unwrap()
1380 .name_mapping
1381 .as_ref()
1382 .expect("name_mapping should be parsed from the table property");
1383 let fields = mapping.fields();
1384 assert_eq!(fields.len(), 1);
1385 assert_eq!(fields[0].field_id(), Some(1));
1386 assert_eq!(fields[0].names(), &[
1387 "id".to_string(),
1388 "record_id".to_string()
1389 ]);
1390 }
1391
1392 #[test]
1393 fn test_table_scan_with_malformed_name_mapping_property() {
1394 let table = table_with_property(DEFAULT_SCHEMA_NAME_MAPPING, "{ not valid json");
1395
1396 let err = table
1397 .scan()
1398 .build()
1399 .expect_err("malformed name mapping should fail to parse");
1400 assert_eq!(err.kind(), ErrorKind::DataInvalid);
1401 }
1402
1403 #[tokio::test]
1404 async fn test_plan_files_carries_name_mapping_into_file_scan_task() {
1405 let mut fixture = TableTestFixture::new();
1406 fixture.setup_manifest_files().await;
1407
1408 let mapping_json = r#"[{"field-id":1,"names":["id","record_id"]}]"#;
1409 let mut metadata = fixture.table.metadata().clone();
1410 metadata.properties.insert(
1411 DEFAULT_SCHEMA_NAME_MAPPING.to_string(),
1412 mapping_json.to_string(),
1413 );
1414 let table = Table::builder()
1415 .metadata(metadata)
1416 .identifier(fixture.table.identifier().clone())
1417 .file_io(fixture.table.file_io().clone())
1418 .metadata_location(fixture.table.metadata_location().unwrap().to_string())
1419 .runtime(test_runtime())
1420 .build()
1421 .unwrap();
1422
1423 let tasks: Vec<_> = table
1424 .scan()
1425 .build()
1426 .unwrap()
1427 .plan_files()
1428 .await
1429 .unwrap()
1430 .try_collect()
1431 .await
1432 .unwrap();
1433
1434 assert!(!tasks.is_empty(), "expected at least one FileScanTask");
1435 for task in &tasks {
1436 let mapping = task
1437 .name_mapping
1438 .as_ref()
1439 .expect("name_mapping should reach the FileScanTask");
1440 assert_eq!(mapping.fields().len(), 1);
1441 assert_eq!(mapping.fields()[0].field_id(), Some(1));
1442 }
1443 }
1444
1445 #[tokio::test]
1446 async fn test_plan_files_on_table_without_any_snapshots() {
1447 let table = TableTestFixture::new_empty().table;
1448 let batch_stream = table.scan().build().unwrap().to_arrow().await.unwrap();
1449 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1450 assert!(batches.is_empty());
1451 }
1452
1453 #[tokio::test]
1454 async fn test_plan_files_no_deletions() {
1455 let mut fixture = TableTestFixture::new();
1456 fixture.setup_manifest_files().await;
1457
1458 let table_scan = fixture
1460 .table
1461 .scan()
1462 .with_row_selection_enabled(true)
1463 .build()
1464 .unwrap();
1465
1466 let mut tasks = table_scan
1467 .plan_files()
1468 .await
1469 .unwrap()
1470 .try_fold(vec![], |mut acc, task| async move {
1471 acc.push(task);
1472 Ok(acc)
1473 })
1474 .await
1475 .unwrap();
1476
1477 assert_eq!(tasks.len(), 2);
1478
1479 tasks.sort_by_key(|t| t.data_file_path.to_string());
1480
1481 assert_eq!(
1483 tasks[0].data_file_path,
1484 format!("{}/1.parquet", &fixture.table_location)
1485 );
1486
1487 assert_eq!(
1489 tasks[1].data_file_path,
1490 format!("{}/3.parquet", &fixture.table_location)
1491 );
1492 }
1493
1494 #[tokio::test]
1495 async fn test_open_parquet_no_deletions() {
1496 let mut fixture = TableTestFixture::new();
1497 fixture.setup_manifest_files().await;
1498
1499 let table_scan = fixture
1501 .table
1502 .scan()
1503 .with_row_selection_enabled(true)
1504 .build()
1505 .unwrap();
1506
1507 let batch_stream = table_scan.to_arrow().await.unwrap();
1508
1509 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1510
1511 let col = batches[0].column_by_name("x").unwrap();
1512
1513 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
1514 assert_eq!(int64_arr.value(0), 1);
1515 }
1516
1517 #[tokio::test]
1518 async fn test_open_parquet_no_deletions_by_separate_reader() {
1519 let mut fixture = TableTestFixture::new();
1520 fixture.setup_manifest_files().await;
1521
1522 let table_scan = fixture
1524 .table
1525 .scan()
1526 .with_row_selection_enabled(true)
1527 .build()
1528 .unwrap();
1529
1530 let mut plan_task: Vec<_> = table_scan
1531 .plan_files()
1532 .await
1533 .unwrap()
1534 .try_collect()
1535 .await
1536 .unwrap();
1537 assert_eq!(plan_task.len(), 2);
1538
1539 let reader = ArrowReaderBuilder::new(
1540 fixture.table.file_io().clone(),
1541 fixture.table.runtime().clone(),
1542 )
1543 .build();
1544 let batch_stream = reader
1545 .clone()
1546 .read(Box::pin(stream::iter(vec![Ok(plan_task.remove(0))])))
1547 .unwrap()
1548 .stream();
1549 let batch_1: Vec<_> = batch_stream.try_collect().await.unwrap();
1550
1551 let reader = ArrowReaderBuilder::new(
1552 fixture.table.file_io().clone(),
1553 fixture.table.runtime().clone(),
1554 )
1555 .build();
1556 let batch_stream = reader
1557 .read(Box::pin(stream::iter(vec![Ok(plan_task.remove(0))])))
1558 .unwrap()
1559 .stream();
1560 let batch_2: Vec<_> = batch_stream.try_collect().await.unwrap();
1561
1562 assert_eq!(batch_1, batch_2);
1563 }
1564
1565 #[tokio::test]
1566 async fn test_open_parquet_with_projection() {
1567 let mut fixture = TableTestFixture::new();
1568 fixture.setup_manifest_files().await;
1569
1570 let table_scan = fixture
1572 .table
1573 .scan()
1574 .select(["x", "z"])
1575 .with_row_selection_enabled(true)
1576 .build()
1577 .unwrap();
1578
1579 let batch_stream = table_scan.to_arrow().await.unwrap();
1580
1581 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1582
1583 assert_eq!(batches[0].num_columns(), 2);
1584
1585 let col1 = batches[0].column_by_name("x").unwrap();
1586 let int64_arr = col1.as_any().downcast_ref::<Int64Array>().unwrap();
1587 assert_eq!(int64_arr.value(0), 1);
1588
1589 let col2 = batches[0].column_by_name("z").unwrap();
1590 let int64_arr = col2.as_any().downcast_ref::<Int64Array>().unwrap();
1591 assert_eq!(int64_arr.value(0), 3);
1592
1593 let table_scan = fixture.table.scan().select_empty().build().unwrap();
1595 let batch_stream = table_scan.to_arrow().await.unwrap();
1596 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1597
1598 assert_eq!(batches[0].num_columns(), 0);
1599 assert_eq!(batches[0].num_rows(), 1024);
1600 }
1601
1602 #[tokio::test]
1603 async fn test_filter_on_arrow_lt() {
1604 let mut fixture = TableTestFixture::new();
1605 fixture.setup_manifest_files().await;
1606
1607 let mut builder = fixture.table.scan();
1609 let predicate = Reference::new("y").less_than(Datum::long(3));
1610 builder = builder
1611 .with_filter(predicate)
1612 .with_row_selection_enabled(true);
1613 let table_scan = builder.build().unwrap();
1614
1615 let batch_stream = table_scan.to_arrow().await.unwrap();
1616
1617 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1618
1619 assert_eq!(batches[0].num_rows(), 512);
1620
1621 let col = batches[0].column_by_name("x").unwrap();
1622 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
1623 assert_eq!(int64_arr.value(0), 1);
1624
1625 let col = batches[0].column_by_name("y").unwrap();
1626 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
1627 assert_eq!(int64_arr.value(0), 2);
1628 }
1629
1630 #[tokio::test]
1631 async fn test_filter_on_arrow_gt_eq() {
1632 let mut fixture = TableTestFixture::new();
1633 fixture.setup_manifest_files().await;
1634
1635 let mut builder = fixture.table.scan();
1637 let predicate = Reference::new("y").greater_than_or_equal_to(Datum::long(5));
1638 builder = builder
1639 .with_filter(predicate)
1640 .with_row_selection_enabled(true);
1641 let table_scan = builder.build().unwrap();
1642
1643 let batch_stream = table_scan.to_arrow().await.unwrap();
1644
1645 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1646
1647 assert_eq!(batches[0].num_rows(), 12);
1648
1649 let col = batches[0].column_by_name("x").unwrap();
1650 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
1651 assert_eq!(int64_arr.value(0), 1);
1652
1653 let col = batches[0].column_by_name("y").unwrap();
1654 let int64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
1655 assert_eq!(int64_arr.value(0), 5);
1656 }
1657
1658 #[tokio::test]
1659 async fn test_filter_double_eq() {
1660 let mut fixture = TableTestFixture::new();
1661 fixture.setup_manifest_files().await;
1662
1663 let mut builder = fixture.table.scan();
1665 let predicate = Reference::new("dbl").equal_to(Datum::double(150.0f64));
1666 builder = builder
1667 .with_filter(predicate)
1668 .with_row_selection_enabled(true);
1669 let table_scan = builder.build().unwrap();
1670
1671 let batch_stream = table_scan.to_arrow().await.unwrap();
1672
1673 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1674
1675 assert_eq!(batches.len(), 2);
1676 assert_eq!(batches[0].num_rows(), 12);
1677
1678 let col = batches[0].column_by_name("dbl").unwrap();
1679 let f64_arr = col.as_any().downcast_ref::<Float64Array>().unwrap();
1680 assert_eq!(f64_arr.value(1), 150.0f64);
1681 }
1682
1683 #[tokio::test]
1684 async fn test_filter_int_eq() {
1685 let mut fixture = TableTestFixture::new();
1686 fixture.setup_manifest_files().await;
1687
1688 let mut builder = fixture.table.scan();
1690 let predicate = Reference::new("i32").equal_to(Datum::int(150i32));
1691 builder = builder
1692 .with_filter(predicate)
1693 .with_row_selection_enabled(true);
1694 let table_scan = builder.build().unwrap();
1695
1696 let batch_stream = table_scan.to_arrow().await.unwrap();
1697
1698 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1699
1700 assert_eq!(batches.len(), 2);
1701 assert_eq!(batches[0].num_rows(), 12);
1702
1703 let col = batches[0].column_by_name("i32").unwrap();
1704 let i32_arr = col.as_any().downcast_ref::<Int32Array>().unwrap();
1705 assert_eq!(i32_arr.value(1), 150i32);
1706 }
1707
1708 #[tokio::test]
1709 async fn test_filter_long_eq() {
1710 let mut fixture = TableTestFixture::new();
1711 fixture.setup_manifest_files().await;
1712
1713 let mut builder = fixture.table.scan();
1715 let predicate = Reference::new("i64").equal_to(Datum::long(150i64));
1716 builder = builder
1717 .with_filter(predicate)
1718 .with_row_selection_enabled(true);
1719 let table_scan = builder.build().unwrap();
1720
1721 let batch_stream = table_scan.to_arrow().await.unwrap();
1722
1723 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1724
1725 assert_eq!(batches.len(), 2);
1726 assert_eq!(batches[0].num_rows(), 12);
1727
1728 let col = batches[0].column_by_name("i64").unwrap();
1729 let i64_arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
1730 assert_eq!(i64_arr.value(1), 150i64);
1731 }
1732
1733 #[tokio::test]
1734 async fn test_filter_bool_eq() {
1735 let mut fixture = TableTestFixture::new();
1736 fixture.setup_manifest_files().await;
1737
1738 let mut builder = fixture.table.scan();
1740 let predicate = Reference::new("bool").equal_to(Datum::bool(true));
1741 builder = builder
1742 .with_filter(predicate)
1743 .with_row_selection_enabled(true);
1744 let table_scan = builder.build().unwrap();
1745
1746 let batch_stream = table_scan.to_arrow().await.unwrap();
1747
1748 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1749
1750 assert_eq!(batches.len(), 2);
1751 assert_eq!(batches[0].num_rows(), 512);
1752
1753 let col = batches[0].column_by_name("bool").unwrap();
1754 let bool_arr = col.as_any().downcast_ref::<BooleanArray>().unwrap();
1755 assert!(bool_arr.value(1));
1756 }
1757
1758 #[tokio::test]
1759 async fn test_filter_on_arrow_is_null() {
1760 let mut fixture = TableTestFixture::new();
1761 fixture.setup_manifest_files().await;
1762
1763 let mut builder = fixture.table.scan();
1765 let predicate = Reference::new("y").is_null();
1766 builder = builder
1767 .with_filter(predicate)
1768 .with_row_selection_enabled(true);
1769 let table_scan = builder.build().unwrap();
1770
1771 let batch_stream = table_scan.to_arrow().await.unwrap();
1772
1773 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1774 assert_eq!(batches.len(), 0);
1775 }
1776
1777 #[tokio::test]
1778 async fn test_filter_on_arrow_is_not_null() {
1779 let mut fixture = TableTestFixture::new();
1780 fixture.setup_manifest_files().await;
1781
1782 let mut builder = fixture.table.scan();
1784 let predicate = Reference::new("y").is_not_null();
1785 builder = builder
1786 .with_filter(predicate)
1787 .with_row_selection_enabled(true);
1788 let table_scan = builder.build().unwrap();
1789
1790 let batch_stream = table_scan.to_arrow().await.unwrap();
1791
1792 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1793 assert_eq!(batches[0].num_rows(), 1024);
1794 }
1795
1796 #[tokio::test]
1797 async fn test_filter_on_arrow_lt_and_gt() {
1798 let mut fixture = TableTestFixture::new();
1799 fixture.setup_manifest_files().await;
1800
1801 let mut builder = fixture.table.scan();
1803 let predicate = Reference::new("y")
1804 .less_than(Datum::long(5))
1805 .and(Reference::new("z").greater_than_or_equal_to(Datum::long(4)));
1806 builder = builder
1807 .with_filter(predicate)
1808 .with_row_selection_enabled(true);
1809 let table_scan = builder.build().unwrap();
1810
1811 let batch_stream = table_scan.to_arrow().await.unwrap();
1812
1813 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1814 assert_eq!(batches[0].num_rows(), 500);
1815
1816 let col = batches[0].column_by_name("x").unwrap();
1817 let expected_x = Arc::new(Int64Array::from_iter_values(vec![1; 500])) as ArrayRef;
1818 assert_eq!(col, &expected_x);
1819
1820 let col = batches[0].column_by_name("y").unwrap();
1821 let mut values = vec![];
1822 values.append(vec![3; 200].as_mut());
1823 values.append(vec![4; 300].as_mut());
1824 let expected_y = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
1825 assert_eq!(col, &expected_y);
1826
1827 let col = batches[0].column_by_name("z").unwrap();
1828 let expected_z = Arc::new(Int64Array::from_iter_values(vec![4; 500])) as ArrayRef;
1829 assert_eq!(col, &expected_z);
1830 }
1831
1832 #[tokio::test]
1833 async fn test_filter_on_arrow_lt_or_gt() {
1834 let mut fixture = TableTestFixture::new();
1835 fixture.setup_manifest_files().await;
1836
1837 let mut builder = fixture.table.scan();
1839 let predicate = Reference::new("y")
1840 .less_than(Datum::long(5))
1841 .or(Reference::new("z").greater_than_or_equal_to(Datum::long(4)));
1842 builder = builder
1843 .with_filter(predicate)
1844 .with_row_selection_enabled(true);
1845 let table_scan = builder.build().unwrap();
1846
1847 let batch_stream = table_scan.to_arrow().await.unwrap();
1848
1849 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1850 assert_eq!(batches[0].num_rows(), 1024);
1851
1852 let col = batches[0].column_by_name("x").unwrap();
1853 let expected_x = Arc::new(Int64Array::from_iter_values(vec![1; 1024])) as ArrayRef;
1854 assert_eq!(col, &expected_x);
1855
1856 let col = batches[0].column_by_name("y").unwrap();
1857 let mut values = vec![2; 512];
1858 values.append(vec![3; 200].as_mut());
1859 values.append(vec![4; 300].as_mut());
1860 values.append(vec![5; 12].as_mut());
1861 let expected_y = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
1862 assert_eq!(col, &expected_y);
1863
1864 let col = batches[0].column_by_name("z").unwrap();
1865 let mut values = vec![3; 512];
1866 values.append(vec![4; 512].as_mut());
1867 let expected_z = Arc::new(Int64Array::from_iter_values(values)) as ArrayRef;
1868 assert_eq!(col, &expected_z);
1869 }
1870
1871 #[tokio::test]
1872 async fn test_filter_on_arrow_startswith() {
1873 let mut fixture = TableTestFixture::new();
1874 fixture.setup_manifest_files().await;
1875
1876 let mut builder = fixture.table.scan();
1878 let predicate = Reference::new("a").starts_with(Datum::string("Ice"));
1879 builder = builder
1880 .with_filter(predicate)
1881 .with_row_selection_enabled(true);
1882 let table_scan = builder.build().unwrap();
1883
1884 let batch_stream = table_scan.to_arrow().await.unwrap();
1885
1886 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1887
1888 assert_eq!(batches[0].num_rows(), 512);
1889
1890 let col = batches[0].column_by_name("a").unwrap();
1891 let string_arr = col.as_any().downcast_ref::<StringArray>().unwrap();
1892 assert_eq!(string_arr.value(0), "Iceberg");
1893 }
1894
1895 #[tokio::test]
1896 async fn test_filter_on_arrow_not_startswith() {
1897 let mut fixture = TableTestFixture::new();
1898 fixture.setup_manifest_files().await;
1899
1900 let mut builder = fixture.table.scan();
1902 let predicate = Reference::new("a").not_starts_with(Datum::string("Ice"));
1903 builder = builder
1904 .with_filter(predicate)
1905 .with_row_selection_enabled(true);
1906 let table_scan = builder.build().unwrap();
1907
1908 let batch_stream = table_scan.to_arrow().await.unwrap();
1909
1910 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1911
1912 assert_eq!(batches[0].num_rows(), 512);
1913
1914 let col = batches[0].column_by_name("a").unwrap();
1915 let string_arr = col.as_any().downcast_ref::<StringArray>().unwrap();
1916 assert_eq!(string_arr.value(0), "Apache");
1917 }
1918
1919 #[tokio::test]
1920 async fn test_filter_on_arrow_in() {
1921 let mut fixture = TableTestFixture::new();
1922 fixture.setup_manifest_files().await;
1923
1924 let mut builder = fixture.table.scan();
1926 let predicate =
1927 Reference::new("a").is_in([Datum::string("Sioux"), Datum::string("Iceberg")]);
1928 builder = builder
1929 .with_filter(predicate)
1930 .with_row_selection_enabled(true);
1931 let table_scan = builder.build().unwrap();
1932
1933 let batch_stream = table_scan.to_arrow().await.unwrap();
1934
1935 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1936
1937 assert_eq!(batches[0].num_rows(), 512);
1938
1939 let col = batches[0].column_by_name("a").unwrap();
1940 let string_arr = col.as_any().downcast_ref::<StringArray>().unwrap();
1941 assert_eq!(string_arr.value(0), "Iceberg");
1942 }
1943
1944 #[tokio::test]
1945 async fn test_filter_on_arrow_not_in() {
1946 let mut fixture = TableTestFixture::new();
1947 fixture.setup_manifest_files().await;
1948
1949 let mut builder = fixture.table.scan();
1951 let predicate =
1952 Reference::new("a").is_not_in([Datum::string("Sioux"), Datum::string("Iceberg")]);
1953 builder = builder
1954 .with_filter(predicate)
1955 .with_row_selection_enabled(true);
1956 let table_scan = builder.build().unwrap();
1957
1958 let batch_stream = table_scan.to_arrow().await.unwrap();
1959
1960 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
1961
1962 assert_eq!(batches[0].num_rows(), 512);
1963
1964 let col = batches[0].column_by_name("a").unwrap();
1965 let string_arr = col.as_any().downcast_ref::<StringArray>().unwrap();
1966 assert_eq!(string_arr.value(0), "Apache");
1967 }
1968
1969 #[test]
1970 fn test_file_scan_task_serialize_deserialize() {
1971 let test_fn = |task: FileScanTask| {
1972 let serialized = serde_json::to_string(&task).unwrap();
1973 let deserialized: FileScanTask = serde_json::from_str(&serialized).unwrap();
1974
1975 assert_eq!(task.data_file_path, deserialized.data_file_path);
1976 assert_eq!(task.start, deserialized.start);
1977 assert_eq!(task.length, deserialized.length);
1978 assert_eq!(task.project_field_ids, deserialized.project_field_ids);
1979 assert_eq!(task.predicate, deserialized.predicate);
1980 assert_eq!(task.schema, deserialized.schema);
1981 };
1982
1983 let schema = Arc::new(
1985 Schema::builder()
1986 .with_fields(vec![Arc::new(NestedField::required(
1987 1,
1988 "x",
1989 Type::Primitive(PrimitiveType::Binary),
1990 ))])
1991 .build()
1992 .unwrap(),
1993 );
1994 let task = FileScanTask::builder()
1995 .with_data_file_path("data_file_path".to_string())
1996 .with_file_size_in_bytes(0)
1997 .with_start(0)
1998 .with_length(100)
1999 .with_project_field_ids(vec![1, 2, 3])
2000 .with_schema(schema.clone())
2001 .with_record_count(Some(100))
2002 .with_data_file_format(DataFileFormat::Parquet)
2003 .with_case_sensitive(false)
2004 .build();
2005 test_fn(task);
2006
2007 let task = FileScanTask::builder()
2009 .with_data_file_path("data_file_path".to_string())
2010 .with_file_size_in_bytes(0)
2011 .with_start(0)
2012 .with_length(100)
2013 .with_project_field_ids(vec![1, 2, 3])
2014 .with_predicate(Some(BoundPredicate::AlwaysTrue))
2015 .with_schema(schema)
2016 .with_data_file_format(DataFileFormat::Avro)
2017 .with_case_sensitive(false)
2018 .build();
2019 test_fn(task);
2020 }
2021
2022 #[tokio::test]
2023 async fn test_select_with_file_column() {
2024 use arrow_array::cast::AsArray;
2025
2026 let mut fixture = TableTestFixture::new();
2027 fixture.setup_manifest_files().await;
2028
2029 let table_scan = fixture
2031 .table
2032 .scan()
2033 .select(["x", RESERVED_COL_NAME_FILE])
2034 .with_row_selection_enabled(true)
2035 .build()
2036 .unwrap();
2037
2038 let batch_stream = table_scan.to_arrow().await.unwrap();
2039 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2040
2041 assert_eq!(batches[0].num_columns(), 2);
2043
2044 let x_col = batches[0].column_by_name("x").unwrap();
2046 let x_arr = x_col.as_primitive::<arrow_array::types::Int64Type>();
2047 assert_eq!(x_arr.value(0), 1);
2048
2049 let file_col = batches[0].column_by_name(RESERVED_COL_NAME_FILE);
2051 assert!(
2052 file_col.is_some(),
2053 "_file column should be present in the batch"
2054 );
2055
2056 let file_col = file_col.unwrap();
2058 assert!(
2059 matches!(
2060 file_col.data_type(),
2061 arrow_schema::DataType::RunEndEncoded(_, _)
2062 ),
2063 "_file column should use RunEndEncoded type"
2064 );
2065
2066 let run_array = file_col
2068 .as_any()
2069 .downcast_ref::<arrow_array::RunArray<arrow_array::types::Int32Type>>()
2070 .expect("_file column should be a RunArray");
2071
2072 let values = run_array.values();
2073 let string_values = values.as_string::<i32>();
2074 assert_eq!(string_values.len(), 1, "Should have a single file path");
2075
2076 let file_path = string_values.value(0);
2077 assert!(
2078 file_path.ends_with(".parquet"),
2079 "File path should end with .parquet, got: {file_path}"
2080 );
2081 }
2082
2083 #[tokio::test]
2084 async fn test_select_file_column_position() {
2085 let mut fixture = TableTestFixture::new();
2086 fixture.setup_manifest_files().await;
2087
2088 let table_scan = fixture
2090 .table
2091 .scan()
2092 .select(["x", RESERVED_COL_NAME_FILE, "z"])
2093 .with_row_selection_enabled(true)
2094 .build()
2095 .unwrap();
2096
2097 let batch_stream = table_scan.to_arrow().await.unwrap();
2098 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2099
2100 assert_eq!(batches[0].num_columns(), 3);
2101
2102 let schema = batches[0].schema();
2104 assert_eq!(schema.field(0).name(), "x");
2105 assert_eq!(schema.field(1).name(), RESERVED_COL_NAME_FILE);
2106 assert_eq!(schema.field(2).name(), "z");
2107
2108 assert!(batches[0].column_by_name("x").is_some());
2110 assert!(batches[0].column_by_name(RESERVED_COL_NAME_FILE).is_some());
2111 assert!(batches[0].column_by_name("z").is_some());
2112 }
2113
2114 #[tokio::test]
2115 async fn test_select_file_column_only() {
2116 let mut fixture = TableTestFixture::new();
2117 fixture.setup_manifest_files().await;
2118
2119 let table_scan = fixture
2121 .table
2122 .scan()
2123 .select([RESERVED_COL_NAME_FILE])
2124 .with_row_selection_enabled(true)
2125 .build()
2126 .unwrap();
2127
2128 let batch_stream = table_scan.to_arrow().await.unwrap();
2129 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2130
2131 assert_eq!(batches[0].num_columns(), 1);
2133
2134 let schema = batches[0].schema();
2136 assert_eq!(schema.field(0).name(), RESERVED_COL_NAME_FILE);
2137
2138 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
2142 assert_eq!(total_rows, 2048);
2143 }
2144
2145 #[tokio::test]
2146 async fn test_file_column_with_multiple_files() {
2147 use std::collections::HashSet;
2148
2149 let mut fixture = TableTestFixture::new();
2150 fixture.setup_manifest_files().await;
2151
2152 let table_scan = fixture
2154 .table
2155 .scan()
2156 .select(["x", RESERVED_COL_NAME_FILE])
2157 .with_row_selection_enabled(true)
2158 .build()
2159 .unwrap();
2160
2161 let batch_stream = table_scan.to_arrow().await.unwrap();
2162 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2163
2164 let mut file_paths = HashSet::new();
2166 for batch in &batches {
2167 let file_col = batch.column_by_name(RESERVED_COL_NAME_FILE).unwrap();
2168 let run_array = file_col
2169 .as_any()
2170 .downcast_ref::<arrow_array::RunArray<arrow_array::types::Int32Type>>()
2171 .expect("_file column should be a RunArray");
2172
2173 let values = run_array.values();
2174 let string_values = values.as_string::<i32>();
2175 for i in 0..string_values.len() {
2176 file_paths.insert(string_values.value(i).to_string());
2177 }
2178 }
2179
2180 assert!(!file_paths.is_empty(), "Should have at least one file path");
2182
2183 for path in &file_paths {
2185 assert!(
2186 path.ends_with(".parquet"),
2187 "All file paths should end with .parquet, got: {path}"
2188 );
2189 }
2190 }
2191
2192 #[tokio::test]
2193 async fn test_file_column_at_start() {
2194 let mut fixture = TableTestFixture::new();
2195 fixture.setup_manifest_files().await;
2196
2197 let table_scan = fixture
2199 .table
2200 .scan()
2201 .select([RESERVED_COL_NAME_FILE, "x", "y"])
2202 .with_row_selection_enabled(true)
2203 .build()
2204 .unwrap();
2205
2206 let batch_stream = table_scan.to_arrow().await.unwrap();
2207 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2208
2209 assert_eq!(batches[0].num_columns(), 3);
2210
2211 let schema = batches[0].schema();
2213 assert_eq!(schema.field(0).name(), RESERVED_COL_NAME_FILE);
2214 assert_eq!(schema.field(1).name(), "x");
2215 assert_eq!(schema.field(2).name(), "y");
2216 }
2217
2218 #[tokio::test]
2219 async fn test_file_column_at_end() {
2220 let mut fixture = TableTestFixture::new();
2221 fixture.setup_manifest_files().await;
2222
2223 let table_scan = fixture
2225 .table
2226 .scan()
2227 .select(["x", "y", RESERVED_COL_NAME_FILE])
2228 .with_row_selection_enabled(true)
2229 .build()
2230 .unwrap();
2231
2232 let batch_stream = table_scan.to_arrow().await.unwrap();
2233 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2234
2235 assert_eq!(batches[0].num_columns(), 3);
2236
2237 let schema = batches[0].schema();
2239 assert_eq!(schema.field(0).name(), "x");
2240 assert_eq!(schema.field(1).name(), "y");
2241 assert_eq!(schema.field(2).name(), RESERVED_COL_NAME_FILE);
2242 }
2243
2244 #[tokio::test]
2245 async fn test_select_with_repeated_column_names() {
2246 let mut fixture = TableTestFixture::new();
2247 fixture.setup_manifest_files().await;
2248
2249 let table_scan = fixture
2252 .table
2253 .scan()
2254 .select([
2255 "x",
2256 RESERVED_COL_NAME_FILE,
2257 "x", "y",
2259 RESERVED_COL_NAME_FILE, "y", ])
2262 .with_row_selection_enabled(true)
2263 .build()
2264 .unwrap();
2265
2266 let batch_stream = table_scan.to_arrow().await.unwrap();
2267 let batches: Vec<_> = batch_stream.try_collect().await.unwrap();
2268
2269 assert_eq!(
2271 batches[0].num_columns(),
2272 6,
2273 "Should have exactly 6 columns with duplicates"
2274 );
2275
2276 let schema = batches[0].schema();
2277
2278 assert_eq!(schema.field(0).name(), "x", "Column 0 should be x");
2280 assert_eq!(
2281 schema.field(1).name(),
2282 RESERVED_COL_NAME_FILE,
2283 "Column 1 should be _file"
2284 );
2285 assert_eq!(
2286 schema.field(2).name(),
2287 "x",
2288 "Column 2 should be x (duplicate)"
2289 );
2290 assert_eq!(schema.field(3).name(), "y", "Column 3 should be y");
2291 assert_eq!(
2292 schema.field(4).name(),
2293 RESERVED_COL_NAME_FILE,
2294 "Column 4 should be _file (duplicate)"
2295 );
2296 assert_eq!(
2297 schema.field(5).name(),
2298 "y",
2299 "Column 5 should be y (duplicate)"
2300 );
2301
2302 assert!(
2304 matches!(schema.field(0).data_type(), arrow_schema::DataType::Int64),
2305 "Column x should be Int64"
2306 );
2307 assert!(
2308 matches!(schema.field(2).data_type(), arrow_schema::DataType::Int64),
2309 "Column x (duplicate) should be Int64"
2310 );
2311 assert!(
2312 matches!(schema.field(3).data_type(), arrow_schema::DataType::Int64),
2313 "Column y should be Int64"
2314 );
2315 assert!(
2316 matches!(schema.field(5).data_type(), arrow_schema::DataType::Int64),
2317 "Column y (duplicate) should be Int64"
2318 );
2319 assert!(
2320 matches!(
2321 schema.field(1).data_type(),
2322 arrow_schema::DataType::RunEndEncoded(_, _)
2323 ),
2324 "_file column should use RunEndEncoded type"
2325 );
2326 assert!(
2327 matches!(
2328 schema.field(4).data_type(),
2329 arrow_schema::DataType::RunEndEncoded(_, _)
2330 ),
2331 "_file column (duplicate) should use RunEndEncoded type"
2332 );
2333 }
2334
2335 #[tokio::test]
2336 async fn test_scan_deadlock() {
2337 let mut fixture = TableTestFixture::new();
2338 fixture.setup_deadlock_manifests().await;
2339
2340 let table_scan = fixture
2347 .table
2348 .scan()
2349 .with_concurrency_limit(1)
2350 .build()
2351 .unwrap();
2352
2353 let result = tokio::time::timeout(std::time::Duration::from_secs(5), async {
2356 table_scan
2357 .plan_files()
2358 .await
2359 .unwrap()
2360 .try_collect::<Vec<_>>()
2361 .await
2362 })
2363 .await;
2364
2365 assert!(result.is_ok(), "Scan timed out - deadlock detected");
2367 }
2368}