|
41 | 41 | ) |
42 | 42 |
|
43 | 43 | from .nanoarrow_bridge import ( |
44 | | - NANOARROW_TIME_UNIT_SECOND, |
45 | | - NANOARROW_TIME_UNIT_MILLI, |
46 | | - NANOARROW_TIME_UNIT_MICRO, |
47 | | - NANOARROW_TIME_UNIT_NANO, |
48 | | - NANOARROW_TYPE_BINARY, |
49 | | - NANOARROW_TYPE_DOUBLE, |
50 | | - NANOARROW_TYPE_FLOAT, |
51 | | - NANOARROW_TYPE_INT64, |
52 | | - NANOARROW_TYPE_LARGE_BINARY, |
53 | | - NANOARROW_TYPE_LARGE_STRING, |
54 | | - NANOARROW_TYPE_STRING, |
55 | | - NANOARROW_TYPE_TIMESTAMP, |
56 | | - NANOARROW_TYPE_DECIMAL128, |
| 44 | + ArrowTimeUnit, |
| 45 | + ArrowType, |
57 | 46 | ) |
58 | 47 |
|
59 | 48 |
|
@@ -92,8 +81,8 @@ def _offsets_buffer(self): |
92 | 81 | size_in_bytes=size_bytes, address=address, buffer_type="offsets" |
93 | 82 | ) |
94 | 83 | if self.ora_arrow_array.arrow_type in ( |
95 | | - NANOARROW_TYPE_LARGE_STRING, |
96 | | - NANOARROW_TYPE_LARGE_BINARY, |
| 84 | + ArrowType.NANOARROW_TYPE_LARGE_STRING, |
| 85 | + ArrowType.NANOARROW_TYPE_LARGE_BINARY, |
97 | 86 | ): |
98 | 87 | dtype = (DtypeKind.INT, 64, "l", "=") |
99 | 88 | else: |
@@ -133,36 +122,38 @@ def dtype(self) -> Dtype: |
133 | 122 | Returns the data type of the column. The returned dtype provides |
134 | 123 | information on the storage format and the type of data in the column. |
135 | 124 | """ |
136 | | - if self.ora_arrow_array.arrow_type == NANOARROW_TYPE_INT64: |
| 125 | + arrow_type = self.ora_arrow_array.arrow_type |
| 126 | + if arrow_type == ArrowType.NANOARROW_TYPE_INT64: |
137 | 127 | return (DtypeKind.INT, 64, "l", "=") |
138 | | - elif self.ora_arrow_array.arrow_type == NANOARROW_TYPE_DOUBLE: |
| 128 | + elif arrow_type == ArrowType.NANOARROW_TYPE_DOUBLE: |
139 | 129 | return (DtypeKind.FLOAT, 64, "g", "=") |
140 | | - elif self.ora_arrow_array.arrow_type == NANOARROW_TYPE_FLOAT: |
| 130 | + elif arrow_type == ArrowType.NANOARROW_TYPE_FLOAT: |
141 | 131 | return (DtypeKind.FLOAT, 64, "g", "=") |
142 | | - elif self.ora_arrow_array.arrow_type == NANOARROW_TYPE_STRING: |
| 132 | + elif arrow_type == ArrowType.NANOARROW_TYPE_STRING: |
143 | 133 | return (DtypeKind.STRING, 8, "u", "=") |
144 | | - elif self.ora_arrow_array.arrow_type == NANOARROW_TYPE_TIMESTAMP: |
145 | | - if self.ora_arrow_array.time_unit == NANOARROW_TIME_UNIT_MICRO: |
| 134 | + elif arrow_type == ArrowType.NANOARROW_TYPE_TIMESTAMP: |
| 135 | + time_unit = self.ora_arrow_array.time_unit |
| 136 | + if time_unit == ArrowTimeUnit.NANOARROW_TIME_UNIT_MICRO: |
146 | 137 | return (DtypeKind.DATETIME, 64, "tsu:", "=") |
147 | | - elif self.ora_arrow_array.time_unit == NANOARROW_TIME_UNIT_SECOND: |
| 138 | + elif time_unit == ArrowTimeUnit.NANOARROW_TIME_UNIT_SECOND: |
148 | 139 | return (DtypeKind.DATETIME, 64, "tss:", "=") |
149 | | - elif self.ora_arrow_array.time_unit == NANOARROW_TIME_UNIT_MILLI: |
| 140 | + elif time_unit == ArrowTimeUnit.NANOARROW_TIME_UNIT_MILLI: |
150 | 141 | return (DtypeKind.DATETIME, 64, "tsm:", "=") |
151 | | - elif self.ora_arrow_array.time_unit == NANOARROW_TIME_UNIT_NANO: |
| 142 | + elif time_unit == ArrowTimeUnit.NANOARROW_TIME_UNIT_NANO: |
152 | 143 | return (DtypeKind.DATETIME, 64, "tsn:", "=") |
153 | | - elif self.ora_arrow_array.arrow_type == NANOARROW_TYPE_DECIMAL128: |
| 144 | + elif arrow_type == ArrowType.NANOARROW_TYPE_DECIMAL128: |
154 | 145 | array = self.ora_arrow_array |
155 | 146 | return ( |
156 | 147 | DtypeKind.DECIMAL, |
157 | 148 | 128, |
158 | 149 | f"d:{array.precision}.{array.scale}", |
159 | 150 | "=", |
160 | 151 | ) |
161 | | - elif self.ora_arrow_array.arrow_type == NANOARROW_TYPE_BINARY: |
| 152 | + elif arrow_type == ArrowType.NANOARROW_TYPE_BINARY: |
162 | 153 | return (DtypeKind.STRING, 8, "z", "=") |
163 | | - elif self.ora_arrow_array.arrow_type == NANOARROW_TYPE_LARGE_BINARY: |
| 154 | + elif arrow_type == ArrowType.NANOARROW_TYPE_LARGE_BINARY: |
164 | 155 | return (DtypeKind.STRING, 8, "Z", "=") |
165 | | - elif self.ora_arrow_array.arrow_type == NANOARROW_TYPE_LARGE_STRING: |
| 156 | + elif arrow_type == ArrowType.NANOARROW_TYPE_LARGE_STRING: |
166 | 157 | return (DtypeKind.STRING, 8, "U", "=") |
167 | 158 |
|
168 | 159 | def get_buffers(self) -> ColumnBuffers: |
|
0 commit comments