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 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#[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 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 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 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 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 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}