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..20709ae7 100644 --- a/test/test_integration.py +++ b/test/test_integration.py @@ -1,10 +1,14 @@ import os +from types import SimpleNamespace +from unittest.mock import Mock, patch import numpy as np import pytest +import datasketch.storage as storage_module 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 = { @@ -29,6 +33,55 @@ DO_TEST_CASSANDRA = os.environ.get("DO_TEST_CASSANDRA") == "true" +@pytest.mark.skipif( + not hasattr(storage_module, "CassandraClient"), + reason="cassandra-driver is unavailable", +) +def test_cassandra_schema_validation_and_set_upsert(): + client = object.__new__(storage_module.CassandraClient) + client._session = Mock(keyspace="lsh_test") + client._session.prepare.return_value = "prepared-schema-query" + client._session.execute.return_value = [ + SimpleNamespace(kind="partition_key", position=0, column_name="key"), + SimpleNamespace(kind="clustering", position=0, column_name="value"), + SimpleNamespace(kind="clustering", position=1, column_name="ts"), + SimpleNamespace(kind="regular", position=-1, column_name="ignored"), + ] + + client._validate_table_schema("lsh_table") + client._session.prepare.assert_called_once_with(client.QUERY_GET_TABLE_SCHEMA) + client._session.execute.assert_called_once_with( + "prepared-schema-query", + ("lsh_test", "lsh_table"), + ) + + client._session.execute.return_value = [ + SimpleNamespace(kind="partition_key", position=0, column_name="key"), + SimpleNamespace(kind="clustering", position=0, column_name="value"), + ] + with pytest.raises(RuntimeError, match="incompatible primary key"): + client._validate_table_schema("legacy_table") + + client._stmt_upsert = "set-upsert" + client._execute = Mock() + client.upsert(b"key", [b"same", b"same"]) + client._execute.assert_called_once_with( + [ + ("set-upsert", (b"key", b"same")), + ("set-upsert", (b"key", b"same")), + ] + ) + + session = Mock() + with ( + patch.object(storage_module.CassandraSharedSession, "get_session", return_value=session), + patch.object(storage_module.c_cluster, "MonotonicTimestampGenerator", return_value=Mock()), + patch.object(storage_module.CassandraClient, "_validate_table_schema") as validate_schema, + ): + storage_module.CassandraClient({}, b"table", 100) + validate_schema.assert_called_once_with("lsh_table") + + def _clear_redis_keys(pattern="lsh_test*"): if not DO_TEST_REDIS: return @@ -108,6 +161,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)