delta_arrow_reader/
direct.rs

1//! Public DataFusion-independent Delta-to-Arrow reader.
2
3use std::{
4    collections::VecDeque,
5    fmt,
6    pin::Pin,
7    sync::Arc,
8    task::{Context, Poll},
9};
10
11use arrow::{datatypes::SchemaRef, record_batch::RecordBatch};
12use futures_util::Stream;
13use snafu::ResultExt;
14
15use crate::{
16    DeltaPredicate, DeltaProtocolInfo, DeltaReadMetrics, DeltaReaderBackend, DeltaReaderError,
17    DeltaReaderExecutionOptions, DeltaSnapshotSelection, DeltaStorageOptions,
18    error::{DataFileReadSnafu, InvalidConfigurationSnafu, ScanPlanningSnafu},
19    kernel::{delta_predicate_kernel_pruning_is_exact, delta_predicate_to_kernel_pruning},
20    planning::{
21        DeltaScanPartitionTargetOptions, DeltaScanPlan, plan_scan, validate_backend_available,
22    },
23    predicate::{evaluate_predicate, referenced_columns, validate_predicate},
24    protocol::validate_protocol,
25    scheduling::{
26        DeltaScanExecution, FileAdmission, FileAdmissionFn, FileBatchStream, FileExecutor,
27        PartitionStream,
28    },
29    snapshot::{
30        LoadedDeltaTableSnapshot, StagedDeltaTableSnapshot, load_delta_table_snapshot_async,
31        load_delta_table_snapshot_blocking, load_staged_delta_table_snapshot_async,
32        load_staged_delta_table_snapshot_blocking,
33    },
34};
35
36const TRACING_TARGET: &str = "delta_arrow_reader";
37
38/// Configures and loads one immutable Delta table snapshot.
39///
40/// The asynchronous path uses the caller's Tokio runtime. Scans return a
41/// pull-driven stream and do not materialize the whole table.
42///
43/// # Example
44///
45/// ```no_run
46/// use delta_arrow_reader::{DeltaComparison, DeltaPredicate, DeltaScalar, DeltaTableBuilder};
47/// use futures_util::TryStreamExt;
48///
49/// # async fn read_table() -> Result<(), Box<dyn std::error::Error>> {
50/// let table = DeltaTableBuilder::new("/tmp/example-delta-table")
51///     .load_async()
52///     .await?;
53/// let scan = table
54///     .scan()
55///     .with_projection(vec!["id".into(), "name".into()])
56///     .with_predicate(DeltaPredicate::Compare {
57///         column: "id".into(),
58///         op: DeltaComparison::GtEq,
59///         value: DeltaScalar::Int64(10),
60///     })
61///     .with_limit(100)
62///     .build()
63///     .await?;
64/// let mut batches = scan.execute().await?;
65///
66/// while let Some(batch) = batches.try_next().await? {
67///     println!("rows={}", batch.num_rows());
68/// }
69/// # Ok(())
70/// # }
71/// ```
72pub struct DeltaTableBuilder {
73    table_uri: String,
74    storage_options: DeltaStorageOptions,
75    snapshot_selection: DeltaSnapshotSelection,
76    execution_options: DeltaReaderExecutionOptions,
77}
78
79impl DeltaTableBuilder {
80    /// Creates a builder for the latest snapshot with default execution settings.
81    pub fn new(table_uri: impl Into<String>) -> Self {
82        Self {
83            table_uri: table_uri.into(),
84            storage_options: DeltaStorageOptions::new(),
85            snapshot_selection: DeltaSnapshotSelection::Latest,
86            execution_options: DeltaReaderExecutionOptions::new(),
87        }
88    }
89
90    /// Replaces the storage options forwarded during table loading.
91    pub fn with_storage_options(mut self, value: DeltaStorageOptions) -> Self {
92        self.storage_options = value;
93        self
94    }
95
96    /// Selects the Delta snapshot to load.
97    pub const fn with_snapshot_selection(mut self, value: DeltaSnapshotSelection) -> Self {
98        self.snapshot_selection = value;
99        self
100    }
101
102    /// Replaces the default execution settings used by scans of this table.
103    pub const fn with_execution_options(mut self, value: DeltaReaderExecutionOptions) -> Self {
104        self.execution_options = value;
105        self
106    }
107
108    /// Loads the snapshot on the calling thread.
109    pub fn load(self) -> Result<DeltaTable, DeltaReaderError> {
110        validate_direct_execution_options(self.execution_options)?;
111        let snapshot = load_delta_table_snapshot_blocking(
112            &self.table_uri,
113            &self.storage_options,
114            self.snapshot_selection,
115        )?;
116        Ok(DeltaTable::new(snapshot, self.execution_options))
117    }
118
119    /// Loads the snapshot through the caller-owned Tokio runtime.
120    pub async fn load_async(self) -> Result<DeltaTable, DeltaReaderError> {
121        validate_direct_execution_options(self.execution_options)?;
122        let snapshot = load_delta_table_snapshot_async(
123            self.table_uri,
124            self.storage_options,
125            self.snapshot_selection,
126        )
127        .await?;
128        Ok(DeltaTable::new(snapshot, self.execution_options))
129    }
130
131    /// Loads snapshot metadata without converting its logical Arrow schema.
132    pub fn load_snapshot(self) -> Result<DeltaTableSnapshot, DeltaReaderError> {
133        validate_direct_execution_options(self.execution_options)?;
134        let snapshot = load_staged_delta_table_snapshot_blocking(
135            &self.table_uri,
136            &self.storage_options,
137            self.snapshot_selection,
138        )?;
139        Ok(DeltaTableSnapshot::new(snapshot, self.execution_options))
140    }
141
142    /// Loads snapshot metadata through the caller-owned Tokio runtime.
143    pub async fn load_snapshot_async(self) -> Result<DeltaTableSnapshot, DeltaReaderError> {
144        validate_direct_execution_options(self.execution_options)?;
145        let snapshot = load_staged_delta_table_snapshot_async(
146            self.table_uri,
147            self.storage_options,
148            self.snapshot_selection,
149        )
150        .await?;
151        Ok(DeltaTableSnapshot::new(snapshot, self.execution_options))
152    }
153}
154
155impl fmt::Debug for DeltaTableBuilder {
156    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
157        formatter
158            .debug_struct("DeltaTableBuilder")
159            .field("table_uri", &"<redacted>")
160            .field("storage_options", &"<redacted>")
161            .field("snapshot_selection", &self.snapshot_selection)
162            .field("execution_options", &self.execution_options)
163            .finish()
164    }
165}
166
167/// Loaded Delta snapshot metadata awaiting logical Arrow schema conversion.
168pub struct DeltaTableSnapshot {
169    snapshot: StagedDeltaTableSnapshot,
170    execution_options: DeltaReaderExecutionOptions,
171}
172
173impl DeltaTableSnapshot {
174    fn new(
175        snapshot: StagedDeltaTableSnapshot,
176        execution_options: DeltaReaderExecutionOptions,
177    ) -> Self {
178        Self {
179            snapshot,
180            execution_options,
181        }
182    }
183
184    /// Returns the loaded Delta snapshot version.
185    pub fn version(&self) -> u64 {
186        self.snapshot.version()
187    }
188
189    /// Returns the loaded Delta protocol metadata.
190    pub fn protocol(&self) -> &DeltaProtocolInfo {
191        self.snapshot.protocol_info()
192    }
193
194    /// Returns the normalized table URI.
195    ///
196    /// This value may contain sensitive caller input. Do not log or expose it.
197    pub fn table_uri(&self) -> &str {
198        self.snapshot.table_uri()
199    }
200
201    /// Validates the loaded snapshot against the supported reader protocol.
202    pub fn validate_protocol(&self) -> Result<(), DeltaReaderError> {
203        validate_protocol(self.protocol())
204    }
205
206    /// Converts the logical Arrow schema and finishes constructing the table.
207    pub fn into_table(self) -> Result<DeltaTable, DeltaReaderError> {
208        Ok(DeltaTable::new(
209            self.snapshot.into_loaded()?,
210            self.execution_options,
211        ))
212    }
213}
214
215impl fmt::Debug for DeltaTableSnapshot {
216    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
217        formatter
218            .debug_struct("DeltaTableSnapshot")
219            .field("version", &self.version())
220            .finish_non_exhaustive()
221    }
222}
223
224/// One immutable loaded Delta table snapshot.
225#[derive(Clone)]
226pub struct DeltaTable {
227    snapshot: Arc<LoadedDeltaTableSnapshot>,
228    version: u64,
229    execution_options: DeltaReaderExecutionOptions,
230}
231
232impl DeltaTable {
233    fn new(
234        snapshot: LoadedDeltaTableSnapshot,
235        execution_options: DeltaReaderExecutionOptions,
236    ) -> Self {
237        let version = snapshot.version();
238        Self {
239            snapshot: Arc::new(snapshot),
240            version,
241            execution_options,
242        }
243    }
244
245    /// Returns the loaded Delta snapshot version.
246    pub const fn version(&self) -> u64 {
247        self.version
248    }
249
250    /// Returns the logical Arrow schema.
251    pub fn schema(&self) -> &SchemaRef {
252        self.snapshot.schema_ref()
253    }
254
255    /// Returns the loaded Delta protocol metadata.
256    pub fn protocol(&self) -> &DeltaProtocolInfo {
257        self.snapshot.protocol_info()
258    }
259
260    /// Returns the normalized table URI.
261    ///
262    /// This value may contain sensitive caller input. Do not log or expose it.
263    pub fn table_uri(&self) -> &str {
264        self.snapshot.table_uri()
265    }
266
267    #[allow(dead_code)]
268    pub(crate) fn partition_columns(&self) -> &[String] {
269        self.snapshot.partition_columns()
270    }
271
272    #[allow(dead_code)]
273    pub(crate) fn snapshot(&self) -> &LoadedDeltaTableSnapshot {
274        self.snapshot.as_ref()
275    }
276
277    /// Validates the loaded snapshot against the supported reader protocol.
278    pub fn validate_protocol(&self) -> Result<(), DeltaReaderError> {
279        validate_protocol(self.protocol())
280    }
281
282    /// Starts configuring a new single-use scan.
283    pub fn scan(&self) -> DeltaScanBuilder<'_> {
284        DeltaScanBuilder {
285            table: self,
286            projection: None,
287            predicate: None,
288            limit: None,
289            target_partitions: None,
290            execution_options: self.execution_options,
291        }
292    }
293}
294
295impl fmt::Debug for DeltaTable {
296    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
297        formatter
298            .debug_struct("DeltaTable")
299            .field("version", &self.version)
300            .finish_non_exhaustive()
301    }
302}
303
304/// Configures one single-use direct Delta scan.
305pub struct DeltaScanBuilder<'table> {
306    table: &'table DeltaTable,
307    projection: Option<Vec<String>>,
308    predicate: Option<DeltaPredicate>,
309    limit: Option<usize>,
310    target_partitions: Option<usize>,
311    execution_options: DeltaReaderExecutionOptions,
312}
313
314impl<'table> DeltaScanBuilder<'table> {
315    /// Selects visible logical columns in caller order.
316    pub fn with_projection(mut self, logical_columns: Vec<String>) -> Self {
317        self.projection = Some(logical_columns);
318        self
319    }
320
321    /// Replaces the exact logical row predicate.
322    pub fn with_predicate(mut self, predicate: DeltaPredicate) -> Self {
323        self.predicate = Some(predicate);
324        self
325    }
326
327    /// Sets the maximum number of output rows.
328    pub const fn with_limit(mut self, limit: usize) -> Self {
329        self.limit = Some(limit);
330        self
331    }
332
333    /// Overrides the number of planned scan partitions.
334    pub fn with_target_partitions(mut self, value: usize) -> Result<Self, DeltaReaderError> {
335        if value == 0 {
336            return InvalidConfigurationSnafu {
337                reason: "scan_partition_target_must_be_positive",
338            }
339            .fail();
340        }
341        self.target_partitions = Some(value);
342        Ok(self)
343    }
344
345    /// Replaces the execution settings for this scan.
346    pub fn with_execution_options(
347        mut self,
348        value: DeltaReaderExecutionOptions,
349    ) -> Result<Self, DeltaReaderError> {
350        validate_direct_execution_options(value)?;
351        self.execution_options = value;
352        Ok(self)
353    }
354
355    /// Builds one immutable single-use scan plan without reading data files.
356    pub async fn build(self) -> Result<DeltaScan, DeltaReaderError> {
357        self.table.validate_protocol()?;
358        validate_direct_execution_options(self.execution_options)?;
359        if let Some(predicate) = self.predicate.as_ref() {
360            validate_predicate(predicate, self.table.schema().as_ref())?;
361        }
362
363        let snapshot_version = self.table.version();
364        let backend = self.execution_options.reader_backend();
365        trace_planning_started(snapshot_version, backend);
366        let snapshot = Arc::clone(&self.table.snapshot);
367        let projection = self.projection;
368        let predicate = self.predicate;
369        let hidden_columns = predicate
370            .as_ref()
371            .map(referenced_columns)
372            .unwrap_or_default();
373        let enforce_physical_predicate_rows = predicate
374            .as_ref()
375            .is_some_and(delta_predicate_kernel_pruning_is_exact);
376        let kernel_predicate = predicate
377            .as_ref()
378            .and_then(delta_predicate_to_kernel_pruning);
379        let include_stats = kernel_predicate.is_some();
380        let execution_options = self.execution_options;
381        let target_partitions = self.target_partitions;
382        let result = tokio::task::spawn_blocking(move || {
383            plan_scan(
384                snapshot.as_ref(),
385                projection.as_deref(),
386                &hidden_columns,
387                kernel_predicate,
388                include_stats,
389                execution_options,
390                DeltaScanPartitionTargetOptions {
391                    explicit_target_partitions: target_partitions,
392                    caller_target_partitions: None,
393                },
394            )
395        })
396        .await
397        .boxed()
398        .context(ScanPlanningSnafu {
399            reason: "scan_planning_task_failed",
400        })
401        .and_then(|result| result);
402
403        match result {
404            Ok(plan) => {
405                trace_planning_completed(snapshot_version, backend, plan.partitions.len());
406                Ok(DeltaScan {
407                    partition_count: plan.partitions.len(),
408                    plan: Arc::new(plan),
409                    predicate,
410                    limit: self.limit,
411                    enforce_physical_predicate_rows,
412                })
413            }
414            Err(error) => {
415                trace_planning_failed(snapshot_version, backend, &error);
416                Err(error)
417            }
418        }
419    }
420}
421
422/// One immutable, single-use direct Delta scan plan.
423///
424/// A scan cannot be cloned or executed twice.
425///
426/// ```compile_fail
427/// use delta_arrow_reader::DeltaScan;
428///
429/// async fn execute_twice(scan: DeltaScan) {
430///     let _ = scan.execute().await;
431///     let _ = scan.execute().await;
432/// }
433/// ```
434///
435/// ```compile_fail
436/// use delta_arrow_reader::DeltaScan;
437///
438/// fn clone_scan(scan: DeltaScan) {
439///     let _ = scan.clone();
440/// }
441/// ```
442pub struct DeltaScan {
443    plan: Arc<DeltaScanPlan>,
444    predicate: Option<DeltaPredicate>,
445    limit: Option<usize>,
446    partition_count: usize,
447    enforce_physical_predicate_rows: bool,
448}
449
450impl DeltaScan {
451    /// Returns the visible logical output schema.
452    pub fn schema(&self) -> &SchemaRef {
453        &self.plan.projected_schema
454    }
455
456    /// Returns the number of planned execution partitions.
457    pub const fn partition_count(&self) -> usize {
458        self.partition_count
459    }
460
461    /// Creates the pull-driven direct Arrow batch stream.
462    pub async fn execute(self) -> Result<DeltaBatchStream, DeltaReaderError> {
463        let metrics = self.plan.metrics.clone();
464        let schema = Arc::clone(&self.plan.projected_schema);
465        let partition_count = self.plan.partitions.len();
466        let snapshot_version = self.plan.snapshot_version;
467        let backend = self.plan.execution_options.reader_backend();
468        let projection = (self.plan.logical_schema.as_ref() != schema.as_ref())
469            .then(|| (0..schema.fields().len()).collect::<Vec<_>>());
470        let mut partitions = VecDeque::new();
471
472        if self.limit != Some(0) {
473            let execution = DeltaScanExecution::new(Arc::clone(&self.plan));
474            let admission: FileAdmissionFn<_> = Arc::new(|_| Ok(FileAdmission::Admit));
475            let executor = match backend {
476                DeltaReaderBackend::NativeAsync => native_async_executor(
477                    &self.plan,
478                    None,
479                    self.enforce_physical_predicate_rows
480                        .then(|| self.plan.physical_predicate.clone())
481                        .flatten(),
482                )?,
483                DeltaReaderBackend::OfficialKernel => official_kernel_executor(&self.plan)?,
484            };
485            for partition in 0..partition_count {
486                partitions.push_back(execution.partition_stream(
487                    partition,
488                    Arc::clone(&admission),
489                    Arc::clone(&executor),
490                )?);
491            }
492        }
493
494        Ok(DeltaBatchStream {
495            schema,
496            metrics,
497            partitions,
498            predicate: self.predicate,
499            projection,
500            remaining: self.limit,
501            snapshot_version,
502            backend,
503            partition_count,
504            started: false,
505            done: false,
506        })
507    }
508}
509
510/// Pull-driven stream of finalized logical Arrow batches from one Delta scan.
511///
512/// The stream has no inherent whole-result collection method. Callers that
513/// intentionally materialize a result must opt into a stream extension trait.
514///
515/// ```compile_fail
516/// use delta_arrow_reader::DeltaBatchStream;
517///
518/// fn collect_without_opt_in(stream: DeltaBatchStream) {
519///     let _ = stream.collect();
520/// }
521/// ```
522pub struct DeltaBatchStream {
523    schema: SchemaRef,
524    metrics: DeltaReadMetrics,
525    partitions: VecDeque<PartitionStream>,
526    predicate: Option<DeltaPredicate>,
527    projection: Option<Vec<usize>>,
528    remaining: Option<usize>,
529    snapshot_version: u64,
530    backend: DeltaReaderBackend,
531    partition_count: usize,
532    started: bool,
533    done: bool,
534}
535
536impl DeltaBatchStream {
537    /// Returns the visible logical output schema.
538    pub fn schema(&self) -> &SchemaRef {
539        &self.schema
540    }
541
542    /// Returns a lightweight shared handle to point-in-time scan metrics.
543    pub fn metrics(&self) -> DeltaReadMetrics {
544        self.metrics.clone()
545    }
546
547    fn start(&mut self) {
548        if self.started {
549            return;
550        }
551        self.started = true;
552        trace_execution_started(self.snapshot_version, self.backend, self.partition_count);
553        for partition in &mut self.partitions {
554            partition.start();
555        }
556    }
557
558    fn complete(&mut self) {
559        if self.done {
560            return;
561        }
562        self.partitions.clear();
563        self.done = true;
564        trace_execution_completed(self.snapshot_version, self.backend, self.partition_count);
565    }
566
567    fn fail(&mut self, error: &DeltaReaderError) {
568        self.partitions.clear();
569        self.done = true;
570        trace_execution_failed(
571            self.snapshot_version,
572            self.backend,
573            self.partition_count,
574            error,
575        );
576    }
577
578    fn finalize_batch(&self, mut batch: RecordBatch) -> Result<RecordBatch, DeltaReaderError> {
579        if let Some(predicate) = self.predicate.as_ref() {
580            batch = evaluate_predicate(&batch, predicate)?;
581        }
582        if let Some(projection) = self.projection.as_ref() {
583            batch = batch
584                .project(projection)
585                .boxed()
586                .context(DataFileReadSnafu {
587                    reason: "direct_projection_failed",
588                })?;
589        }
590        Ok(batch)
591    }
592}
593
594impl Stream for DeltaBatchStream {
595    type Item = Result<RecordBatch, DeltaReaderError>;
596
597    fn poll_next(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
598        let this = self.get_mut();
599        if this.done {
600            return Poll::Ready(None);
601        }
602        this.start();
603
604        loop {
605            let Some(partition) = this.partitions.front_mut() else {
606                this.complete();
607                return Poll::Ready(None);
608            };
609            match Pin::new(partition).poll_next(context) {
610                Poll::Ready(Some(Ok(batch))) => {
611                    let mut batch = match this.finalize_batch(batch) {
612                        Ok(batch) => batch,
613                        Err(error) => {
614                            this.fail(&error);
615                            return Poll::Ready(Some(Err(error)));
616                        }
617                    };
618                    if let Some(remaining) = this.remaining.as_mut() {
619                        if batch.num_rows() >= *remaining {
620                            batch = batch.slice(0, *remaining);
621                            *remaining = 0;
622                            this.complete();
623                        } else {
624                            *remaining -= batch.num_rows();
625                        }
626                    }
627                    return Poll::Ready(Some(Ok(batch)));
628                }
629                Poll::Ready(Some(Err(error))) => {
630                    this.fail(&error);
631                    return Poll::Ready(Some(Err(error)));
632                }
633                Poll::Ready(None) => {
634                    this.partitions.pop_front();
635                }
636                Poll::Pending => return Poll::Pending,
637            }
638        }
639    }
640}
641
642impl Drop for DeltaBatchStream {
643    fn drop(&mut self) {
644        if self.done {
645            return;
646        }
647        self.partitions.clear();
648        self.done = true;
649        trace_execution_dropped(self.snapshot_version, self.backend, self.partition_count);
650    }
651}
652
653fn validate_direct_execution_options(
654    options: DeltaReaderExecutionOptions,
655) -> Result<(), DeltaReaderError> {
656    options.validate()?;
657    validate_backend_available(options)?;
658    Ok(())
659}
660
661#[cfg(feature = "native-async")]
662pub(crate) fn native_async_executor(
663    plan: &Arc<DeltaScanPlan>,
664    output_batch_size: Option<usize>,
665    row_predicate: Option<crate::kernel::DeltaKernelPredicate>,
666) -> Result<FileExecutor<crate::planning::DeltaScanFileTask, FileBatchStream>, DeltaReaderError> {
667    Ok(crate::native_async_reader::native_async_file_executor(
668        plan,
669        output_batch_size,
670        row_predicate,
671    ))
672}
673
674#[cfg(not(feature = "native-async"))]
675pub(crate) fn native_async_executor(
676    _plan: &Arc<DeltaScanPlan>,
677    _output_batch_size: Option<usize>,
678    _row_predicate: Option<crate::kernel::DeltaKernelPredicate>,
679) -> Result<FileExecutor<crate::planning::DeltaScanFileTask, FileBatchStream>, DeltaReaderError> {
680    crate::error::UnsupportedBackendSnafu {
681        reason: "native_async_feature_disabled",
682    }
683    .fail()
684}
685
686#[cfg(feature = "official-kernel")]
687pub(crate) fn official_kernel_executor(
688    plan: &Arc<DeltaScanPlan>,
689) -> Result<FileExecutor<crate::planning::DeltaScanFileTask, FileBatchStream>, DeltaReaderError> {
690    Ok(crate::official_kernel_reader::official_kernel_file_executor(plan))
691}
692
693#[cfg(not(feature = "official-kernel"))]
694pub(crate) fn official_kernel_executor(
695    _plan: &Arc<DeltaScanPlan>,
696) -> Result<FileExecutor<crate::planning::DeltaScanFileTask, FileBatchStream>, DeltaReaderError> {
697    crate::error::UnsupportedBackendSnafu {
698        reason: "official_kernel_feature_disabled",
699    }
700    .fail()
701}
702
703fn trace_planning_started(snapshot_version: u64, backend: DeltaReaderBackend) {
704    tracing::debug!(
705        target: TRACING_TARGET,
706        event = "scan_planning.started",
707        snapshot_version,
708        backend = ?backend,
709        partition_count = tracing::field::Empty,
710        outcome = "started"
711    );
712}
713
714fn trace_planning_completed(
715    snapshot_version: u64,
716    backend: DeltaReaderBackend,
717    partition_count: usize,
718) {
719    tracing::debug!(
720        target: TRACING_TARGET,
721        event = "scan_planning.completed",
722        snapshot_version,
723        backend = ?backend,
724        partition_count,
725        outcome = "completed"
726    );
727}
728
729fn trace_planning_failed(
730    snapshot_version: u64,
731    backend: DeltaReaderBackend,
732    error: &DeltaReaderError,
733) {
734    tracing::debug!(
735        target: TRACING_TARGET,
736        event = "scan_planning.failed",
737        snapshot_version,
738        backend = ?backend,
739        partition_count = tracing::field::Empty,
740        outcome = "failed",
741        error_variant = error.as_str(),
742        error_phase = error.phase().as_str()
743    );
744}
745
746fn trace_execution_started(
747    snapshot_version: u64,
748    backend: DeltaReaderBackend,
749    partition_count: usize,
750) {
751    tracing::debug!(
752        target: TRACING_TARGET,
753        event = "scan_execution.started",
754        snapshot_version,
755        backend = ?backend,
756        partition_count,
757        outcome = "started"
758    );
759}
760
761fn trace_execution_completed(
762    snapshot_version: u64,
763    backend: DeltaReaderBackend,
764    partition_count: usize,
765) {
766    tracing::debug!(
767        target: TRACING_TARGET,
768        event = "scan_execution.completed",
769        snapshot_version,
770        backend = ?backend,
771        partition_count,
772        outcome = "completed"
773    );
774}
775
776fn trace_execution_failed(
777    snapshot_version: u64,
778    backend: DeltaReaderBackend,
779    partition_count: usize,
780    error: &DeltaReaderError,
781) {
782    tracing::debug!(
783        target: TRACING_TARGET,
784        event = "scan_execution.failed",
785        snapshot_version,
786        backend = ?backend,
787        partition_count,
788        outcome = "failed",
789        error_variant = error.as_str(),
790        error_phase = error.phase().as_str()
791    );
792}
793
794fn trace_execution_dropped(
795    snapshot_version: u64,
796    backend: DeltaReaderBackend,
797    partition_count: usize,
798) {
799    tracing::debug!(
800        target: TRACING_TARGET,
801        event = "scan_execution.dropped",
802        snapshot_version,
803        backend = ?backend,
804        partition_count,
805        outcome = "dropped"
806    );
807}
808
809#[cfg(test)]
810mod tests {
811    use std::{
812        collections::VecDeque,
813        future::pending,
814        sync::{Arc, Mutex},
815        time::Duration,
816    };
817
818    use arrow::{
819        array::Int32Array,
820        datatypes::{DataType, Field, Schema, SchemaRef},
821        record_batch::RecordBatch,
822    };
823    use futures_util::{FutureExt, StreamExt, stream};
824    use tokio::{sync::Notify, time::timeout};
825    use tracing::{
826        Event, Level, Metadata, Subscriber,
827        span::{Attributes, Id, Record},
828        subscriber::{Interest, with_default},
829    };
830
831    use super::{
832        DeltaBatchStream, trace_execution_completed, trace_execution_dropped,
833        trace_execution_failed, trace_execution_started, trace_planning_completed,
834        trace_planning_failed, trace_planning_started,
835    };
836    use crate::{
837        DeltaReadMetrics, DeltaReaderBackend, DeltaReaderExecutionOptions,
838        error::InvalidConfigurationSnafu,
839        metrics::DeltaReadMetricsConfig,
840        scheduling::{
841            FileAdmission, FileAdmissionFn, FileBatchStream, FileExecutor, FileReadPermit,
842            PartitionStream, ScanCancellation, ScanReadLimiter,
843        },
844    };
845
846    #[derive(Clone, Default)]
847    struct EventFields(Arc<Mutex<Vec<Vec<String>>>>);
848
849    impl Subscriber for EventFields {
850        fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
851            if metadata.target() == "delta_arrow_reader" && *metadata.level() == Level::DEBUG {
852                Interest::always()
853            } else {
854                Interest::sometimes()
855            }
856        }
857
858        fn enabled(&self, metadata: &Metadata<'_>) -> bool {
859            metadata.target() == "delta_arrow_reader" && *metadata.level() == Level::DEBUG
860        }
861
862        fn new_span(&self, _attributes: &Attributes<'_>) -> Id {
863            Id::from_u64(1)
864        }
865
866        fn record(&self, _span: &Id, _values: &Record<'_>) {}
867
868        fn record_follows_from(&self, _span: &Id, _follows: &Id) {}
869
870        fn event(&self, event: &Event<'_>) {
871            let metadata = event.metadata();
872            assert_eq!(metadata.target(), "delta_arrow_reader");
873            self.0.lock().expect("event lock").push(
874                metadata
875                    .fields()
876                    .iter()
877                    .map(|field| field.name().to_owned())
878                    .collect(),
879            );
880        }
881
882        fn enter(&self, _span: &Id) {}
883
884        fn exit(&self, _span: &Id) {}
885    }
886
887    struct ControlledMerge {
888        stream: DeltaBatchStream,
889        limiter: Arc<ScanReadLimiter>,
890        cancellation: ScanCancellation,
891        metrics: DeltaReadMetrics,
892        first_partition_gate: Arc<Notify>,
893    }
894
895    fn schema() -> SchemaRef {
896        Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]))
897    }
898
899    fn batch(id: i32) -> RecordBatch {
900        RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(vec![id]))])
901            .expect("valid test batch")
902    }
903
904    fn batch_id(batch: &RecordBatch) -> i32 {
905        batch
906            .column(0)
907            .as_any()
908            .downcast_ref::<Int32Array>()
909            .expect("Int32 id")
910            .value(0)
911    }
912
913    fn execution_options() -> Result<DeltaReaderExecutionOptions, crate::DeltaReaderError> {
914        DeltaReaderExecutionOptions::new()
915            .with_native_async_prefetch_file_count_per_partition(0)?
916            .with_max_concurrent_file_reads_per_partition(1)?
917            .with_max_concurrent_file_reads_per_scan(Some(2))?
918            .with_output_buffer_capacity_per_partition(1)
919    }
920
921    fn metrics() -> DeltaReadMetrics {
922        DeltaReadMetrics::new(DeltaReadMetricsConfig {
923            snapshot_version: 7,
924            reader_backend: DeltaReaderBackend::NativeAsync,
925            scan_metadata_exhausted: Some(true),
926            scan_partitions_planned: 2,
927            files_planned: 2,
928            files_filtered_during_planning: Some(0),
929            estimated_rows: Some(4),
930            estimated_bytes: Some(4),
931        })
932    }
933
934    fn file_stream(permit: FileReadPermit, batches: Vec<RecordBatch>) -> FileBatchStream {
935        Box::pin(stream::unfold(
936            (VecDeque::from(batches), permit),
937            |(mut batches, permit)| async move {
938                batches
939                    .pop_front()
940                    .map(|batch| (Ok(batch), (batches, permit)))
941            },
942        ))
943    }
944
945    fn gated_file_stream(
946        permit: FileReadPermit,
947        batches: Vec<RecordBatch>,
948        gate: Arc<Notify>,
949    ) -> FileBatchStream {
950        Box::pin(stream::unfold(
951            (false, VecDeque::from(batches), permit, gate),
952            |(wait, mut batches, permit, gate)| async move {
953                let batch = batches.pop_front()?;
954                if wait {
955                    gate.notified().await;
956                }
957                Some((Ok(batch), (true, batches, permit, gate)))
958            },
959        ))
960    }
961
962    fn direct_stream(
963        partitions: VecDeque<PartitionStream>,
964        metrics: DeltaReadMetrics,
965    ) -> DeltaBatchStream {
966        DeltaBatchStream {
967            schema: schema(),
968            metrics,
969            partitions,
970            predicate: None,
971            projection: None,
972            remaining: None,
973            snapshot_version: 7,
974            backend: DeltaReaderBackend::NativeAsync,
975            partition_count: 2,
976            started: false,
977            done: false,
978        }
979    }
980
981    fn controlled_merge() -> Result<ControlledMerge, Box<dyn std::error::Error>> {
982        let options = execution_options()?;
983        let limiter = ScanReadLimiter::new(options, 2, 2);
984        let cancellation = ScanCancellation::new();
985        let metrics = metrics();
986        let first_partition_gate = Arc::new(Notify::new());
987        let executor: FileExecutor<i32, FileBatchStream> = {
988            let gate = Arc::clone(&first_partition_gate);
989            Arc::new(move |task, permit, _| {
990                let gate = Arc::clone(&gate);
991                async move {
992                    let batches = vec![batch(task), batch(task * 2)];
993                    Ok(if task == 1 {
994                        gated_file_stream(permit, batches, gate)
995                    } else {
996                        file_stream(permit, batches)
997                    })
998                }
999                .boxed()
1000            })
1001        };
1002        let admission: FileAdmissionFn<i32> = Arc::new(|_: &i32| Ok(FileAdmission::Admit));
1003        let first = PartitionStream::new(
1004            vec![1],
1005            limiter.partition(0)?,
1006            options,
1007            admission.clone(),
1008            Arc::clone(&executor),
1009            metrics.clone(),
1010            cancellation.clone(),
1011        );
1012        let second = PartitionStream::new(
1013            vec![10],
1014            limiter.partition(1)?,
1015            options,
1016            admission,
1017            executor,
1018            metrics.clone(),
1019            cancellation.clone(),
1020        );
1021
1022        Ok(ControlledMerge {
1023            stream: direct_stream(VecDeque::from([first, second]), metrics.clone()),
1024            limiter,
1025            cancellation,
1026            metrics,
1027            first_partition_gate,
1028        })
1029    }
1030
1031    async fn wait_for_batches(metrics: &DeltaReadMetrics, expected: u64) {
1032        timeout(Duration::from_secs(5), async {
1033            while metrics.snapshot().batches_produced < expected {
1034                tokio::task::yield_now().await;
1035            }
1036        })
1037        .await
1038        .expect("batch production reached expected bound");
1039    }
1040
1041    #[test]
1042    fn lifecycle_tracing_has_only_bounded_fields() {
1043        let events = Arc::new(Mutex::new(Vec::new()));
1044        let subscriber = EventFields(Arc::clone(&events));
1045        let error = InvalidConfigurationSnafu { reason: "test" }.build();
1046
1047        let _ = tracing::subscriber::set_global_default(EventFields::default());
1048        with_default(subscriber, || {
1049            tracing::callsite::rebuild_interest_cache();
1050            trace_planning_started(7, DeltaReaderBackend::NativeAsync);
1051            trace_planning_completed(7, DeltaReaderBackend::NativeAsync, 2);
1052            trace_planning_failed(7, DeltaReaderBackend::NativeAsync, &error);
1053            trace_execution_started(7, DeltaReaderBackend::NativeAsync, 2);
1054            trace_execution_completed(7, DeltaReaderBackend::NativeAsync, 2);
1055            trace_execution_failed(7, DeltaReaderBackend::NativeAsync, 2, &error);
1056            trace_execution_dropped(7, DeltaReaderBackend::NativeAsync, 2);
1057        });
1058        tracing::callsite::rebuild_interest_cache();
1059
1060        let events = events.lock().expect("event lock");
1061        assert_eq!(events.len(), 7);
1062        let allowed = [
1063            "backend",
1064            "error_phase",
1065            "error_variant",
1066            "event",
1067            "outcome",
1068            "partition_count",
1069            "snapshot_version",
1070        ];
1071        for fields in events.iter() {
1072            assert!(fields.iter().all(|field| allowed.contains(&field.as_str())));
1073            assert!(fields.contains(&"event".to_owned()));
1074            assert!(fields.contains(&"snapshot_version".to_owned()));
1075            assert!(fields.contains(&"backend".to_owned()));
1076            assert!(fields.contains(&"partition_count".to_owned()));
1077            assert!(fields.contains(&"outcome".to_owned()));
1078        }
1079    }
1080
1081    #[tokio::test]
1082    async fn merged_stream_is_ordered_and_bounds_later_partition_queues()
1083    -> Result<(), Box<dyn std::error::Error>> {
1084        let ControlledMerge {
1085            mut stream,
1086            limiter,
1087            metrics,
1088            first_partition_gate,
1089            ..
1090        } = controlled_merge()?;
1091
1092        let first = stream.next().await.ok_or("first batch missing")??;
1093        assert_eq!(batch_id(&first), 1);
1094        wait_for_batches(&metrics, 2).await;
1095        for _ in 0..32 {
1096            tokio::task::yield_now().await;
1097        }
1098        assert_eq!(metrics.snapshot().batches_produced, 2);
1099        assert_eq!(limiter.active_file_reads(), 2);
1100
1101        first_partition_gate.notify_one();
1102        let mut ids = vec![batch_id(
1103            &stream.next().await.ok_or("second batch missing")??,
1104        )];
1105        while let Some(batch) = stream.next().await {
1106            ids.push(batch_id(&batch?));
1107        }
1108        assert_eq!(ids, [2, 10, 20]);
1109        assert_eq!(metrics.snapshot().batches_produced, 4);
1110        assert_eq!(metrics.snapshot().scan_partitions_completed, 2);
1111        assert_eq!(limiter.active_file_reads(), 0);
1112        Ok(())
1113    }
1114
1115    #[tokio::test]
1116    async fn merged_stream_drop_cancels_blocked_partitions_and_releases_permits()
1117    -> Result<(), Box<dyn std::error::Error>> {
1118        let ControlledMerge {
1119            mut stream,
1120            limiter,
1121            cancellation,
1122            metrics,
1123            ..
1124        } = controlled_merge()?;
1125
1126        let first = stream.next().await.ok_or("first batch missing")??;
1127        assert_eq!(batch_id(&first), 1);
1128        wait_for_batches(&metrics, 2).await;
1129        assert_eq!(limiter.active_file_reads(), 2);
1130        drop(stream);
1131
1132        assert!(cancellation.is_cancelled());
1133        timeout(Duration::from_secs(5), async {
1134            while limiter.active_file_reads() != 0 {
1135                tokio::task::yield_now().await;
1136            }
1137        })
1138        .await?;
1139        assert_eq!(metrics.snapshot().batches_produced, 2);
1140        assert_eq!(metrics.snapshot().scan_partitions_completed, 0);
1141        Ok(())
1142    }
1143
1144    #[tokio::test]
1145    async fn merged_stream_forwards_one_concurrent_error_and_releases_permits()
1146    -> Result<(), Box<dyn std::error::Error>> {
1147        let options = execution_options()?;
1148        let limiter = ScanReadLimiter::new(options, 2, 2);
1149        let cancellation = ScanCancellation::new();
1150        let metrics = metrics();
1151        let executor: FileExecutor<i32, FileBatchStream> = Arc::new(|task, permit, _| {
1152            async move {
1153                Ok(if task == 1 {
1154                    Box::pin(stream::once(async move {
1155                        let _permit = permit;
1156                        pending::<Result<RecordBatch, crate::DeltaReaderError>>().await
1157                    })) as FileBatchStream
1158                } else {
1159                    Box::pin(stream::once(async move {
1160                        let _permit = permit;
1161                        Err(InvalidConfigurationSnafu {
1162                            reason: "controlled_partition_failure",
1163                        }
1164                        .build())
1165                    })) as FileBatchStream
1166                })
1167            }
1168            .boxed()
1169        });
1170        let admission = Arc::new(|_: &i32| Ok(FileAdmission::Admit));
1171        let first = PartitionStream::new(
1172            vec![1],
1173            limiter.partition(0)?,
1174            options,
1175            admission.clone(),
1176            Arc::clone(&executor),
1177            metrics.clone(),
1178            cancellation.clone(),
1179        );
1180        let second = PartitionStream::new(
1181            vec![2],
1182            limiter.partition(1)?,
1183            options,
1184            admission,
1185            executor,
1186            metrics.clone(),
1187            cancellation.clone(),
1188        );
1189        let mut stream = direct_stream(VecDeque::from([first, second]), metrics);
1190
1191        let error = timeout(Duration::from_secs(5), stream.next())
1192            .await?
1193            .ok_or("error item missing")?
1194            .expect_err("controlled partition must fail");
1195        assert_eq!(error.as_str(), "invalid_configuration");
1196        assert!(stream.next().await.is_none());
1197        assert!(cancellation.is_cancelled());
1198        timeout(Duration::from_secs(5), async {
1199            while limiter.active_file_reads() != 0 {
1200                tokio::task::yield_now().await;
1201            }
1202        })
1203        .await?;
1204        Ok(())
1205    }
1206}
delta_arrow_reader/direct.rs

delta_arrow_reader/
direct.rs
