@@ -28,9 +28,13 @@ use arrow::compute::kernels;
2828use arrow:: datatypes:: { DataType , Field , Schema } ;
2929use arrow:: error:: ArrowError ;
3030use arrow:: ffi_stream:: ArrowArrayStreamReader ;
31- use arrow:: pyarrow:: PyArrowConvert ;
31+ use arrow:: pyarrow:: { PyArrowConvert , PyArrowException , PyArrowType } ;
3232use arrow:: record_batch:: RecordBatch ;
3333
34+ fn to_py_err ( err : ArrowError ) -> PyErr {
35+ PyArrowException :: new_err ( err. to_string ( ) )
36+ }
37+
3438/// Returns `array + array` of an int64 array.
3539#[ pyfunction]
3640fn double ( array : & PyAny , py : Python ) -> PyResult < PyObject > {
@@ -41,8 +45,10 @@ fn double(array: &PyAny, py: Python) -> PyResult<PyObject> {
4145 let array = array
4246 . as_any ( )
4347 . downcast_ref :: < Int64Array > ( )
44- . ok_or ( ArrowError :: ParseError ( "Expects an int64" . to_string ( ) ) ) ?;
45- let array = kernels:: arithmetic:: add ( array, array) ?;
48+ . ok_or_else ( || ArrowError :: ParseError ( "Expects an int64" . to_string ( ) ) )
49+ . map_err ( to_py_err) ?;
50+
51+ let array = kernels:: arithmetic:: add ( array, array) . map_err ( to_py_err) ?;
4652
4753 // export
4854 array. to_pyarrow ( py)
@@ -66,56 +72,61 @@ fn double_py(lambda: &PyAny, py: Python) -> PyResult<bool> {
6672
6773/// Returns the substring
6874#[ pyfunction]
69- fn substring ( array : ArrayData , start : i64 ) -> PyResult < ArrayData > {
75+ fn substring (
76+ array : PyArrowType < ArrayData > ,
77+ start : i64 ,
78+ ) -> PyResult < PyArrowType < ArrayData > > {
7079 // import
71- let array = ArrayRef :: from ( array) ;
80+ let array = ArrayRef :: from ( array. 0 ) ;
7281
7382 // substring
74- let array = kernels:: substring:: substring ( array. as_ref ( ) , start, None ) ?;
83+ let array = kernels:: substring:: substring ( array. as_ref ( ) , start, None ) . map_err ( to_py_err ) ?;
7584
76- Ok ( array. data ( ) . to_owned ( ) )
85+ Ok ( array. data ( ) . to_owned ( ) . into ( ) )
7786}
7887
7988/// Returns the concatenate
8089#[ pyfunction]
81- fn concatenate ( array : ArrayData , py : Python ) -> PyResult < PyObject > {
82- let array = ArrayRef :: from ( array) ;
90+ fn concatenate ( array : PyArrowType < ArrayData > , py : Python ) -> PyResult < PyObject > {
91+ let array = ArrayRef :: from ( array. 0 ) ;
8392
8493 // concat
85- let array = kernels:: concat:: concat ( & [ array. as_ref ( ) , array. as_ref ( ) ] ) ?;
94+ let array = kernels:: concat:: concat ( & [ array. as_ref ( ) , array. as_ref ( ) ] ) . map_err ( to_py_err ) ?;
8695
8796 array. to_pyarrow ( py)
8897}
8998
9099#[ pyfunction]
91- fn round_trip_type ( obj : DataType ) -> PyResult < DataType > {
100+ fn round_trip_type ( obj : PyArrowType < DataType > ) -> PyResult < PyArrowType < DataType > > {
92101 Ok ( obj)
93102}
94103
95104#[ pyfunction]
96- fn round_trip_field ( obj : Field ) -> PyResult < Field > {
105+ fn round_trip_field ( obj : PyArrowType < Field > ) -> PyResult < PyArrowType < Field > > {
97106 Ok ( obj)
98107}
99108
100109#[ pyfunction]
101- fn round_trip_schema ( obj : Schema ) -> PyResult < Schema > {
110+ fn round_trip_schema ( obj : PyArrowType < Schema > ) -> PyResult < PyArrowType < Schema > > {
102111 Ok ( obj)
103112}
104113
105114#[ pyfunction]
106- fn round_trip_array ( obj : ArrayData ) -> PyResult < ArrayData > {
115+ fn round_trip_array ( obj : PyArrowType < ArrayData > ) -> PyResult < PyArrowType < ArrayData > > {
107116 Ok ( obj)
108117}
109118
110119#[ pyfunction]
111- fn round_trip_record_batch ( obj : RecordBatch ) -> PyResult < RecordBatch > {
120+ fn round_trip_record_batch (
121+ obj : PyArrowType < RecordBatch > ,
122+ ) -> PyResult < PyArrowType < RecordBatch > > {
112123 Ok ( obj)
113124}
114125
115126#[ pyfunction]
116127fn round_trip_record_batch_reader (
117- obj : ArrowArrayStreamReader ,
118- ) -> PyResult < ArrowArrayStreamReader > {
128+ obj : PyArrowType < ArrowArrayStreamReader > ,
129+ ) -> PyResult < PyArrowType < ArrowArrayStreamReader > > {
119130 Ok ( obj)
120131}
121132
0 commit comments