def test_string_sliced(self):
a = pyarrow.array(["a", None, "ccc"]).slice(1, 2)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_fixed_binary(self):
a = pyarrow.array([b"aa", None, b"cc"], pyarrow.binary(2))
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def _test_struct_sliced(self):
fields = [
("f1", pyarrow.int32()),
("f2", pyarrow.string()),
]
a = pyarrow.array(
[
{
"f1": 1,
"f2": "a"
},
None,
{
"f1": 3,
"f2": None
},
{
"f1": None,
"f2": "d"
},
{
"f1": None,
"f2": None
},
],
pyarrow.struct(fields),
).slice(1, 2)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_boolean_sliced(self):
a = pyarrow.array([True, None, False, True, False]).slice(1, 2)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_primitive_sliced(self):
a = pyarrow.array([0, None, 2, 3, 4]).slice(1, 2)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_list_sliced(self):
a = pyarrow.array([[], None, [1, 2], [4, 5, 6]],
pyarrow.list_(pyarrow.int64())).slice(1, 2)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_dictionary_python():
"""
Python -> Rust -> Python
"""
a = pa.array(["a", None, "b", None, "a"], type=pa.dictionary(pa.int8(), pa.string()))
b = rust.round_trip_array(a)
assert a == b
del a
del b
def test_null(self):
"""
Python -> Rust -> Python
"""
a = pyarrow.array([None], type=pyarrow.null())
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_fixed_len_list_array():
"""
Python -> Rust -> Python
"""
a = pa.array([[1, 2], None, [3, 4], [5, 6]], pa.list_(pa.int64(), 2))
b = rust.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
del a
del b
def test_fixed_len_binary_array():
"""
Python -> Rust -> Python
"""
a = pa.array(["aaa", None, "bbb", "ccc"], pa.binary(3))
b = rust.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
del a
del b
def test_decimal_roundtrip(self):
"""
Python -> Rust -> Python
"""
data = [
round(decimal.Decimal(722.82), 2),
round(decimal.Decimal(-934.11), 2),
None,
]
a = pyarrow.array(data, pyarrow.decimal128(5, 2))
b = arrow_pyarrow_integration_testing.round_trip_array(a)
self.assertEqual(a, b)
def test_list_array(self):
"""
Python -> Rust -> Python
"""
for _ in range(2):
a = pyarrow.array([[], None, [1, 2], [4, 5, 6]],
pyarrow.list_(pyarrow.int64()))
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_decimal_python():
"""
Python -> Rust -> Python
"""
data = [
round(decimal.Decimal(123.45), 2),
round(decimal.Decimal(-123.45), 2), None
]
a = pa.array(data, pa.decimal128(6, 2))
b = rust.round_trip_array(a)
assert a == b
del a
del b
def _test_fixed_list_sliced(self):
"""
Python -> Rust -> Python
"""
a = pyarrow.array(
[None, [1, 2], [4, 5]],
pyarrow.list_(pyarrow.int64(), 2),
).slice(1, 2)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_dict(self):
"""
Python -> Rust -> Python
"""
a = pyarrow.array(
["a", "a", "b", None, "c"],
pyarrow.dictionary(pyarrow.int64(), pyarrow.utf8()),
)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_map(self):
"""
Python -> Rust -> Python
"""
offsets = [0, None, 2, 6]
pykeys = [b"a", b"b", b"c", b"d", b"e", b"f"]
pyitems = [1, 2, 3, None, 4, 5]
keys = pyarrow.array(pykeys, type="binary")
items = pyarrow.array(pyitems, type="i4")
a = pyarrow.MapArray.from_arrays(offsets, keys, items)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_sparse_union(self):
"""
Python -> Rust -> Python
"""
a = pyarrow.UnionArray.from_sparse(
pyarrow.array([0, 1, 1, 0, 1], pyarrow.int8()),
[
pyarrow.array(["a", "", "", "", "c"], pyarrow.utf8()),
pyarrow.array([0, 1, 2, None, 0], pyarrow.int64()),
],
)
b = arrow_pyarrow_integration_testing.round_trip_array(a)
b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
def test_string(self):
a = pyarrow.array(["a", None, "ccc"])
b = arrow_pyarrow_integration_testing.round_trip_array(a)
c = pyarrow.array(["a", None, "ccc"])
self.assertEqual(b, c)
