From 7e97a8935579f430ddae500db4df0a65072396fd Mon Sep 17 00:00:00 2001 From: Dipesh Babu Date: Sat, 18 Jul 2026 23:33:24 -0400 Subject: [PATCH 1/3] Preserve duplicate Cassandra list values --- datasketch/storage.py | 46 +++++++++++++++++++++++++++++++++------- test/test_integration.py | 20 +++++++++++++++++ 2 files changed, 58 insertions(+), 8 deletions(-) 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) From 74377da288e96e61a0dd1330120e8cd1309788d0 Mon Sep 17 00:00:00 2001 From: Dipesh Babu Date: Sun, 19 Jul 2026 00:58:31 -0400 Subject: [PATCH 2/3] Cover Cassandra schema validation --- test/test_integration.py | 43 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/test/test_integration.py b/test/test_integration.py index 6f2ce88f..c5f05e34 100644 --- a/test/test_integration.py +++ b/test/test_integration.py @@ -1,8 +1,11 @@ import os +from types import SimpleNamespace +from unittest.mock import Mock 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 @@ -30,6 +33,46 @@ 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")), + ] + ) + + def _clear_redis_keys(pattern="lsh_test*"): if not DO_TEST_REDIS: return From df2acb62f5bd19bcc03fe4244cd8d3f8155823ef Mon Sep 17 00:00:00 2001 From: Dipesh Babu Date: Sun, 19 Jul 2026 01:10:19 -0400 Subject: [PATCH 3/3] Cover Cassandra schema validation wiring --- test/test_integration.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/test/test_integration.py b/test/test_integration.py index c5f05e34..20709ae7 100644 --- a/test/test_integration.py +++ b/test/test_integration.py @@ -1,6 +1,6 @@ import os from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch import numpy as np import pytest @@ -72,6 +72,15 @@ def test_cassandra_schema_validation_and_set_upsert(): ] ) + 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: