Cassandra multi-table for performance

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'
+        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

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
+        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

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 = {{ 
+                with replication = {{
-                   'class' : 'SimpleStrategy', 
+                   'class' : 'SimpleStrategy',
-                   'replication_factor' : 1 
+                   '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(
+        if self.use_legacy:
-            f"insert into {self.table} (collection, s, p, o) values (%s, %s, %s, %s)",
+            self.session.execute(
-            (collection, s, p, o)
+                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(
+        if self.use_legacy:
-            f"select s, p, o from {self.table} where collection = %s limit {limit}",
+            return self.session.execute(
-            (collection,)
+                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(
+        if self.use_legacy:
-            f"select p, o from {self.table} where collection = %s and s = %s limit {limit}",
+            return self.session.execute(
-            (collection, s)
+                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(
+        if self.use_legacy:
-            f"select s, o from {self.table} where collection = %s and p = %s limit {limit}",
+            return self.session.execute(
-            (collection, p)
+                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(
+        if self.use_legacy:
-            f"select s, p from {self.table} where collection = %s and o = %s limit {limit}",
+            return self.session.execute(
-            (collection, o)
+                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(
+        if self.use_legacy:
-            f"select o from {self.table} where collection = %s and s = %s and p = %s limit {limit}",
+            return self.session.execute(
-            (collection, s, p)
+                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(
+        if self.use_legacy:
-            f"select s from {self.table} where collection = %s and p = %s and o = %s limit {limit} allow filtering",
+            return self.session.execute(
-            (collection, p, o)
+                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(
+        if self.use_legacy:
-            f"select p from {self.table} where collection = %s and o = %s and s = %s limit {limit} allow filtering",
+            return self.session.execute(
-            (collection, o, s)
+                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(
+        if self.use_legacy:
-            f"""select s as x from {self.table} where collection = %s and s = %s and p = %s and o = %s limit {limit}""",
+            return self.session.execute(
-            (collection, s, p, o)
+                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(
+        if self.use_legacy:
-            f"delete from {self.table} where collection = %s",
+            self.session.execute(
-            (collection,)
+                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"""