Skip to main content

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        if let Some(provider) = provider.as_any().downcast_ref::<ForeignTableProvider>() {
394            return provider.0.clone();
395        }
396        let private_data = Box::new(ProviderPrivateData { provider, runtime });
397
398        Self {
399            schema: schema_fn_wrapper,
400            scan: scan_fn_wrapper,
401            table_type: table_type_fn_wrapper,
402            supports_filters_pushdown: match can_support_pushdown_filters {
403                true => Some(supports_filters_pushdown_fn_wrapper),
404                false => None,
405            },
406            insert_into: insert_into_fn_wrapper,
407            logical_codec,
408            clone: clone_fn_wrapper,
409            release: release_fn_wrapper,
410            version: super::version,
411            private_data: Box::into_raw(private_data) as *mut c_void,
412            library_marker_id: crate::get_library_marker_id,
413        }
414    }
415}
416
417/// This wrapper struct exists on the receiver side of the FFI interface, so it has
418/// no guarantees about being able to access the data in `private_data`. Any functions
419/// defined on this struct must only use the stable functions provided in
420/// FFI_TableProvider to interact with the foreign table provider.
421#[derive(Debug)]
422pub struct ForeignTableProvider(pub FFI_TableProvider);
423
424unsafe impl Send for ForeignTableProvider {}
425unsafe impl Sync for ForeignTableProvider {}
426
427impl From<&FFI_TableProvider> for Arc<dyn TableProvider> {
428    fn from(provider: &FFI_TableProvider) -> Self {
429        if (provider.library_marker_id)() == crate::get_library_marker_id() {
430            Arc::clone(provider.inner()) as Arc<dyn TableProvider>
431        } else {
432            Arc::new(ForeignTableProvider(provider.clone()))
433        }
434    }
435}
436
437impl Clone for FFI_TableProvider {
438    fn clone(&self) -> Self {
439        unsafe { (self.clone)(self) }
440    }
441}
442
443#[async_trait]
444impl TableProvider for ForeignTableProvider {
445    fn as_any(&self) -> &dyn Any {
446        self
447    }
448
449    fn schema(&self) -> SchemaRef {
450        let wrapped_schema = unsafe { (self.0.schema)(&self.0) };
451        wrapped_schema.into()
452    }
453
454    fn table_type(&self) -> TableType {
455        unsafe { (self.0.table_type)(&self.0).into() }
456    }
457
458    async fn scan(
459        &self,
460        session: &dyn Session,
461        projection: Option<&Vec<usize>>,
462        filters: &[Expr],
463        limit: Option<usize>,
464    ) -> Result<Arc<dyn ExecutionPlan>> {
465        let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone());
466
467        let projections: Option<RVec<usize>> =
468            projection.map(|p| p.iter().map(|v| v.to_owned()).collect());
469
470        let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
471        let filter_list = LogicalExprList {
472            expr: serialize_exprs(filters, codec.as_ref())?,
473        };
474        let filters_serialized = filter_list.encode_to_vec().into();
475
476        let plan = unsafe {
477            let maybe_plan = (self.0.scan)(
478                &self.0,
479                session,
480                projections.unwrap_or_default(),
481                filters_serialized,
482                limit.into(),
483            )
484            .await;
485
486            <Arc<dyn ExecutionPlan>>::try_from(&df_result!(maybe_plan)?)?
487        };
488
489        Ok(plan)
490    }
491
492    /// Tests whether the table provider can make use of a filter expression
493    /// to optimize data retrieval.
494    fn supports_filters_pushdown(
495        &self,
496        filters: &[&Expr],
497    ) -> Result<Vec<TableProviderFilterPushDown>> {
498        unsafe {
499            let pushdown_fn = match self.0.supports_filters_pushdown {
500                Some(func) => func,
501                None => {
502                    return Ok(vec![
503                        TableProviderFilterPushDown::Unsupported;
504                        filters.len()
505                    ]);
506                }
507            };
508
509            let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
510
511            let expr_list = LogicalExprList {
512                expr: serialize_exprs(
513                    filters.iter().map(|f| f.to_owned()),
514                    codec.as_ref(),
515                )?,
516            };
517            let serialized_filters = expr_list.encode_to_vec();
518
519            let pushdowns = df_result!(pushdown_fn(&self.0, serialized_filters.into()))?;
520
521            Ok(pushdowns.iter().map(|v| v.into()).collect())
522        }
523    }
524
525    async fn insert_into(
526        &self,
527        session: &dyn Session,
528        input: Arc<dyn ExecutionPlan>,
529        insert_op: InsertOp,
530    ) -> Result<Arc<dyn ExecutionPlan>> {
531        let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone());
532
533        let rc = Handle::try_current().ok();
534        let input = FFI_ExecutionPlan::new(input, rc);
535        let insert_op: FFI_InsertOp = insert_op.into();
536
537        let plan = unsafe {
538            let maybe_plan =
539                (self.0.insert_into)(&self.0, session, &input, insert_op).await;
540
541            <Arc<dyn ExecutionPlan>>::try_from(&df_result!(maybe_plan)?)?
542        };
543
544        Ok(plan)
545    }
546}
547
548#[cfg(test)]
549mod tests {
550    use arrow::datatypes::Schema;
551    use datafusion::prelude::{SessionContext, col, lit};
552    use datafusion_execution::TaskContextProvider;
553
554    use super::*;
555
556    fn create_test_table_provider() -> Result<Arc<dyn TableProvider>> {
557        use arrow::datatypes::Field;
558        use datafusion::arrow::array::Float32Array;
559        use datafusion::arrow::datatypes::DataType;
560        use datafusion::arrow::record_batch::RecordBatch;
561        use datafusion::datasource::MemTable;
562
563        let schema =
564            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
565
566        // define data in two partitions
567        let batch1 = RecordBatch::try_new(
568            Arc::clone(&schema),
569            vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))],
570        )?;
571        let batch2 = RecordBatch::try_new(
572            Arc::clone(&schema),
573            vec![Arc::new(Float32Array::from(vec![64.0]))],
574        )?;
575
576        Ok(Arc::new(MemTable::try_new(
577            schema,
578            vec![vec![batch1], vec![batch2]],
579        )?))
580    }
581
582    #[tokio::test]
583    async fn test_round_trip_ffi_table_provider_scan() -> Result<()> {
584        let provider = create_test_table_provider()?;
585        let ctx = Arc::new(SessionContext::new());
586        let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
587        let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider);
588
589        let mut ffi_provider =
590            FFI_TableProvider::new(provider, true, None, task_ctx_provider, None);
591        ffi_provider.library_marker_id = crate::mock_foreign_marker_id;
592
593        let foreign_table_provider: Arc<dyn TableProvider> = (&ffi_provider).into();
594
595        ctx.register_table("t", foreign_table_provider)?;
596
597        let df = ctx.table("t").await?;
598
599        df.select(vec![col("a")])?
600            .filter(col("a").gt(lit(3.0)))?
601            .show()
602            .await?;
603
604        Ok(())
605    }
606
607    #[tokio::test]
608    async fn test_round_trip_ffi_table_provider_insert_into() -> Result<()> {
609        let provider = create_test_table_provider()?;
610        let ctx = Arc::new(SessionContext::new());
611        let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
612        let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider);
613
614        let mut ffi_provider =
615            FFI_TableProvider::new(provider, true, None, task_ctx_provider, None);
616        ffi_provider.library_marker_id = crate::mock_foreign_marker_id;
617
618        let foreign_table_provider: Arc<dyn TableProvider> = (&ffi_provider).into();
619
620        ctx.register_table("t", foreign_table_provider)?;
621
622        let result = ctx
623            .sql("INSERT INTO t VALUES (128.0);")
624            .await?
625            .collect()
626            .await?;
627
628        assert!(result.len() == 1 && result[0].num_rows() == 1);
629
630        ctx.table("t")
631            .await?
632            .select(vec![col("a")])?
633            .filter(col("a").gt(lit(3.0)))?
634            .show()
635            .await?;
636
637        Ok(())
638    }
639
640    #[tokio::test]
641    async fn test_aggregation() -> Result<()> {
642        use arrow::datatypes::Field;
643        use datafusion::arrow::array::Float32Array;
644        use datafusion::arrow::datatypes::DataType;
645        use datafusion::arrow::record_batch::RecordBatch;
646        use datafusion::common::assert_batches_eq;
647        use datafusion::datasource::MemTable;
648
649        let schema =
650            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
651
652        // define data in two partitions
653        let batch1 = RecordBatch::try_new(
654            Arc::clone(&schema),
655            vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))],
656        )?;
657
658        let ctx = Arc::new(SessionContext::new());
659        let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
660        let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider);
661
662        let provider = Arc::new(MemTable::try_new(schema, vec![vec![batch1]])?);
663
664        let ffi_provider =
665            FFI_TableProvider::new(provider, true, None, task_ctx_provider, None);
666
667        let foreign_table_provider: Arc<dyn TableProvider> = (&ffi_provider).into();
668
669        ctx.register_table("t", foreign_table_provider)?;
670
671        let result = ctx
672            .sql("SELECT COUNT(*) as cnt FROM t")
673            .await?
674            .collect()
675            .await?;
676        #[rustfmt::skip]
677        let expected = [
678            "+-----+",
679            "| cnt |",
680            "+-----+",
681            "| 3   |",
682            "+-----+"
683        ];
684        assert_batches_eq!(expected, &result);
685        Ok(())
686    }
687
688    #[test]
689    fn test_ffi_table_provider_local_bypass() -> Result<()> {
690        let table_provider = create_test_table_provider()?;
691
692        let ctx = Arc::new(SessionContext::new()) as Arc<dyn TaskContextProvider>;
693        let task_ctx_provider = FFI_TaskContextProvider::from(&ctx);
694        let mut ffi_table =
695            FFI_TableProvider::new(table_provider, false, None, task_ctx_provider, None);
696
697        // Verify local libraries can be downcast to their original
698        let foreign_table: Arc<dyn TableProvider> = (&ffi_table).into();
699        assert!(
700            foreign_table
701                .as_any()
702                .downcast_ref::<datafusion::datasource::MemTable>()
703                .is_some()
704        );
705
706        // Verify different library markers generate foreign providers
707        ffi_table.library_marker_id = crate::mock_foreign_marker_id;
708        let foreign_table: Arc<dyn TableProvider> = (&ffi_table).into();
709        assert!(
710            foreign_table
711                .as_any()
712                .downcast_ref::<ForeignTableProvider>()
713                .is_some()
714        );
715
716        Ok(())
717    }
718}