def test_list_of_numpy_roundtrip(self):
x = [np.array([i, i * 2]) for i in range(5)]
arr = ArrowTensorArray.from_numpy(x)
batch = pa.RecordBatch.from_arrays([arr], ["batched_tensor"])
result_batch = _roundtrip_batch(batch)
result_arr = result_batch.column(0)
result = result_arr.to_numpy()
expected = np.stack(x)
npt.assert_array_equal(expected, result)
def test_numpy_roundtrip(self):
x = np.array([[1, 2], [3, 4], [5, 6]])
arr = ArrowTensorArray.from_numpy(x)
self.assertEqual(len(arr), 3)
batch = pa.RecordBatch.from_arrays([arr], ["batched_tensor"])
result_batch = _roundtrip_batch(batch)
result_arr = result_batch.column(0)
result = result_arr.to_numpy()
npt.assert_array_equal(x, result)
def test_batch_size(self):
x = [np.array([i, i * 2]) for i in range(6)]
arr_iter = ArrowTensorArray.from_numpy(x, batch_size=3)
result_obj_list = []
for arr in arr_iter:
batch = pa.RecordBatch.from_arrays([arr], ["batched_tensor"])
result_batch = _roundtrip_batch(batch)
result_arr = result_batch.column(0)
result_obj_list.append(result_arr.to_numpy())
self.assertEqual(len(result_obj_list), 2)
result = np.concatenate(result_obj_list)
expected = np.stack(x)
npt.assert_array_equal(expected, result)
def test_ndarray_dict(self):
obj = {
'a': [np.array([i, i * 2]) for i in range(10)],
'b': [np.array([i, i * i]) for i in range(10)]
}
table = ArrowTensorArray.from_numpy(obj)
result_table = _roundtrip_table(table)
results_a = [
chunk.to_numpy()
for chunk in result_table.column('a').iterchunks()
]
results_b = [
chunk.to_numpy()
for chunk in result_table.column('b').iterchunks()
]
result_a = np.concatenate(results_a)
result_b = np.concatenate(results_b)
self.assertTrue(np.array_equal(obj['a'], result_a))
self.assertTrue(np.array_equal(obj['b'], result_b))
def __from_arrow__(self, extension_array):
from text_extensions_for_pandas.array.arrow_conversion import ArrowTensorArray
values = ArrowTensorArray.to_numpy(extension_array)
return TensorArray(values)
def __arrow_array__(self, type=None):
from text_extensions_for_pandas.array.arrow_conversion import ArrowTensorArray
return ArrowTensorArray.from_numpy(self._tensor)