diff --git a/datasketch/storage.py b/datasketch/storage.py index 54b1c0d3..2b4a95dd 100644 --- a/datasketch/storage.py +++ b/datasketch/storage.py @@ -326,8 +326,14 @@ class CassandraClient: key blob, value blob, ts bigint, - PRIMARY KEY (key, value) - ) WITH CLUSTERING ORDER BY (value DESC) + PRIMARY KEY (key, value, ts) + ) WITH CLUSTERING ORDER BY (value DESC, ts ASC) + """ + + QUERY_GET_TABLE_SCHEMA = """ + SELECT column_name, kind, position + FROM system_schema.columns + WHERE keyspace_name = ? AND table_name = ? """ QUERY_DROP_TABLE = "DROP TABLE IF EXISTS {}" @@ -366,11 +372,7 @@ class CassandraClient: WHERE key = ? AND value = ? """ - QUERY_UPSERT = """ - UPDATE {} - SET ts = ? - WHERE key = ? AND value = ? - """ + QUERY_UPSERT = "INSERT INTO {} (key, value, ts) VALUES (?, ?, 0)" QUERY_INSERT = "INSERT INTO {} (key, value, ts) VALUES (?, ?, ?)" @@ -406,6 +408,7 @@ def __init__(self, cassandra_params, name, buffer_size): if cassandra_params.get("drop_tables", False): self._session.execute(self.QUERY_DROP_TABLE.format(table_name)) self._session.execute(self.QUERY_CREATE_TABLE.format(table_name)) + self._validate_table_schema(table_name) # Prepare all the statements for this table self._stmt_insert = self._session.prepare(self.QUERY_INSERT.format(table_name)) @@ -417,6 +420,30 @@ def __init__(self, cassandra_params, name, buffer_size): self._stmt_delete_key = self._session.prepare(self.QUERY_DELETE_KEY.format(table_name)) self._stmt_delete_val = self._session.prepare(self.QUERY_DELETE_VAL.format(table_name)) + def _validate_table_schema(self, table_name): + """Reject tables created with the old duplicate-collapsing key.""" + statement = self._session.prepare(self.QUERY_GET_TABLE_SCHEMA) + rows = self._session.execute( + statement, + (self._session.keyspace, table_name), + ) + primary_key = { + (row.kind, row.position): row.column_name + for row in rows + if row.kind in ("partition_key", "clustering") + } + expected = { + ("partition_key", 0): "key", + ("clustering", 0): "value", + ("clustering", 1): "ts", + } + if primary_key != expected: + raise RuntimeError( + "Cassandra table %r uses an incompatible primary key. " + "Rebuild the LSH table (or set drop_tables=True while rebuilding) " + "so duplicate ordered values can be preserved." % table_name + ) + @property def buffer_size(self): """Get the buffer size. @@ -527,7 +554,10 @@ def upsert(self, key, vals, buffer=False): :param iterable[byte|str] vals: the iterable of values :param boolean buffer: whether the upsert statements should be buffered """ - statements_and_parameters = [(self._stmt_upsert, (self._ts(), key, val)) for val in vals] + # The timestamp is part of the primary key so ordered storage can + # preserve duplicates. Sets use the fixed timestamp in + # QUERY_UPSERT, making repeated (key, value) inserts idempotent. + statements_and_parameters = [(self._stmt_upsert, (key, val)) for val in vals] if buffer: self._buffer(statements_and_parameters) else: diff --git a/test/test_integration.py b/test/test_integration.py index aa1ea91c..6f2ce88f 100644 --- a/test/test_integration.py +++ b/test/test_integration.py @@ -5,6 +5,7 @@ from datasketch.lsh import MinHashLSH from datasketch.minhash import MinHash +from datasketch.storage import ordered_storage from datasketch.weighted_minhash import WeightedMinHashGenerator STORAGE_CONFIG_REDIS = { @@ -108,6 +109,25 @@ def test_insert(self, storage_config): for i, H in enumerate(lsh.keys[b"a"]): assert b"a" in lsh.hashtables[i][H] + def test_ordered_storage_preserves_duplicate_values(self, storage_config): + storage = ordered_storage(storage_config, name=b"lsh_test_duplicate_values") + storage.insert(b"key", b"same", b"same") + assert list(storage.get(b"key")) == [b"same", b"same"] + + def test_remove_empty_minhash_clears_every_band(self, storage_config): + lsh = MinHashLSH(threshold=0.5, num_perm=16, storage_config=storage_config, prepickle=False) + empty = MinHash(16) + lsh.insert(b"empty", empty) + + band_hashes = list(lsh.keys[b"empty"]) + assert len(band_hashes) == lsh.b + assert len(set(band_hashes)) == 1 + + lsh.remove(b"empty") + assert b"empty" not in lsh.keys + for band_hash, hashtable in zip(band_hashes, lsh.hashtables): + assert b"empty" not in hashtable.get(band_hash) + def test_insert_non_bytes_key_raises_error(self, storage_config): """Test that inserting non-bytes keys with prepickle=False raises TypeError.""" lsh = MinHashLSH(threshold=0.5, num_perm=16, storage_config=storage_config, prepickle=False)