Add equality methods.

This commit is contained in:
Rok authored and Antoine Pitrou committed 2019-07-01 11:39:30 +02:00
1 parent 4a30487fce
commit 3fcc1929e8
2 files changed
+51

No files matched your search

+24
View File
@@ -936,6 +936,18 @@ shape: {0.shape}""".format(self)
&out_data, &out_coords))
return PyObject_to_object(out_data), PyObject_to_object(out_coords)
def equals(self, SparseTensorCOO other):
"""
Return true if sparse tensors contains exactly equal data
"""
return self.stp.Equals(deref(other.stp))
def __eq__(self, other):
if isinstance(other, SparseTensorCOO):
return self.equals(other)
else:
return NotImplemented
@property
def is_mutable(self):
return self.stp.is_mutable()
@@ -1034,6 +1046,18 @@ shape: {0.shape}""".format(self)
return (PyObject_to_object(out_data), PyObject_to_object(out_indptr),
PyObject_to_object(out_indices))
def equals(self, SparseTensorCSR other):
"""
Return true if sparse tensors contains exactly equal data
"""
return self.stp.Equals(deref(other.stp))
def __eq__(self, other):
if isinstance(other, SparseTensorCSR):
return self.equals(other)
else:
return NotImplemented
@property
def is_mutable(self):
return self.stp.is_mutable()
@@ -59,6 +59,32 @@ def test_sparse_tensor_csr_base_object():
assert sys.getrefcount(sparse_tensor) == n + 1
@pytest.mark.skip
@pytest.mark.parametrize('sparse_tensor_type', [
pa.SparseTensorCSR,
pa.SparseTensorCOO,
])
def test_sparse_tensor_equals(sparse_tensor_type):
def eq(a, b):
assert a.equals(b)
assert a == b
assert not (a != b)
def ne(a, b):
assert not a.equals(b)
assert not (a == b)
assert a != b
data = np.random.randn(10, 6)[::, ::2]
sparse_tensor1 = sparse_tensor_type.from_dense_numpy(data)
sparse_tensor2 = sparse_tensor_type.from_dense_numpy(np.ascontiguousarray(data))
eq(sparse_tensor1, sparse_tensor2)
data = data.copy()
data[9, 0] = 1.0
sparse_tensor2 = sparse_tensor_type.from_dense_numpy(np.ascontiguousarray(data))
ne(sparse_tensor1, sparse_tensor2)
@pytest.mark.parametrize('dtype_str,arrow_type', [
('i1', pa.int8()),
('i2', pa.int16()),
@@ -141,6 +167,7 @@ def test_sparse_tensor_csr_from_dense(dtype_str, arrow_type):
assert (indices == result_indices).all()
@pytest.mark.skip
@pytest.mark.parametrize('dtype_str,arrow_type', [
('i1', pa.int8()),
('i2', pa.int16()),