From 0683a31e30ce21bbd5a5df78c808cf46dd59ac72 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Thu, 18 Sep 2025 10:33:32 +0100 Subject: [PATCH] Cassandra multi-table for performance --- .../test_triples_cassandra_query.py | 201 +++++++++++- .../test_triples_cassandra_storage.py | 97 +++++- .../trustgraph/direct/cassandra_kg.py | 291 +++++++++++++++--- 3 files changed, 541 insertions(+), 48 deletions(-) diff --git a/tests/unit/test_query/test_triples_cassandra_query.py b/tests/unit/test_query/test_triples_cassandra_query.py index 72871456..f5be4961 100644 --- a/tests/unit/test_query/test_triples_cassandra_query.py +++ b/tests/unit/test_query/test_triples_cassandra_query.py @@ -534,4 +534,203 @@ class TestCassandraQueryProcessor: assert len(result) == 2 assert result[0].o.value == 'object1' - assert result[1].o.value == 'object2' \ No newline at end of file + assert result[1].o.value == 'object2' + + +class TestCassandraQueryPerformanceOptimizations: + """Test cases for multi-table performance optimizations in query service""" + + @pytest.mark.asyncio + @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') + async def test_get_po_query_optimization(self, mock_trustgraph): + """Test that get_po queries use optimized table (no ALLOW FILTERING)""" + from trustgraph.schema import TriplesQueryRequest, Value + + mock_tg_instance = MagicMock() + mock_trustgraph.return_value = mock_tg_instance + + mock_result = MagicMock() + mock_result.s = 'result_subject' + mock_tg_instance.get_po.return_value = [mock_result] + + processor = Processor(taskgroup=MagicMock()) + + # PO query pattern (predicate + object, find subjects) + query = TriplesQueryRequest( + user='test_user', + collection='test_collection', + s=None, + p=Value(value='test_predicate', is_uri=False), + o=Value(value='test_object', is_uri=False), + limit=50 + ) + + result = await processor.query_triples(query) + + # Verify get_po was called (should use optimized po_table) + mock_tg_instance.get_po.assert_called_once_with( + 'test_collection', 'test_predicate', 'test_object', limit=50 + ) + + assert len(result) == 1 + assert result[0].s.value == 'result_subject' + assert result[0].p.value == 'test_predicate' + assert result[0].o.value == 'test_object' + + @pytest.mark.asyncio + @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') + async def test_get_os_query_optimization(self, mock_trustgraph): + """Test that get_os queries use optimized table (no ALLOW FILTERING)""" + from trustgraph.schema import TriplesQueryRequest, Value + + mock_tg_instance = MagicMock() + mock_trustgraph.return_value = mock_tg_instance + + mock_result = MagicMock() + mock_result.p = 'result_predicate' + mock_tg_instance.get_os.return_value = [mock_result] + + processor = Processor(taskgroup=MagicMock()) + + # OS query pattern (object + subject, find predicates) + query = TriplesQueryRequest( + user='test_user', + collection='test_collection', + s=Value(value='test_subject', is_uri=False), + p=None, + o=Value(value='test_object', is_uri=False), + limit=25 + ) + + result = await processor.query_triples(query) + + # Verify get_os was called (should use optimized subject_table with clustering) + mock_tg_instance.get_os.assert_called_once_with( + 'test_collection', 'test_object', 'test_subject', limit=25 + ) + + assert len(result) == 1 + assert result[0].s.value == 'test_subject' + assert result[0].p.value == 'result_predicate' + assert result[0].o.value == 'test_object' + + @pytest.mark.asyncio + @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') + async def test_all_query_patterns_use_correct_tables(self, mock_trustgraph): + """Test that all query patterns route to their optimal tables""" + from trustgraph.schema import TriplesQueryRequest, Value + + mock_tg_instance = MagicMock() + mock_trustgraph.return_value = mock_tg_instance + + # Mock empty results for all queries + mock_tg_instance.get_all.return_value = [] + mock_tg_instance.get_s.return_value = [] + mock_tg_instance.get_p.return_value = [] + mock_tg_instance.get_o.return_value = [] + mock_tg_instance.get_sp.return_value = [] + mock_tg_instance.get_po.return_value = [] + mock_tg_instance.get_os.return_value = [] + mock_tg_instance.get_spo.return_value = [] + + processor = Processor(taskgroup=MagicMock()) + + # Test each query pattern + test_patterns = [ + # (s, p, o, expected_method) + (None, None, None, 'get_all'), # All triples + ('s1', None, None, 'get_s'), # Subject only + (None, 'p1', None, 'get_p'), # Predicate only + (None, None, 'o1', 'get_o'), # Object only + ('s1', 'p1', None, 'get_sp'), # Subject + Predicate + (None, 'p1', 'o1', 'get_po'), # Predicate + Object (CRITICAL OPTIMIZATION) + ('s1', None, 'o1', 'get_os'), # Object + Subject + ('s1', 'p1', 'o1', 'get_spo'), # All three + ] + + for s, p, o, expected_method in test_patterns: + # Reset mock call counts + mock_tg_instance.reset_mock() + + query = TriplesQueryRequest( + user='test_user', + collection='test_collection', + s=Value(value=s, is_uri=False) if s else None, + p=Value(value=p, is_uri=False) if p else None, + o=Value(value=o, is_uri=False) if o else None, + limit=10 + ) + + await processor.query_triples(query) + + # Verify the correct method was called + method = getattr(mock_tg_instance, expected_method) + assert method.called, f"Expected {expected_method} to be called for pattern s={s}, p={p}, o={o}" + + def test_legacy_vs_optimized_mode_configuration(self): + """Test that environment variable controls query optimization mode""" + taskgroup_mock = MagicMock() + + # Test optimized mode (default) + with patch.dict('os.environ', {}, clear=True): + processor = Processor(taskgroup=taskgroup_mock) + # Mode is determined in KnowledgeGraph initialization + + # Test legacy mode + with patch.dict('os.environ', {'CASSANDRA_USE_LEGACY': 'true'}): + processor = Processor(taskgroup=taskgroup_mock) + # Mode is determined in KnowledgeGraph initialization + + # Test explicit optimized mode + with patch.dict('os.environ', {'CASSANDRA_USE_LEGACY': 'false'}): + processor = Processor(taskgroup=taskgroup_mock) + # Mode is determined in KnowledgeGraph initialization + + @pytest.mark.asyncio + @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') + async def test_performance_critical_po_query_no_filtering(self, mock_trustgraph): + """Test the performance-critical PO query that eliminates ALLOW FILTERING""" + from trustgraph.schema import TriplesQueryRequest, Value + + mock_tg_instance = MagicMock() + mock_trustgraph.return_value = mock_tg_instance + + # Mock multiple subjects for the same predicate-object pair + mock_results = [] + for i in range(5): + mock_result = MagicMock() + mock_result.s = f'subject_{i}' + mock_results.append(mock_result) + + mock_tg_instance.get_po.return_value = mock_results + + processor = Processor(taskgroup=MagicMock()) + + # This is the query pattern that was slow with ALLOW FILTERING + query = TriplesQueryRequest( + user='large_dataset_user', + collection='massive_collection', + s=None, + p=Value(value='http://www.w3.org/1999/02/22-rdf-syntax-ns#type', is_uri=True), + o=Value(value='http://example.com/Person', is_uri=True), + limit=1000 + ) + + result = await processor.query_triples(query) + + # Verify optimized get_po was used (no ALLOW FILTERING needed!) + mock_tg_instance.get_po.assert_called_once_with( + 'massive_collection', + 'http://www.w3.org/1999/02/22-rdf-syntax-ns#type', + 'http://example.com/Person', + limit=1000 + ) + + # Verify all results were returned + assert len(result) == 5 + for i, triple in enumerate(result): + assert triple.s.value == f'subject_{i}' + assert triple.p.value == 'http://www.w3.org/1999/02/22-rdf-syntax-ns#type' + assert triple.p.is_uri is True + assert triple.o.value == 'http://example.com/Person' + assert triple.o.is_uri is True \ No newline at end of file diff --git a/tests/unit/test_storage/test_triples_cassandra_storage.py b/tests/unit/test_storage/test_triples_cassandra_storage.py index a6a6a539..54ea1a95 100644 --- a/tests/unit/test_storage/test_triples_cassandra_storage.py +++ b/tests/unit/test_storage/test_triples_cassandra_storage.py @@ -415,4 +415,99 @@ class TestCassandraStorageProcessor: # Table should remain unchanged since self.table = table happens after try/except assert processor.table == ('old_user', 'old_collection') # TrustGraph should be set to None though - assert processor.tg is None \ No newline at end of file + assert processor.tg is None + + +class TestCassandraPerformanceOptimizations: + """Test cases for multi-table performance optimizations""" + + @pytest.mark.asyncio + @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') + async def test_legacy_mode_uses_single_table(self, mock_trustgraph): + """Test that legacy mode still works with single table""" + taskgroup_mock = MagicMock() + mock_tg_instance = MagicMock() + mock_trustgraph.return_value = mock_tg_instance + + with patch.dict('os.environ', {'CASSANDRA_USE_LEGACY': 'true'}): + processor = Processor(taskgroup=taskgroup_mock) + + mock_message = MagicMock() + mock_message.metadata.user = 'user1' + mock_message.metadata.collection = 'collection1' + mock_message.triples = [] + + await processor.store_triples(mock_message) + + # Verify KnowledgeGraph instance uses legacy mode + kg_instance = mock_trustgraph.return_value + assert kg_instance is not None + + @pytest.mark.asyncio + @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') + async def test_optimized_mode_uses_multi_table(self, mock_trustgraph): + """Test that optimized mode uses multi-table schema""" + taskgroup_mock = MagicMock() + mock_tg_instance = MagicMock() + mock_trustgraph.return_value = mock_tg_instance + + with patch.dict('os.environ', {'CASSANDRA_USE_LEGACY': 'false'}): + processor = Processor(taskgroup=taskgroup_mock) + + mock_message = MagicMock() + mock_message.metadata.user = 'user1' + mock_message.metadata.collection = 'collection1' + mock_message.triples = [] + + await processor.store_triples(mock_message) + + # Verify KnowledgeGraph instance is in optimized mode + kg_instance = mock_trustgraph.return_value + assert kg_instance is not None + + @pytest.mark.asyncio + @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') + async def test_batch_write_consistency(self, mock_trustgraph): + """Test that all tables stay consistent during batch writes""" + taskgroup_mock = MagicMock() + mock_tg_instance = MagicMock() + mock_trustgraph.return_value = mock_tg_instance + + processor = Processor(taskgroup=taskgroup_mock) + + # Create test triple + triple = MagicMock() + triple.s.value = 'test_subject' + triple.p.value = 'test_predicate' + triple.o.value = 'test_object' + + mock_message = MagicMock() + mock_message.metadata.user = 'user1' + mock_message.metadata.collection = 'collection1' + mock_message.triples = [triple] + + await processor.store_triples(mock_message) + + # Verify insert was called for the triple (implementation details tested in KnowledgeGraph) + mock_tg_instance.insert.assert_called_once_with( + 'collection1', 'test_subject', 'test_predicate', 'test_object' + ) + + def test_environment_variable_controls_mode(self): + """Test that CASSANDRA_USE_LEGACY environment variable controls operation mode""" + taskgroup_mock = MagicMock() + + # Test legacy mode + with patch.dict('os.environ', {'CASSANDRA_USE_LEGACY': 'true'}): + processor = Processor(taskgroup=taskgroup_mock) + # Mode is determined in KnowledgeGraph initialization + + # Test optimized mode + with patch.dict('os.environ', {'CASSANDRA_USE_LEGACY': 'false'}): + processor = Processor(taskgroup=taskgroup_mock) + # Mode is determined in KnowledgeGraph initialization + + # Test default mode (optimized when env var not set) + with patch.dict('os.environ', {}, clear=True): + processor = Processor(taskgroup=taskgroup_mock) + # Mode is determined in KnowledgeGraph initialization \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/direct/cassandra_kg.py b/trustgraph-flow/trustgraph/direct/cassandra_kg.py index 93e41230..a4cf12b4 100644 --- a/trustgraph-flow/trustgraph/direct/cassandra_kg.py +++ b/trustgraph-flow/trustgraph/direct/cassandra_kg.py @@ -1,11 +1,16 @@ from cassandra.cluster import Cluster from cassandra.auth import PlainTextAuthProvider +from cassandra.query import BatchStatement, SimpleStatement from ssl import SSLContext, PROTOCOL_TLSv1_2 +import os +import logging # Global list to track clusters for cleanup _active_clusters = [] +logger = logging.getLogger(__name__) + class KnowledgeGraph: def __init__( @@ -17,9 +22,19 @@ class KnowledgeGraph: hosts = ["localhost"] self.keyspace = keyspace - self.table = "triples" # Fixed table name for unified schema self.username = username - + + # Multi-table schema design for optimal performance + self.use_legacy = os.getenv('CASSANDRA_USE_LEGACY', 'false').lower() == 'true' + + if self.use_legacy: + self.table = "triples" # Legacy single table + else: + # New optimized tables + self.subject_table = "triples_by_subject" + self.po_table = "triples_by_po" + self.object_table = "triples_by_object" + if username and password: ssl_context = SSLContext(PROTOCOL_TLSv1_2) auth_provider = PlainTextAuthProvider(username=username, password=password) @@ -27,12 +42,15 @@ class KnowledgeGraph: else: self.cluster = Cluster(hosts) self.session = self.cluster.connect() - + # Track this cluster globally _active_clusters.append(self.cluster) self.init() + if not self.use_legacy: + self.prepare_statements() + def clear(self): self.session.execute(f""" @@ -45,14 +63,21 @@ class KnowledgeGraph: self.session.execute(f""" create keyspace if not exists {self.keyspace} - with replication = {{ - 'class' : 'SimpleStrategy', - 'replication_factor' : 1 + with replication = {{ + 'class' : 'SimpleStrategy', + 'replication_factor' : 1 }}; """); self.session.set_keyspace(self.keyspace) + if self.use_legacy: + self.init_legacy_schema() + else: + self.init_optimized_schema() + + def init_legacy_schema(self): + """Initialize legacy single-table schema for backward compatibility""" self.session.execute(f""" create table if not exists {self.table} ( collection text, @@ -78,67 +103,241 @@ class KnowledgeGraph: ON {self.table} (o); """); + def init_optimized_schema(self): + """Initialize optimized multi-table schema for performance""" + # Table 1: Subject-centric queries (get_s, get_sp, get_spo, get_os) + self.session.execute(f""" + CREATE TABLE IF NOT EXISTS {self.subject_table} ( + collection text, + s text, + p text, + o text, + PRIMARY KEY ((collection, s), p, o) + ); + """); + + # Table 2: Predicate-Object queries (get_p, get_po) - eliminates ALLOW FILTERING! + self.session.execute(f""" + CREATE TABLE IF NOT EXISTS {self.po_table} ( + collection text, + p text, + o text, + s text, + PRIMARY KEY ((collection, p), o, s) + ); + """); + + # Table 3: Object-centric queries (get_o) + self.session.execute(f""" + CREATE TABLE IF NOT EXISTS {self.object_table} ( + collection text, + o text, + s text, + p text, + PRIMARY KEY ((collection, o), s, p) + ); + """); + + logger.info("Optimized multi-table schema initialized") + + def prepare_statements(self): + """Prepare statements for optimal performance""" + # Insert statements for batch operations + self.insert_subject_stmt = self.session.prepare( + f"INSERT INTO {self.subject_table} (collection, s, p, o) VALUES (?, ?, ?, ?)" + ) + + self.insert_po_stmt = self.session.prepare( + f"INSERT INTO {self.po_table} (collection, p, o, s) VALUES (?, ?, ?, ?)" + ) + + self.insert_object_stmt = self.session.prepare( + f"INSERT INTO {self.object_table} (collection, o, s, p) VALUES (?, ?, ?, ?)" + ) + + # Query statements for optimized access + self.get_all_stmt = self.session.prepare( + f"SELECT s, p, o FROM {self.subject_table} WHERE collection = ? LIMIT ?" + ) + + self.get_s_stmt = self.session.prepare( + f"SELECT p, o FROM {self.subject_table} WHERE collection = ? AND s = ? LIMIT ?" + ) + + self.get_p_stmt = self.session.prepare( + f"SELECT s, o FROM {self.po_table} WHERE collection = ? AND p = ? LIMIT ?" + ) + + self.get_o_stmt = self.session.prepare( + f"SELECT s, p FROM {self.object_table} WHERE collection = ? AND o = ? LIMIT ?" + ) + + self.get_sp_stmt = self.session.prepare( + f"SELECT o FROM {self.subject_table} WHERE collection = ? AND s = ? AND p = ? LIMIT ?" + ) + + # The critical optimization: get_po without ALLOW FILTERING! + self.get_po_stmt = self.session.prepare( + f"SELECT s FROM {self.po_table} WHERE collection = ? AND p = ? AND o = ? LIMIT ?" + ) + + self.get_os_stmt = self.session.prepare( + f"SELECT p FROM {self.subject_table} WHERE collection = ? AND s = ? AND o = ? LIMIT ?" + ) + + self.get_spo_stmt = self.session.prepare( + f"SELECT s as x FROM {self.subject_table} WHERE collection = ? AND s = ? AND p = ? AND o = ? LIMIT ?" + ) + + logger.info("Prepared statements initialized for optimal performance") + def insert(self, collection, s, p, o): - self.session.execute( - f"insert into {self.table} (collection, s, p, o) values (%s, %s, %s, %s)", - (collection, s, p, o) - ) + if self.use_legacy: + self.session.execute( + f"insert into {self.table} (collection, s, p, o) values (%s, %s, %s, %s)", + (collection, s, p, o) + ) + else: + # Batch write to all three tables for consistency + batch = BatchStatement() + + # Insert into subject table + batch.add(self.insert_subject_stmt, (collection, s, p, o)) + + # Insert into predicate-object table (column order: collection, p, o, s) + batch.add(self.insert_po_stmt, (collection, p, o, s)) + + # Insert into object table (column order: collection, o, s, p) + batch.add(self.insert_object_stmt, (collection, o, s, p)) + + self.session.execute(batch) def get_all(self, collection, limit=50): - return self.session.execute( - f"select s, p, o from {self.table} where collection = %s limit {limit}", - (collection,) - ) + if self.use_legacy: + return self.session.execute( + f"select s, p, o from {self.table} where collection = %s limit {limit}", + (collection,) + ) + else: + # Use subject table for get_all queries + return self.session.execute( + self.get_all_stmt, + (collection, limit) + ) def get_s(self, collection, s, limit=10): - return self.session.execute( - f"select p, o from {self.table} where collection = %s and s = %s limit {limit}", - (collection, s) - ) + if self.use_legacy: + return self.session.execute( + f"select p, o from {self.table} where collection = %s and s = %s limit {limit}", + (collection, s) + ) + else: + # Optimized: Direct partition access with (collection, s) + return self.session.execute( + self.get_s_stmt, + (collection, s, limit) + ) def get_p(self, collection, p, limit=10): - return self.session.execute( - f"select s, o from {self.table} where collection = %s and p = %s limit {limit}", - (collection, p) - ) + if self.use_legacy: + return self.session.execute( + f"select s, o from {self.table} where collection = %s and p = %s limit {limit}", + (collection, p) + ) + else: + # Optimized: Use po_table for direct partition access + return self.session.execute( + self.get_p_stmt, + (collection, p, limit) + ) def get_o(self, collection, o, limit=10): - return self.session.execute( - f"select s, p from {self.table} where collection = %s and o = %s limit {limit}", - (collection, o) - ) + if self.use_legacy: + return self.session.execute( + f"select s, p from {self.table} where collection = %s and o = %s limit {limit}", + (collection, o) + ) + else: + # Optimized: Use object_table for direct partition access + return self.session.execute( + self.get_o_stmt, + (collection, o, limit) + ) def get_sp(self, collection, s, p, limit=10): - return self.session.execute( - f"select o from {self.table} where collection = %s and s = %s and p = %s limit {limit}", - (collection, s, p) - ) + if self.use_legacy: + return self.session.execute( + f"select o from {self.table} where collection = %s and s = %s and p = %s limit {limit}", + (collection, s, p) + ) + else: + # Optimized: Use subject_table with clustering key access + return self.session.execute( + self.get_sp_stmt, + (collection, s, p, limit) + ) def get_po(self, collection, p, o, limit=10): - return self.session.execute( - f"select s from {self.table} where collection = %s and p = %s and o = %s limit {limit} allow filtering", - (collection, p, o) - ) + if self.use_legacy: + return self.session.execute( + f"select s from {self.table} where collection = %s and p = %s and o = %s limit {limit} allow filtering", + (collection, p, o) + ) + else: + # CRITICAL OPTIMIZATION: Use po_table - NO MORE ALLOW FILTERING! + return self.session.execute( + self.get_po_stmt, + (collection, p, o, limit) + ) def get_os(self, collection, o, s, limit=10): - return self.session.execute( - f"select p from {self.table} where collection = %s and o = %s and s = %s limit {limit} allow filtering", - (collection, o, s) - ) + if self.use_legacy: + return self.session.execute( + f"select p from {self.table} where collection = %s and o = %s and s = %s limit {limit} allow filtering", + (collection, o, s) + ) + else: + # Optimized: Use subject_table with clustering access (no more ALLOW FILTERING) + return self.session.execute( + self.get_os_stmt, + (collection, s, o, limit) + ) def get_spo(self, collection, s, p, o, limit=10): - return self.session.execute( - f"""select s as x from {self.table} where collection = %s and s = %s and p = %s and o = %s limit {limit}""", - (collection, s, p, o) - ) + if self.use_legacy: + return self.session.execute( + f"""select s as x from {self.table} where collection = %s and s = %s and p = %s and o = %s limit {limit}""", + (collection, s, p, o) + ) + else: + # Optimized: Use subject_table for exact key lookup + return self.session.execute( + self.get_spo_stmt, + (collection, s, p, o, limit) + ) def delete_collection(self, collection): """Delete all triples for a specific collection""" - self.session.execute( - f"delete from {self.table} where collection = %s", - (collection,) - ) + if self.use_legacy: + self.session.execute( + f"delete from {self.table} where collection = %s", + (collection,) + ) + else: + # Delete from all three tables + self.session.execute( + f"delete from {self.subject_table} where collection = %s", + (collection,) + ) + self.session.execute( + f"delete from {self.po_table} where collection = %s", + (collection,) + ) + self.session.execute( + f"delete from {self.object_table} where collection = %s", + (collection,) + ) def close(self): """Close the Cassandra session and cluster connections properly"""