ARROW-439: [Python] Add option in "to_pandas" conversions to yield Categorical from...
[arrow.git] / python / pyarrow / array.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
19 def array(object sequence, DataType type=None, MemoryPool memory_pool=None,
20           size=None):
21     """
22     Create pyarrow.Array instance from a Python sequence
23
24     Parameters
25     ----------
26     sequence : sequence-like or iterable object of Python objects.
27         If both type and size are specified may be a single use iterable.
28     type : pyarrow.DataType, optional
29         If not passed, will be inferred from the data
30     memory_pool : pyarrow.MemoryPool, optional
31         If not passed, will allocate memory from the currently-set default
32         memory pool
33     size : int64, optional
34         Size of the elements. If the imput is larger than size bail at this
35         length. For iterators, if size is larger than the input iterator this
36         will be treated as a "max size", but will involve an initial allocation
37         of size followed by a resize to the actual size (so if you know the
38         exact size specifying it correctly will give you better performance).
39
40     Returns
41     -------
42     array : pyarrow.Array
43     """
44     cdef:
45         shared_ptr[CArray] sp_array
46         CMemoryPool* pool
47         int64_t c_size
48
49     pool = maybe_unbox_memory_pool(memory_pool)
50     if type is None:
51         with nogil:
52             check_status(ConvertPySequence(sequence, pool, &sp_array))
53     else:
54         if size is None:
55             with nogil:
56                 check_status(
57                     ConvertPySequence(
58                         sequence, pool, &sp_array, type.sp_type
59                     )
60                 )
61         else:
62             c_size = size
63             with nogil:
64                 check_status(
65                     ConvertPySequence(
66                         sequence, pool, &sp_array, type.sp_type, c_size
67                     )
68                 )
69
70     return pyarrow_wrap_array(sp_array)
71
72
73 def _normalize_slice(object arrow_obj, slice key):
74     cdef Py_ssize_t n = len(arrow_obj)
75
76     start = key.start or 0
77     while start < 0:
78         start += n
79
80     stop = key.stop if key.stop is not None else n
81     while stop < 0:
82         stop += n
83
84     step = key.step or 1
85     if step != 1:
86         raise IndexError('only slices with step 1 supported')
87     else:
88         return arrow_obj.slice(start, stop - start)
89
90
91 cdef class Array:
92
93     cdef void init(self, const shared_ptr[CArray]& sp_array):
94         self.sp_array = sp_array
95         self.ap = sp_array.get()
96         self.type = pyarrow_wrap_data_type(self.sp_array.get().type())
97
98     def _debug_print(self):
99         with nogil:
100             check_status(DebugPrint(deref(self.ap), 0))
101
102     @staticmethod
103     def from_pandas(obj, mask=None, DataType type=None,
104                     timestamps_to_ms=False,
105                     MemoryPool memory_pool=None):
106         """
107         Convert pandas.Series to an Arrow Array.
108
109         Parameters
110         ----------
111         series : pandas.Series or numpy.ndarray
112
113         mask : pandas.Series or numpy.ndarray, optional
114             boolean mask if the object is null (True) or valid (False)
115
116         type : pyarrow.DataType
117             Explicit type to attempt to coerce to
118
119         timestamps_to_ms : bool, optional
120             Convert datetime columns to ms resolution. This is needed for
121             compatibility with other functionality like Parquet I/O which
122             only supports milliseconds.
123
124         memory_pool: MemoryPool, optional
125             Specific memory pool to use to allocate the resulting Arrow array.
126
127         Notes
128         -----
129         Localized timestamps will currently be returned as UTC (pandas's native
130         representation).  Timezone-naive data will be implicitly interpreted as
131         UTC.
132
133         Examples
134         --------
135
136         >>> import pandas as pd
137         >>> import pyarrow as pa
138         >>> pa.Array.from_pandas(pd.Series([1, 2]))
139         <pyarrow.array.Int64Array object at 0x7f674e4c0e10>
140         [
141           1,
142           2
143         ]
144
145         >>> import numpy as np
146         >>> pa.Array.from_pandas(pd.Series([1, 2]), np.array([0, 1],
147         ... dtype=bool))
148         <pyarrow.array.Int64Array object at 0x7f9019e11208>
149         [
150           1,
151           NA
152         ]
153
154         Returns
155         -------
156         array : pyarrow.Array or pyarrow.ChunkedArray (if object data
157         overflowed binary storage)
158         """
159         cdef:
160             shared_ptr[CArray] out
161             shared_ptr[CChunkedArray] chunked_out
162             shared_ptr[CDataType] c_type
163             CMemoryPool* pool
164
165         if mask is not None:
166             mask = get_series_values(mask)
167
168         values = get_series_values(obj)
169         pool = maybe_unbox_memory_pool(memory_pool)
170
171         if isinstance(values, Categorical):
172             return DictionaryArray.from_arrays(
173                 values.codes, values.categories.values,
174                 mask=mask, ordered=values.ordered,
175                 memory_pool=memory_pool)
176         elif values.dtype == object:
177             # Object dtype undergoes a different conversion path as more type
178             # inference may be needed
179             if type is not None:
180                 c_type = type.sp_type
181             with nogil:
182                 check_status(PandasObjectsToArrow(
183                     pool, values, mask, c_type, &chunked_out))
184
185             if chunked_out.get().num_chunks() > 1:
186                 return pyarrow_wrap_chunked_array(chunked_out)
187             else:
188                 out = chunked_out.get().chunk(0)
189         else:
190             values, type = pdcompat.maybe_coerce_datetime64(
191                 values, obj.dtype, type, timestamps_to_ms=timestamps_to_ms)
192
193             if type is None:
194                 dtype = values.dtype
195                 with nogil:
196                     check_status(NumPyDtypeToArrow(dtype, &c_type))
197             else:
198                 c_type = type.sp_type
199
200             with nogil:
201                 check_status(PandasToArrow(
202                     pool, values, mask, c_type, &out))
203
204         return pyarrow_wrap_array(out)
205
206     property null_count:
207
208         def __get__(self):
209             return self.sp_array.get().null_count()
210
211     def __iter__(self):
212         for i in range(len(self)):
213             yield self.getitem(i)
214         raise StopIteration
215
216     def __repr__(self):
217         from pyarrow.formatting import array_format
218         type_format = object.__repr__(self)
219         values = array_format(self, window=10)
220         return '{0}\n{1}'.format(type_format, values)
221
222     def equals(Array self, Array other):
223         return self.ap.Equals(deref(other.ap))
224
225     def __len__(self):
226         if self.sp_array.get():
227             return self.sp_array.get().length()
228         else:
229             return 0
230
231     def isnull(self):
232         raise NotImplemented
233
234     def __getitem__(self, key):
235         cdef Py_ssize_t n = len(self)
236
237         if PySlice_Check(key):
238             return _normalize_slice(self, key)
239
240         while key < 0:
241             key += len(self)
242
243         return self.getitem(key)
244
245     cdef getitem(self, int64_t i):
246         return box_scalar(self.type, self.sp_array, i)
247
248     def slice(self, offset=0, length=None):
249         """
250         Compute zero-copy slice of this array
251
252         Parameters
253         ----------
254         offset : int, default 0
255             Offset from start of array to slice
256         length : int, default None
257             Length of slice (default is until end of Array starting from
258             offset)
259
260         Returns
261         -------
262         sliced : RecordBatch
263         """
264         cdef:
265             shared_ptr[CArray] result
266
267         if offset < 0:
268             raise IndexError('Offset must be non-negative')
269
270         if length is None:
271             result = self.ap.Slice(offset)
272         else:
273             result = self.ap.Slice(offset, length)
274
275         return pyarrow_wrap_array(result)
276
277     def to_pandas(self, c_bool strings_to_categorical=False):
278         """
279         Convert to an array object suitable for use in pandas
280
281         Parameters
282         ----------
283         strings_to_categorical : boolean, default False
284             Encode string (UTF8) and binary types to pandas.Categorical
285
286         See also
287         --------
288         Column.to_pandas
289         Table.to_pandas
290         RecordBatch.to_pandas
291         """
292         cdef:
293             PyObject* out
294             PandasOptions options
295
296         options = PandasOptions(strings_to_categorical=strings_to_categorical)
297         with nogil:
298             check_status(ConvertArrayToPandas(options, self.sp_array,
299                                               self, &out))
300         return wrap_array_output(out)
301
302     def to_pylist(self):
303         """
304         Convert to an list of native Python objects.
305         """
306         return [x.as_py() for x in self]
307
308     def validate(self):
309         """
310         Perform any validation checks implemented by
311         arrow::ValidateArray. Raises exception with error message if array does
312         not validate
313
314         Raises
315         ------
316         ArrowInvalid
317         """
318         with nogil:
319             check_status(ValidateArray(deref(self.ap)))
320
321
322 cdef class Tensor:
323
324     cdef void init(self, const shared_ptr[CTensor]& sp_tensor):
325         self.sp_tensor = sp_tensor
326         self.tp = sp_tensor.get()
327         self.type = pyarrow_wrap_data_type(self.tp.type())
328
329     def __repr__(self):
330         return """<pyarrow.Tensor>
331 type: {0}
332 shape: {1}
333 strides: {2}""".format(self.type, self.shape, self.strides)
334
335     @staticmethod
336     def from_numpy(obj):
337         cdef shared_ptr[CTensor] ctensor
338         with nogil:
339             check_status(NdarrayToTensor(c_default_memory_pool(), obj,
340                                          &ctensor))
341         return pyarrow_wrap_tensor(ctensor)
342
343     def to_numpy(self):
344         """
345         Convert arrow::Tensor to numpy.ndarray with zero copy
346         """
347         cdef:
348             PyObject* out
349
350         with nogil:
351             check_status(TensorToNdarray(deref(self.tp), self, &out))
352         return PyObject_to_object(out)
353
354     def equals(self, Tensor other):
355         """
356         Return true if the tensors contains exactly equal data
357         """
358         return self.tp.Equals(deref(other.tp))
359
360     property is_mutable:
361
362         def __get__(self):
363             return self.tp.is_mutable()
364
365     property is_contiguous:
366
367         def __get__(self):
368             return self.tp.is_contiguous()
369
370     property ndim:
371
372         def __get__(self):
373             return self.tp.ndim()
374
375     property size:
376
377         def __get__(self):
378             return self.tp.size()
379
380     property shape:
381
382         def __get__(self):
383             cdef size_t i
384             py_shape = []
385             for i in range(self.tp.shape().size()):
386                 py_shape.append(self.tp.shape()[i])
387             return py_shape
388
389     property strides:
390
391         def __get__(self):
392             cdef size_t i
393             py_strides = []
394             for i in range(self.tp.strides().size()):
395                 py_strides.append(self.tp.strides()[i])
396             return py_strides
397
398
399 cdef wrap_array_output(PyObject* output):
400     cdef object obj = PyObject_to_object(output)
401
402     if isinstance(obj, dict):
403         return Categorical(obj['indices'],
404                            categories=obj['dictionary'],
405                            fastpath=True)
406     else:
407         return obj
408
409
410 cdef class NullArray(Array):
411     pass
412
413
414 cdef class BooleanArray(Array):
415     pass
416
417
418 cdef class NumericArray(Array):
419     pass
420
421
422 cdef class IntegerArray(NumericArray):
423     pass
424
425
426 cdef class FloatingPointArray(NumericArray):
427     pass
428
429
430 cdef class Int8Array(IntegerArray):
431     pass
432
433
434 cdef class UInt8Array(IntegerArray):
435     pass
436
437
438 cdef class Int16Array(IntegerArray):
439     pass
440
441
442 cdef class UInt16Array(IntegerArray):
443     pass
444
445
446 cdef class Int32Array(IntegerArray):
447     pass
448
449
450 cdef class UInt32Array(IntegerArray):
451     pass
452
453
454 cdef class Int64Array(IntegerArray):
455     pass
456
457
458 cdef class UInt64Array(IntegerArray):
459     pass
460
461
462 cdef class Date32Array(NumericArray):
463     pass
464
465
466 cdef class Date64Array(NumericArray):
467     pass
468
469
470 cdef class TimestampArray(NumericArray):
471     pass
472
473
474 cdef class Time32Array(NumericArray):
475     pass
476
477
478 cdef class Time64Array(NumericArray):
479     pass
480
481
482 cdef class FloatArray(FloatingPointArray):
483     pass
484
485
486 cdef class DoubleArray(FloatingPointArray):
487     pass
488
489
490 cdef class FixedSizeBinaryArray(Array):
491     pass
492
493
494 cdef class DecimalArray(FixedSizeBinaryArray):
495     pass
496
497
498 cdef class ListArray(Array):
499
500     @staticmethod
501     def from_arrays(Array offsets, Array values, MemoryPool pool=None):
502         """
503         Construct ListArray from arrays of int32 offsets and values
504
505         Parameters
506         ----------
507         offset : Array (int32 type)
508         values : Array (any type)
509
510         Returns
511         -------
512         list_array : ListArray
513         """
514         cdef shared_ptr[CArray] out
515         cdef CMemoryPool* cpool = maybe_unbox_memory_pool(pool)
516         with nogil:
517             check_status(CListArray.FromArrays(
518                 deref(offsets.ap), deref(values.ap), cpool, &out))
519         return pyarrow_wrap_array(out)
520
521
522 cdef class StringArray(Array):
523     pass
524
525
526 cdef class BinaryArray(Array):
527     pass
528
529
530 cdef class DictionaryArray(Array):
531
532     cdef getitem(self, int64_t i):
533         cdef Array dictionary = self.dictionary
534         index = self.indices[i]
535         if index is NA:
536             return index
537         else:
538             return box_scalar(dictionary.type, dictionary.sp_array,
539                               index.as_py())
540
541     property dictionary:
542
543         def __get__(self):
544             cdef CDictionaryArray* darr = <CDictionaryArray*>(self.ap)
545
546             if self._dictionary is None:
547                 self._dictionary = pyarrow_wrap_array(darr.dictionary())
548
549             return self._dictionary
550
551     property indices:
552
553         def __get__(self):
554             cdef CDictionaryArray* darr = <CDictionaryArray*>(self.ap)
555
556             if self._indices is None:
557                 self._indices = pyarrow_wrap_array(darr.indices())
558
559             return self._indices
560
561     @staticmethod
562     def from_arrays(indices, dictionary, mask=None, ordered=False,
563                     MemoryPool memory_pool=None):
564         """
565         Construct Arrow DictionaryArray from array of indices (must be
566         non-negative integers) and corresponding array of dictionary values
567
568         Parameters
569         ----------
570         indices : ndarray or pandas.Series, integer type
571         dictionary : ndarray or pandas.Series
572         mask : ndarray or pandas.Series, boolean type
573             True values indicate that indices are actually null
574         ordered : boolean, default False
575             Set to True if the category values are ordered
576
577         Returns
578         -------
579         dict_array : DictionaryArray
580         """
581         cdef:
582             Array arrow_indices, arrow_dictionary
583             DictionaryArray result
584             shared_ptr[CDataType] c_type
585             shared_ptr[CArray] c_result
586
587         if isinstance(indices, Array):
588             if mask is not None:
589                 raise NotImplementedError(
590                     "mask not implemented with Arrow array inputs yet")
591             arrow_indices = indices
592         else:
593             if mask is None:
594                 mask = indices == -1
595             else:
596                 mask = mask | (indices == -1)
597             arrow_indices = Array.from_pandas(indices, mask=mask,
598                                               memory_pool=memory_pool)
599
600         if isinstance(dictionary, Array):
601             arrow_dictionary = dictionary
602         else:
603             arrow_dictionary = Array.from_pandas(dictionary,
604                                                  memory_pool=memory_pool)
605
606         if not isinstance(arrow_indices, IntegerArray):
607             raise ValueError('Indices must be integer type')
608
609         cdef c_bool c_ordered = ordered
610
611         c_type.reset(new CDictionaryType(arrow_indices.type.sp_type,
612                                          arrow_dictionary.sp_array, c_ordered))
613         c_result.reset(new CDictionaryArray(c_type, arrow_indices.sp_array))
614
615         result = DictionaryArray()
616         result.init(c_result)
617         return result
618
619
620 cdef class StructArray(Array):
621     @staticmethod
622     def from_arrays(field_names, arrays):
623         cdef:
624             Array array
625             shared_ptr[CArray] c_array
626             vector[shared_ptr[CArray]] c_arrays
627             shared_ptr[CArray] c_result
628             ssize_t num_arrays
629             ssize_t length
630             ssize_t i
631
632         num_arrays = len(arrays)
633         if num_arrays == 0:
634             raise ValueError("arrays list is empty")
635
636         length = len(arrays[0])
637
638         c_arrays.resize(num_arrays)
639         for i in range(num_arrays):
640             array = arrays[i]
641             if len(array) != length:
642                 raise ValueError("All arrays must have the same length")
643             c_arrays[i] = array.sp_array
644
645         cdef DataType struct_type = struct([
646             field(name, array.type)
647             for name, array in
648             zip(field_names, arrays)
649         ])
650
651         c_result.reset(new CStructArray(struct_type.sp_type, length, c_arrays))
652         result = StructArray()
653         result.init(c_result)
654         return result
655
656
657 cdef dict _array_classes = {
658     _Type_NA: NullArray,
659     _Type_BOOL: BooleanArray,
660     _Type_UINT8: UInt8Array,
661     _Type_UINT16: UInt16Array,
662     _Type_UINT32: UInt32Array,
663     _Type_UINT64: UInt64Array,
664     _Type_INT8: Int8Array,
665     _Type_INT16: Int16Array,
666     _Type_INT32: Int32Array,
667     _Type_INT64: Int64Array,
668     _Type_DATE32: Date32Array,
669     _Type_DATE64: Date64Array,
670     _Type_TIMESTAMP: TimestampArray,
671     _Type_TIME32: Time32Array,
672     _Type_TIME64: Time64Array,
673     _Type_FLOAT: FloatArray,
674     _Type_DOUBLE: DoubleArray,
675     _Type_LIST: ListArray,
676     _Type_BINARY: BinaryArray,
677     _Type_STRING: StringArray,
678     _Type_DICTIONARY: DictionaryArray,
679     _Type_FIXED_SIZE_BINARY: FixedSizeBinaryArray,
680     _Type_DECIMAL: DecimalArray,
681     _Type_STRUCT: StructArray,
682 }
683
684
685 cdef object get_series_values(object obj):
686     if isinstance(obj, PandasSeries):
687         result = obj.values
688     elif isinstance(obj, np.ndarray):
689         result = obj
690     else:
691         result = PandasSeries(obj).values
692
693     return result