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

nautilus_persistence/common/
datafusion.rs
