]>
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 | # --------------------------------------------------------------------- | |
19 | # Implement Feather file format | |
20 | ||
21 | # cython: profile=False | |
22 | # distutils: language = c++ | |
23 | # cython: language_level=3 | |
24 | ||
25 | from cython.operator cimport dereference as deref | |
26 | from pyarrow.includes.common cimport * | |
27 | from pyarrow.includes.libarrow cimport * | |
28 | from pyarrow.includes.libarrow_feather cimport * | |
29 | from pyarrow.lib cimport (check_status, Table, _Weakrefable, | |
30 | get_writer, get_reader, pyarrow_wrap_table) | |
31 | from pyarrow.lib import tobytes | |
32 | ||
33 | ||
34 | class FeatherError(Exception): | |
35 | pass | |
36 | ||
37 | ||
38 | def write_feather(Table table, object dest, compression=None, | |
39 | compression_level=None, chunksize=None, version=2): | |
40 | cdef shared_ptr[COutputStream] sink | |
41 | get_writer(dest, &sink) | |
42 | ||
43 | cdef CFeatherProperties properties | |
44 | if version == 2: | |
45 | properties.version = kFeatherV2Version | |
46 | else: | |
47 | properties.version = kFeatherV1Version | |
48 | ||
49 | if compression == 'zstd': | |
50 | properties.compression = CCompressionType_ZSTD | |
51 | elif compression == 'lz4': | |
52 | properties.compression = CCompressionType_LZ4_FRAME | |
53 | else: | |
54 | properties.compression = CCompressionType_UNCOMPRESSED | |
55 | ||
56 | if chunksize is not None: | |
57 | properties.chunksize = chunksize | |
58 | ||
59 | if compression_level is not None: | |
60 | properties.compression_level = compression_level | |
61 | ||
62 | with nogil: | |
63 | check_status(WriteFeather(deref(table.table), sink.get(), | |
64 | properties)) | |
65 | ||
66 | ||
67 | cdef class FeatherReader(_Weakrefable): | |
68 | cdef: | |
69 | shared_ptr[CFeatherReader] reader | |
70 | ||
71 | def __cinit__(self, source, c_bool use_memory_map): | |
72 | cdef shared_ptr[CRandomAccessFile] reader | |
73 | get_reader(source, use_memory_map, &reader) | |
74 | with nogil: | |
75 | self.reader = GetResultValue(CFeatherReader.Open(reader)) | |
76 | ||
77 | @property | |
78 | def version(self): | |
79 | return self.reader.get().version() | |
80 | ||
81 | def read(self): | |
82 | cdef shared_ptr[CTable] sp_table | |
83 | with nogil: | |
84 | check_status(self.reader.get() | |
85 | .Read(&sp_table)) | |
86 | ||
87 | return pyarrow_wrap_table(sp_table) | |
88 | ||
89 | def read_indices(self, indices): | |
90 | cdef: | |
91 | shared_ptr[CTable] sp_table | |
92 | vector[int] c_indices | |
93 | ||
94 | for index in indices: | |
95 | c_indices.push_back(index) | |
96 | with nogil: | |
97 | check_status(self.reader.get() | |
98 | .Read(c_indices, &sp_table)) | |
99 | ||
100 | return pyarrow_wrap_table(sp_table) | |
101 | ||
102 | def read_names(self, names): | |
103 | cdef: | |
104 | shared_ptr[CTable] sp_table | |
105 | vector[c_string] c_names | |
106 | ||
107 | for name in names: | |
108 | c_names.push_back(tobytes(name)) | |
109 | with nogil: | |
110 | check_status(self.reader.get() | |
111 | .Read(c_names, &sp_table)) | |
112 | ||
113 | return pyarrow_wrap_table(sp_table) |