datafusion_ffi/
table_provider.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::any::Any;
19use std::ffi::c_void;
20use std::sync::Arc;
21
22use abi_stable::StableAbi;
23use abi_stable::std_types::{ROption, RResult, RVec};
24use arrow::datatypes::SchemaRef;
25use async_ffi::{FfiFuture, FutureExt};
26use async_trait::async_trait;
27use datafusion_catalog::{Session, TableProvider};
28use datafusion_common::error::{DataFusionError, Result};
29use datafusion_execution::TaskContext;
30use datafusion_expr::dml::InsertOp;
31use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType};
32use datafusion_physical_plan::ExecutionPlan;
33use datafusion_proto::logical_plan::from_proto::parse_exprs;
34use datafusion_proto::logical_plan::to_proto::serialize_exprs;
35use datafusion_proto::logical_plan::{
36    DefaultLogicalExtensionCodec, LogicalExtensionCodec,
37};
38use datafusion_proto::protobuf::LogicalExprList;
39use prost::Message;
40use tokio::runtime::Handle;
41
42use super::execution_plan::FFI_ExecutionPlan;
43use super::insert_op::FFI_InsertOp;
44use crate::arrow_wrappers::WrappedSchema;
45use crate::execution::FFI_TaskContextProvider;
46use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
47use crate::session::{FFI_SessionRef, ForeignSession};
48use crate::table_source::{FFI_TableProviderFilterPushDown, FFI_TableType};
49use crate::util::FFIResult;
50use crate::{df_result, rresult_return};
51
52/// A stable struct for sharing [`TableProvider`] across FFI boundaries.
53///
54/// # Struct Layout
55///
56/// The following description applies to all structs provided in this crate.
57///
58/// Each of the exposed structs in this crate is provided with a variant prefixed
59/// with `Foreign`. This variant is designed to be used by the consumer of the
60/// foreign code. The `Foreign` structs should _never_ access the `private_data`
61/// fields. Instead they should only access the data returned through the function
62/// calls defined on the `FFI_` structs. The second purpose of the `Foreign`
63/// structs is to contain additional data that may be needed by the traits that
64/// are implemented on them. Some of these traits require borrowing data which
65/// can be far more convenient to be locally stored.
66///
67/// For example, we have a struct `FFI_TableProvider` to give access to the
68/// `TableProvider` functions like `table_type()` and `scan()`. If we write a
69/// library that wishes to expose it's `TableProvider`, then we can access the
70/// private data that contains the Arc reference to the `TableProvider` via
71/// `FFI_TableProvider`. This data is local to the library.
72///
73/// If we have a program that accesses a `TableProvider` via FFI, then it
74/// will use `ForeignTableProvider`. When using `ForeignTableProvider` we **must**
75/// not attempt to access the `private_data` field in `FFI_TableProvider`. If a
76/// user is testing locally, you may be able to successfully access this field, but
77/// it will only work if you are building against the exact same version of
78/// `DataFusion` for both libraries **and** the same compiler. It will not work
79/// in general.
80///
81/// It is worth noting that which library is the `local` and which is `foreign`
82/// depends on which interface we are considering. For example, suppose we have a
83/// Python library called `my_provider` that exposes a `TableProvider` called
84/// `MyProvider` via `FFI_TableProvider`. Within the library `my_provider` we can
85/// access the `private_data` via `FFI_TableProvider`. We connect this to
86/// `datafusion-python`, where we access it as a `ForeignTableProvider`. Now when
87/// we call `scan()` on this interface, we have to pass it a `FFI_SessionConfig`.
88/// The `SessionConfig` is local to `datafusion-python` and **not** `my_provider`.
89/// It is important to be careful when expanding these functions to be certain which
90/// side of the interface each object refers to.
91#[repr(C)]
92#[derive(Debug, StableAbi)]
93pub struct FFI_TableProvider {
94    /// Return the table schema
95    schema: unsafe extern "C" fn(provider: &Self) -> WrappedSchema,
96
97    /// Perform a scan on the table. See [`TableProvider`] for detailed usage information.
98    ///
99    /// # Arguments
100    ///
101    /// * `provider` - the table provider
102    /// * `session` - session
103    /// * `projections` - if specified, only a subset of the columns are returned
104    /// * `filters_serialized` - filters to apply to the scan, which are a
105    ///   [`LogicalExprList`] protobuf message serialized into bytes to pass
106    ///   across the FFI boundary.
107    /// * `limit` - if specified, limit the number of rows returned
108    scan: unsafe extern "C" fn(
109        provider: &Self,
110        session: FFI_SessionRef,
111        projections: RVec<usize>,
112        filters_serialized: RVec<u8>,
113        limit: ROption<usize>,
114    ) -> FfiFuture<FFIResult<FFI_ExecutionPlan>>,
115
116    /// Return the type of table. See [`TableType`] for options.
117    table_type: unsafe extern "C" fn(provider: &Self) -> FFI_TableType,
118
119    /// Based upon the input filters, identify which are supported. The filters
120    /// are a [`LogicalExprList`] protobuf message serialized into bytes to pass
121    /// across the FFI boundary.
122    supports_filters_pushdown: Option<
123        unsafe extern "C" fn(
124            provider: &FFI_TableProvider,
125            filters_serialized: RVec<u8>,
126        ) -> FFIResult<RVec<FFI_TableProviderFilterPushDown>>,
127    >,
128
129    insert_into: unsafe extern "C" fn(
130        provider: &Self,
131        session: FFI_SessionRef,
132        input: &FFI_ExecutionPlan,
133        insert_op: FFI_InsertOp,
134    ) -> FfiFuture<FFIResult<FFI_ExecutionPlan>>,
135
136    pub logical_codec: FFI_LogicalExtensionCodec,
137
138    /// Used to create a clone on the provider of the execution plan. This should
139    /// only need to be called by the receiver of the plan.
140    clone: unsafe extern "C" fn(plan: &Self) -> Self,
141
142    /// Release the memory of the private data when it is no longer being used.
143    release: unsafe extern "C" fn(arg: &mut Self),
144
145    /// Return the major DataFusion version number of this provider.
146    pub version: unsafe extern "C" fn() -> u64,
147
148    /// Internal data. This is only to be accessed by the provider of the plan.
149    /// A [`ForeignTableProvider`] should never attempt to access this data.
150    private_data: *mut c_void,
151
152    /// Utility to identify when FFI objects are accessed locally through
153    /// the foreign interface. See [`crate::get_library_marker_id`] and
154    /// the crate's `README.md` for more information.
155    pub library_marker_id: extern "C" fn() -> usize,
156}
157
158unsafe impl Send for FFI_TableProvider {}
159unsafe impl Sync for FFI_TableProvider {}
160
161struct ProviderPrivateData {
162    provider: Arc<dyn TableProvider + Send>,
163    runtime: Option<Handle>,
164}
165
166impl FFI_TableProvider {
167    fn inner(&self) -> &Arc<dyn TableProvider + Send> {
168        let private_data = self.private_data as *const ProviderPrivateData;
169        unsafe { &(*private_data).provider }
170    }
171
172    fn runtime(&self) -> &Option<Handle> {
173        let private_data = self.private_data as *const ProviderPrivateData;
174        unsafe { &(*private_data).runtime }
175    }
176}
177
178unsafe extern "C" fn schema_fn_wrapper(provider: &FFI_TableProvider) -> WrappedSchema {
179    provider.inner().schema().into()
180}
181
182unsafe extern "C" fn table_type_fn_wrapper(
183    provider: &FFI_TableProvider,
184) -> FFI_TableType {
185    provider.inner().table_type().into()
186}
187
188fn supports_filters_pushdown_internal(
189    provider: &Arc<dyn TableProvider + Send>,
190    filters_serialized: &[u8],
191    task_ctx: &Arc<TaskContext>,
192    codec: &dyn LogicalExtensionCodec,
193) -> Result<RVec<FFI_TableProviderFilterPushDown>> {
194    let filters = match filters_serialized.is_empty() {
195        true => vec![],
196        false => {
197            let proto_filters = LogicalExprList::decode(filters_serialized)
198                .map_err(|e| DataFusionError::Plan(e.to_string()))?;
199
200            parse_exprs(proto_filters.expr.iter(), task_ctx.as_ref(), codec)?
201        }
202    };
203    let filters_borrowed: Vec<&Expr> = filters.iter().collect();
204
205    let results: RVec<_> = provider
206        .supports_filters_pushdown(&filters_borrowed)?
207        .iter()
208        .map(|v| v.into())
209        .collect();
210
211    Ok(results)
212}
213
214unsafe extern "C" fn supports_filters_pushdown_fn_wrapper(
215    provider: &FFI_TableProvider,
216    filters_serialized: RVec<u8>,
217) -> FFIResult<RVec<FFI_TableProviderFilterPushDown>> {
218    let logical_codec: Arc<dyn LogicalExtensionCodec> = (&provider.logical_codec).into();
219    let task_ctx = rresult_return!(<Arc<TaskContext>>::try_from(
220        &provider.logical_codec.task_ctx_provider
221    ));
222    supports_filters_pushdown_internal(
223        provider.inner(),
224        &filters_serialized,
225        &task_ctx,
226        logical_codec.as_ref(),
227    )
228    .map_err(|e| e.to_string().into())
229    .into()
230}
231
232unsafe extern "C" fn scan_fn_wrapper(
233    provider: &FFI_TableProvider,
234    session: FFI_SessionRef,
235    projections: RVec<usize>,
236    filters_serialized: RVec<u8>,
237    limit: ROption<usize>,
238) -> FfiFuture<FFIResult<FFI_ExecutionPlan>> {
239    let task_ctx: Result<Arc<TaskContext>, DataFusionError> =
240        (&provider.logical_codec.task_ctx_provider).try_into();
241    let runtime = provider.runtime().clone();
242    let logical_codec: Arc<dyn LogicalExtensionCodec> = (&provider.logical_codec).into();
243    let internal_provider = Arc::clone(provider.inner());
244
245    async move {
246        let mut foreign_session = None;
247        let session = rresult_return!(
248            session
249                .as_local()
250                .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
251                .unwrap_or_else(|| {
252                    foreign_session = Some(ForeignSession::try_from(&session)?);
253                    Ok(foreign_session.as_ref().unwrap())
254                })
255        );
256
257        let task_ctx = rresult_return!(task_ctx);
258        let filters = match filters_serialized.is_empty() {
259            true => vec![],
260            false => {
261                let proto_filters =
262                    rresult_return!(LogicalExprList::decode(filters_serialized.as_ref()));
263
264                rresult_return!(parse_exprs(
265                    proto_filters.expr.iter(),
266                    task_ctx.as_ref(),
267                    logical_codec.as_ref(),
268                ))
269            }
270        };
271
272        let projections: Vec<_> = projections.into_iter().collect();
273
274        let plan = rresult_return!(
275            internal_provider
276                .scan(session, Some(&projections), &filters, limit.into())
277                .await
278        );
279
280        RResult::ROk(FFI_ExecutionPlan::new(plan, runtime.clone()))
281    }
282    .into_ffi()
283}
284
285unsafe extern "C" fn insert_into_fn_wrapper(
286    provider: &FFI_TableProvider,
287    session: FFI_SessionRef,
288    input: &FFI_ExecutionPlan,
289    insert_op: FFI_InsertOp,
290) -> FfiFuture<FFIResult<FFI_ExecutionPlan>> {
291    let runtime = provider.runtime().clone();
292    let internal_provider = Arc::clone(provider.inner());
293    let input = input.clone();
294
295    async move {
296        let mut foreign_session = None;
297        let session = rresult_return!(
298            session
299                .as_local()
300                .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
301                .unwrap_or_else(|| {
302                    foreign_session = Some(ForeignSession::try_from(&session)?);
303                    Ok(foreign_session.as_ref().unwrap())
304                })
305        );
306
307        let input = rresult_return!(<Arc<dyn ExecutionPlan>>::try_from(&input));
308
309        let insert_op = InsertOp::from(insert_op);
310
311        let plan = rresult_return!(
312            internal_provider
313                .insert_into(session, input, insert_op)
314                .await
315        );
316
317        RResult::ROk(FFI_ExecutionPlan::new(plan, runtime.clone()))
318    }
319    .into_ffi()
320}
321
322unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_TableProvider) {
323    unsafe {
324        debug_assert!(!provider.private_data.is_null());
325        let private_data =
326            Box::from_raw(provider.private_data as *mut ProviderPrivateData);
327        drop(private_data);
328        provider.private_data = std::ptr::null_mut();
329    }
330}
331
332unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_TableProvider) -> FFI_TableProvider {
333    let runtime = provider.runtime().clone();
334    let old_provider = Arc::clone(provider.inner());
335
336    let private_data = Box::into_raw(Box::new(ProviderPrivateData {
337        provider: old_provider,
338        runtime,
339    })) as *mut c_void;
340
341    FFI_TableProvider {
342        schema: schema_fn_wrapper,
343        scan: scan_fn_wrapper,
344        table_type: table_type_fn_wrapper,
345        supports_filters_pushdown: provider.supports_filters_pushdown,
346        insert_into: provider.insert_into,
347        logical_codec: provider.logical_codec.clone(),
348        clone: clone_fn_wrapper,
349        release: release_fn_wrapper,
350        version: super::version,
351        private_data,
352        library_marker_id: crate::get_library_marker_id,
353    }
354}
355
356impl Drop for FFI_TableProvider {
357    fn drop(&mut self) {
358        unsafe { (self.release)(self) }
359    }
360}
361
362impl FFI_TableProvider {
363    /// Creates a new [`FFI_TableProvider`].
364    pub fn new(
365        provider: Arc<dyn TableProvider + Send>,
366        can_support_pushdown_filters: bool,
367        runtime: Option<Handle>,
368        task_ctx_provider: impl Into<FFI_TaskContextProvider>,
369        logical_codec: Option<Arc<dyn LogicalExtensionCodec>>,
370    ) -> Self {
371        let task_ctx_provider = task_ctx_provider.into();
372        let logical_codec =
373            logical_codec.unwrap_or_else(|| Arc::new(DefaultLogicalExtensionCodec {}));
374        let logical_codec = FFI_LogicalExtensionCodec::new(
375            logical_codec,
376            runtime.clone(),
377            task_ctx_provider.clone(),
378        );
379        Self::new_with_ffi_codec(
380            provider,
381            can_support_pushdown_filters,
382            runtime,
383            logical_codec,
384        )
385    }
386
387    pub fn new_with_ffi_codec(
388        provider: Arc<dyn TableProvider + Send>,
389        can_support_pushdown_filters: bool,
390        runtime: Option<Handle>,
391        logical_codec: FFI_LogicalExtensionCodec,
392    ) -> Self {
393        let private_data = Box::new(ProviderPrivateData { provider, runtime });
394
395        Self {
396            schema: schema_fn_wrapper,
397            scan: scan_fn_wrapper,
398            table_type: table_type_fn_wrapper,
399            supports_filters_pushdown: match can_support_pushdown_filters {
400                true => Some(supports_filters_pushdown_fn_wrapper),
401                false => None,
402            },
403            insert_into: insert_into_fn_wrapper,
404            logical_codec,
405            clone: clone_fn_wrapper,
406            release: release_fn_wrapper,
407            version: super::version,
408            private_data: Box::into_raw(private_data) as *mut c_void,
409            library_marker_id: crate::get_library_marker_id,
410        }
411    }
412}
413
414/// This wrapper struct exists on the receiver side of the FFI interface, so it has
415/// no guarantees about being able to access the data in `private_data`. Any functions
416/// defined on this struct must only use the stable functions provided in
417/// FFI_TableProvider to interact with the foreign table provider.
418#[derive(Debug)]
419pub struct ForeignTableProvider(pub FFI_TableProvider);
420
421unsafe impl Send for ForeignTableProvider {}
422unsafe impl Sync for ForeignTableProvider {}
423
424impl From<&FFI_TableProvider> for Arc<dyn TableProvider> {
425    fn from(provider: &FFI_TableProvider) -> Self {
426        if (provider.library_marker_id)() == crate::get_library_marker_id() {
427            Arc::clone(provider.inner()) as Arc<dyn TableProvider>
428        } else {
429            Arc::new(ForeignTableProvider(provider.clone()))
430        }
431    }
432}
433
434impl Clone for FFI_TableProvider {
435    fn clone(&self) -> Self {
436        unsafe { (self.clone)(self) }
437    }
438}
439
440#[async_trait]
441impl TableProvider for ForeignTableProvider {
442    fn as_any(&self) -> &dyn Any {
443        self
444    }
445
446    fn schema(&self) -> SchemaRef {
447        let wrapped_schema = unsafe { (self.0.schema)(&self.0) };
448        wrapped_schema.into()
449    }
450
451    fn table_type(&self) -> TableType {
452        unsafe { (self.0.table_type)(&self.0).into() }
453    }
454
455    async fn scan(
456        &self,
457        session: &dyn Session,
458        projection: Option<&Vec<usize>>,
459        filters: &[Expr],
460        limit: Option<usize>,
461    ) -> Result<Arc<dyn ExecutionPlan>> {
462        let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone());
463
464        let projections: Option<RVec<usize>> =
465            projection.map(|p| p.iter().map(|v| v.to_owned()).collect());
466
467        let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
468        let filter_list = LogicalExprList {
469            expr: serialize_exprs(filters, codec.as_ref())?,
470        };
471        let filters_serialized = filter_list.encode_to_vec().into();
472
473        let plan = unsafe {
474            let maybe_plan = (self.0.scan)(
475                &self.0,
476                session,
477                projections.unwrap_or_default(),
478                filters_serialized,
479                limit.into(),
480            )
481            .await;
482
483            <Arc<dyn ExecutionPlan>>::try_from(&df_result!(maybe_plan)?)?
484        };
485
486        Ok(plan)
487    }
488
489    /// Tests whether the table provider can make use of a filter expression
490    /// to optimize data retrieval.
491    fn supports_filters_pushdown(
492        &self,
493        filters: &[&Expr],
494    ) -> Result<Vec<TableProviderFilterPushDown>> {
495        unsafe {
496            let pushdown_fn = match self.0.supports_filters_pushdown {
497                Some(func) => func,
498                None => {
499                    return Ok(vec![
500                        TableProviderFilterPushDown::Unsupported;
501                        filters.len()
502                    ]);
503                }
504            };
505
506            let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
507
508            let expr_list = LogicalExprList {
509                expr: serialize_exprs(
510                    filters.iter().map(|f| f.to_owned()),
511                    codec.as_ref(),
512                )?,
513            };
514            let serialized_filters = expr_list.encode_to_vec();
515
516            let pushdowns = df_result!(pushdown_fn(&self.0, serialized_filters.into()))?;
517
518            Ok(pushdowns.iter().map(|v| v.into()).collect())
519        }
520    }
521
522    async fn insert_into(
523        &self,
524        session: &dyn Session,
525        input: Arc<dyn ExecutionPlan>,
526        insert_op: InsertOp,
527    ) -> Result<Arc<dyn ExecutionPlan>> {
528        let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone());
529
530        let rc = Handle::try_current().ok();
531        let input = FFI_ExecutionPlan::new(input, rc);
532        let insert_op: FFI_InsertOp = insert_op.into();
533
534        let plan = unsafe {
535            let maybe_plan =
536                (self.0.insert_into)(&self.0, session, &input, insert_op).await;
537
538            <Arc<dyn ExecutionPlan>>::try_from(&df_result!(maybe_plan)?)?
539        };
540
541        Ok(plan)
542    }
543}
544
545#[cfg(test)]
546mod tests {
547    use arrow::datatypes::Schema;
548    use datafusion::prelude::{SessionContext, col, lit};
549    use datafusion_execution::TaskContextProvider;
550
551    use super::*;
552
553    fn create_test_table_provider() -> Result<Arc<dyn TableProvider>> {
554        use arrow::datatypes::Field;
555        use datafusion::arrow::array::Float32Array;
556        use datafusion::arrow::datatypes::DataType;
557        use datafusion::arrow::record_batch::RecordBatch;
558        use datafusion::datasource::MemTable;
559
560        let schema =
561            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
562
563        // define data in two partitions
564        let batch1 = RecordBatch::try_new(
565            Arc::clone(&schema),
566            vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))],
567        )?;
568        let batch2 = RecordBatch::try_new(
569            Arc::clone(&schema),
570            vec![Arc::new(Float32Array::from(vec![64.0]))],
571        )?;
572
573        Ok(Arc::new(MemTable::try_new(
574            schema,
575            vec![vec![batch1], vec![batch2]],
576        )?))
577    }
578
579    #[tokio::test]
580    async fn test_round_trip_ffi_table_provider_scan() -> Result<()> {
581        let provider = create_test_table_provider()?;
582        let ctx = Arc::new(SessionContext::new());
583        let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
584        let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider);
585
586        let mut ffi_provider =
587            FFI_TableProvider::new(provider, true, None, task_ctx_provider, None);
588        ffi_provider.library_marker_id = crate::mock_foreign_marker_id;
589
590        let foreign_table_provider: Arc<dyn TableProvider> = (&ffi_provider).into();
591
592        ctx.register_table("t", foreign_table_provider)?;
593
594        let df = ctx.table("t").await?;
595
596        df.select(vec![col("a")])?
597            .filter(col("a").gt(lit(3.0)))?
598            .show()
599            .await?;
600
601        Ok(())
602    }
603
604    #[tokio::test]
605    async fn test_round_trip_ffi_table_provider_insert_into() -> Result<()> {
606        let provider = create_test_table_provider()?;
607        let ctx = Arc::new(SessionContext::new());
608        let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
609        let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider);
610
611        let mut ffi_provider =
612            FFI_TableProvider::new(provider, true, None, task_ctx_provider, None);
613        ffi_provider.library_marker_id = crate::mock_foreign_marker_id;
614
615        let foreign_table_provider: Arc<dyn TableProvider> = (&ffi_provider).into();
616
617        ctx.register_table("t", foreign_table_provider)?;
618
619        let result = ctx
620            .sql("INSERT INTO t VALUES (128.0);")
621            .await?
622            .collect()
623            .await?;
624
625        assert!(result.len() == 1 && result[0].num_rows() == 1);
626
627        ctx.table("t")
628            .await?
629            .select(vec![col("a")])?
630            .filter(col("a").gt(lit(3.0)))?
631            .show()
632            .await?;
633
634        Ok(())
635    }
636
637    #[tokio::test]
638    async fn test_aggregation() -> Result<()> {
639        use arrow::datatypes::Field;
640        use datafusion::arrow::array::Float32Array;
641        use datafusion::arrow::datatypes::DataType;
642        use datafusion::arrow::record_batch::RecordBatch;
643        use datafusion::common::assert_batches_eq;
644        use datafusion::datasource::MemTable;
645
646        let schema =
647            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
648
649        // define data in two partitions
650        let batch1 = RecordBatch::try_new(
651            Arc::clone(&schema),
652            vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))],
653        )?;
654
655        let ctx = Arc::new(SessionContext::new());
656        let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
657        let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider);
658
659        let provider = Arc::new(MemTable::try_new(schema, vec![vec![batch1]])?);
660
661        let ffi_provider =
662            FFI_TableProvider::new(provider, true, None, task_ctx_provider, None);
663
664        let foreign_table_provider: Arc<dyn TableProvider> = (&ffi_provider).into();
665
666        ctx.register_table("t", foreign_table_provider)?;
667
668        let result = ctx
669            .sql("SELECT COUNT(*) as cnt FROM t")
670            .await?
671            .collect()
672            .await?;
673        #[rustfmt::skip]
674        let expected = [
675            "+-----+",
676            "| cnt |",
677            "+-----+",
678            "| 3   |",
679            "+-----+"
680        ];
681        assert_batches_eq!(expected, &result);
682        Ok(())
683    }
684
685    #[test]
686    fn test_ffi_table_provider_local_bypass() -> Result<()> {
687        let table_provider = create_test_table_provider()?;
688
689        let ctx = Arc::new(SessionContext::new()) as Arc<dyn TaskContextProvider>;
690        let task_ctx_provider = FFI_TaskContextProvider::from(&ctx);
691        let mut ffi_table =
692            FFI_TableProvider::new(table_provider, false, None, task_ctx_provider, None);
693
694        // Verify local libraries can be downcast to their original
695        let foreign_table: Arc<dyn TableProvider> = (&ffi_table).into();
696        assert!(
697            foreign_table
698                .as_any()
699                .downcast_ref::<datafusion::datasource::MemTable>()
700                .is_some()
701        );
702
703        // Verify different library markers generate foreign providers
704        ffi_table.library_marker_id = crate::mock_foreign_marker_id;
705        let foreign_table: Arc<dyn TableProvider> = (&ffi_table).into();
706        assert!(
707            foreign_table
708                .as_any()
709                .downcast_ref::<ForeignTableProvider>()
710                .is_some()
711        );
712
713        Ok(())
714    }
715}