1use 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#[repr(C)]
92#[derive(Debug, StableAbi)]
93pub struct FFI_TableProvider {
94 schema: unsafe extern "C" fn(provider: &Self) -> WrappedSchema,
96
97 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 table_type: unsafe extern "C" fn(provider: &Self) -> FFI_TableType,
118
119 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 clone: unsafe extern "C" fn(plan: &Self) -> Self,
141
142 release: unsafe extern "C" fn(arg: &mut Self),
144
145 pub version: unsafe extern "C" fn() -> u64,
147
148 private_data: *mut c_void,
151
152 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 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#[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 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 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 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 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 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}