ARROW-7261: [Python] Add Python support for Fixed Size List type

cpp/src/arrow/array.cc

@@ -491,6 +491,24 @@ std::shared_ptr<DataType> FixedSizeListArray::value_type() const {
 
 std::shared_ptr<Array> FixedSizeListArray::values() const { return values_; }
 
+Status FixedSizeListArray::FromArrays(const std::shared_ptr<Array>& values,
+                                      int64_t list_size, std::shared_ptr<Array>* out) {
+  if (list_size < 0) {
+    return Status::Invalid("list_size needs to be a positive integer");
+  }
+  if ((values->length() % list_size) != 0) {
+    return Status::Invalid(
+        "The length of the values Array needs to be a multiple of the list_size");
+  }
+  int64_t length = values->length() / list_size;
+  auto list_type = std::make_shared<FixedSizeListType>(values->type(), list_size);
+  std::shared_ptr<Buffer> validity_buf;
+
+  *out = std::make_shared<FixedSizeListArray>(list_type, length, values, validity_buf, -1,
+                                              0);
+  return Status::OK();
+}
+
 // ----------------------------------------------------------------------
 // String and binary
 

cpp/src/arrow/array.h

@@ -736,6 +736,14 @@ class ARROW_EXPORT FixedSizeListArray : public Array {
     return values_->Slice(value_offset(i), value_length(i));
   }
 
+  /// \brief Construct FixedSizeListArray from child value array and value_length
+  ///
+  /// \param[in] values Array containing list values
+  /// \param[in] list_size The fixed length of each list
+  /// \param[out] out Will have length equal to values.length() / list_size
+  static Status FromArrays(const std::shared_ptr<Array>& values, int64_t list_size,
+                           std::shared_ptr<Array>* out);
+
  protected:
   void SetData(const std::shared_ptr<ArrayData>& data);
   int32_t list_size_;

cpp/src/arrow/python/python_to_arrow.cc

@@ -1149,6 +1149,15 @@ Status GetConverter(const std::shared_ptr<DataType>& type, bool from_pandas,
       } else {
         *out = std::unique_ptr<SeqConverter>(
             new MapConverter<NullCoding::NONE_ONLY>(from_pandas, strict_conversions));
+    case Type::FIXED_SIZE_LIST:
+      if (from_pandas) {
+        *out = std::unique_ptr<SeqConverter>(
+            new ListConverter<FixedSizeListType, NullCoding::PANDAS_SENTINELS>(
+                from_pandas, strict_conversions));
+      } else {
+        *out = std::unique_ptr<SeqConverter>(
+            new ListConverter<FixedSizeListType, NullCoding::NONE_ONLY>(
+                from_pandas, strict_conversions));
       }
       return Status::OK();
     case Type::STRUCT:

python/pyarrow/__init__.py

@@ -59,7 +59,8 @@ from pyarrow.lib import (null, bool_,
                          field,
                          type_for_alias,
                          DataType, DictionaryType, StructType,
-                         ListType, LargeListType, MapType, UnionType,
+                         ListType, LargeListType, MapType, FixedSizeListType,
+                         UnionType,
                          TimestampType, Time32Type, Time64Type, DurationType,
                          FixedSizeBinaryType, Decimal128Type,
                          BaseExtensionType, ExtensionType,
@@ -80,7 +81,8 @@ from pyarrow.lib import (null, bool_,
                          Int16Array, UInt16Array,
                          Int32Array, UInt32Array,
                          Int64Array, UInt64Array,
-                         ListArray, LargeListArray, MapArray, UnionArray,
+                         ListArray, LargeListArray, MapArray,
+                         FixedSizeListArray, UnionArray,
                          BinaryArray, StringArray,
                          LargeBinaryArray, LargeStringArray,
                          FixedSizeBinaryArray,
@@ -93,7 +95,7 @@ from pyarrow.lib import (null, bool_,
                          Int8Value, Int16Value, Int32Value, Int64Value,
                          UInt8Value, UInt16Value, UInt32Value, UInt64Value,
                          HalfFloatValue, FloatValue, DoubleValue,
-                         ListValue, LargeListValue, MapValue,
+                         ListValue, LargeListValue, MapValue, FixedSizeListValue,
                          BinaryValue, StringValue,
                          LargeBinaryValue, LargeStringValue,
                          FixedSizeBinaryValue,

python/pyarrow/array.pxi

@@ -1438,6 +1438,54 @@ cdef class MapArray(Array):
         return pyarrow_wrap_array((<CMapArray*> self.ap).items())
 
 
+cdef class FixedSizeListArray(Array):
+    """
+    Concrete class for Arrow arrays of a fixed size list data type.
+    """
+
+    @staticmethod
+    def from_arrays(values, int64_t length, MemoryPool pool=None):
+        """
+        Construct ListArray from arrays of int32 offsets and values
+
+        Parameters
+        ----------
+        offsets : Array (int32 type)
+        values : Array (any type)
+
+        Returns
+        -------
+        list_array : ListArray
+        """
+        cdef:
+            Array _values
+            shared_ptr[CArray] out
+
+        _values = asarray(values)
+
+        with nogil:
+            check_status(CFixedSizeListArray.FromArrays(
+                _values.sp_array, length, &out))
+        cdef Array result = pyarrow_wrap_array(out)
+        result.validate()
+        return result
+
+    @property
+    def values(self):
+        return self.flatten()
+
+    def flatten(self):
+        """
+        Unnest this ListArray by one level
+
+        Returns
+        -------
+        result : Array
+        """
+        cdef CFixedSizeListArray* arr = <CFixedSizeListArray*> self.ap
+        return pyarrow_wrap_array(arr.values())
+
+
 cdef class UnionArray(Array):
     """
     Concrete class for Arrow arrays of a Union data type.
@@ -1890,6 +1938,7 @@ cdef dict _array_classes = {
     _Type_LIST: ListArray,
     _Type_LARGE_LIST: LargeListArray,
     _Type_MAP: MapArray,
+    _Type_FIXED_SIZE_LIST: FixedSizeListArray,
     _Type_UNION: UnionArray,
     _Type_BINARY: BinaryArray,
     _Type_STRING: StringArray,

python/pyarrow/includes/libarrow.pxd

@@ -73,6 +73,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
         _Type_LIST" arrow::Type::LIST"
         _Type_LARGE_LIST" arrow::Type::LARGE_LIST"
+        _Type_FIXED_SIZE_LIST" arrow::Type::FIXED_SIZE_LIST"
         _Type_STRUCT" arrow::Type::STRUCT"
         _Type_UNION" arrow::Type::UNION"
         _Type_DICTIONARY" arrow::Type::DICTIONARY"
@@ -293,6 +294,14 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         shared_ptr[CDataType] item_type()
         c_bool keys_sorted()
 
+    cdef cppclass CFixedSizeListType" arrow::FixedSizeListType"(CDataType):
+        CFixedSizeListType(const shared_ptr[CDataType]& value_type,
+                           int32_t list_size)
+        CFixedSizeListType(const shared_ptr[CField]& field, int32_t list_size)
+        shared_ptr[CDataType] value_type()
+        shared_ptr[CField] value_field()
+        int32_t list_size()
+
     cdef cppclass CStringType" arrow::StringType"(CDataType):
         pass
 
@@ -473,7 +482,16 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         CStatus FromArrays(const CArray& offsets, const CArray& values,
                            CMemoryPool* pool, shared_ptr[CArray]* out)
 
-        const int64_t* raw_value_offsets()
+        int64_t value_offset(int i)
+        int64_t value_length(int i)
+        shared_ptr[CArray] values()
+        shared_ptr[CDataType] value_type()
+
+    cdef cppclass CFixedSizeListArray" arrow::FixedSizeListArray"(CArray):
+        @staticmethod
+        CStatus FromArrays(const shared_ptr[CArray]& values, int64_t list_size,
+                           shared_ptr[CArray]* out)
+
         int64_t value_offset(int i)
         int64_t value_length(int i)
         shared_ptr[CArray] values()

python/pyarrow/lib.pxd

@@ -74,6 +74,11 @@ cdef class MapType(DataType):
         const CMapType* map_type
 
 
+cdef class FixedSizeListType(DataType):
+    cdef:
+        const CFixedSizeListType* list_type
+
+
 cdef class StructType(DataType):
     cdef:
         const CStructType* struct_type
@@ -226,6 +231,17 @@ cdef class MapValue(ArrayValue):
     cdef int64_t length(self)
 
 
+cdef class FixedSizeListValue(ArrayValue):
+    cdef readonly:
+        DataType value_type
+
+    cdef:
+        CFixedSizeListArray* ap
+
+    cdef getitem(self, int64_t i)
+    cdef int64_t length(self)
+
+
 cdef class StructValue(ArrayValue):
     cdef:
         CStructArray* ap
@@ -388,6 +404,10 @@ cdef class MapArray(Array):
     pass
 
 
+cdef class FixedSizeListArray(Array):
+    pass
+
+
 cdef class UnionArray(Array):
     pass
 

python/pyarrow/lib.pyx

@@ -93,10 +93,11 @@ Type_LARGE_STRING = _Type_LARGE_STRING
 Type_FIXED_SIZE_BINARY = _Type_FIXED_SIZE_BINARY
 Type_LIST = _Type_LIST
 Type_LARGE_LIST = _Type_LARGE_LIST
+Type_MAP = _Type_MAP
+Type_FIXED_SIZE_LIST = _Type_FIXED_SIZE_LIST
 Type_STRUCT = _Type_STRUCT
 Type_UNION = _Type_UNION
 Type_DICTIONARY = _Type_DICTIONARY
-Type_MAP = _Type_MAP
 
 UnionMode_SPARSE = _UnionMode_SPARSE
 UnionMode_DENSE = _UnionMode_DENSE

python/pyarrow/public-api.pxi

@@ -88,6 +88,8 @@ cdef api object pyarrow_wrap_data_type(
         out = LargeListType.__new__(LargeListType)
     elif type.get().id() == _Type_MAP:
         out = MapType.__new__(MapType)
+    elif type.get().id() == _Type_FIXED_SIZE_LIST:
+        out = FixedSizeListType.__new__(FixedSizeListType)
     elif type.get().id() == _Type_STRUCT:
         out = StructType.__new__(StructType)
     elif type.get().id() == _Type_UNION:

python/pyarrow/scalar.pxi

@@ -786,6 +786,57 @@ cdef class MapValue(ArrayValue):
         return result
 
 
+cdef class FixedSizeListValue(ArrayValue):
+    """
+    Concrete class for fixed size list array elements.
+    """
+
+    def __len__(self):
+        """
+        Return the number of values.
+        """
+        return self.length()
+
+    def __getitem__(self, i):
+        """
+        Return the value at the given index.
+        """
+        return self.getitem(_normalize_index(i, self.length()))
+
+    def __iter__(self):
+        """
+        Iterate over this element's values.
+        """
+        for i in range(len(self)):
+            yield self.getitem(i)
+        raise StopIteration
+
+    cdef void _set_array(self, const shared_ptr[CArray]& sp_array):
+        self.sp_array = sp_array
+        self.ap = <CFixedSizeListArray*> sp_array.get()
+        self.value_type = pyarrow_wrap_data_type(self.ap.value_type())
+
+    cdef getitem(self, int64_t i):
+        cdef int64_t j = self.ap.value_offset(self.index) + i
+        return box_scalar(self.value_type, self.ap.values(), j)
+
+    cdef int64_t length(self):
+        return self.ap.value_length(self.index)
+
+    def as_py(self):
+        """
+        Return this value as a Python list.
+        """
+        cdef:
+            int64_t j
+            list result = []
+
+        for j in range(len(self)):
+            result.append(self.getitem(j).as_py())
+
+        return result
+
+
 cdef class UnionValue(ArrayValue):
     """
     Concrete class for union array elements.
@@ -932,6 +983,7 @@ cdef dict _array_value_classes = {
     _Type_LIST: ListValue,
     _Type_LARGE_LIST: LargeListValue,
     _Type_MAP: MapValue,
+    _Type_FIXED_SIZE_LIST: FixedSizeListValue,
     _Type_UNION: UnionValue,
     _Type_BINARY: BinaryValue,
     _Type_STRING: StringValue,

python/pyarrow/tests/test_array.py

@@ -678,6 +678,21 @@ def test_map_from_arrays():
         pa.MapArray.from_arrays(offsets, keys_with_null, items)
 
 
+def test_fixed_size_list_from_arrays():
+    values = pa.array(range(12), pa.int64())
+    result = pa.FixedSizeListArray.from_arrays(values, 4)
+    assert result.to_pylist() == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]
+    assert result.type.equals(pa.list_(pa.int64(), 4))
+
+    # raise on invalid values / list_size
+    with pytest.raises(ValueError):
+        pa.FixedSizeListArray.from_arrays(values, -4)
+
+    with pytest.raises(ValueError):
+        # length of values not multiple of 5
+        pa.FixedSizeListArray.from_arrays(values, 5)
+
+
 def test_union_from_dense():
     binary = pa.array([b'a', b'b', b'c', b'd'], type='binary')
     int64 = pa.array([1, 2, 3], type='int64')
@@ -967,6 +982,7 @@ def test_cast_from_null():
         pa.binary(),
         pa.binary(10),
         pa.list_(pa.int16()),
+        pa.list_(pa.int32(), 4),
         pa.large_list(pa.uint8()),
         pa.decimal128(19, 4),
         pa.timestamp('us'),
@@ -1804,6 +1820,43 @@ def test_list_array_values_offsets_sliced(klass):
     assert arr2[0].as_py() == arr2.values[i:j].to_pylist() == [4]
 
 
+def test_fixed_size_list_array_flatten():
+    typ2 = pa.list_(pa.list_(pa.int64(), 2), 3)
+    arr2 = pa.array([
+        [
+            [1, 2],
+            [3, 4],
+            [5, 6],
+        ],
+        None,
+        [
+            [7, None],
+            None,
+            [8, 9]
+        ],
+    ], type=typ2)
+    assert arr2.type.equals(typ2)
+
+    typ1 = pa.list_(pa.int64(), 2)
+    arr1 = pa.array([
+        [1, 2], [3, 4], [5, 6],
+        None, None, None,
+        [7, None], None, [8, 9]
+    ], type=typ1)
+    assert arr1.type.equals(typ1)
+    assert arr2.flatten().equals(arr1)
+
+    typ0 = pa.int64()
+    arr0 = pa.array([
+        1, 2, 3, 4, 5, 6,
+        None, None, None, None, None, None,
+        7, None, None, None, 8, 9,
+    ], type=typ0)
+    assert arr0.type.equals(typ0)
+    assert arr1.flatten().equals(arr0)
+    assert arr2.flatten().flatten().equals(arr0)
+
+
 def test_struct_array_flatten():
     ty = pa.struct([pa.field('x', pa.int16()),
                     pa.field('y', pa.float32())])

python/pyarrow/tests/test_convert_builtin.py

@@ -234,6 +234,18 @@ def test_nested_arrays(seq):
     assert arr.to_pylist() == [[], [1, 2], None]
 
 
+@parametrize_with_iterable_types
+def test_nested_fixed_size_list(seq):
+    data = [[1, 2], [3, None], None]
+    arr = pa.array(seq(data), type=pa.list_(pa.int32(), 2))
+    assert len(arr) == 3
+    assert arr.null_count == 1
+    assert arr.type == pa.list_(pa.int32(), 2)
+    assert arr.to_pylist() == data
+
+    # TODO check for invalid data (non-equal length nested lists)
+
+
 @parametrize_with_iterable_types
 def test_sequence_all_none(seq):
     arr = pa.array(seq([None, None]))

python/pyarrow/tests/test_misc.py

@@ -53,6 +53,8 @@ def test_cpu_count():
     pa.Tensor,
     pa.lib.DataType,
     pa.lib.ListType,
+    pa.lib.LargeListType,
+    pa.lib.FixedSizeListType,
     pa.lib.UnionType,
     pa.lib.StructType,
     pa.lib.Time32Type,
@@ -77,6 +79,7 @@ def test_cpu_count():
     pa.ListArray,
     pa.LargeListArray,
     pa.MapArray,
+    pa.FixedSizeListArray,
     pa.UnionArray,
     pa.BinaryArray,
     pa.StringArray,
@@ -116,6 +119,7 @@ def test_cpu_count():
     pa.ListValue,
     pa.LargeListValue,
     pa.MapValue,
+    pa.FixedSizeListValue,
     pa.UnionValue,
     pa.StructValue,
     pa.DictionaryValue,

python/pyarrow/tests/test_scalars.py

@@ -249,6 +249,22 @@ class TestScalars(unittest.TestCase):
         assert v[0][0] == u'c'
         assert v[0][1] is pa.NA
 
+    def test_fixed_size_list(self):
+        arr = pa.array([[1, None, 3], [4, 5, 6]], pa.list_(pa.int64(), 3))
+
+        v = arr[0]
+        assert len(v) == 3
+        assert isinstance(v, pa.FixedSizeListValue)
+        assert repr(v) == "[1, None, 3]"
+        assert v.as_py() == [1, None, 3]
+        assert v[0].as_py() == 1
+        assert v[1] is pa.NA
+        assert v[-1] == v[2]
+        with pytest.raises(IndexError):
+            v[-4]
+        with pytest.raises(IndexError):
+            v[3]
+
     def test_date(self):
         # ARROW-5125
         d1, d2 = datetime.date(3200, 1, 1), datetime.date(1960, 1, 1),

python/pyarrow/tests/test_schema.py

@@ -64,6 +64,7 @@ def test_type_to_pandas_dtype():
         (pa.binary(12), np.object_),
         (pa.string(), np.object_),
         (pa.list_(pa.int8()), np.object_),
+        # (pa.list_(pa.int8(), 2), np.object_),  # TODO needs pandas conversion
     ]
     for arrow_type, numpy_type in cases:
         assert arrow_type.to_pandas_dtype() == numpy_type

python/pyarrow/tests/test_types.py

@@ -53,6 +53,7 @@ def get_many_types():
         pa.large_string(),
         pa.large_binary(),
         pa.list_(pa.int32()),
+        pa.list_(pa.int32(), 2),
         pa.large_list(pa.uint16()),
         pa.map_(pa.string(), pa.int32()),
         pa.struct([pa.field('a', pa.int32()),
@@ -118,11 +119,17 @@ def test_is_decimal():
 def test_is_list():
     a = pa.list_(pa.int32())
     b = pa.large_list(pa.int32())
+    c = pa.list_(pa.int32(), 3)
 
     assert types.is_list(a)
     assert not types.is_large_list(a)
+    assert not types.is_fixed_size_list(a)
     assert types.is_large_list(b)
     assert not types.is_list(b)
+    assert not types.is_fixed_size_list(b)
+    assert types.is_fixed_size_list(c)
+    assert not types.is_list(c)
+    assert not types.is_large_list(c)
 
     assert not types.is_list(pa.int32())
 
@@ -311,6 +318,13 @@ def test_map_type():
         pa.map_(pa.int32(), None)
 
 
+def test_fixed_size_list_type():
+    ty = pa.list_(pa.float64(), 2)
+    assert isinstance(ty, pa.FixedSizeListType)
+    assert ty.value_type == pa.float64()
+    assert ty.list_size == 2
+
+
 def test_struct_type():
     fields = [
         # Duplicate field name on purpose

python/pyarrow/types.pxi

@@ -309,6 +309,34 @@ cdef class MapType(DataType):
         return pyarrow_wrap_data_type(self.map_type.item_type())
 
 
+
+cdef class FixedSizeListType(DataType):
+    """
+    Concrete class for fixed size list data types.
+    """
+
+    cdef void init(self, const shared_ptr[CDataType]& type) except *:
+        DataType.init(self, type)
+        self.list_type = <const CFixedSizeListType*> type.get()
+
+    def __reduce__(self):
+        return list_, (self.value_type, self.list_size)
+
+    @property
+    def value_type(self):
+        """
+        The data type of large list values.
+        """
+        return pyarrow_wrap_data_type(self.list_type.value_type())
+
+    @property
+    def list_size(self):
+        """
+        The size of the fixed size lists.
+        """
+        return self.list_type.list_size()
+
+
 cdef class StructType(DataType):
     """
     Concrete class for struct data types.
@@ -1822,13 +1850,16 @@ def large_utf8():
     return large_string()
 
 
-cpdef ListType list_(value_type):
+def list_(value_type, int list_size=-1):
     """
     Create ListType instance from child data type or field
 
     Parameters
     ----------
     value_type : DataType or Field
+    list_size : int, optional, default -1
+        If length == -1 then return a variable length list type. If length is
+        greater than or equal to 0 then return a fixed size list type.
 
     Returns
     -------
@@ -1837,7 +1868,6 @@ cpdef ListType list_(value_type):
     cdef:
         Field _field
         shared_ptr[CDataType] list_type
-        ListType out = ListType.__new__(ListType)
 
     if isinstance(value_type, DataType):
         _field = field('item', value_type)
@@ -1846,9 +1876,12 @@ cpdef ListType list_(value_type):
     else:
         raise TypeError('List requires DataType or Field')
 
-    list_type.reset(new CListType(_field.sp_field))
-    out.init(list_type)
-    return out
+    if list_size == -1:
+        list_type.reset(new CListType(_field.sp_field))
+    else:
+        list_type.reset(new CFixedSizeListType(_field.sp_field, list_size))
+
+    return pyarrow_wrap_data_type(list_type)
 
 
 cpdef LargeListType large_list(value_type):

python/pyarrow/types.py

@@ -173,6 +173,13 @@ def is_large_list(t):
     return t.id == lib.Type_LARGE_LIST
 
 
+def is_fixed_size_list(t):
+    """
+    Return True if value is an instance of a fixed size list type
+    """
+    return t.id == lib.Type_FIXED_SIZE_LIST
+
+
 def is_struct(t):
     """
     Return True if value is an instance of a struct type