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

delta_arrow_reader/
direct.rs
