nautilus_persistence/common/
datafusion.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! DataFusion planning and blocking execution shared by catalog backends.
17#![expect(
18    clippy::missing_errors_doc,
19    reason = "DataFusion session methods forward query and Arrow errors for controlled schemas"
20)]
21
22use std::sync::Arc;
23
24use ahash::AHashSet;
25use async_trait::async_trait;
26use datafusion::{
27    arrow::{
28        array::{
29            Array, ArrayRef, BinaryViewArray, FixedSizeBinaryBuilder, FixedSizeListArray,
30            ListArray, StringArray, StringViewArray, UInt32Array, new_empty_array, new_null_array,
31        },
32        buffer::{OffsetBuffer, ScalarBuffer},
33        compute::{cast, concat, take_record_batch},
34        datatypes::{DataType, Schema, SchemaRef},
35        error::ArrowError,
36        record_batch::RecordBatch,
37    },
38    catalog::{Session, TableProvider},
39    common::{Statistics, stats::Precision},
40    datasource::{
41        file_format::{
42            FileFormat, FileMeta, file_compression_type::FileCompressionType, options::ReadOptions,
43        },
44        listing::{ListingTable, ListingTableConfig, ListingTableUrl},
45        physical_plan::{FileScanConfig, FileSource},
46        table_schema::TableSchema,
47    },
48    error::{DataFusionError, Result},
49    physical_expr::LexOrdering,
50    physical_plan::{EmptyRecordBatchStream, ExecutionPlan, SendableRecordBatchStream},
51    prelude::*,
52};
53use futures::{Stream, StreamExt, TryStreamExt};
54use nautilus_common::live::{block_on_nautilus_with, get_runtime};
55use nautilus_core::UnixNanos;
56use nautilus_serialization::arrow::{KEY_IDENTIFIER, StringColumnRef};
57use object_store::{ObjectMeta, ObjectStore};
58use tokio::{
59    sync::mpsc::{self, Receiver},
60    task::JoinHandle,
61};
62use url::Url;
63
64use crate::common::{arrow::validate_catalog_schema, storage::StorageBackend};
65
66/// Batches buffered ahead of a blocking consumer.
67///
68/// A depth of one would stall the producing task on every item until the consumer takes it, so the
69/// object-store read of the next batches cannot overlap with decoding the current one.
70const BATCH_STREAM_BUFFER: usize = 4;
71
72pub(crate) struct BlockingBatchStream<T> {
73    receiver: Receiver<T>,
74    task: JoinHandle<()>,
75}
76
77impl<T> BlockingBatchStream<T> {
78    pub(crate) fn from_stream_with_runtime<S>(stream: S, runtime: &tokio::runtime::Handle) -> Self
79    where
80        S: Stream<Item = T> + Send + 'static,
81        T: Send + 'static,
82    {
83        let (sender, receiver) = mpsc::channel(BATCH_STREAM_BUFFER);
84
85        let task = runtime.spawn(async move {
86            futures::pin_mut!(stream);
87
88            while let Some(item) = stream.next().await {
89                if sender.send(item).await.is_err() {
90                    break;
91                }
92            }
93        });
94
95        Self { receiver, task }
96    }
97}
98
99impl<T: Send> Iterator for BlockingBatchStream<T> {
100    type Item = T;
101
102    fn next(&mut self) -> Option<Self::Item> {
103        block_on_nautilus_with(|| self.receiver.recv())
104    }
105}
106
107impl<T> Drop for BlockingBatchStream<T> {
108    fn drop(&mut self) {
109        self.receiver.close();
110        self.task.abort();
111    }
112}
113
114/// Provides a DataFusion session for registering and querying catalog table sources.
115pub struct DataBackendSession {
116    pub chunk_size: usize,
117    pub runtime: tokio::runtime::Handle,
118    pub(crate) session_ctx: SessionContext,
119    pub(crate) registered_tables: AHashSet<String>,
120}
121
122impl DataBackendSession {
123    /// Creates a new [`DataBackendSession`] instance.
124    #[must_use]
125    pub fn new(chunk_size: usize) -> Self {
126        let runtime = get_runtime().handle().clone();
127        let session_ctx = SessionContext::new_with_config(session_config());
128
129        Self {
130            chunk_size,
131            runtime,
132            session_ctx,
133            registered_tables: AHashSet::new(),
134        }
135    }
136
137    /// Register an object store with the session context
138    pub fn register_object_store(&mut self, url: &Url, object_store: Arc<dyn ObjectStore>) {
139        self.session_ctx.register_object_store(url, object_store);
140    }
141
142    /// Registers a storage backend's object store with the session context.
143    ///
144    /// External catalog implementations can call this before adding native table providers or
145    /// object-store-relative file paths to the session.
146    ///
147    /// # Errors
148    ///
149    /// Returns an error if the storage URI cannot be converted into a DataFusion root URL.
150    pub fn register_storage_backend(&mut self, storage: &StorageBackend) -> anyhow::Result<()> {
151        let root_url = storage.datafusion_root_url()?;
152        self.register_object_store(&root_url, storage.object_store.clone());
153        Ok(())
154    }
155
156    /// Registers a table provider with the session context.
157    ///
158    /// This supports non-Parquet table formats such as Delta Lake while keeping the Parquet
159    /// registration path unchanged.
160    pub fn register_table_provider(
161        &mut self,
162        table_name: &str,
163        provider: Arc<dyn TableProvider>,
164    ) -> Result<()> {
165        if !self.registered_tables.contains(table_name) {
166            self.session_ctx.register_table(table_name, provider)?;
167            self.registered_tables.insert(table_name.to_string());
168        }
169
170        Ok(())
171    }
172
173    pub(crate) fn collect_parquet_files_batches(
174        &mut self,
175        table_name: &str,
176        file_paths: Vec<String>,
177        sql_query: Option<&str>,
178    ) -> anyhow::Result<Vec<RecordBatch>> {
179        if file_paths.is_empty() {
180            return Ok(Vec::new());
181        }
182
183        let batch_stream = self.parquet_files_batch_stream(table_name, file_paths, sql_query)?;
184        let schema = batch_stream.schema();
185        let mut batches = block_on_nautilus_with(|| batch_stream.try_collect::<Vec<_>>())?;
186        if batches.is_empty() {
187            batches.push(RecordBatch::new_empty(schema));
188        }
189
190        Ok(batches)
191    }
192
193    pub(crate) fn parquet_files_batch_stream(
194        &mut self,
195        table_name: &str,
196        file_paths: Vec<String>,
197        sql_query: Option<&str>,
198    ) -> anyhow::Result<SendableRecordBatchStream> {
199        if file_paths.is_empty() {
200            return Ok(Box::pin(EmptyRecordBatchStream::new(Arc::new(
201                Schema::empty(),
202            ))));
203        }
204
205        self.register_parquet_files_table(table_name, file_paths)?;
206        Ok(self.execute_registered(table_name, sql_query)?)
207    }
208
209    // Returns `false` without registering when the files' schemas cannot merge, such as files
210    // written at different precisions.
211    pub(crate) fn try_register_parquet_files_table(
212        &mut self,
213        table_name: &str,
214        file_paths: Vec<String>,
215    ) -> anyhow::Result<bool> {
216        match self.register_parquet_files_table(table_name, file_paths) {
217            Err(e) if is_schema_merge_error(&e) => Ok(false),
218            result => result.map(|()| true),
219        }
220    }
221
222    fn register_parquet_files_table(
223        &mut self,
224        table_name: &str,
225        file_paths: Vec<String>,
226    ) -> anyhow::Result<()> {
227        if self.registered_tables.contains(table_name) {
228            return Ok(());
229        }
230
231        // Catalog files are written in ascending ts_init order, declaring it lets files without
232        // sorting metadata skip a sort that can reorder rows with equal ts_init.
233        let parquet_options = ParquetReadOptions::<'_> {
234            skip_metadata: Some(false),
235            file_sort_order: vec![vec![col("ts_init").sort(true, false)]],
236            ..Default::default()
237        };
238
239        let mut listing_options = parquet_options.to_listing_options(
240            &self.session_ctx.copied_config(),
241            self.session_ctx.copied_table_options(),
242        );
243        listing_options.format = Arc::new(CatalogParquetFormat(listing_options.format));
244        let table_paths = file_paths
245            .into_iter()
246            .map(ListingTableUrl::parse)
247            .collect::<Result<Vec<_>>>()?;
248        let state = self.session_ctx.state();
249        let config = block_on_nautilus_with(|| {
250            ListingTableConfig::new_with_multi_paths(table_paths)
251                .with_listing_options(listing_options)
252                .infer_schema(&state)
253        })?;
254        let table = ListingTable::try_new(config)?.with_cache(
255            self.session_ctx
256                .runtime_env()
257                .cache_manager
258                .get_file_statistic_cache(),
259        );
260        let dataframe = self.session_ctx.read_table(Arc::new(table))?;
261
262        validate_catalog_schema(dataframe.schema().as_arrow())?;
263        self.session_ctx
264            .register_table(table_name, dataframe.into_view())?;
265        self.registered_tables.insert(table_name.to_string());
266        Ok(())
267    }
268
269    pub(crate) fn execute_registered(
270        &self,
271        table_name: &str,
272        sql_query: Option<&str>,
273    ) -> Result<SendableRecordBatchStream> {
274        let default_query = format!("SELECT * FROM {table_name} ORDER BY ts_init");
275        let sql_query = sql_query.unwrap_or(&default_query);
276        let query = block_on_nautilus_with(|| self.session_ctx.sql(sql_query))?;
277        block_on_nautilus_with(|| query.execute_stream())
278    }
279
280    /// Clears all registered tables.
281    ///
282    /// This is useful when the underlying files have changed and we need to
283    /// re-register tables with updated data.
284    pub fn clear_registered_tables(&mut self) {
285        self.registered_tables.clear();
286
287        // Create a new session context to completely reset the DataFusion state
288        self.session_ctx = SessionContext::new_with_config(session_config());
289    }
290}
291
292fn is_schema_merge_error(error: &anyhow::Error) -> bool {
293    error.chain().any(|cause| {
294        matches!(
295            cause.downcast_ref::<ArrowError>(),
296            Some(ArrowError::SchemaError(_))
297        )
298    })
299}
300
301pub(crate) fn session_config() -> SessionConfig {
302    SessionConfig::new()
303        .set_str("datafusion.optimizer.repartition_file_scans", "false")
304        // Repartitioning filtered batches can reorder rows with equal ts_init
305        .set_str(
306            "datafusion.optimizer.enable_round_robin_repartition",
307            "false",
308        )
309        .set_str("datafusion.optimizer.prefer_existing_sort", "true")
310}
311
312// Parquet format whose file statistics omit float bounds. Parquet float statistics skip NaN, so a
313// column holding NaN and one other value reports equal bounds, and DataFusion then reads that value
314// for every row (apache/datafusion#26490). The other bounds stay, so DataFusion can still prove the
315// files of one scan group are ordered by `ts_init` and stream them without sorting.
316#[derive(Debug)]
317struct CatalogParquetFormat(Arc<dyn FileFormat>);
318
319#[async_trait]
320impl FileFormat for CatalogParquetFormat {
321    fn get_ext(&self) -> String {
322        self.0.get_ext()
323    }
324
325    fn get_ext_with_compression(&self, compression: &FileCompressionType) -> Result<String> {
326        self.0.get_ext_with_compression(compression)
327    }
328
329    fn compression_type(&self) -> Option<FileCompressionType> {
330        self.0.compression_type()
331    }
332
333    async fn infer_schema(
334        &self,
335        state: &dyn Session,
336        store: &Arc<dyn ObjectStore>,
337        objects: &[ObjectMeta],
338    ) -> Result<SchemaRef> {
339        self.0.infer_schema(state, store, objects).await
340    }
341
342    async fn infer_stats(
343        &self,
344        state: &dyn Session,
345        store: &Arc<dyn ObjectStore>,
346        table_schema: SchemaRef,
347        object: &ObjectMeta,
348    ) -> Result<Statistics> {
349        let statistics = self
350            .0
351            .infer_stats(state, store, Arc::clone(&table_schema), object)
352            .await?;
353        Ok(statistics_without_float_bounds(statistics, &table_schema))
354    }
355
356    async fn infer_ordering(
357        &self,
358        state: &dyn Session,
359        store: &Arc<dyn ObjectStore>,
360        table_schema: SchemaRef,
361        object: &ObjectMeta,
362    ) -> Result<Option<LexOrdering>> {
363        self.0
364            .infer_ordering(state, store, table_schema, object)
365            .await
366    }
367
368    async fn infer_stats_and_ordering(
369        &self,
370        state: &dyn Session,
371        store: &Arc<dyn ObjectStore>,
372        table_schema: SchemaRef,
373        object: &ObjectMeta,
374    ) -> Result<FileMeta> {
375        let meta = self
376            .0
377            .infer_stats_and_ordering(state, store, Arc::clone(&table_schema), object)
378            .await?;
379        let statistics = statistics_without_float_bounds(meta.statistics, &table_schema);
380        Ok(FileMeta::new(statistics).with_ordering(meta.ordering))
381    }
382
383    async fn create_physical_plan(
384        &self,
385        state: &dyn Session,
386        conf: FileScanConfig,
387    ) -> Result<Arc<dyn ExecutionPlan>> {
388        self.0.create_physical_plan(state, conf).await
389    }
390
391    fn file_source(&self, table_schema: TableSchema) -> Arc<dyn FileSource> {
392        self.0.file_source(table_schema)
393    }
394}
395
396fn statistics_without_float_bounds(mut statistics: Statistics, schema: &Schema) -> Statistics {
397    for (column, field) in statistics.column_statistics.iter_mut().zip(schema.fields()) {
398        if field.data_type().is_floating() {
399            column.min_value = Precision::Absent;
400            column.max_value = Precision::Absent;
401        }
402    }
403    statistics
404}
405
406pub(crate) fn cast_record_batch_to_schema(
407    batch: &RecordBatch,
408    schema: &Schema,
409) -> Result<RecordBatch> {
410    let batch_schema = batch.schema();
411    let mut fields = Vec::with_capacity(batch_schema.fields().len());
412    let mut columns = Vec::with_capacity(batch.columns().len());
413
414    for (column, batch_field) in batch.columns().iter().zip(batch_schema.fields()) {
415        let field = schema
416            .field_with_name(batch_field.name())
417            .unwrap_or(batch_field.as_ref());
418
419        fields.push(Arc::new(field.clone()));
420        columns.push(cast_column_to_data_type(column, field.data_type())?);
421    }
422
423    let schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()));
424    Ok(RecordBatch::try_new(schema, columns)?)
425}
426
427fn cast_column_to_data_type(column: &ArrayRef, data_type: &DataType) -> Result<ArrayRef> {
428    if column.data_type() == data_type {
429        return Ok(column.clone());
430    }
431
432    if let (DataType::BinaryView, DataType::FixedSizeBinary(width)) =
433        (column.data_type(), data_type)
434    {
435        let array = column
436            .as_any()
437            .downcast_ref::<BinaryViewArray>()
438            .expect("BinaryView column should downcast to BinaryViewArray");
439        let mut builder = FixedSizeBinaryBuilder::with_capacity(array.len(), *width);
440        for row in 0..array.len() {
441            if array.is_null(row) {
442                builder.append_null();
443            } else {
444                builder.append_value(array.value(row))?;
445            }
446        }
447
448        return Ok(Arc::new(builder.finish()));
449    }
450
451    if let (DataType::FixedSizeList(_, size), DataType::List(field)) =
452        (column.data_type(), data_type)
453    {
454        let array = column
455            .as_any()
456            .downcast_ref::<FixedSizeListArray>()
457            .expect("FixedSizeList column should downcast to FixedSizeListArray");
458
459        let size = usize::try_from(*size).map_err(|e| {
460            DataFusionError::Execution(format!("Invalid fixed-size list length {size}: {e}"))
461        })?;
462
463        let mut parts = Vec::with_capacity(array.len());
464        let mut offsets = Vec::with_capacity(array.len() + 1);
465        offsets.push(0_i32);
466
467        for row in 0..array.len() {
468            if !array.is_null(row) {
469                let value = array.value(row);
470                parts.push(cast_column_to_data_type(&value, field.data_type())?);
471            }
472
473            let offset = parts
474                .len()
475                .checked_mul(size)
476                .and_then(|offset| i32::try_from(offset).ok())
477                .ok_or_else(|| {
478                    DataFusionError::Execution(
479                        "List offset exceeds the supported i32 range".to_string(),
480                    )
481                })?;
482
483            offsets.push(offset);
484        }
485
486        let values = concat_list_values(&parts, field.data_type())?;
487        let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets));
488        return Ok(Arc::new(ListArray::try_new(
489            field.clone(),
490            offsets,
491            values,
492            array.nulls().cloned(),
493        )?));
494    }
495
496    if let (DataType::List(_), DataType::FixedSizeList(field, size)) =
497        (column.data_type(), data_type)
498    {
499        let array = column
500            .as_any()
501            .downcast_ref::<ListArray>()
502            .expect("List column should downcast to ListArray");
503
504        let size_usize = usize::try_from(*size).map_err(|e| {
505            DataFusionError::Execution(format!("Invalid fixed-size list length {size}: {e}"))
506        })?;
507
508        let mut parts = Vec::with_capacity(array.len());
509        for row in 0..array.len() {
510            if array.is_null(row) {
511                parts.push(new_null_array(field.data_type(), size_usize));
512                continue;
513            }
514
515            let value = array.value(row);
516            if value.len() != size_usize {
517                return Err(DataFusionError::Execution(format!(
518                    "List row {row} has length {}, expected {size}",
519                    value.len(),
520                )));
521            }
522
523            parts.push(cast_column_to_data_type(&value, field.data_type())?);
524        }
525
526        let values = concat_list_values(&parts, field.data_type())?;
527        return Ok(Arc::new(FixedSizeListArray::try_new(
528            field.clone(),
529            *size,
530            values,
531            array.nulls().cloned(),
532        )?));
533    }
534
535    Ok(cast(column, data_type)?)
536}
537
538fn concat_list_values(parts: &[ArrayRef], data_type: &DataType) -> Result<ArrayRef> {
539    if parts.is_empty() {
540        return Ok(new_empty_array(data_type));
541    }
542
543    let parts = parts.iter().map(AsRef::as_ref).collect::<Vec<_>>();
544    Ok(concat(&parts)?)
545}
546
547#[must_use]
548pub fn build_query(
549    table: &str,
550    start: Option<UnixNanos>,
551    end: Option<UnixNanos>,
552    where_clause: Option<&str>,
553) -> String {
554    let filter = where_sql(start, end, where_clause);
555    format!("SELECT * FROM {table}{filter} ORDER BY ts_init")
556}
557
558#[must_use]
559pub fn build_identifier_query(
560    table: &str,
561    start: Option<UnixNanos>,
562    end: Option<UnixNanos>,
563    where_clause: Option<&str>,
564) -> String {
565    let filter = where_sql(start, end, where_clause);
566    format!("SELECT DISTINCT identifier FROM {table}{filter} ORDER BY identifier")
567}
568
569fn where_sql(
570    start: Option<UnixNanos>,
571    end: Option<UnixNanos>,
572    where_clause: Option<&str>,
573) -> String {
574    let mut conditions = Vec::new();
575
576    if let Some(clause) = where_clause {
577        // Parenthesized so a caller clause containing a top-level `OR` still binds as one
578        // condition once the timestamp bounds are appended with `AND`.
579        conditions.push(format!("({clause})"));
580    }
581
582    if let Some(start_ts) = start {
583        conditions.push(format!("CAST(ts_init AS BIGINT) >= {start_ts}"));
584    }
585
586    if let Some(end_ts) = end {
587        conditions.push(format!("CAST(ts_init AS BIGINT) <= {end_ts}"));
588    }
589
590    if conditions.is_empty() {
591        return String::new();
592    }
593
594    format!(" WHERE {}", conditions.join(" AND "))
595}
596
597pub fn identifiers_from_record_batches(batches: &[RecordBatch]) -> anyhow::Result<Vec<String>> {
598    let mut identifiers = AHashSet::new();
599
600    for batch in batches {
601        let column = batch
602            .column_by_name("identifier")
603            .ok_or_else(|| anyhow::anyhow!("identifier column not found"))?
604            .as_any();
605
606        if let Some(array) = column.downcast_ref::<StringArray>() {
607            for row in 0..array.len() {
608                if !array.is_null(row) {
609                    identifiers.insert(array.value(row).to_string());
610                }
611            }
612
613            continue;
614        }
615
616        if let Some(array) = column.downcast_ref::<StringViewArray>() {
617            for row in 0..array.len() {
618                if !array.is_null(row) {
619                    identifiers.insert(array.value(row).to_string());
620                }
621            }
622
623            continue;
624        }
625
626        anyhow::bail!("identifier column must be Utf8 or Utf8View");
627    }
628
629    let mut identifiers = identifiers.into_iter().collect::<Vec<_>>();
630    identifiers.sort();
631    Ok(identifiers)
632}
633
634/// Keeps the rows whose catalog identifier satisfies `matches`.
635///
636/// Each row's identifier comes from the `identifier` column, or `fallback_identifier` when the
637/// batch has no such column or the row value is null. Returns `None` when no row matches.
638pub(crate) fn filter_record_batch_by_identifier(
639    batch: &RecordBatch,
640    fallback_identifier: Option<&str>,
641    matches: impl Fn(Option<&str>) -> bool,
642) -> anyhow::Result<Option<RecordBatch>> {
643    let Some(column) = batch.column_by_name(KEY_IDENTIFIER) else {
644        return Ok(matches(fallback_identifier).then(|| batch.clone()));
645    };
646
647    let identifiers = StringColumnRef::try_from_array(column.as_ref())
648        .ok_or_else(|| anyhow::anyhow!("Identifier column must be an Arrow string type"))?;
649
650    let indices = (0..batch.num_rows())
651        .filter_map(|row| {
652            let identifier = (!identifiers.is_null(row))
653                .then(|| identifiers.value(row))
654                .or(fallback_identifier);
655            matches(identifier).then(|| u32::try_from(row))
656        })
657        .collect::<Result<Vec<_>, _>>()?;
658
659    if indices.is_empty() {
660        return Ok(None);
661    }
662
663    if indices.len() == batch.num_rows() {
664        return Ok(Some(batch.clone()));
665    }
666
667    take_record_batch(batch, &UInt32Array::from(indices))
668        .map(Some)
669        .map_err(|e| anyhow::anyhow!("Failed to filter catalog batch by identifier: {e}"))
670}
671
672#[cfg(test)]
673mod tests {
674    use std::{
675        fs::File,
676        ops::Range,
677        path::Path,
678        sync::{Arc, mpsc},
679        time::Duration,
680    };
681
682    use datafusion::{
683        arrow::{
684            array::{AsArray, BooleanArray, Decimal128Array, FixedSizeBinaryArray, UInt32Array},
685            buffer::NullBuffer,
686            compute::{concat_batches, filter_record_batch},
687            datatypes::{Field, TimestampNanosecondType},
688        },
689        datasource::MemTable,
690        execution::object_store::ObjectStoreUrl,
691        physical_plan::displayable,
692    };
693    use nautilus_testkit::common::get_test_data_file_path;
694    use parquet::{
695        arrow::{ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder},
696        file::{metadata::SortingColumn, properties::WriterProperties},
697    };
698    use rstest::rstest;
699    use tempfile::TempDir;
700
701    use super::*;
702    use crate::common::storage::create_storage_backend_from_path;
703
704    #[rstest]
705    fn register_storage_backend_accepts_memory_backend() {
706        let storage = create_storage_backend_from_path("memory://", None).unwrap();
707        let mut session = DataBackendSession::new(10);
708
709        session.register_storage_backend(&storage).unwrap();
710
711        assert_registered_object_store(&session, &storage);
712    }
713
714    #[rstest]
715    fn register_storage_backend_accepts_local_backend() {
716        let temp_dir = TempDir::new().unwrap();
717        let storage =
718            create_storage_backend_from_path(temp_dir.path().to_str().unwrap(), None).unwrap();
719        let mut session = DataBackendSession::new(10);
720
721        session.register_storage_backend(&storage).unwrap();
722
723        assert_registered_object_store(&session, &storage);
724    }
725
726    fn assert_registered_object_store(session: &DataBackendSession, storage: &StorageBackend) {
727        let root_url = ObjectStoreUrl::parse(storage.datafusion_root_url().unwrap()).unwrap();
728        let registered = session
729            .session_ctx
730            .runtime_env()
731            .object_store(root_url)
732            .unwrap();
733
734        assert!(std::ptr::addr_eq(
735            Arc::as_ptr(&registered),
736            Arc::as_ptr(&storage.object_store),
737        ));
738    }
739
740    fn memory_table(column: &str) -> Arc<dyn TableProvider> {
741        let schema = Arc::new(Schema::new(vec![Field::new(column, DataType::Utf8, true)]));
742        Arc::new(MemTable::try_new(schema, vec![Vec::new()]).unwrap())
743    }
744
745    #[rstest]
746    fn register_table_provider_keeps_first_registration() {
747        let mut session = DataBackendSession::new(10);
748
749        session
750            .register_table_provider("records", memory_table("first"))
751            .unwrap();
752        session
753            .register_table_provider("records", memory_table("second"))
754            .unwrap();
755        let provider =
756            futures::executor::block_on(session.session_ctx.table_provider("records")).unwrap();
757
758        assert_eq!(provider.schema().field(0).name(), "first");
759    }
760
761    #[rstest]
762    fn session_config_keeps_file_scan_order() {
763        let config = session_config();
764        let optimizer = &config.options().optimizer;
765
766        assert!(!optimizer.repartition_file_scans);
767        assert!(!optimizer.enable_round_robin_repartition);
768        assert!(optimizer.prefer_existing_sort);
769    }
770
771    #[rstest]
772    fn identifiers_from_record_batches_reads_utf8_view_columns() {
773        let schema = Arc::new(Schema::new(vec![Field::new(
774            "identifier",
775            DataType::Utf8View,
776            true,
777        )]));
778        let batch = RecordBatch::try_new(
779            schema,
780            vec![Arc::new(StringViewArray::from(vec![
781                Some("ESZ4.XCME"),
782                None,
783                Some("ESM4.XCME"),
784            ])) as ArrayRef],
785        )
786        .unwrap();
787
788        let identifiers = identifiers_from_record_batches(&[batch]).unwrap();
789
790        assert_eq!(identifiers, vec!["ESM4.XCME", "ESZ4.XCME"]);
791    }
792
793    #[rstest]
794    #[case::missing_column("instrument_id", DataType::Utf8, "identifier column not found")]
795    #[case::unsupported_type(
796        "identifier",
797        DataType::Int32,
798        "identifier column must be Utf8 or Utf8View"
799    )]
800    fn identifiers_from_record_batches_rejects_invalid_columns(
801        #[case] name: &str,
802        #[case] data_type: DataType,
803        #[case] expected: &str,
804    ) {
805        let column = new_null_array(&data_type, 1);
806        let schema = Arc::new(Schema::new(vec![Field::new(name, data_type, true)]));
807        let batch = RecordBatch::try_new(schema, vec![column]).unwrap();
808
809        let error = identifiers_from_record_batches(&[batch]).unwrap_err();
810
811        assert_eq!(error.to_string(), expected);
812    }
813
814    #[rstest]
815    fn binary_view_cast_to_fixed_size_binary_preserves_nulls() {
816        let values = vec![Some(b"ab".as_slice()), None, Some(b"cd".as_slice())];
817        let column = Arc::new(BinaryViewArray::from(values.clone())) as ArrayRef;
818
819        let cast = cast_column_to_data_type(&column, &DataType::FixedSizeBinary(2)).unwrap();
820
821        let expected =
822            FixedSizeBinaryArray::try_from_sparse_iter_with_size(values.into_iter(), 2).unwrap();
823        assert_eq!(cast.to_data(), expected.to_data());
824    }
825
826    #[rstest]
827    fn list_cast_to_fixed_size_list_rejects_wrong_row_length() {
828        let field = Arc::new(Field::new("element", DataType::UInt32, false));
829        let list = Arc::new(
830            ListArray::try_new(
831                field.clone(),
832                OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 3])),
833                Arc::new(UInt32Array::from(vec![1_u32, 2, 3])),
834                None,
835            )
836            .unwrap(),
837        ) as ArrayRef;
838
839        let error =
840            cast_column_to_data_type(&list, &DataType::FixedSizeList(field, 2)).unwrap_err();
841
842        assert_eq!(
843            error.to_string(),
844            "Execution error: List row 0 has length 3, expected 2"
845        );
846    }
847
848    struct DropSignal(mpsc::Sender<()>);
849
850    impl Drop for DropSignal {
851        fn drop(&mut self) {
852            let _ = self.0.send(());
853        }
854    }
855
856    #[rstest]
857    fn blocking_batch_stream_drop_stops_producer() {
858        let runtime = tokio::runtime::Runtime::new().unwrap();
859        let (dropped_sender, dropped_receiver) = mpsc::channel();
860        let signal = DropSignal(dropped_sender);
861
862        let stream = futures::stream::pending::<i32>().chain(futures::stream::once(async move {
863            drop(signal);
864            0
865        }));
866
867        let stream = BlockingBatchStream::from_stream_with_runtime(stream, runtime.handle());
868
869        drop(stream);
870
871        assert_eq!(
872            dropped_receiver.recv_timeout(Duration::from_secs(5)),
873            Ok(())
874        );
875    }
876
877    #[rstest]
878    fn build_identifier_query_projects_distinct_identifier_with_filters() {
879        let query = build_identifier_query(
880            "quotes",
881            Some(UnixNanos::from(10)),
882            Some(UnixNanos::from(20)),
883            Some("instrument_id LIKE 'ES%'"),
884        );
885
886        assert_eq!(
887            query,
888            "SELECT DISTINCT identifier FROM quotes WHERE (instrument_id LIKE 'ES%') \
889             AND CAST(ts_init AS BIGINT) >= 10 AND CAST(ts_init AS BIGINT) <= 20 ORDER BY identifier"
890        );
891    }
892
893    #[rstest]
894    fn build_query_groups_a_disjunctive_where_clause_against_timestamp_bounds() {
895        let query = build_query(
896            "bars",
897            Some(UnixNanos::from(10)),
898            Some(UnixNanos::from(20)),
899            Some("identifier = 'ES.GLBX' OR starts_with(identifier, 'ES.GLBX-')"),
900        );
901
902        assert_eq!(
903            query,
904            "SELECT * FROM bars WHERE (identifier = 'ES.GLBX' OR starts_with(identifier, 'ES.GLBX-')) \
905             AND CAST(ts_init AS BIGINT) >= 10 AND CAST(ts_init AS BIGINT) <= 20 ORDER BY ts_init"
906        );
907    }
908
909    #[rstest]
910    fn build_query_without_bounds_keeps_the_where_clause_alone() {
911        let query = build_query("bars", None, None, Some("identifier = 'ES.GLBX'"));
912
913        assert_eq!(
914            query,
915            "SELECT * FROM bars WHERE (identifier = 'ES.GLBX') ORDER BY ts_init"
916        );
917    }
918
919    #[rstest]
920    fn build_query_without_conditions_omits_the_where_keyword() {
921        let query = build_query("bars", None, None, None);
922
923        assert_eq!(query, "SELECT * FROM bars ORDER BY ts_init");
924    }
925
926    #[rstest]
927    fn identifiers_from_record_batches_returns_sorted_unique_non_null_values() {
928        let schema = Arc::new(Schema::new(vec![Field::new(
929            "identifier",
930            DataType::Utf8,
931            true,
932        )]));
933        let batch = RecordBatch::try_new(
934            schema,
935            vec![Arc::new(StringArray::from(vec![
936                Some("ESZ4.XCME"),
937                None,
938                Some("ESM4.XCME"),
939                Some("ESZ4.XCME"),
940            ])) as ArrayRef],
941        )
942        .unwrap();
943
944        let identifiers = identifiers_from_record_batches(&[batch]).unwrap();
945
946        assert_eq!(identifiers, vec!["ESM4.XCME", "ESZ4.XCME"]);
947    }
948
949    #[rstest]
950    fn fixed_size_list_cast_round_trips_through_delta_list_type() {
951        let field = Arc::new(Field::new("element", DataType::Decimal128(38, 16), true));
952        let values = Decimal128Array::from(vec![1_i128, 2, 3, 4])
953            .with_precision_and_scale(38, 16)
954            .unwrap();
955        let fixed_type = DataType::FixedSizeList(field.clone(), 2);
956        let fixed = Arc::new(
957            FixedSizeListArray::try_new(field.clone(), 2, Arc::new(values), None).unwrap(),
958        ) as ArrayRef;
959
960        let list_type = DataType::List(field);
961        let list = cast_column_to_data_type(&fixed, &list_type).unwrap();
962        let restored = cast_column_to_data_type(&list, &fixed_type).unwrap();
963
964        assert_eq!(list.data_type(), &list_type);
965        assert_eq!(restored.data_type(), &fixed_type);
966        assert_eq!(restored.to_data(), fixed.to_data());
967
968        let field = Arc::new(Field::new("element", DataType::UInt32, false));
969        let fixed = Arc::new(
970            FixedSizeListArray::try_new(
971                field.clone(),
972                2,
973                Arc::new(UInt32Array::from(vec![1_u32, 2])),
974                Some(NullBuffer::from(vec![false])),
975            )
976            .unwrap(),
977        ) as ArrayRef;
978        let list_type = DataType::List(field);
979        let list = cast_column_to_data_type(&fixed, &list_type).unwrap();
980        let list = list.as_any().downcast_ref::<ListArray>().unwrap();
981
982        assert!(list.is_null(0));
983        assert_eq!(list.values().len(), 0);
984    }
985
986    #[rstest]
987    fn data_backend_sessions_share_global_runtime() {
988        let first = DataBackendSession::new(10);
989        let second = DataBackendSession::new(10);
990
991        assert_eq!(first.runtime.id(), second.runtime.id());
992        assert_eq!(first.runtime.id(), get_runtime().handle().id());
993    }
994
995    #[rstest]
996    fn blocking_batch_stream_prefetches_first_item() {
997        let runtime = tokio::runtime::Runtime::new().unwrap();
998        let (polled_sender, polled_receiver) = mpsc::channel();
999
1000        let stream = futures::stream::once(async move {
1001            polled_sender.send(()).unwrap();
1002            42
1003        });
1004
1005        let mut stream = BlockingBatchStream::from_stream_with_runtime(stream, runtime.handle());
1006
1007        assert_eq!(polled_receiver.recv_timeout(Duration::from_secs(1)), Ok(()),);
1008        assert_eq!(stream.next(), Some(42));
1009    }
1010
1011    #[rstest]
1012    fn blocking_batch_stream_works_inside_current_thread_runtime() {
1013        let runtime = tokio::runtime::Builder::new_current_thread()
1014            .enable_all()
1015            .build()
1016            .unwrap();
1017
1018        runtime.block_on(async {
1019            let mut stream = BlockingBatchStream::from_stream_with_runtime(
1020                futures::stream::iter([42]),
1021                get_runtime().handle(),
1022            );
1023
1024            assert_eq!(stream.next(), Some(42));
1025        });
1026    }
1027
1028    // Rows sharing a ts_init sit on both sides of the second row group boundary in this file
1029    const FILTER_QUERY_FIXTURE: &str = "nautilus/arrow/quotes-3-groups-filter-query.parquet";
1030    const FILTER_QUERY_START: i64 = 1_701_388_832_486_000_000;
1031
1032    #[derive(Clone, Copy, Debug)]
1033    enum Registration {
1034        File,
1035        Directory,
1036        DisjointFiles,
1037    }
1038
1039    #[rstest]
1040    fn query_preserves_equal_timestamp_order(
1041        #[values(
1042            Registration::File,
1043            Registration::Directory,
1044            Registration::DisjointFiles
1045        )]
1046        registration: Registration,
1047        #[values(false, true)] sorting_metadata: bool,
1048        #[values(false, true)] filtered: bool,
1049        #[values(1, 2, 4, 8)] partitions: usize,
1050    ) {
1051        let (fixture, row_groups) = read_filter_query_fixture();
1052        let temp_dir = TempDir::new().unwrap();
1053        let file_count = match registration {
1054            Registration::File | Registration::Directory => 1,
1055            Registration::DisjointFiles => 3,
1056        };
1057
1058        for (index, rows) in disjoint_file_ranges(&fixture, file_count)
1059            .into_iter()
1060            .enumerate()
1061        {
1062            let path = temp_dir.path().join(format!("part-{index}.parquet"));
1063            write_quotes(&path, &fixture, rows, &row_groups, sorting_metadata);
1064        }
1065
1066        // Directory tables register with a trailing separator, like catalog directory tables
1067        let path = match registration {
1068            Registration::File => temp_dir.path().join("part-0.parquet"),
1069            Registration::Directory | Registration::DisjointFiles => temp_dir.path().join(""),
1070        };
1071        let path = path.to_str().unwrap().to_string();
1072
1073        let start = filtered.then(|| UnixNanos::from(FILTER_QUERY_START as u64));
1074        let expected = if filtered {
1075            let ts_init = fixture
1076                .column_by_name("ts_init")
1077                .unwrap()
1078                .as_primitive::<TimestampNanosecondType>();
1079            let mask = ts_init
1080                .values()
1081                .iter()
1082                .map(|value| Some(*value >= FILTER_QUERY_START))
1083                .collect::<BooleanArray>();
1084            filter_record_batch(&fixture, &mask).unwrap()
1085        } else {
1086            fixture
1087        };
1088
1089        let mut session = DataBackendSession::new(10);
1090        session.session_ctx =
1091            SessionContext::new_with_config(session_config().with_target_partitions(partitions));
1092        let query = build_query("quotes", start, None, None);
1093        let batches = session
1094            .collect_parquet_files_batches("quotes", vec![path], Some(&query))
1095            .unwrap();
1096        let actual = concat_batches(&batches[0].schema(), &batches).unwrap();
1097        let actual = cast_record_batch_to_schema(&actual, &expected.schema()).unwrap();
1098        let first_mismatch = (0..expected.num_rows().min(actual.num_rows()))
1099            .find(|&row| actual.slice(row, 1) != expected.slice(row, 1));
1100
1101        assert_eq!(expected.num_rows(), if filtered { 10_001 } else { 15_000 });
1102        assert_eq!(actual.num_rows(), expected.num_rows());
1103        assert_eq!(first_mismatch, None);
1104    }
1105
1106    // Streamed reads rely on DataFusion proving the files of one scan group are ordered, which
1107    // needs their ts_init statistics; a SortExec would buffer the whole input instead.
1108    #[rstest]
1109    fn query_of_disjoint_files_keeps_ordered_scan(#[values(1, 2, 8)] partitions: usize) {
1110        let (fixture, row_groups) = read_filter_query_fixture();
1111        let temp_dir = TempDir::new().unwrap();
1112
1113        for (index, rows) in disjoint_file_ranges(&fixture, 3).into_iter().enumerate() {
1114            let path = temp_dir.path().join(format!("part-{index}.parquet"));
1115            write_quotes(&path, &fixture, rows, &row_groups, true);
1116        }
1117        let path = temp_dir.path().join("").to_str().unwrap().to_string();
1118        let mut session = DataBackendSession::new(10);
1119        session.session_ctx =
1120            SessionContext::new_with_config(session_config().with_target_partitions(partitions));
1121        session
1122            .register_parquet_files_table("quotes", vec![path])
1123            .unwrap();
1124        let query = build_query("quotes", None, None, None);
1125        let plan = block_on_nautilus_with(|| async {
1126            session
1127                .session_ctx
1128                .sql(&query)
1129                .await?
1130                .create_physical_plan()
1131                .await
1132        })
1133        .unwrap();
1134        let plan = displayable(plan.as_ref()).indent(true).to_string();
1135
1136        assert!(!plan.contains("SortExec"), "{plan}");
1137        assert!(plan.contains("output_ordering=[ts_init@"), "{plan}");
1138    }
1139
1140    fn read_filter_query_fixture() -> (RecordBatch, Vec<usize>) {
1141        let path = get_test_data_file_path(FILTER_QUERY_FIXTURE);
1142        let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(path).unwrap()).unwrap();
1143        let row_groups = builder
1144            .metadata()
1145            .row_groups()
1146            .iter()
1147            .map(|row_group| usize::try_from(row_group.num_rows()).unwrap())
1148            .collect();
1149        let schema = builder.schema().clone();
1150        let batches = builder
1151            .build()
1152            .unwrap()
1153            .collect::<std::result::Result<Vec<_>, _>>()
1154            .unwrap();
1155
1156        (concat_batches(&schema, &batches).unwrap(), row_groups)
1157    }
1158
1159    // Splits rows into contiguous ranges that never divide rows sharing a ts_init
1160    fn disjoint_file_ranges(batch: &RecordBatch, count: usize) -> Vec<Range<usize>> {
1161        let ts_init = batch
1162            .column_by_name("ts_init")
1163            .unwrap()
1164            .as_primitive::<TimestampNanosecondType>()
1165            .values();
1166        let mut ranges = Vec::with_capacity(count);
1167        let mut start = 0;
1168
1169        for index in 1..count {
1170            let mut end = (batch.num_rows() * index / count).max(start + 1);
1171            while end < batch.num_rows() && ts_init[end] == ts_init[end - 1] {
1172                end += 1;
1173            }
1174            ranges.push(start..end);
1175            start = end;
1176        }
1177
1178        ranges.push(start..batch.num_rows());
1179        ranges
1180    }
1181
1182    // Writes the rows keeping the fixture's row group boundaries, with or without the
1183    // ts_init sorting metadata the catalog writer stores.
1184    fn write_quotes(
1185        path: &Path,
1186        batch: &RecordBatch,
1187        rows: Range<usize>,
1188        row_groups: &[usize],
1189        sorting_metadata: bool,
1190    ) {
1191        let sorting_columns = sorting_metadata.then(|| {
1192            vec![SortingColumn {
1193                column_idx: i32::try_from(batch.schema().index_of("ts_init").unwrap()).unwrap(),
1194                descending: false,
1195                nulls_first: false,
1196            }]
1197        });
1198        let properties = WriterProperties::builder()
1199            .set_sorting_columns(sorting_columns)
1200            .build();
1201        let mut writer = ArrowWriter::try_new(
1202            File::create(path).unwrap(),
1203            batch.schema(),
1204            Some(properties),
1205        )
1206        .unwrap();
1207        let mut boundary = 0;
1208
1209        for row_group in row_groups {
1210            let group = boundary..boundary + row_group;
1211            boundary = group.end;
1212            let start = group.start.max(rows.start);
1213            let end = group.end.min(rows.end);
1214            if start < end {
1215                writer.write(&batch.slice(start, end - start)).unwrap();
1216                writer.flush().unwrap();
1217            }
1218        }
1219
1220        writer.close().unwrap();
1221    }
1222}
nautilus_persistence/common/datafusion.rs

nautilus_persistence/common/
datafusion.rs
