def testCooFromListArray(self, list_array, expected_coo,
expected_dense_shape):
for input_array in [
pa.array(list_array),
# it should work for sliced arrays.
pa.array(list_array + list_array).slice(0, len(list_array)),
pa.array(list_array + list_array).slice(len(list_array)),
]:
coo, dense_shape = array_util.CooFromListArray(input_array)
self.assertTrue(coo.type.equals(pa.int64()))
self.assertTrue(dense_shape.type.equals(pa.int64()))
self.assertEqual(expected_coo, coo.to_pylist())
self.assertEqual(expected_dense_shape, dense_shape.to_pylist())
def GetTensor(self, record_batch: pa.RecordBatch,
produce_eager_tensors: bool) -> Any:
array = record_batch.column(self._column_index)
coo_array, dense_shape_array = array_util.CooFromListArray(array)
dense_shape_np = dense_shape_array.to_numpy()
values_np = np.asarray(array.flatten())
coo_np = coo_array.to_numpy().reshape(values_np.size, 2)
if produce_eager_tensors:
return tf.sparse.SparseTensor(
indices=tf.convert_to_tensor(coo_np),
dense_shape=tf.convert_to_tensor(dense_shape_np),
values=tf.convert_to_tensor(values_np))
return tf.compat.v1.SparseTensorValue(
indices=coo_np, dense_shape=dense_shape_np, values=values_np)