datafusion_ffi/tests/
mod.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::sync::Arc;
19
20use arrow::array::{RecordBatch, record_batch};
21use arrow_schema::{DataType, Field, Schema};
22use async_provider::create_async_table_provider;
23use async_trait::async_trait;
24use catalog::create_catalog_provider;
25use datafusion_catalog::MemTable;
26use datafusion_catalog::{Session, TableProvider};
27use datafusion_common::stats::Precision;
28use datafusion_common::{ColumnStatistics, Statistics};
29use datafusion_common::{Result, ScalarValue};
30use datafusion_expr::{Expr, TableType};
31use datafusion_physical_expr::PhysicalExpr;
32use datafusion_physical_plan::ExecutionPlan;
33use sync_provider::create_sync_table_provider;
34use udf_udaf_udwf::{
35    create_ffi_abs_func, create_ffi_first_value_func, create_ffi_random_func,
36    create_ffi_rank_func, create_ffi_stddev_func, create_ffi_sum_func,
37    create_ffi_table_func,
38};
39
40use crate::catalog_provider::FFI_CatalogProvider;
41use crate::catalog_provider_list::FFI_CatalogProviderList;
42use crate::config::extension_options::FFI_ExtensionOptions;
43use crate::execution_plan::FFI_ExecutionPlan;
44use crate::execution_plan::tests::{EmptyExec, create_dynamic_filter};
45use crate::physical_optimizer::FFI_PhysicalOptimizerRule;
46use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
47use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
48use crate::query_planner::FFI_QueryPlanner;
49use crate::table_provider::FFI_TableProvider;
50use crate::table_provider_factory::FFI_TableProviderFactory;
51use crate::tests::catalog::create_catalog_provider_list;
52use crate::udaf::FFI_AggregateUDF;
53use crate::udf::FFI_ScalarUDF;
54use crate::udtf::FFI_TableFunction;
55use crate::udwf::FFI_WindowUDF;
56use crate::util::FFI_Option;
57
58mod async_provider;
59pub mod catalog;
60pub mod config;
61mod physical_optimizer;
62mod query_planner;
63mod sync_provider;
64mod table_provider_factory;
65mod udf_udaf_udwf;
66pub mod utils;
67
68#[repr(C)]
69/// This struct defines the module interfaces. It is to be shared by
70/// both the module loading program and library that implements the
71/// module.
72pub struct ForeignLibraryModule {
73    /// Construct an opinionated catalog provider
74    pub create_catalog:
75        extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_CatalogProvider,
76
77    /// Construct an opinionated catalog provider list
78    pub create_catalog_list:
79        extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_CatalogProviderList,
80
81    /// Constructs the table provider
82    pub create_table: extern "C" fn(
83        synchronous: bool,
84        codec: FFI_LogicalExtensionCodec,
85    ) -> FFI_TableProvider,
86
87    /// Constructs the table provider factory
88    pub create_table_factory:
89        extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProviderFactory,
90
91    /// Create a scalar UDF
92    pub create_scalar_udf: extern "C" fn() -> FFI_ScalarUDF,
93
94    pub create_nullary_udf: extern "C" fn() -> FFI_ScalarUDF,
95
96    pub create_timezone_udf: extern "C" fn() -> FFI_ScalarUDF,
97
98    pub create_placement_udf: extern "C" fn() -> FFI_ScalarUDF,
99
100    pub create_table_function:
101        extern "C" fn(FFI_LogicalExtensionCodec) -> FFI_TableFunction,
102
103    /// Create an aggregate UDAF using sum
104    pub create_sum_udaf: extern "C" fn() -> FFI_AggregateUDF,
105
106    /// Create  grouping UDAF using stddev
107    pub create_stddev_udaf: extern "C" fn() -> FFI_AggregateUDF,
108
109    pub create_rank_udwf: extern "C" fn() -> FFI_WindowUDF,
110
111    /// Create extension options, for either ConfigOptions or TableOptions
112    pub create_extension_options: extern "C" fn() -> FFI_ExtensionOptions,
113
114    pub create_empty_exec: extern "C" fn() -> FFI_ExecutionPlan,
115
116    pub create_exec_with_expressions: extern "C" fn() -> FFI_ExecutionPlan,
117
118    pub create_exec_with_dynamic_expressions: extern "C" fn() -> FFI_ExecutionPlan,
119
120    pub create_exec_with_statistics: extern "C" fn() -> FFI_ExecutionPlan,
121
122    pub create_table_with_statistics:
123        extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider,
124
125    pub create_physical_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule,
126
127    pub create_context_aware_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule,
128
129    /// Construct a query planner. When `library_a_planner` is provided the
130    /// planner delegates to it, as library C does after library A swaps planners.
131    pub create_query_planner: extern "C" fn(
132        logical_codec: FFI_LogicalExtensionCodec,
133        physical_codec: FFI_PhysicalExtensionCodec,
134        library_a_planner: FFI_Option<FFI_QueryPlanner>,
135    ) -> FFI_QueryPlanner,
136
137    pub version: extern "C" fn() -> u64,
138
139    /// Create an aggregate UDAF using first_value
140    pub create_first_value_udaf: extern "C" fn() -> FFI_AggregateUDF,
141}
142
143pub fn create_test_schema() -> Arc<Schema> {
144    Arc::new(Schema::new(vec![
145        Field::new("a", DataType::Int32, true),
146        Field::new("b", DataType::Float64, true),
147    ]))
148}
149
150pub fn create_record_batch(start_value: i32, num_values: usize) -> RecordBatch {
151    let end_value = start_value + num_values as i32;
152    let a_vals: Vec<i32> = (start_value..end_value).collect();
153    let b_vals: Vec<f64> = a_vals.iter().map(|v| *v as f64).collect();
154
155    record_batch!(("a", Int32, a_vals), ("b", Float64, b_vals)).unwrap()
156}
157
158/// Here we only wish to create a simple table provider as an example.
159/// We create an in-memory table and convert it to it's FFI counterpart.
160extern "C" fn construct_table_provider(
161    synchronous: bool,
162    codec: FFI_LogicalExtensionCodec,
163) -> FFI_TableProvider {
164    match synchronous {
165        true => create_sync_table_provider(codec),
166        false => create_async_table_provider(codec),
167    }
168}
169
170/// Here we only wish to create a simple table provider as an example.
171/// We create an in-memory table and convert it to it's FFI counterpart.
172extern "C" fn construct_table_provider_factory(
173    codec: FFI_LogicalExtensionCodec,
174) -> FFI_TableProviderFactory {
175    table_provider_factory::create(codec)
176}
177
178pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan {
179    let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
180
181    let plan = Arc::new(EmptyExec::new(schema));
182    FFI_ExecutionPlan::new(plan, None)
183}
184
185pub(crate) extern "C" fn create_exec_with_expressions() -> FFI_ExecutionPlan {
186    let schema = Arc::new(Schema::empty());
187    let expression: Arc<dyn PhysicalExpr> = create_dynamic_filter();
188    let plan = Arc::new(EmptyExec::new(schema).with_expressions(vec![expression]));
189    FFI_ExecutionPlan::new(plan, None)
190}
191
192pub(crate) extern "C" fn create_exec_with_dynamic_expressions() -> FFI_ExecutionPlan {
193    let schema = Arc::new(Schema::empty());
194    let expression: Arc<dyn PhysicalExpr> = create_dynamic_filter();
195    let plan =
196        Arc::new(EmptyExec::new(schema).with_dynamic_expressions(vec![expression]));
197    FFI_ExecutionPlan::new(plan, None)
198}
199
200/// Returns canonical statistics used by both the producer and consumer sides of
201/// the integration tests so round-trips can be asserted without hard-coding
202/// the values in two places.
203pub fn make_test_statistics() -> Statistics {
204    Statistics {
205        num_rows: Precision::Exact(42),
206        total_byte_size: Precision::Exact(672),
207        column_statistics: vec![
208            ColumnStatistics {
209                null_count: Precision::Exact(0),
210                max_value: Precision::Exact(ScalarValue::Int32(Some(100))),
211                min_value: Precision::Exact(ScalarValue::Int32(Some(-10))),
212                sum_value: Precision::Exact(ScalarValue::Int64(Some(1890))),
213                distinct_count: Precision::Inexact(40),
214                byte_size: Precision::Exact(168),
215            },
216            ColumnStatistics {
217                null_count: Precision::Exact(1),
218                max_value: Precision::Exact(ScalarValue::Float64(Some(99.5))),
219                min_value: Precision::Exact(ScalarValue::Float64(Some(-1.5))),
220                sum_value: Precision::Absent,
221                distinct_count: Precision::Absent,
222                byte_size: Precision::Exact(328),
223            },
224        ],
225    }
226}
227
228pub(crate) extern "C" fn create_exec_with_statistics() -> FFI_ExecutionPlan {
229    let schema = create_test_schema();
230    let plan = Arc::new(EmptyExec::new(schema).with_statistics(make_test_statistics()));
231    FFI_ExecutionPlan::new(plan, None)
232}
233
234/// Thin wrapper that attaches a fixed [`Statistics`] snapshot to any inner
235/// [`TableProvider`] without changing its scan behaviour.
236#[derive(Debug)]
237struct TableWithStats {
238    inner: Arc<dyn TableProvider>,
239    stats: Statistics,
240}
241
242#[async_trait]
243impl TableProvider for TableWithStats {
244    fn schema(&self) -> arrow_schema::SchemaRef {
245        self.inner.schema()
246    }
247
248    fn table_type(&self) -> TableType {
249        self.inner.table_type()
250    }
251
252    fn statistics(&self) -> Option<Statistics> {
253        Some(self.stats.clone())
254    }
255
256    async fn scan(
257        &self,
258        session: &dyn Session,
259        projection: Option<&Vec<usize>>,
260        filters: &[Expr],
261        limit: Option<usize>,
262    ) -> Result<Arc<dyn ExecutionPlan>> {
263        self.inner.scan(session, projection, filters, limit).await
264    }
265}
266
267pub(crate) extern "C" fn create_table_with_statistics(
268    codec: FFI_LogicalExtensionCodec,
269) -> FFI_TableProvider {
270    let schema = create_test_schema();
271    let batch = create_record_batch(1, 5);
272    let inner = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
273    let provider = Arc::new(TableWithStats {
274        inner,
275        stats: make_test_statistics(),
276    });
277    FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec)
278}
279
280/// This defines the entry point for using the module.
281#[unsafe(no_mangle)]
282pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule {
283    ForeignLibraryModule {
284        create_catalog: create_catalog_provider,
285        create_catalog_list: create_catalog_provider_list,
286        create_table: construct_table_provider,
287        create_table_factory: construct_table_provider_factory,
288        create_scalar_udf: create_ffi_abs_func,
289        create_nullary_udf: create_ffi_random_func,
290        create_timezone_udf: udf_udaf_udwf::create_timezone_func,
291        create_placement_udf: udf_udaf_udwf::create_placement_func,
292        create_table_function: create_ffi_table_func,
293        create_sum_udaf: create_ffi_sum_func,
294        create_stddev_udaf: create_ffi_stddev_func,
295        create_rank_udwf: create_ffi_rank_func,
296        create_extension_options: config::create_extension_options,
297        create_empty_exec,
298        create_exec_with_expressions,
299        create_exec_with_dynamic_expressions,
300        create_exec_with_statistics,
301        create_table_with_statistics,
302        create_physical_optimizer_rule:
303            physical_optimizer::create_physical_optimizer_rule,
304        create_context_aware_optimizer_rule:
305            physical_optimizer::create_context_aware_optimizer_rule,
306        create_query_planner: query_planner::create_query_planner,
307        version: super::version,
308        create_first_value_udaf: create_ffi_first_value_func,
309    }
310}
datafusion_ffi/tests/mod.rs

datafusion_ffi/tests/
mod.rs
