renato2099 commented on code in PR #19079: URL: https://github.com/apache/datafusion/pull/19079#discussion_r2595363784
########## datafusion/ffi/src/proto/physical_extension_codec.rs: ########## @@ -0,0 +1,701 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::ffi::c_void; +use std::sync::Arc; + +use abi_stable::std_types::{RResult, RSlice, RStr, RVec}; +use abi_stable::StableAbi; +use async_trait::async_trait; +use datafusion_common::error::Result; +use datafusion_execution::TaskContext; +use datafusion_expr::{ + AggregateUDF, AggregateUDFImpl, ScalarUDF, ScalarUDFImpl, WindowUDF, WindowUDFImpl, +}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use tokio::runtime::Handle; + +use crate::execution::FFI_TaskContextProvider; +use crate::execution_plan::FFI_ExecutionPlan; +use crate::udaf::FFI_AggregateUDF; +use crate::udf::FFI_ScalarUDF; +use crate::udwf::FFI_WindowUDF; +use crate::util::FFIResult; +use crate::{df_result, rresult_return}; + +/// A stable struct for sharing [`PhysicalExtensionCodec`] across FFI boundaries. +#[repr(C)] +#[derive(Debug, StableAbi)] +#[allow(non_camel_case_types)] +pub struct FFI_PhysicalExtensionCodec { + /// Decode bytes into an execution plan. + try_decode: unsafe extern "C" fn( + &Self, + buf: RSlice<u8>, + inputs: RVec<FFI_ExecutionPlan>, + ) -> FFIResult<FFI_ExecutionPlan>, + + /// Encode an execution plan into bytes. + try_encode: + unsafe extern "C" fn(&Self, node: FFI_ExecutionPlan) -> FFIResult<RVec<u8>>, + + /// Decode bytes into a user defined scalar function. + try_decode_udf: unsafe extern "C" fn( + &Self, + name: RStr, + buf: RSlice<u8>, + ) -> FFIResult<FFI_ScalarUDF>, + + /// Encode a user defined scalar function into bytes. + try_encode_udf: + unsafe extern "C" fn(&Self, node: FFI_ScalarUDF) -> FFIResult<RVec<u8>>, + + /// Decode bytes into a user defined aggregate function. + try_decode_udaf: unsafe extern "C" fn( + &Self, + name: RStr, + buf: RSlice<u8>, + ) -> FFIResult<FFI_AggregateUDF>, + + /// Encode a user defined aggregate function into bytes. + try_encode_udaf: + unsafe extern "C" fn(&Self, node: FFI_AggregateUDF) -> FFIResult<RVec<u8>>, + + /// Decode bytes into a user defined window function. + try_decode_udwf: unsafe extern "C" fn( + &Self, + name: RStr, + buf: RSlice<u8>, + ) -> FFIResult<FFI_WindowUDF>, + + /// Encode a user defined window function into bytes. + try_encode_udwf: + unsafe extern "C" fn(&Self, node: FFI_WindowUDF) -> FFIResult<RVec<u8>>, + + /// Access the current [`TaskContext`]. + task_ctx_provider: FFI_TaskContextProvider, + + /// Used to create a clone on the provider of the execution plan. This should + /// only need to be called by the receiver of the plan. + pub clone: unsafe extern "C" fn(plan: &Self) -> Self, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Return the major DataFusion version number of this provider. + pub version: unsafe extern "C" fn() -> u64, + + /// Internal data. This is only to be accessed by the provider of the plan. + /// A [`ForeignPhysicalExtensionCodec`] should never attempt to access this data. + pub private_data: *mut c_void, + + /// Utility to identify when FFI objects are accessed locally through + /// the foreign interface. + pub library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_PhysicalExtensionCodec {} +unsafe impl Sync for FFI_PhysicalExtensionCodec {} + +struct PhysicalExtensionCodecPrivateData { + provider: Arc<dyn PhysicalExtensionCodec>, + runtime: Option<Handle>, +} + +impl FFI_PhysicalExtensionCodec { + fn inner(&self) -> &Arc<dyn PhysicalExtensionCodec> { + let private_data = self.private_data as *const PhysicalExtensionCodecPrivateData; + unsafe { &(*private_data).provider } + } + + fn runtime(&self) -> &Option<Handle> { + let private_data = self.private_data as *const PhysicalExtensionCodecPrivateData; + unsafe { &(*private_data).runtime } + } +} + +unsafe extern "C" fn try_decode_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + buf: RSlice<u8>, + inputs: RVec<FFI_ExecutionPlan>, +) -> FFIResult<FFI_ExecutionPlan> { + let task_ctx: Arc<TaskContext> = + rresult_return!((&codec.task_ctx_provider).try_into()); + let codec = codec.inner(); + let inputs = inputs + .into_iter() + .map(|plan| <Arc<dyn ExecutionPlan>>::try_from(&plan)) + .collect::<Result<Vec<_>>>(); + let inputs = rresult_return!(inputs); + + let plan = + rresult_return!(codec.try_decode(buf.as_ref(), &inputs, task_ctx.as_ref())); + + RResult::ROk(FFI_ExecutionPlan::new(plan, task_ctx, None)) +} + +unsafe extern "C" fn try_encode_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + node: FFI_ExecutionPlan, +) -> FFIResult<RVec<u8>> { + let codec = codec.inner(); + + let plan: Arc<dyn ExecutionPlan> = rresult_return!((&node).try_into()); + + let mut bytes = Vec::new(); + rresult_return!(codec.try_encode(plan, &mut bytes)); + + RResult::ROk(bytes.into()) +} + +unsafe extern "C" fn try_decode_udf_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + name: RStr, + buf: RSlice<u8>, +) -> FFIResult<FFI_ScalarUDF> { + let codec = codec.inner(); + + let udf = rresult_return!(codec.try_decode_udf(name.as_str(), buf.as_ref())); + let udf = FFI_ScalarUDF::from(udf); + + RResult::ROk(udf) +} + +unsafe extern "C" fn try_encode_udf_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + node: FFI_ScalarUDF, +) -> FFIResult<RVec<u8>> { + let codec = codec.inner(); + let node: Arc<dyn ScalarUDFImpl> = (&node).into(); + let node = ScalarUDF::new_from_shared_impl(node); + + let mut bytes = Vec::new(); + rresult_return!(codec.try_encode_udf(&node, &mut bytes)); + + RResult::ROk(bytes.into()) +} + +unsafe extern "C" fn try_decode_udaf_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + name: RStr, + buf: RSlice<u8>, +) -> FFIResult<FFI_AggregateUDF> { + let codec_inner = codec.inner(); + let udaf = rresult_return!(codec_inner.try_decode_udaf(name.into(), buf.as_ref())); + let udaf = FFI_AggregateUDF::from(udaf); + + RResult::ROk(udaf) +} + +unsafe extern "C" fn try_encode_udaf_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + node: FFI_AggregateUDF, +) -> FFIResult<RVec<u8>> { + let codec = codec.inner(); + let udaf: Arc<dyn AggregateUDFImpl> = (&node).into(); + let udaf = AggregateUDF::new_from_shared_impl(udaf); + + let mut bytes = Vec::new(); + rresult_return!(codec.try_encode_udaf(&udaf, &mut bytes)); + + RResult::ROk(bytes.into()) +} + +unsafe extern "C" fn try_decode_udwf_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + name: RStr, + buf: RSlice<u8>, +) -> FFIResult<FFI_WindowUDF> { + let codec = codec.inner(); + let udwf = rresult_return!(codec.try_decode_udwf(name.into(), buf.as_ref())); + let udwf = FFI_WindowUDF::from(udwf); + + RResult::ROk(udwf) +} + +unsafe extern "C" fn try_encode_udwf_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + node: FFI_WindowUDF, +) -> FFIResult<RVec<u8>> { + let codec = codec.inner(); + let udwf: Arc<dyn WindowUDFImpl> = (&node).into(); + let udwf = WindowUDF::new_from_shared_impl(udwf); + + let mut bytes = Vec::new(); + rresult_return!(codec.try_encode_udwf(&udwf, &mut bytes)); + + RResult::ROk(bytes.into()) +} + +unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_PhysicalExtensionCodec) { + let private_data = + Box::from_raw(provider.private_data as *mut PhysicalExtensionCodecPrivateData); + drop(private_data); +} + +unsafe extern "C" fn clone_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, +) -> FFI_PhysicalExtensionCodec { + let old_codec = Arc::clone(codec.inner()); + let runtime = codec.runtime().clone(); + + FFI_PhysicalExtensionCodec::new(old_codec, runtime, codec.task_ctx_provider.clone()) +} + +impl Drop for FFI_PhysicalExtensionCodec { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl FFI_PhysicalExtensionCodec { + /// Creates a new [`FFI_PhysicalExtensionCodec`]. + pub fn new( + provider: Arc<dyn PhysicalExtensionCodec + Send>, + runtime: Option<Handle>, + task_ctx_provider: impl Into<FFI_TaskContextProvider>, + ) -> Self { + let task_ctx_provider = task_ctx_provider.into(); + let private_data = + Box::new(PhysicalExtensionCodecPrivateData { provider, runtime }); + + Self { + try_decode: try_decode_fn_wrapper, + try_encode: try_encode_fn_wrapper, + try_decode_udf: try_decode_udf_fn_wrapper, + try_encode_udf: try_encode_udf_fn_wrapper, + try_decode_udaf: try_decode_udaf_fn_wrapper, + try_encode_udaf: try_encode_udaf_fn_wrapper, + try_decode_udwf: try_decode_udwf_fn_wrapper, + try_encode_udwf: try_encode_udwf_fn_wrapper, + task_ctx_provider, + + clone: clone_fn_wrapper, + release: release_fn_wrapper, + version: crate::version, + private_data: Box::into_raw(private_data) as *mut c_void, + library_marker_id: crate::get_library_marker_id, + } + } +} + +/// This wrapper struct exists on the receiver side of the FFI interface, so it has +/// no guarantees about being able to access the data in `private_data`. Any functions +/// defined on this struct must only use the stable functions provided in +/// FFI_PhysicalExtensionCodec to interact with the foreign table provider. +#[derive(Debug)] +pub struct ForeignPhysicalExtensionCodec(pub FFI_PhysicalExtensionCodec); + +unsafe impl Send for ForeignPhysicalExtensionCodec {} +unsafe impl Sync for ForeignPhysicalExtensionCodec {} + +impl From<&FFI_PhysicalExtensionCodec> for Arc<dyn PhysicalExtensionCodec> { + fn from(provider: &FFI_PhysicalExtensionCodec) -> Self { + if (provider.library_marker_id)() == crate::get_library_marker_id() { + Arc::clone(provider.inner()) + } else { + Arc::new(ForeignPhysicalExtensionCodec(provider.clone())) + } + } +} + +impl Clone for FFI_PhysicalExtensionCodec { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +#[async_trait] +impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc<dyn ExecutionPlan>], + _ctx: &TaskContext, + ) -> Result<Arc<dyn ExecutionPlan>> { + let task_ctx = (&self.0.task_ctx_provider).try_into()?; + let inputs = inputs + .iter() + .map(|plan| { + FFI_ExecutionPlan::new(Arc::clone(plan), Arc::clone(&task_ctx), None) + }) + .collect(); + + let plan = + df_result!(unsafe { (self.0.try_decode)(&self.0, buf.into(), inputs) })?; + let plan: Arc<dyn ExecutionPlan> = (&plan).try_into()?; + + Ok(plan) + } + + fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> { + let task_ctx = (&self.0.task_ctx_provider).try_into()?; + let plan = FFI_ExecutionPlan::new(node, task_ctx, None); + let bytes = df_result!(unsafe { (self.0.try_encode)(&self.0, plan) })?; + + buf.extend(bytes); + Ok(()) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> { + let udf = unsafe { + df_result!((self.0.try_decode_udf)(&self.0, name.into(), buf.into())) + }?; + let udf: Arc<dyn ScalarUDFImpl> = (&udf).into(); + + Ok(Arc::new(ScalarUDF::new_from_shared_impl(udf))) + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> { + let node = FFI_ScalarUDF::from(Arc::new(node.clone())); + let bytes = df_result!(unsafe { (self.0.try_encode_udf)(&self.0, node) })?; + + buf.extend(bytes); + + Ok(()) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> { + let udaf = unsafe { + df_result!((self.0.try_decode_udaf)(&self.0, name.into(), buf.into())) + }?; + let udaf: Arc<dyn AggregateUDFImpl> = (&udaf).into(); + + Ok(Arc::new(AggregateUDF::new_from_shared_impl(udaf))) + } + + fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> { + let node = Arc::new(node.clone()); + let node = FFI_AggregateUDF::from(node); + let bytes = df_result!(unsafe { (self.0.try_encode_udaf)(&self.0, node) })?; + + buf.extend(bytes); + + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> { + let udwf = unsafe { + df_result!((self.0.try_decode_udwf)(&self.0, name.into(), buf.into())) + }?; + let udwf: Arc<dyn WindowUDFImpl> = (&udwf).into(); + + Ok(Arc::new(WindowUDF::new_from_shared_impl(udwf))) + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> { + let node = Arc::new(node.clone()); + let node = FFI_WindowUDF::from(node); + let bytes = df_result!(unsafe { (self.0.try_encode_udwf)(&self.0, node) })?; + + buf.extend(bytes); + + Ok(()) + } +} + +#[cfg(test)] +pub(crate) mod tests { + use std::sync::Arc; + + use arrow_schema::{DataType, Field, Schema}; + use datafusion_common::exec_err; + use datafusion_execution::TaskContext; + use datafusion_expr::ptr_eq::arc_ptr_eq; + use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF, WindowUDFImpl}; + use datafusion_functions::math::abs::AbsFunc; + use datafusion_functions_aggregate::sum::Sum; + use datafusion_functions_window::rank::{Rank, RankType}; + use datafusion_physical_plan::ExecutionPlan; + use datafusion_proto::physical_plan::PhysicalExtensionCodec; + + use crate::execution_plan::tests::EmptyExec; + use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; + + #[derive(Debug)] + pub(crate) struct TestExtensionCodec; + + impl TestExtensionCodec { + pub(crate) const MAGIC_NUMBER: u8 = 127; + } + + impl PhysicalExtensionCodec for TestExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + _inputs: &[Arc<dyn ExecutionPlan>], + _ctx: &TaskContext, + ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> { + if buf[0] != Self::MAGIC_NUMBER { + return exec_err!( + "TestExtensionCodec input buffer does not start with magic number" + ); + } + + if buf.len() != 2 || buf[1] != Self::MAGIC_NUMBER { + return exec_err!("TestExtensionCodec unable to decode udf"); + } Review Comment: I think it is just because of the way how the tests are setup , we could endup with a case with `[MAGIC_NUMBER, MAGIC_NUMBER]`, but I agree it could be simplified. @timsaucer feel free to correct me if I got it wrong 😅 ########## datafusion/ffi/src/proto/logical_extension_codec.rs: ########## @@ -0,0 +1,739 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::ffi::c_void; +use std::sync::Arc; + +use abi_stable::std_types::{RResult, RSlice, RStr, RVec}; +use abi_stable::StableAbi; +use arrow::datatypes::SchemaRef; +use async_trait::async_trait; +use datafusion_catalog::TableProvider; +use datafusion_common::error::Result; +use datafusion_common::{not_impl_err, TableReference}; +use datafusion_datasource::file_format::FileFormatFactory; +use datafusion_execution::TaskContext; +use datafusion_expr::{ + AggregateUDF, AggregateUDFImpl, Extension, LogicalPlan, ScalarUDF, ScalarUDFImpl, + WindowUDF, WindowUDFImpl, +}; +use datafusion_proto::logical_plan::LogicalExtensionCodec; +use tokio::runtime::Handle; + +use crate::arrow_wrappers::WrappedSchema; +use crate::execution::FFI_TaskContextProvider; +use crate::table_provider::FFI_TableProvider; +use crate::udaf::FFI_AggregateUDF; +use crate::udf::FFI_ScalarUDF; +use crate::udwf::FFI_WindowUDF; +use crate::util::FFIResult; +use crate::{df_result, rresult_return}; + +/// A stable struct for sharing [`LogicalExtensionCodec`] across FFI boundaries. +#[repr(C)] +#[derive(Debug, StableAbi)] +#[allow(non_camel_case_types)] +pub struct FFI_LogicalExtensionCodec { + /// Decode bytes into a table provider. + try_decode_table_provider: unsafe extern "C" fn( + &Self, + buf: RSlice<u8>, + table_ref: RStr, + schema: WrappedSchema, + ) -> FFIResult<FFI_TableProvider>, + + /// Encode a table provider into bytes. + try_encode_table_provider: unsafe extern "C" fn( + &Self, + table_ref: RStr, + node: FFI_TableProvider, + ) -> FFIResult<RVec<u8>>, + + /// Decode bytes into a user defined scalar function. + try_decode_udf: unsafe extern "C" fn( + &Self, + name: RStr, + buf: RSlice<u8>, + ) -> FFIResult<FFI_ScalarUDF>, + + /// Encode a user defined scalar function into bytes. + try_encode_udf: + unsafe extern "C" fn(&Self, node: FFI_ScalarUDF) -> FFIResult<RVec<u8>>, + + /// Decode bytes into a user defined aggregate function. + try_decode_udaf: unsafe extern "C" fn( + &Self, + name: RStr, + buf: RSlice<u8>, + ) -> FFIResult<FFI_AggregateUDF>, + + /// Encode a user defined aggregate function into bytes. + try_encode_udaf: + unsafe extern "C" fn(&Self, node: FFI_AggregateUDF) -> FFIResult<RVec<u8>>, + + /// Decode bytes into a user defined window function. + try_decode_udwf: unsafe extern "C" fn( + &Self, + name: RStr, + buf: RSlice<u8>, + ) -> FFIResult<FFI_WindowUDF>, + + /// Encode a user defined window function into bytes. + try_encode_udwf: + unsafe extern "C" fn(&Self, node: FFI_WindowUDF) -> FFIResult<RVec<u8>>, + + task_ctx_provider: FFI_TaskContextProvider, + + /// Used to create a clone on the provider of the execution plan. This should + /// only need to be called by the receiver of the plan. + pub clone: unsafe extern "C" fn(plan: &Self) -> Self, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Return the major DataFusion version number of this provider. + pub version: unsafe extern "C" fn() -> u64, + + /// Internal data. This is only to be accessed by the provider of the plan. + /// A [`ForeignLogicalExtensionCodec`] should never attempt to access this data. + pub private_data: *mut c_void, + + /// Utility to identify when FFI objects are accessed locally through + /// the foreign interface. + pub library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_LogicalExtensionCodec {} +unsafe impl Sync for FFI_LogicalExtensionCodec {} + +struct LogicalExtensionCodecPrivateData { + provider: Arc<dyn LogicalExtensionCodec>, + runtime: Option<Handle>, +} + +impl FFI_LogicalExtensionCodec { + fn inner(&self) -> &Arc<dyn LogicalExtensionCodec> { + let private_data = self.private_data as *const LogicalExtensionCodecPrivateData; + unsafe { &(*private_data).provider } + } + + fn runtime(&self) -> &Option<Handle> { + let private_data = self.private_data as *const LogicalExtensionCodecPrivateData; + unsafe { &(*private_data).runtime } + } + + fn task_ctx(&self) -> Result<Arc<TaskContext>> { + (&self.task_ctx_provider).try_into() + } +} + +unsafe extern "C" fn try_decode_table_provider_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, + buf: RSlice<u8>, + table_ref: RStr, + schema: WrappedSchema, +) -> FFIResult<FFI_TableProvider> { + let ctx = rresult_return!(codec.task_ctx()); + let runtime = codec.runtime().clone(); + let codec = codec.inner(); + let table_ref = TableReference::from(table_ref.as_str()); + let schema: SchemaRef = schema.into(); + + let table_provider = rresult_return!(codec.try_decode_table_provider( + buf.as_ref(), + &table_ref, + schema, + ctx.as_ref() + )); + + RResult::ROk(FFI_TableProvider::new(table_provider, true, runtime)) +} + +unsafe extern "C" fn try_encode_table_provider_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, + table_ref: RStr, + node: FFI_TableProvider, +) -> FFIResult<RVec<u8>> { + let table_ref = TableReference::from(table_ref.as_str()); + let table_provider: Arc<dyn TableProvider> = (&node).into(); + let codec = codec.inner(); + + let mut bytes = Vec::new(); + rresult_return!(codec.try_encode_table_provider( + &table_ref, + table_provider, + &mut bytes + )); + + RResult::ROk(bytes.into()) +} + +unsafe extern "C" fn try_decode_udf_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, + name: RStr, + buf: RSlice<u8>, +) -> FFIResult<FFI_ScalarUDF> { + let codec = codec.inner(); + + let udf = rresult_return!(codec.try_decode_udf(name.as_str(), buf.as_ref())); + let udf = FFI_ScalarUDF::from(udf); + + RResult::ROk(udf) +} + +unsafe extern "C" fn try_encode_udf_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, + node: FFI_ScalarUDF, +) -> FFIResult<RVec<u8>> { + let codec = codec.inner(); + let node: Arc<dyn ScalarUDFImpl> = (&node).into(); + let node = ScalarUDF::new_from_shared_impl(node); + + let mut bytes = Vec::new(); + rresult_return!(codec.try_encode_udf(&node, &mut bytes)); + + RResult::ROk(bytes.into()) +} + +unsafe extern "C" fn try_decode_udaf_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, + name: RStr, + buf: RSlice<u8>, +) -> FFIResult<FFI_AggregateUDF> { + let codec_inner = codec.inner(); + let udaf = rresult_return!(codec_inner.try_decode_udaf(name.into(), buf.as_ref())); + let udaf = FFI_AggregateUDF::from(udaf); + + RResult::ROk(udaf) +} + +unsafe extern "C" fn try_encode_udaf_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, + node: FFI_AggregateUDF, +) -> FFIResult<RVec<u8>> { + let codec = codec.inner(); + let udaf: Arc<dyn AggregateUDFImpl> = (&node).into(); + let udaf = AggregateUDF::new_from_shared_impl(udaf); + + let mut bytes = Vec::new(); + rresult_return!(codec.try_encode_udaf(&udaf, &mut bytes)); + + RResult::ROk(bytes.into()) +} + +unsafe extern "C" fn try_decode_udwf_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, + name: RStr, + buf: RSlice<u8>, +) -> FFIResult<FFI_WindowUDF> { + let codec = codec.inner(); + let udwf = rresult_return!(codec.try_decode_udwf(name.into(), buf.as_ref())); + let udwf = FFI_WindowUDF::from(udwf); + + RResult::ROk(udwf) +} + +unsafe extern "C" fn try_encode_udwf_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, + node: FFI_WindowUDF, +) -> FFIResult<RVec<u8>> { + let codec = codec.inner(); + let udwf: Arc<dyn WindowUDFImpl> = (&node).into(); + let udwf = WindowUDF::new_from_shared_impl(udwf); + + let mut bytes = Vec::new(); + rresult_return!(codec.try_encode_udwf(&udwf, &mut bytes)); + + RResult::ROk(bytes.into()) +} + +unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_LogicalExtensionCodec) { + let private_data = + Box::from_raw(provider.private_data as *mut LogicalExtensionCodecPrivateData); + drop(private_data); +} + +unsafe extern "C" fn clone_fn_wrapper( + codec: &FFI_LogicalExtensionCodec, +) -> FFI_LogicalExtensionCodec { + let old_codec = Arc::clone(codec.inner()); + let runtime = codec.runtime().clone(); + + FFI_LogicalExtensionCodec::new(old_codec, runtime, codec.task_ctx_provider.clone()) +} + +impl Drop for FFI_LogicalExtensionCodec { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl FFI_LogicalExtensionCodec { + /// Creates a new [`FFI_LogicalExtensionCodec`]. + pub fn new( + provider: Arc<dyn LogicalExtensionCodec + Send>, + runtime: Option<Handle>, + task_ctx_provider: impl Into<FFI_TaskContextProvider>, + ) -> Self { + let task_ctx_provider = task_ctx_provider.into(); + let private_data = + Box::new(LogicalExtensionCodecPrivateData { provider, runtime }); + + Self { + try_decode_table_provider: try_decode_table_provider_fn_wrapper, + try_encode_table_provider: try_encode_table_provider_fn_wrapper, + try_decode_udf: try_decode_udf_fn_wrapper, + try_encode_udf: try_encode_udf_fn_wrapper, + try_decode_udaf: try_decode_udaf_fn_wrapper, + try_encode_udaf: try_encode_udaf_fn_wrapper, + try_decode_udwf: try_decode_udwf_fn_wrapper, + try_encode_udwf: try_encode_udwf_fn_wrapper, + task_ctx_provider, + + clone: clone_fn_wrapper, + release: release_fn_wrapper, + version: crate::version, + private_data: Box::into_raw(private_data) as *mut c_void, + library_marker_id: crate::get_library_marker_id, + } + } +} + +/// This wrapper struct exists on the receiver side of the FFI interface, so it has +/// no guarantees about being able to access the data in `private_data`. Any functions +/// defined on this struct must only use the stable functions provided in +/// FFI_LogicalExtensionCodec to interact with the foreign table provider. +#[derive(Debug)] +pub struct ForeignLogicalExtensionCodec(pub FFI_LogicalExtensionCodec); + +unsafe impl Send for ForeignLogicalExtensionCodec {} +unsafe impl Sync for ForeignLogicalExtensionCodec {} + +impl From<&FFI_LogicalExtensionCodec> for Arc<dyn LogicalExtensionCodec> { + fn from(provider: &FFI_LogicalExtensionCodec) -> Self { + if (provider.library_marker_id)() == crate::get_library_marker_id() { + Arc::clone(provider.inner()) + } else { + Arc::new(ForeignLogicalExtensionCodec(provider.clone())) + } + } +} + +impl Clone for FFI_LogicalExtensionCodec { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +#[async_trait] +impl LogicalExtensionCodec for ForeignLogicalExtensionCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[LogicalPlan], + _ctx: &TaskContext, + ) -> Result<Extension> { + not_impl_err!("FFI does not support decode of Extensions") + } + + fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> { + not_impl_err!("FFI does not support encode of Extensions") + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &TableReference, + schema: SchemaRef, + _ctx: &TaskContext, + ) -> Result<Arc<dyn TableProvider>> { + let table_ref = table_ref.to_string(); + let schema: WrappedSchema = schema.into(); + + let ffi_table_provider = unsafe { + df_result!((self.0.try_decode_table_provider)( + &self.0, + buf.into(), + table_ref.as_str().into(), + schema + )) + }?; + + Ok((&ffi_table_provider).into()) + } + + fn try_encode_table_provider( + &self, + table_ref: &TableReference, + node: Arc<dyn TableProvider>, + buf: &mut Vec<u8>, + ) -> Result<()> { + let table_ref = table_ref.to_string(); + let node = FFI_TableProvider::new(node, true, None); + + let bytes = df_result!(unsafe { + (self.0.try_encode_table_provider)(&self.0, table_ref.as_str().into(), node) + })?; + + buf.extend(bytes); + + Ok(()) + } + + fn try_decode_file_format( + &self, + _buf: &[u8], + _ctx: &TaskContext, + ) -> Result<Arc<dyn FileFormatFactory>> { + not_impl_err!("FFI does not support decode_file_format") + } + + fn try_encode_file_format( + &self, + _buf: &mut Vec<u8>, + _node: Arc<dyn FileFormatFactory>, + ) -> Result<()> { + not_impl_err!("FFI does not support encode_file_format") + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> { + let udf = unsafe { + df_result!((self.0.try_decode_udf)(&self.0, name.into(), buf.into())) + }?; + let udf: Arc<dyn ScalarUDFImpl> = (&udf).into(); + + Ok(Arc::new(ScalarUDF::new_from_shared_impl(udf))) + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> { + let node = FFI_ScalarUDF::from(Arc::new(node.clone())); + let bytes = df_result!(unsafe { (self.0.try_encode_udf)(&self.0, node) })?; + + buf.extend(bytes); + + Ok(()) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> { + let udaf = unsafe { + df_result!((self.0.try_decode_udaf)(&self.0, name.into(), buf.into())) + }?; + let udaf: Arc<dyn AggregateUDFImpl> = (&udaf).into(); + + Ok(Arc::new(AggregateUDF::new_from_shared_impl(udaf))) + } + + fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> { + let node = Arc::new(node.clone()); + let node = FFI_AggregateUDF::from(node); + let bytes = df_result!(unsafe { (self.0.try_encode_udaf)(&self.0, node) })?; + + buf.extend(bytes); + + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> { + let udwf = unsafe { + df_result!((self.0.try_decode_udwf)(&self.0, name.into(), buf.into())) + }?; + let udwf: Arc<dyn WindowUDFImpl> = (&udwf).into(); + + Ok(Arc::new(WindowUDF::new_from_shared_impl(udwf))) + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> { + let node = Arc::new(node.clone()); + let node = FFI_WindowUDF::from(node); + let bytes = df_result!(unsafe { (self.0.try_encode_udwf)(&self.0, node) })?; + + buf.extend(bytes); + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::record_batch; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use datafusion_catalog::{MemTable, TableProvider}; + use datafusion_common::{exec_err, TableReference}; + use datafusion_datasource::file_format::FileFormatFactory; + use datafusion_execution::TaskContext; + use datafusion_expr::ptr_eq::arc_ptr_eq; + use datafusion_expr::{AggregateUDF, Extension, LogicalPlan, ScalarUDF, WindowUDF}; + use datafusion_functions::math::abs::AbsFunc; + use datafusion_functions_aggregate::sum::Sum; + use datafusion_functions_window::rank::{Rank, RankType}; + use datafusion_proto::logical_plan::LogicalExtensionCodec; + use datafusion_proto::physical_plan::PhysicalExtensionCodec; + + use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; + use crate::proto::physical_extension_codec::tests::TestExtensionCodec; + + fn create_test_table() -> MemTable { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let rb = record_batch!(("a", Int32, [1, 2, 3])) + .expect("should be able to create a record batch"); + MemTable::try_new(schema, vec![vec![rb]]) + .expect("should be able to create an in memory table") + } + + impl LogicalExtensionCodec for TestExtensionCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[LogicalPlan], + _ctx: &TaskContext, + ) -> datafusion_common::Result<Extension> { + unimplemented!() + } + + fn try_encode( + &self, + _node: &Extension, + _buf: &mut Vec<u8>, + ) -> datafusion_common::Result<()> { + unimplemented!() + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + _table_ref: &TableReference, + schema: SchemaRef, + _ctx: &TaskContext, + ) -> datafusion_common::Result<Arc<dyn TableProvider>> { + if buf[0] != Self::MAGIC_NUMBER { + return exec_err!( + "TestExtensionCodec input buffer does not start with magic number" + ); + } + + if schema != create_test_table().schema() { + return exec_err!("Incorrect test table schema"); + } + + if buf.len() != 2 || buf[1] != Self::MAGIC_NUMBER { + return exec_err!("TestExtensionCodec unable to decode table provider"); + } + + Ok(Arc::new(create_test_table()) as Arc<dyn TableProvider>) + } + + fn try_encode_table_provider( + &self, + _table_ref: &TableReference, + node: Arc<dyn TableProvider>, + buf: &mut Vec<u8>, + ) -> datafusion_common::Result<()> { + buf.push(Self::MAGIC_NUMBER); + + if !node.as_any().is::<MemTable>() { + return exec_err!("TestExtensionCodec only expects Abs UDF"); + }; + + if node.schema() != create_test_table().schema() { + return exec_err!("Unexpected schema for encoding."); + } + + buf.push(Self::MAGIC_NUMBER); + + Ok(()) + } + + fn try_decode_file_format( + &self, + _buf: &[u8], + _ctx: &TaskContext, + ) -> datafusion_common::Result<Arc<dyn FileFormatFactory>> { + unimplemented!() + } + + fn try_encode_file_format( + &self, + _buf: &mut Vec<u8>, + _node: Arc<dyn FileFormatFactory>, + ) -> datafusion_common::Result<()> { + unimplemented!() + } + + fn try_decode_udf( + &self, + name: &str, + buf: &[u8], + ) -> datafusion_common::Result<Arc<ScalarUDF>> { + PhysicalExtensionCodec::try_decode_udf(self, name, buf) + } + + fn try_encode_udf( + &self, + node: &ScalarUDF, + buf: &mut Vec<u8>, + ) -> datafusion_common::Result<()> { + PhysicalExtensionCodec::try_encode_udf(self, node, buf) + } + + fn try_decode_udaf( + &self, + name: &str, + buf: &[u8], + ) -> datafusion_common::Result<Arc<AggregateUDF>> { + PhysicalExtensionCodec::try_decode_udaf(self, name, buf) + } + + fn try_encode_udaf( + &self, + node: &AggregateUDF, + buf: &mut Vec<u8>, + ) -> datafusion_common::Result<()> { + PhysicalExtensionCodec::try_encode_udaf(self, node, buf) + } + + fn try_decode_udwf( + &self, + name: &str, + buf: &[u8], + ) -> datafusion_common::Result<Arc<WindowUDF>> { + PhysicalExtensionCodec::try_decode_udwf(self, name, buf) + } + + fn try_encode_udwf( + &self, + node: &WindowUDF, + buf: &mut Vec<u8>, + ) -> datafusion_common::Result<()> { + PhysicalExtensionCodec::try_encode_udwf(self, node, buf) + } + } + + #[test] + fn roundtrip_ffi_logical_extension_codec_table_provider( + ) -> datafusion_common::Result<()> { + let codec = Arc::new(TestExtensionCodec {}); + let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + + let mut ffi_codec = + FFI_LogicalExtensionCodec::new(codec, None, task_ctx_provider); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + let foreign_codec: Arc<dyn LogicalExtensionCodec> = (&ffi_codec).into(); + + let table = Arc::new(create_test_table()) as Arc<dyn TableProvider>; + let mut bytes = Vec::new(); + foreign_codec.try_encode_table_provider(&"my_table".into(), table, &mut bytes)?; + + let returned_table = foreign_codec.try_decode_table_provider( + &bytes, + &"my_table".into(), + create_test_table().schema(), + ctx.task_ctx().as_ref(), + )?; + + assert!(returned_table.as_any().is::<MemTable>()); + + Ok(()) + } + + #[test] + fn roundtrip_ffi_logical_extension_codec_udf() -> datafusion_common::Result<()> { + let codec = Arc::new(TestExtensionCodec {}); + let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + + let mut ffi_codec = + FFI_LogicalExtensionCodec::new(codec, None, task_ctx_provider); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + let foreign_codec: Arc<dyn LogicalExtensionCodec> = (&ffi_codec).into(); + + let udf = Arc::new(ScalarUDF::from(AbsFunc::new())); + let mut bytes = Vec::new(); + foreign_codec.try_encode_udf(udf.as_ref(), &mut bytes)?; Review Comment: I am wondering if having a unit test that shows error handling (or expected behavior) when trying the round trip but the data we are reading is corrupted/incomplete somehow. That was at least it is clear for the consumers what the intent would be -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
