forked from apache/arrow
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathfeather.pxi
More file actions
140 lines (106 loc) · 3.97 KB
/
feather.pxi
File metadata and controls
140 lines (106 loc) · 3.97 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# ---------------------------------------------------------------------
# Implement legacy Feather file format
class FeatherError(Exception):
pass
cdef class FeatherWriter:
cdef:
unique_ptr[CFeatherWriter] writer
cdef public:
int64_t num_rows
def __cinit__(self):
self.num_rows = -1
def open(self, object dest):
cdef shared_ptr[COutputStream] sink
get_writer(dest, &sink)
with nogil:
check_status(CFeatherWriter.Open(sink, &self.writer))
def close(self):
if self.num_rows < 0:
self.num_rows = 0
self.writer.get().SetNumRows(self.num_rows)
with nogil:
check_status(self.writer.get().Finalize())
def write_array(self, object name, object col, object mask=None):
cdef Array arr
if self.num_rows >= 0:
if len(col) != self.num_rows:
raise ValueError('prior column had a different number of rows')
else:
self.num_rows = len(col)
if isinstance(col, Array):
arr = col
else:
arr = Array.from_pandas(col, mask=mask)
cdef c_string c_name = tobytes(name)
with nogil:
check_status(
self.writer.get().Append(c_name, deref(arr.sp_array)))
cdef class FeatherReader:
cdef:
unique_ptr[CFeatherReader] reader
def __cinit__(self):
pass
def open(self, source, c_bool use_memory_map=True):
cdef shared_ptr[CRandomAccessFile] reader
get_reader(source, use_memory_map, &reader)
with nogil:
check_status(CFeatherReader.Open(reader, &self.reader))
@property
def num_rows(self):
return self.reader.get().num_rows()
@property
def num_columns(self):
return self.reader.get().num_columns()
def get_column_name(self, int i):
cdef c_string name = self.reader.get().GetColumnName(i)
return frombytes(name)
def get_column(self, int i):
if i < 0 or i >= self.num_columns:
raise IndexError(i)
cdef shared_ptr[CChunkedArray] sp_chunked_array
with nogil:
check_status(self.reader.get()
.GetColumn(i, &sp_chunked_array))
return pyarrow_wrap_chunked_array(sp_chunked_array)
def _read(self):
cdef shared_ptr[CTable] sp_table
with nogil:
check_status(self.reader.get()
.Read(&sp_table))
return pyarrow_wrap_table(sp_table)
def _read_indices(self, indices):
cdef:
shared_ptr[CTable] sp_table
vector[int] c_indices
for index in indices:
c_indices.push_back(index)
with nogil:
check_status(self.reader.get()
.Read(c_indices, &sp_table))
return pyarrow_wrap_table(sp_table)
def _read_names(self, names):
cdef:
shared_ptr[CTable] sp_table
vector[c_string] c_names
for name in names:
c_names.push_back(tobytes(name))
with nogil:
check_status(self.reader.get()
.Read(c_names, &sp_table))
return pyarrow_wrap_table(sp_table)