Preserve duplicate Cassandra list values

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:

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)