ceph/src/arrow/python/pyarrow/public-api.pxi

   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
  18 from libcpp.memory cimport shared_ptr
  19 from pyarrow.includes.libarrow cimport (CArray, CDataType, CField,
  20                                         CRecordBatch, CSchema,
  21                                         CTable, CTensor,
  22                                         CSparseCOOTensor, CSparseCSRMatrix,
  23                                         CSparseCSCMatrix, CSparseCSFTensor)
  24
  25 # You cannot assign something to a dereferenced pointer in Cython thus these
  26 # methods don't use Status to indicate a successful operation.
  27
  28
  29 cdef api bint pyarrow_is_buffer(object buffer):
  30     return isinstance(buffer, Buffer)
  31
  32
  33 cdef api shared_ptr[CBuffer] pyarrow_unwrap_buffer(object buffer):
  34     cdef Buffer buf
  35     if pyarrow_is_buffer(buffer):
  36         buf = <Buffer>(buffer)
  37         return buf.buffer
  38
  39     return shared_ptr[CBuffer]()
  40
  41
  42 cdef api object pyarrow_wrap_buffer(const shared_ptr[CBuffer]& buf):
  43     cdef Buffer result = Buffer.__new__(Buffer)
  44     result.init(buf)
  45     return result
  46
  47
  48 cdef api object pyarrow_wrap_resizable_buffer(
  49         const shared_ptr[CResizableBuffer]& buf):
  50     cdef ResizableBuffer result = ResizableBuffer.__new__(ResizableBuffer)
  51     result.init_rz(buf)
  52     return result
  53
  54
  55 cdef api bint pyarrow_is_data_type(object type_):
  56     return isinstance(type_, DataType)
  57
  58
  59 cdef api shared_ptr[CDataType] pyarrow_unwrap_data_type(
  60         object data_type):
  61     cdef DataType type_
  62     if pyarrow_is_data_type(data_type):
  63         type_ = <DataType>(data_type)
  64         return type_.sp_type
  65
  66     return shared_ptr[CDataType]()
  67
  68
  69 # Workaround for Cython parsing bug
  70 # https://github.com/cython/cython/issues/2143
  71 ctypedef const CPyExtensionType* _CPyExtensionTypePtr
  72
  73
  74 cdef api object pyarrow_wrap_data_type(
  75         const shared_ptr[CDataType]& type):
  76     cdef:
  77         const CExtensionType* ext_type
  78         const CPyExtensionType* cpy_ext_type
  79         DataType out
  80
  81     if type.get() == NULL:
  82         return None
  83
  84     if type.get().id() == _Type_DICTIONARY:
  85         out = DictionaryType.__new__(DictionaryType)
  86     elif type.get().id() == _Type_LIST:
  87         out = ListType.__new__(ListType)
  88     elif type.get().id() == _Type_LARGE_LIST:
  89         out = LargeListType.__new__(LargeListType)
  90     elif type.get().id() == _Type_MAP:
  91         out = MapType.__new__(MapType)
  92     elif type.get().id() == _Type_FIXED_SIZE_LIST:
  93         out = FixedSizeListType.__new__(FixedSizeListType)
  94     elif type.get().id() == _Type_STRUCT:
  95         out = StructType.__new__(StructType)
  96     elif type.get().id() == _Type_SPARSE_UNION:
  97         out = SparseUnionType.__new__(SparseUnionType)
  98     elif type.get().id() == _Type_DENSE_UNION:
  99         out = DenseUnionType.__new__(DenseUnionType)
 100     elif type.get().id() == _Type_TIMESTAMP:
 101         out = TimestampType.__new__(TimestampType)
 102     elif type.get().id() == _Type_DURATION:
 103         out = DurationType.__new__(DurationType)
 104     elif type.get().id() == _Type_FIXED_SIZE_BINARY:
 105         out = FixedSizeBinaryType.__new__(FixedSizeBinaryType)
 106     elif type.get().id() == _Type_DECIMAL128:
 107         out = Decimal128Type.__new__(Decimal128Type)
 108     elif type.get().id() == _Type_DECIMAL256:
 109         out = Decimal256Type.__new__(Decimal256Type)
 110     elif type.get().id() == _Type_EXTENSION:
 111         ext_type = <const CExtensionType*> type.get()
 112         cpy_ext_type = dynamic_cast[_CPyExtensionTypePtr](ext_type)
 113         if cpy_ext_type != nullptr:
 114             return cpy_ext_type.GetInstance()
 115         else:
 116             out = BaseExtensionType.__new__(BaseExtensionType)
 117     else:
 118         out = DataType.__new__(DataType)
 119
 120     out.init(type)
 121     return out
 122
 123
 124 cdef object pyarrow_wrap_metadata(
 125         const shared_ptr[const CKeyValueMetadata]& meta):
 126     if meta.get() == nullptr:
 127         return None
 128     else:
 129         return KeyValueMetadata.wrap(meta)
 130
 131
 132 cdef api bint pyarrow_is_metadata(object metadata):
 133     return isinstance(metadata, KeyValueMetadata)
 134
 135
 136 cdef shared_ptr[const CKeyValueMetadata] pyarrow_unwrap_metadata(object meta):
 137     cdef shared_ptr[const CKeyValueMetadata] c_meta
 138     if pyarrow_is_metadata(meta):
 139         c_meta = (<KeyValueMetadata>meta).unwrap()
 140     return c_meta
 141
 142
 143 cdef api bint pyarrow_is_field(object field):
 144     return isinstance(field, Field)
 145
 146
 147 cdef api shared_ptr[CField] pyarrow_unwrap_field(object field):
 148     cdef Field field_
 149     if pyarrow_is_field(field):
 150         field_ = <Field>(field)
 151         return field_.sp_field
 152
 153     return shared_ptr[CField]()
 154
 155
 156 cdef api object pyarrow_wrap_field(const shared_ptr[CField]& field):
 157     if field.get() == NULL:
 158         return None
 159     cdef Field out = Field.__new__(Field)
 160     out.init(field)
 161     return out
 162
 163
 164 cdef api bint pyarrow_is_schema(object schema):
 165     return isinstance(schema, Schema)
 166
 167
 168 cdef api shared_ptr[CSchema] pyarrow_unwrap_schema(object schema):
 169     cdef Schema sch
 170     if pyarrow_is_schema(schema):
 171         sch = <Schema>(schema)
 172         return sch.sp_schema
 173
 174     return shared_ptr[CSchema]()
 175
 176
 177 cdef api object pyarrow_wrap_schema(const shared_ptr[CSchema]& schema):
 178     cdef Schema out = Schema.__new__(Schema)
 179     out.init_schema(schema)
 180     return out
 181
 182
 183 cdef api bint pyarrow_is_array(object array):
 184     return isinstance(array, Array)
 185
 186
 187 cdef api shared_ptr[CArray] pyarrow_unwrap_array(object array):
 188     cdef Array arr
 189     if pyarrow_is_array(array):
 190         arr = <Array>(array)
 191         return arr.sp_array
 192
 193     return shared_ptr[CArray]()
 194
 195
 196 cdef api object pyarrow_wrap_array(const shared_ptr[CArray]& sp_array):
 197     if sp_array.get() == NULL:
 198         raise ValueError('Array was NULL')
 199
 200     klass = get_array_class_from_type(sp_array.get().type())
 201
 202     cdef Array arr = klass.__new__(klass)
 203     arr.init(sp_array)
 204     return arr
 205
 206
 207 cdef api bint pyarrow_is_chunked_array(object array):
 208     return isinstance(array, ChunkedArray)
 209
 210
 211 cdef api shared_ptr[CChunkedArray] pyarrow_unwrap_chunked_array(object array):
 212     cdef ChunkedArray arr
 213     if pyarrow_is_chunked_array(array):
 214         arr = <ChunkedArray>(array)
 215         return arr.sp_chunked_array
 216
 217     return shared_ptr[CChunkedArray]()
 218
 219
 220 cdef api object pyarrow_wrap_chunked_array(
 221         const shared_ptr[CChunkedArray]& sp_array):
 222     if sp_array.get() == NULL:
 223         raise ValueError('ChunkedArray was NULL')
 224
 225     cdef CDataType* data_type = sp_array.get().type().get()
 226
 227     if data_type == NULL:
 228         raise ValueError('ChunkedArray data type was NULL')
 229
 230     cdef ChunkedArray arr = ChunkedArray.__new__(ChunkedArray)
 231     arr.init(sp_array)
 232     return arr
 233
 234
 235 cdef api bint pyarrow_is_scalar(object value):
 236     return isinstance(value, Scalar)
 237
 238
 239 cdef api shared_ptr[CScalar] pyarrow_unwrap_scalar(object scalar):
 240     if pyarrow_is_scalar(scalar):
 241         return (<Scalar> scalar).unwrap()
 242     return shared_ptr[CScalar]()
 243
 244
 245 cdef api object pyarrow_wrap_scalar(const shared_ptr[CScalar]& sp_scalar):
 246     if sp_scalar.get() == NULL:
 247         raise ValueError('Scalar was NULL')
 248
 249     cdef CDataType* data_type = sp_scalar.get().type.get()
 250
 251     if data_type == NULL:
 252         raise ValueError('Scalar data type was NULL')
 253
 254     if data_type.id() == _Type_NA:
 255         return _NULL
 256
 257     if data_type.id() not in _scalar_classes:
 258         raise ValueError('Scalar type not supported')
 259
 260     klass = _scalar_classes[data_type.id()]
 261
 262     cdef Scalar scalar = klass.__new__(klass)
 263     scalar.init(sp_scalar)
 264     return scalar
 265
 266
 267 cdef api bint pyarrow_is_tensor(object tensor):
 268     return isinstance(tensor, Tensor)
 269
 270
 271 cdef api shared_ptr[CTensor] pyarrow_unwrap_tensor(object tensor):
 272     cdef Tensor ten
 273     if pyarrow_is_tensor(tensor):
 274         ten = <Tensor>(tensor)
 275         return ten.sp_tensor
 276
 277     return shared_ptr[CTensor]()
 278
 279
 280 cdef api object pyarrow_wrap_tensor(
 281         const shared_ptr[CTensor]& sp_tensor):
 282     if sp_tensor.get() == NULL:
 283         raise ValueError('Tensor was NULL')
 284
 285     cdef Tensor tensor = Tensor.__new__(Tensor)
 286     tensor.init(sp_tensor)
 287     return tensor
 288
 289
 290 cdef api bint pyarrow_is_sparse_coo_tensor(object sparse_tensor):
 291     return isinstance(sparse_tensor, SparseCOOTensor)
 292
 293 cdef api shared_ptr[CSparseCOOTensor] pyarrow_unwrap_sparse_coo_tensor(
 294         object sparse_tensor):
 295     cdef SparseCOOTensor sten
 296     if pyarrow_is_sparse_coo_tensor(sparse_tensor):
 297         sten = <SparseCOOTensor>(sparse_tensor)
 298         return sten.sp_sparse_tensor
 299
 300     return shared_ptr[CSparseCOOTensor]()
 301
 302 cdef api object pyarrow_wrap_sparse_coo_tensor(
 303         const shared_ptr[CSparseCOOTensor]& sp_sparse_tensor):
 304     if sp_sparse_tensor.get() == NULL:
 305         raise ValueError('SparseCOOTensor was NULL')
 306
 307     cdef SparseCOOTensor sparse_tensor = SparseCOOTensor.__new__(
 308         SparseCOOTensor)
 309     sparse_tensor.init(sp_sparse_tensor)
 310     return sparse_tensor
 311
 312
 313 cdef api bint pyarrow_is_sparse_csr_matrix(object sparse_tensor):
 314     return isinstance(sparse_tensor, SparseCSRMatrix)
 315
 316 cdef api shared_ptr[CSparseCSRMatrix] pyarrow_unwrap_sparse_csr_matrix(
 317         object sparse_tensor):
 318     cdef SparseCSRMatrix sten
 319     if pyarrow_is_sparse_csr_matrix(sparse_tensor):
 320         sten = <SparseCSRMatrix>(sparse_tensor)
 321         return sten.sp_sparse_tensor
 322
 323     return shared_ptr[CSparseCSRMatrix]()
 324
 325 cdef api object pyarrow_wrap_sparse_csr_matrix(
 326         const shared_ptr[CSparseCSRMatrix]& sp_sparse_tensor):
 327     if sp_sparse_tensor.get() == NULL:
 328         raise ValueError('SparseCSRMatrix was NULL')
 329
 330     cdef SparseCSRMatrix sparse_tensor = SparseCSRMatrix.__new__(
 331         SparseCSRMatrix)
 332     sparse_tensor.init(sp_sparse_tensor)
 333     return sparse_tensor
 334
 335
 336 cdef api bint pyarrow_is_sparse_csc_matrix(object sparse_tensor):
 337     return isinstance(sparse_tensor, SparseCSCMatrix)
 338
 339 cdef api shared_ptr[CSparseCSCMatrix] pyarrow_unwrap_sparse_csc_matrix(
 340         object sparse_tensor):
 341     cdef SparseCSCMatrix sten
 342     if pyarrow_is_sparse_csc_matrix(sparse_tensor):
 343         sten = <SparseCSCMatrix>(sparse_tensor)
 344         return sten.sp_sparse_tensor
 345
 346     return shared_ptr[CSparseCSCMatrix]()
 347
 348 cdef api object pyarrow_wrap_sparse_csc_matrix(
 349         const shared_ptr[CSparseCSCMatrix]& sp_sparse_tensor):
 350     if sp_sparse_tensor.get() == NULL:
 351         raise ValueError('SparseCSCMatrix was NULL')
 352
 353     cdef SparseCSCMatrix sparse_tensor = SparseCSCMatrix.__new__(
 354         SparseCSCMatrix)
 355     sparse_tensor.init(sp_sparse_tensor)
 356     return sparse_tensor
 357
 358
 359 cdef api bint pyarrow_is_sparse_csf_tensor(object sparse_tensor):
 360     return isinstance(sparse_tensor, SparseCSFTensor)
 361
 362 cdef api shared_ptr[CSparseCSFTensor] pyarrow_unwrap_sparse_csf_tensor(
 363         object sparse_tensor):
 364     cdef SparseCSFTensor sten
 365     if pyarrow_is_sparse_csf_tensor(sparse_tensor):
 366         sten = <SparseCSFTensor>(sparse_tensor)
 367         return sten.sp_sparse_tensor
 368
 369     return shared_ptr[CSparseCSFTensor]()
 370
 371 cdef api object pyarrow_wrap_sparse_csf_tensor(
 372         const shared_ptr[CSparseCSFTensor]& sp_sparse_tensor):
 373     if sp_sparse_tensor.get() == NULL:
 374         raise ValueError('SparseCSFTensor was NULL')
 375
 376     cdef SparseCSFTensor sparse_tensor = SparseCSFTensor.__new__(
 377         SparseCSFTensor)
 378     sparse_tensor.init(sp_sparse_tensor)
 379     return sparse_tensor
 380
 381
 382 cdef api bint pyarrow_is_table(object table):
 383     return isinstance(table, Table)
 384
 385
 386 cdef api shared_ptr[CTable] pyarrow_unwrap_table(object table):
 387     cdef Table tab
 388     if pyarrow_is_table(table):
 389         tab = <Table>(table)
 390         return tab.sp_table
 391
 392     return shared_ptr[CTable]()
 393
 394
 395 cdef api object pyarrow_wrap_table(const shared_ptr[CTable]& ctable):
 396     cdef Table table = Table.__new__(Table)
 397     table.init(ctable)
 398     return table
 399
 400
 401 cdef api bint pyarrow_is_batch(object batch):
 402     return isinstance(batch, RecordBatch)
 403
 404
 405 cdef api shared_ptr[CRecordBatch] pyarrow_unwrap_batch(object batch):
 406     cdef RecordBatch bat
 407     if pyarrow_is_batch(batch):
 408         bat = <RecordBatch>(batch)
 409         return bat.sp_batch
 410
 411     return shared_ptr[CRecordBatch]()
 412
 413
 414 cdef api object pyarrow_wrap_batch(
 415         const shared_ptr[CRecordBatch]& cbatch):
 416     cdef RecordBatch batch = RecordBatch.__new__(RecordBatch)
 417     batch.init(cbatch)
 418     return batch