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, new_empty_array, new_null_array,
30        },
31        buffer::{OffsetBuffer, ScalarBuffer},
32        compute::{cast, concat},
33        datatypes::{DataType, Field, Schema},
34        record_batch::RecordBatch,
35    },
36    catalog::TableProvider,
37    error::{DataFusionError, Result},
38    physical_plan::{EmptyRecordBatchStream, SendableRecordBatchStream},
39    prelude::*,
40};
41use futures::{Stream, StreamExt, TryStreamExt};
42use nautilus_common::live::{block_on_nautilus_with, get_runtime};
43use nautilus_core::UnixNanos;
44use object_store::ObjectStore;
45use tokio::{
46    sync::mpsc::{self, Receiver},
47    task::JoinHandle,
48};
49use url::Url;
50
51use crate::common::{arrow::validate_catalog_schema, storage::StorageBackend};
52
53/// Batches buffered ahead of a blocking consumer.
54///
55/// A depth of one would stall the producing task on every item until the consumer takes it, so the
56/// object-store read of the next batches cannot overlap with decoding the current one.
57const BATCH_STREAM_BUFFER: usize = 4;
58
59pub(crate) struct BlockingBatchStream<T> {
60    receiver: Receiver<T>,
61    task: JoinHandle<()>,
62}
63
64impl<T> BlockingBatchStream<T> {
65    pub(crate) fn from_stream_with_runtime<S>(stream: S, runtime: &tokio::runtime::Handle) -> Self
66    where
67        S: Stream<Item = T> + Send + 'static,
68        T: Send + 'static,
69    {
70        let (sender, receiver) = mpsc::channel(BATCH_STREAM_BUFFER);
71
72        let task = runtime.spawn(async move {
73            futures::pin_mut!(stream);
74            while let Some(item) = stream.next().await {
75                if sender.send(item).await.is_err() {
76                    break;
77                }
78            }
79        });
80
81        Self { receiver, task }
82    }
83}
84
85impl<T: Send> Iterator for BlockingBatchStream<T> {
86    type Item = T;
87
88    fn next(&mut self) -> Option<Self::Item> {
89        block_on_nautilus_with(|| self.receiver.recv())
90    }
91}
92
93impl<T> Drop for BlockingBatchStream<T> {
94    fn drop(&mut self) {
95        self.receiver.close();
96        self.task.abort();
97    }
98}
99
100/// Provides a DataFusion session for registering and querying catalog table sources.
101pub struct DataBackendSession {
102    pub chunk_size: usize,
103    pub runtime: tokio::runtime::Handle,
104    pub(crate) session_ctx: SessionContext,
105    pub(crate) registered_tables: AHashSet<String>,
106}
107
108impl DataBackendSession {
109    /// Creates a new [`DataBackendSession`] instance.
110    #[must_use]
111    pub fn new(chunk_size: usize) -> Self {
112        let session_ctx = SessionContext::new_with_config(session_config());
113        Self {
114            session_ctx,
115            chunk_size,
116            runtime: get_runtime().handle().clone(),
117            registered_tables: AHashSet::new(),
118        }
119    }
120
121    /// Register an object store with the session context
122    pub fn register_object_store(&mut self, url: &Url, object_store: Arc<dyn ObjectStore>) {
123        self.session_ctx.register_object_store(url, object_store);
124    }
125
126    /// Registers an OpenDAL-backed storage backend with the session context.
127    ///
128    /// External catalog implementations can call this before adding native table providers or
129    /// object-store-relative file paths to the session.
130    ///
131    /// # Errors
132    ///
133    /// Returns an error if the storage URI cannot be converted into a DataFusion root URL.
134    pub fn register_storage_backend(&mut self, storage: &StorageBackend) -> anyhow::Result<()> {
135        let root_url = storage.datafusion_root_url()?;
136        self.register_object_store(&root_url, storage.object_store.clone());
137        Ok(())
138    }
139
140    /// Registers a table provider with the session context.
141    ///
142    /// This supports non-Parquet table formats such as Delta Lake while keeping the Parquet
143    /// registration path unchanged.
144    pub fn register_table_provider(
145        &mut self,
146        table_name: &str,
147        provider: Arc<dyn TableProvider>,
148    ) -> Result<()> {
149        if !self.registered_tables.contains(table_name) {
150            self.session_ctx.register_table(table_name, provider)?;
151            self.registered_tables.insert(table_name.to_string());
152        }
153
154        Ok(())
155    }
156
157    pub(crate) fn collect_parquet_files_batches(
158        &mut self,
159        table_name: &str,
160        file_paths: Vec<String>,
161        sql_query: Option<&str>,
162    ) -> anyhow::Result<Vec<RecordBatch>> {
163        if file_paths.is_empty() {
164            return Ok(Vec::new());
165        }
166
167        let batch_stream = self.parquet_files_batch_stream(table_name, file_paths, sql_query)?;
168        let schema = batch_stream.schema();
169        let mut batches = block_on_nautilus_with(|| batch_stream.try_collect::<Vec<_>>())?;
170        if batches.is_empty() {
171            batches.push(RecordBatch::new_empty(schema));
172        }
173        Ok(batches)
174    }
175
176    pub(crate) fn parquet_files_batch_stream(
177        &mut self,
178        table_name: &str,
179        file_paths: Vec<String>,
180        sql_query: Option<&str>,
181    ) -> anyhow::Result<SendableRecordBatchStream> {
182        if file_paths.is_empty() {
183            return Ok(Box::pin(EmptyRecordBatchStream::new(Arc::new(
184                Schema::empty(),
185            ))));
186        }
187
188        self.register_parquet_files_table(table_name, file_paths)?;
189        Ok(self.execute_registered(table_name, sql_query)?)
190    }
191
192    fn register_parquet_files_table(
193        &mut self,
194        table_name: &str,
195        file_paths: Vec<String>,
196    ) -> anyhow::Result<()> {
197        if !self.registered_tables.contains(table_name) {
198            let parquet_options = ParquetReadOptions::<'_> {
199                skip_metadata: Some(false),
200                ..Default::default()
201            };
202            let dataframe = block_on_nautilus_with(|| {
203                self.session_ctx.read_parquet(file_paths, parquet_options)
204            })?;
205            validate_catalog_schema(dataframe.schema().as_arrow())?;
206            self.session_ctx
207                .register_table(table_name, dataframe.into_view())?;
208            self.registered_tables.insert(table_name.to_string());
209        }
210
211        Ok(())
212    }
213
214    pub(crate) fn execute_registered(
215        &self,
216        table_name: &str,
217        sql_query: Option<&str>,
218    ) -> Result<SendableRecordBatchStream> {
219        let default_query = format!("SELECT * FROM {table_name} ORDER BY ts_init");
220        let sql_query = sql_query.unwrap_or(&default_query);
221        let query = block_on_nautilus_with(|| self.session_ctx.sql(sql_query))?;
222        block_on_nautilus_with(|| query.execute_stream())
223    }
224
225    /// Clears all registered tables.
226    ///
227    /// This is useful when the underlying files have changed and we need to
228    /// re-register tables with updated data.
229    pub fn clear_registered_tables(&mut self) {
230        self.registered_tables.clear();
231
232        // Create a new session context to completely reset the DataFusion state
233        self.session_ctx = SessionContext::new_with_config(session_config());
234    }
235}
236
237fn session_config() -> SessionConfig {
238    SessionConfig::new()
239        .set_str("datafusion.optimizer.repartition_file_scans", "false")
240        .set_str("datafusion.optimizer.prefer_existing_sort", "true")
241}
242
243pub(crate) fn cast_record_batch_to_schema(
244    batch: &RecordBatch,
245    schema: &Schema,
246) -> Result<RecordBatch> {
247    let batch_schema = batch.schema();
248    let mut fields = Vec::with_capacity(batch_schema.fields().len());
249    let mut columns = Vec::with_capacity(batch.columns().len());
250
251    for (column, batch_field) in batch.columns().iter().zip(batch_schema.fields()) {
252        let field = schema
253            .field_with_name(batch_field.name())
254            .unwrap_or(batch_field.as_ref());
255
256        fields.push(Arc::new(field.clone()));
257        columns.push(cast_column_to_field(column, field)?);
258    }
259
260    let schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()));
261    Ok(RecordBatch::try_new(schema, columns)?)
262}
263
264fn cast_column_to_field(column: &ArrayRef, field: &Field) -> Result<ArrayRef> {
265    cast_column_to_data_type(column, field.data_type())
266}
267
268#[expect(
269    clippy::too_many_lines,
270    reason = "the function keeps recursive Arrow cast rules in one exhaustive type dispatcher"
271)]
272pub(crate) fn cast_column_to_data_type(
273    column: &ArrayRef,
274    data_type: &DataType,
275) -> Result<ArrayRef> {
276    if column.data_type() == data_type {
277        return Ok(column.clone());
278    }
279
280    if let (DataType::BinaryView, DataType::FixedSizeBinary(width)) =
281        (column.data_type(), data_type)
282    {
283        let array = column
284            .as_any()
285            .downcast_ref::<BinaryViewArray>()
286            .expect("BinaryView column should downcast to BinaryViewArray");
287        let mut builder = FixedSizeBinaryBuilder::with_capacity(array.len(), *width);
288        for row in 0..array.len() {
289            if array.is_null(row) {
290                builder.append_null();
291            } else {
292                builder.append_value(array.value(row))?;
293            }
294        }
295        return Ok(Arc::new(builder.finish()));
296    }
297
298    if let (DataType::FixedSizeList(_, size), DataType::List(field)) =
299        (column.data_type(), data_type)
300    {
301        let array = column
302            .as_any()
303            .downcast_ref::<FixedSizeListArray>()
304            .expect("FixedSizeList column should downcast to FixedSizeListArray");
305        let size = usize::try_from(*size).map_err(|e| {
306            DataFusionError::Execution(format!("Invalid fixed-size list length {size}: {e}"))
307        })?;
308        let mut parts = Vec::with_capacity(array.len());
309        let mut offsets = Vec::with_capacity(array.len() + 1);
310        offsets.push(0_i32);
311
312        for row in 0..array.len() {
313            if !array.is_null(row) {
314                let value = array.value(row);
315                parts.push(cast_column_to_data_type(&value, field.data_type())?);
316            }
317            let offset = parts
318                .len()
319                .checked_mul(size)
320                .and_then(|offset| i32::try_from(offset).ok())
321                .ok_or_else(|| {
322                    DataFusionError::Execution(
323                        "List offset exceeds the supported i32 range".to_string(),
324                    )
325                })?;
326            offsets.push(offset);
327        }
328        let values = if parts.is_empty() {
329            new_empty_array(field.data_type())
330        } else {
331            let parts = parts
332                .iter()
333                .map(std::convert::AsRef::as_ref)
334                .collect::<Vec<_>>();
335            concat(&parts)?
336        };
337        let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets));
338        return Ok(Arc::new(ListArray::try_new(
339            field.clone(),
340            offsets,
341            values,
342            array.nulls().cloned(),
343        )?));
344    }
345
346    if let (DataType::List(_), DataType::FixedSizeList(field, size)) =
347        (column.data_type(), data_type)
348    {
349        let array = column
350            .as_any()
351            .downcast_ref::<ListArray>()
352            .expect("List column should downcast to ListArray");
353        let size_usize = usize::try_from(*size).map_err(|e| {
354            DataFusionError::Execution(format!("Invalid fixed-size list length {size}: {e}"))
355        })?;
356        let mut parts = Vec::with_capacity(array.len());
357        for row in 0..array.len() {
358            if array.is_null(row) {
359                parts.push(new_null_array(field.data_type(), size_usize));
360                continue;
361            }
362            let value = array.value(row);
363            if value.len() != size_usize {
364                return Err(DataFusionError::Execution(format!(
365                    "List row {row} has length {}, expected {size}",
366                    value.len(),
367                )));
368            }
369            parts.push(cast_column_to_data_type(&value, field.data_type())?);
370        }
371        let values = if parts.is_empty() {
372            new_empty_array(field.data_type())
373        } else {
374            let parts = parts
375                .iter()
376                .map(std::convert::AsRef::as_ref)
377                .collect::<Vec<_>>();
378            concat(&parts)?
379        };
380        return Ok(Arc::new(FixedSizeListArray::try_new(
381            field.clone(),
382            *size,
383            values,
384            array.nulls().cloned(),
385        )?));
386    }
387
388    Ok(cast(column, data_type)?)
389}
390
391#[must_use]
392pub fn build_query(
393    table: &str,
394    start: Option<UnixNanos>,
395    end: Option<UnixNanos>,
396    where_clause: Option<&str>,
397) -> String {
398    let conditions = query_conditions(start, end, where_clause);
399    let mut query = format!("SELECT * FROM {table}");
400
401    if !conditions.is_empty() {
402        query.push_str(" WHERE ");
403        query.push_str(&conditions.join(" AND "));
404    }
405
406    query.push_str(" ORDER BY ts_init");
407    query
408}
409
410#[must_use]
411pub fn build_identifier_query(
412    table: &str,
413    start: Option<UnixNanos>,
414    end: Option<UnixNanos>,
415    where_clause: Option<&str>,
416) -> String {
417    let conditions = query_conditions(start, end, where_clause);
418    let mut query = format!("SELECT DISTINCT identifier FROM {table}");
419
420    if !conditions.is_empty() {
421        query.push_str(" WHERE ");
422        query.push_str(&conditions.join(" AND "));
423    }
424
425    query.push_str(" ORDER BY identifier");
426    query
427}
428
429fn query_conditions(
430    start: Option<UnixNanos>,
431    end: Option<UnixNanos>,
432    where_clause: Option<&str>,
433) -> Vec<String> {
434    let mut conditions = Vec::new();
435
436    if let Some(clause) = where_clause {
437        // Parenthesized so a caller clause containing a top-level `OR` still binds as one
438        // condition once the timestamp bounds are appended with `AND`.
439        conditions.push(format!("({clause})"));
440    }
441
442    if let Some(start_ts) = start {
443        conditions.push(format!("CAST(ts_init AS BIGINT) >= {start_ts}"));
444    }
445
446    if let Some(end_ts) = end {
447        conditions.push(format!("CAST(ts_init AS BIGINT) <= {end_ts}"));
448    }
449
450    conditions
451}
452
453pub fn identifiers_from_record_batches(batches: &[RecordBatch]) -> anyhow::Result<Vec<String>> {
454    let mut identifiers = AHashSet::new();
455
456    for batch in batches {
457        let column = batch
458            .column_by_name("identifier")
459            .ok_or_else(|| anyhow::anyhow!("identifier column not found"))?
460            .as_any();
461
462        if let Some(array) = column.downcast_ref::<StringArray>() {
463            for row in 0..array.len() {
464                if !array.is_null(row) {
465                    identifiers.insert(array.value(row).to_string());
466                }
467            }
468            continue;
469        }
470
471        if let Some(array) = column.downcast_ref::<StringViewArray>() {
472            for row in 0..array.len() {
473                if !array.is_null(row) {
474                    identifiers.insert(array.value(row).to_string());
475                }
476            }
477            continue;
478        }
479
480        anyhow::bail!("identifier column must be Utf8 or Utf8View");
481    }
482
483    let mut identifiers = identifiers.into_iter().collect::<Vec<_>>();
484    identifiers.sort();
485    Ok(identifiers)
486}
487
488#[cfg(test)]
489mod tests {
490    use std::{
491        sync::{Arc, mpsc},
492        time::Duration,
493    };
494
495    use nautilus_common::live::get_runtime;
496    use nautilus_model::{
497        data::{DataBatch, QuoteTick},
498        identifiers::InstrumentId,
499        types::{Price, Quantity},
500    };
501    use rstest::rstest;
502    use tempfile::TempDir;
503
504    use super::*;
505    use crate::{
506        catalog::session::{DataBatchQuery, TypedDataBatchSession},
507        common::storage::create_storage_backend_from_path,
508    };
509
510    fn typed_quote(ts_init: u64) -> QuoteTick {
511        QuoteTick::new(
512            InstrumentId::from("AUD/USD.SIM"),
513            Price::from("1.0"),
514            Price::from("1.1"),
515            Quantity::from("1000"),
516            Quantity::from("1000"),
517            UnixNanos::from(ts_init),
518            UnixNanos::from(ts_init),
519        )
520    }
521
522    fn batch_ts(batch: &DataBatch) -> Vec<u64> {
523        match batch {
524            DataBatch::Quote(quotes) => quotes
525                .as_ref()
526                .iter()
527                .map(|quote| quote.ts_init.as_u64())
528                .collect(),
529            other => panic!("expected quote batch, found {other:?}"),
530        }
531    }
532
533    #[rstest]
534    fn typed_session_chunks_pages_with_carry_across_pulls() {
535        let pages: Vec<anyhow::Result<Vec<QuoteTick>>> = vec![
536            Ok(vec![typed_quote(1), typed_quote(2), typed_quote(3)]),
537            Ok(Vec::new()),
538            Ok(vec![typed_quote(4), typed_quote(5)]),
539        ];
540        let mut session = TypedDataBatchSession::new(Box::new(pages.into_iter()), Some(2));
541
542        assert_eq!(batch_ts(&session.next_batch().unwrap().unwrap()), [1, 2]);
543        assert_eq!(batch_ts(&session.next_batch().unwrap().unwrap()), [3, 4]);
544        assert_eq!(batch_ts(&session.next_batch().unwrap().unwrap()), [5]);
545        assert!(session.next_batch().unwrap().is_none());
546    }
547
548    #[rstest]
549    fn typed_session_extends_chunk_across_equal_boundary_ts() {
550        let data = vec![
551            typed_quote(1),
552            typed_quote(2),
553            typed_quote(2),
554            typed_quote(2),
555            typed_quote(3),
556        ];
557        let mut session = TypedDataBatchSession::from_vec(data, Some(2));
558
559        assert_eq!(
560            batch_ts(&session.next_batch().unwrap().unwrap()),
561            [1, 2, 2, 2]
562        );
563        assert_eq!(batch_ts(&session.next_batch().unwrap().unwrap()), [3]);
564        assert!(session.next_batch().unwrap().is_none());
565    }
566
567    #[rstest]
568    fn typed_session_empty_source_yields_none() {
569        let mut session = TypedDataBatchSession::<QuoteTick>::from_vec(Vec::new(), None);
570
571        assert!(session.next_batch().unwrap().is_none());
572    }
573
574    #[rstest]
575    fn typed_session_propagates_page_error() {
576        let pages: Vec<anyhow::Result<Vec<QuoteTick>>> = vec![
577            Ok(vec![typed_quote(1)]),
578            Err(anyhow::anyhow!("page failed")),
579        ];
580        let mut session = TypedDataBatchSession::new(Box::new(pages.into_iter()), Some(4));
581
582        assert_eq!(session.next_batch().unwrap_err().to_string(), "page failed");
583    }
584
585    #[rstest]
586    fn register_storage_backend_accepts_memory_backend() {
587        let storage = create_storage_backend_from_path("memory://", None).unwrap();
588        let mut session = DataBackendSession::new(10);
589
590        session.register_storage_backend(&storage).unwrap();
591    }
592
593    #[rstest]
594    fn register_storage_backend_accepts_local_backend() {
595        let temp_dir = TempDir::new().unwrap();
596        let storage =
597            create_storage_backend_from_path(temp_dir.path().to_str().unwrap(), None).unwrap();
598        let mut session = DataBackendSession::new(10);
599
600        session.register_storage_backend(&storage).unwrap();
601    }
602
603    #[rstest]
604    fn build_identifier_query_projects_distinct_identifier_with_filters() {
605        let query = build_identifier_query(
606            "quotes",
607            Some(UnixNanos::from(10)),
608            Some(UnixNanos::from(20)),
609            Some("instrument_id LIKE 'ES%'"),
610        );
611
612        assert_eq!(
613            query,
614            "SELECT DISTINCT identifier FROM quotes WHERE (instrument_id LIKE 'ES%') \
615             AND CAST(ts_init AS BIGINT) >= 10 AND CAST(ts_init AS BIGINT) <= 20 ORDER BY identifier"
616        );
617    }
618
619    #[rstest]
620    fn build_query_groups_a_disjunctive_where_clause_against_timestamp_bounds() {
621        let query = build_query(
622            "bars",
623            Some(UnixNanos::from(10)),
624            Some(UnixNanos::from(20)),
625            Some("identifier = 'ES.GLBX' OR starts_with(identifier, 'ES.GLBX-')"),
626        );
627
628        assert_eq!(
629            query,
630            "SELECT * FROM bars WHERE (identifier = 'ES.GLBX' OR starts_with(identifier, 'ES.GLBX-')) \
631             AND CAST(ts_init AS BIGINT) >= 10 AND CAST(ts_init AS BIGINT) <= 20 ORDER BY ts_init"
632        );
633    }
634
635    #[rstest]
636    fn build_query_without_bounds_keeps_the_where_clause_alone() {
637        let query = build_query("bars", None, None, Some("identifier = 'ES.GLBX'"));
638
639        assert_eq!(
640            query,
641            "SELECT * FROM bars WHERE (identifier = 'ES.GLBX') ORDER BY ts_init"
642        );
643    }
644
645    #[rstest]
646    fn build_query_without_conditions_omits_the_where_keyword() {
647        let query = build_query("bars", None, None, None);
648
649        assert_eq!(query, "SELECT * FROM bars ORDER BY ts_init");
650    }
651
652    #[rstest]
653    fn identifiers_from_record_batches_returns_sorted_unique_non_null_values() {
654        let schema = Arc::new(Schema::new(vec![Field::new(
655            "identifier",
656            DataType::Utf8,
657            true,
658        )]));
659        let batch = RecordBatch::try_new(
660            schema,
661            vec![Arc::new(StringArray::from(vec![
662                Some("ESZ4.XCME"),
663                None,
664                Some("ESM4.XCME"),
665                Some("ESZ4.XCME"),
666            ])) as ArrayRef],
667        )
668        .unwrap();
669
670        let identifiers = identifiers_from_record_batches(&[batch]).unwrap();
671
672        assert_eq!(identifiers, vec!["ESM4.XCME", "ESZ4.XCME"]);
673    }
674
675    #[rstest]
676    fn fixed_size_list_cast_round_trips_through_delta_list_type() {
677        use arrow::{
678            array::{Decimal128Array, UInt32Array},
679            buffer::NullBuffer,
680        };
681
682        let field = Arc::new(Field::new("element", DataType::Decimal128(38, 16), true));
683        let values = Decimal128Array::from(vec![1_i128, 2, 3, 4])
684            .with_precision_and_scale(38, 16)
685            .unwrap();
686        let fixed_type = DataType::FixedSizeList(field.clone(), 2);
687        let fixed = Arc::new(
688            FixedSizeListArray::try_new(field.clone(), 2, Arc::new(values), None).unwrap(),
689        ) as ArrayRef;
690
691        let list_type = DataType::List(field);
692        let list = cast_column_to_data_type(&fixed, &list_type).unwrap();
693        let restored = cast_column_to_data_type(&list, &fixed_type).unwrap();
694
695        assert_eq!(list.data_type(), &list_type);
696        assert_eq!(restored.data_type(), &fixed_type);
697        assert_eq!(restored.to_data(), fixed.to_data());
698
699        let field = Arc::new(Field::new("element", DataType::UInt32, false));
700        let fixed = Arc::new(
701            FixedSizeListArray::try_new(
702                field.clone(),
703                2,
704                Arc::new(UInt32Array::from(vec![1_u32, 2])),
705                Some(NullBuffer::from(vec![false])),
706            )
707            .unwrap(),
708        ) as ArrayRef;
709        let list_type = DataType::List(field);
710        let list = cast_column_to_data_type(&fixed, &list_type).unwrap();
711        let list = list.as_any().downcast_ref::<ListArray>().unwrap();
712
713        assert!(list.is_null(0));
714        assert_eq!(list.values().len(), 0);
715    }
716
717    #[rstest]
718    fn data_backend_sessions_share_global_runtime() {
719        let first = DataBackendSession::new(10);
720        let second = DataBackendSession::new(10);
721
722        assert_eq!(first.runtime.id(), second.runtime.id());
723        assert_eq!(first.runtime.id(), get_runtime().handle().id());
724    }
725
726    #[rstest]
727    fn blocking_batch_stream_prefetches_first_item() {
728        let runtime = tokio::runtime::Runtime::new().unwrap();
729        let (polled_sender, polled_receiver) = mpsc::channel();
730        let stream = futures::stream::once(async move {
731            polled_sender.send(()).unwrap();
732            42
733        });
734        let mut stream = BlockingBatchStream::from_stream_with_runtime(stream, runtime.handle());
735
736        assert_eq!(polled_receiver.recv_timeout(Duration::from_secs(1)), Ok(()),);
737        assert_eq!(stream.next(), Some(42));
738    }
739
740    #[rstest]
741    fn blocking_batch_stream_works_inside_current_thread_runtime() {
742        let runtime = tokio::runtime::Builder::new_current_thread()
743            .enable_all()
744            .build()
745            .unwrap();
746
747        runtime.block_on(async {
748            let mut stream = BlockingBatchStream::from_stream_with_runtime(
749                futures::stream::iter([42]),
750                get_runtime().handle(),
751            );
752
753            assert_eq!(stream.next(), Some(42));
754        });
755    }
756}
nautilus_persistence/common/datafusion.rs

nautilus_persistence/common/
datafusion.rs
