]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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 | #pragma once | |
19 | ||
20 | #include <cstdint> | |
21 | #include <memory> | |
22 | #include <vector> | |
23 | ||
24 | #include "arrow/python/serialize.h" | |
25 | #include "arrow/python/visibility.h" | |
26 | #include "arrow/status.h" | |
27 | ||
28 | namespace arrow { | |
29 | ||
30 | class RecordBatch; | |
31 | class Tensor; | |
32 | ||
33 | namespace io { | |
34 | ||
35 | class RandomAccessFile; | |
36 | ||
37 | } // namespace io | |
38 | ||
39 | namespace py { | |
40 | ||
41 | struct ARROW_PYTHON_EXPORT SparseTensorCounts { | |
42 | int coo; | |
43 | int csr; | |
44 | int csc; | |
45 | int csf; | |
46 | int ndim_csf; | |
47 | ||
48 | int num_total_tensors() const { return coo + csr + csc + csf; } | |
49 | int num_total_buffers() const { | |
50 | return coo * 3 + csr * 4 + csc * 4 + 2 * ndim_csf + csf; | |
51 | } | |
52 | }; | |
53 | ||
54 | /// \brief Read serialized Python sequence from file interface using Arrow IPC | |
55 | /// \param[in] src a RandomAccessFile | |
56 | /// \param[out] out the reconstructed data | |
57 | /// \return Status | |
58 | ARROW_PYTHON_EXPORT | |
59 | Status ReadSerializedObject(io::RandomAccessFile* src, SerializedPyObject* out); | |
60 | ||
61 | /// \brief Reconstruct SerializedPyObject from representation produced by | |
62 | /// SerializedPyObject::GetComponents. | |
63 | /// | |
64 | /// \param[in] num_tensors number of tensors in the object | |
65 | /// \param[in] num_sparse_tensors number of sparse tensors in the object | |
66 | /// \param[in] num_ndarrays number of numpy Ndarrays in the object | |
67 | /// \param[in] num_buffers number of buffers in the object | |
68 | /// \param[in] data a list containing pyarrow.Buffer instances. It must be 1 + | |
69 | /// num_tensors * 2 + num_coo_tensors * 3 + num_csr_tensors * 4 + num_csc_tensors * 4 + | |
70 | /// num_csf_tensors * (2 * ndim_csf + 3) + num_buffers in length | |
71 | /// \param[out] out the reconstructed object | |
72 | /// \return Status | |
73 | ARROW_PYTHON_EXPORT | |
74 | Status GetSerializedFromComponents(int num_tensors, | |
75 | const SparseTensorCounts& num_sparse_tensors, | |
76 | int num_ndarrays, int num_buffers, PyObject* data, | |
77 | SerializedPyObject* out); | |
78 | ||
79 | /// \brief Reconstruct Python object from Arrow-serialized representation | |
80 | /// \param[in] context Serialization context which contains custom serialization | |
81 | /// and deserialization callbacks. Can be any Python object with a | |
82 | /// _serialize_callback method for serialization and a _deserialize_callback | |
83 | /// method for deserialization. If context is None, no custom serialization | |
84 | /// will be attempted. | |
85 | /// \param[in] object Object to deserialize | |
86 | /// \param[in] base a Python object holding the underlying data that any NumPy | |
87 | /// arrays will reference, to avoid premature deallocation | |
88 | /// \param[out] out The returned object | |
89 | /// \return Status | |
90 | /// This acquires the GIL | |
91 | ARROW_PYTHON_EXPORT | |
92 | Status DeserializeObject(PyObject* context, const SerializedPyObject& object, | |
93 | PyObject* base, PyObject** out); | |
94 | ||
95 | /// \brief Reconstruct Ndarray from Arrow-serialized representation | |
96 | /// \param[in] object Object to deserialize | |
97 | /// \param[out] out The deserialized tensor | |
98 | /// \return Status | |
99 | ARROW_PYTHON_EXPORT | |
100 | Status DeserializeNdarray(const SerializedPyObject& object, std::shared_ptr<Tensor>* out); | |
101 | ||
102 | ARROW_PYTHON_EXPORT | |
103 | Status NdarrayFromBuffer(std::shared_ptr<Buffer> src, std::shared_ptr<Tensor>* out); | |
104 | ||
105 | } // namespace py | |
106 | } // namespace arrow |