1717
1818use crate :: dataset:: Dataset ;
1919use crate :: errors:: { py_datafusion_err, to_datafusion_err, PyDataFusionError , PyDataFusionResult } ;
20- use crate :: table:: PyTableProvider ;
2120use crate :: utils:: { validate_pycapsule, wait_for_future} ;
2221use async_trait:: async_trait;
2322use datafusion:: catalog:: { MemoryCatalogProvider , MemorySchemaProvider } ;
@@ -52,7 +51,7 @@ pub struct PySchema {
5251#[ pyclass( name = "RawTable" , module = "datafusion.catalog" , subclass) ]
5352#[ derive( Clone ) ]
5453pub struct PyTable {
55- pub table : Arc < dyn TableProvider + Send > ,
54+ pub table : Arc < dyn TableProvider > ,
5655}
5756
5857impl From < Arc < dyn CatalogProvider > > for PyCatalog {
@@ -68,11 +67,11 @@ impl From<Arc<dyn SchemaProvider>> for PySchema {
6867}
6968
7069impl PyTable {
71- pub fn new ( table : Arc < dyn TableProvider + Send > ) -> Self {
70+ pub fn new ( table : Arc < dyn TableProvider > ) -> Self {
7271 Self { table }
7372 }
7473
75- pub fn table ( & self ) -> Arc < dyn TableProvider + Send > {
74+ pub fn table ( & self ) -> Arc < dyn TableProvider > {
7675 self . table . clone ( )
7776 }
7877}
@@ -206,18 +205,15 @@ impl PySchema {
206205
207206 let provider = unsafe { capsule. reference :: < FFI_TableProvider > ( ) } ;
208207 let provider: ForeignTableProvider = provider. into ( ) ;
209- Arc :: new ( provider) as Arc < dyn TableProvider + Send >
208+ Arc :: new ( provider) as Arc < dyn TableProvider >
210209 } else {
211210 match table_provider. extract :: < PyTable > ( ) {
212211 Ok ( py_table) => py_table. table ,
213- Err ( _) => match table_provider. extract :: < PyTableProvider > ( ) {
214- Ok ( py_provider) => py_provider. into_inner ( ) ,
215- Err ( _) => {
216- let py = table_provider. py ( ) ;
217- let provider = Dataset :: new ( & table_provider, py) ?;
218- Arc :: new ( provider) as Arc < dyn TableProvider + Send >
219- }
220- } ,
212+ Err ( _) => {
213+ let py = table_provider. py ( ) ;
214+ let provider = Dataset :: new ( & table_provider, py) ?;
215+ Arc :: new ( provider) as Arc < dyn TableProvider >
216+ }
221217 }
222218 } ;
223219
@@ -298,7 +294,7 @@ impl RustWrappedPySchemaProvider {
298294 }
299295 }
300296
301- fn table_inner ( & self , name : & str ) -> PyResult < Option < Arc < dyn TableProvider + Send > > > {
297+ fn table_inner ( & self , name : & str ) -> PyResult < Option < Arc < dyn TableProvider > > > {
302298 Python :: with_gil ( |py| {
303299 let provider = self . schema_provider . bind ( py) ;
304300 let py_table_method = provider. getattr ( "table" ) ?;
@@ -309,30 +305,26 @@ impl RustWrappedPySchemaProvider {
309305 }
310306
311307 if py_table. hasattr ( "__datafusion_table_provider__" ) ? {
312- let capsule = py_table . getattr ( "__datafusion_table_provider__" ) ?. call0 ( ) ?;
308+ let capsule = provider . getattr ( "__datafusion_table_provider__" ) ?. call0 ( ) ?;
313309 let capsule = capsule. downcast :: < PyCapsule > ( ) . map_err ( py_datafusion_err) ?;
314310 validate_pycapsule ( capsule, "datafusion_table_provider" ) ?;
315311
316312 let provider = unsafe { capsule. reference :: < FFI_TableProvider > ( ) } ;
317313 let provider: ForeignTableProvider = provider. into ( ) ;
318314
319- Ok ( Some ( Arc :: new ( provider) as Arc < dyn TableProvider + Send > ) )
315+ Ok ( Some ( Arc :: new ( provider) as Arc < dyn TableProvider > ) )
320316 } else {
321317 if let Ok ( inner_table) = py_table. getattr ( "table" ) {
322318 if let Ok ( inner_table) = inner_table. extract :: < PyTable > ( ) {
323319 return Ok ( Some ( inner_table. table ) ) ;
324320 }
325321 }
326322
327- if let Ok ( py_provider) = py_table. extract :: < PyTableProvider > ( ) {
328- return Ok ( Some ( py_provider. into_inner ( ) ) ) ;
329- }
330-
331323 match py_table. extract :: < PyTable > ( ) {
332324 Ok ( py_table) => Ok ( Some ( py_table. table ) ) ,
333325 Err ( _) => {
334326 let ds = Dataset :: new ( & py_table, py) . map_err ( py_datafusion_err) ?;
335- Ok ( Some ( Arc :: new ( ds) as Arc < dyn TableProvider + Send > ) )
327+ Ok ( Some ( Arc :: new ( ds) as Arc < dyn TableProvider > ) )
336328 }
337329 }
338330 }
@@ -368,32 +360,15 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
368360 & self ,
369361 name : & str ,
370362 ) -> datafusion:: common:: Result < Option < Arc < dyn TableProvider > > , DataFusionError > {
371- // Convert from our internal Send type to the trait expected type
372- match self . table_inner ( name) . map_err ( to_datafusion_err) ? {
373- Some ( table) => {
374- // Safe conversion: we're widening the bounds (removing Send)
375- let raw = Arc :: into_raw ( table) ;
376- let wide: * const dyn TableProvider = raw as * const _ ;
377- let arc = unsafe { Arc :: from_raw ( wide) } ;
378- Ok ( Some ( arc) )
379- }
380- None => Ok ( None ) ,
381- }
363+ self . table_inner ( name) . map_err ( to_datafusion_err)
382364 }
383365
384366 fn register_table (
385367 & self ,
386368 name : String ,
387369 table : Arc < dyn TableProvider > ,
388370 ) -> datafusion:: common:: Result < Option < Arc < dyn TableProvider > > > {
389- // Convert from trait type to our internal Send type
390- let send_table = {
391- let raw = Arc :: into_raw ( table) ;
392- let send: * const ( dyn TableProvider + Send ) = raw as * const _ ;
393- unsafe { Arc :: from_raw ( send) }
394- } ;
395-
396- let py_table = PyTable :: new ( send_table) ;
371+ let py_table = PyTable :: new ( table) ;
397372 Python :: with_gil ( |py| {
398373 let provider = self . schema_provider . bind ( py) ;
399374 let _ = provider
@@ -422,14 +397,7 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
422397 // If we can turn this table provider into a `Dataset`, return it.
423398 // Otherwise, return None.
424399 let dataset = match Dataset :: new ( & table, py) {
425- Ok ( dataset) => {
426- // Convert from our internal Send type to trait expected type
427- let send_table = Arc :: new ( dataset) as Arc < dyn TableProvider + Send > ;
428- let raw = Arc :: into_raw ( send_table) ;
429- let wide: * const dyn TableProvider = raw as * const _ ;
430- let arc = unsafe { Arc :: from_raw ( wide) } ;
431- Some ( arc)
432- }
400+ Ok ( dataset) => Some ( Arc :: new ( dataset) as Arc < dyn TableProvider > ) ,
433401 Err ( _) => None ,
434402 } ;
435403
0 commit comments