[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

engine/clients/cassandra/README.md

@@ -9,9 +9,9 @@ Run the following command to start the server (alternatively, run `docker compos
 $ docker compose up
 
 [+] Running 1/1
- ✔ cassandra Pulled                                                                                                                                                                           1.4s 
+ ✔ cassandra Pulled                                                                                                                                                                           1.4s
 [+] Running 1/1
- ✔ Container cassandra-benchmark  Recreated                                                                                                                                                   0.1s 
+ ✔ Container cassandra-benchmark  Recreated                                                                                                                                                   0.1s
 Attaching to cassandra-benchmark
 cassandra-benchmark  | CompileCommand: dontinline org/apache/cassandra/db/Columns$Serializer.deserializeLargeSubset(Lorg/apache/cassandra/io/util/DataInputPlus;Lorg/apache/cassandra/db/Columns;I)Lorg/apache/cassandra/db/Columns; bool dontinline = true
 ...
@@ -29,7 +29,7 @@ cassandra-benchmark  | INFO  [main] 2025-04-04 22:28:25,091 StorageService.java:
 ### Start up the client benchmark
 Run the following command to start the client benchmark using `glove-25-angular` dataset as an example:
 ```bash
-% python3 -m run --engines cassandra-single-node --datasets glove-25-angular 
+% python3 -m run --engines cassandra-single-node --datasets glove-25-angular
 ```
 and you'll see the following output:
 ```bash

engine/clients/cassandra/__init__.py

@@ -2,8 +2,4 @@
 from engine.clients.cassandra.search import CassandraSearcher
 from engine.clients.cassandra.upload import CassandraUploader
 
-__all__ = [
-    "CassandraConfigurator",
-    "CassandraSearcher",
-    "CassandraUploader"
-]
+__all__ = ["CassandraConfigurator", "CassandraSearcher", "CassandraUploader"]

engine/clients/cassandra/config.py

@@ -4,4 +4,4 @@
 CASSANDRA_TABLE = os.getenv("CASSANDRA_TABLE", "vectors")
 ASTRA_API_ENDPOINT = os.getenv("ASTRA_API_ENDPOINT", None)
 ASTRA_API_KEY = os.getenv("ASTRA_API_KEY", None)
-ASTRA_SCB_PATH = os.getenv("ASTRA_SCB_PATH", None)
+ASTRA_SCB_PATH = os.getenv("ASTRA_SCB_PATH", None)

engine/clients/cassandra/configure.py

@@ -1,6 +1,10 @@
-from cassandra.cluster import Cluster, ExecutionProfile, EXEC_PROFILE_DEFAULT
-from cassandra.policies import DCAwareRoundRobinPolicy, TokenAwarePolicy, ExponentialReconnectionPolicy
 from cassandra import ConsistencyLevel, ProtocolVersion
+from cassandra.cluster import EXEC_PROFILE_DEFAULT, Cluster, ExecutionProfile
+from cassandra.policies import (
+    DCAwareRoundRobinPolicy,
+    ExponentialReconnectionPolicy,
+    TokenAwarePolicy,
+)
 
 from benchmark.dataset import Dataset
 from engine.base_client.configure import BaseConfigurator
@@ -13,26 +17,28 @@ class CassandraConfigurator(BaseConfigurator):
     DISTANCE_MAPPING = {
         Distance.L2: "euclidean",
         Distance.COSINE: "cosine",
-        Distance.DOT: "dot_product"
+        Distance.DOT: "dot_product",
     }
 
     def __init__(self, host, collection_params: dict, connection_params: dict):
         super().__init__(host, collection_params, connection_params)
-        
+
         # Set up execution profiles for consistency and performance
         profile = ExecutionProfile(
             load_balancing_policy=TokenAwarePolicy(DCAwareRoundRobinPolicy()),
             consistency_level=ConsistencyLevel.LOCAL_QUORUM,
-            request_timeout=60
+            request_timeout=60,
         )
-        
+
         # Initialize Cassandra cluster connection
         self.cluster = Cluster(
             contact_points=[host],
             execution_profiles={EXEC_PROFILE_DEFAULT: profile},
             protocol_version=ProtocolVersion.V4,
-            reconnection_policy=ExponentialReconnectionPolicy(base_delay=1, max_delay=60),
-            **connection_params
+            reconnection_policy=ExponentialReconnectionPolicy(
+                base_delay=1, max_delay=60
+            ),
+            **connection_params,
         )
         self.session = self.cluster.connect()
 
@@ -44,17 +50,17 @@ def recreate(self, dataset: Dataset, collection_params):
         """Create keyspace and table for vector search"""
         # Create keyspace if not exists
         self.session.execute(
-            f"""CREATE KEYSPACE IF NOT EXISTS {CASSANDRA_KEYSPACE} 
+            f"""CREATE KEYSPACE IF NOT EXISTS {CASSANDRA_KEYSPACE}
             WITH REPLICATION = {{ 'class': 'SimpleStrategy', 'replication_factor': 1 }}"""
         )
-        
+
         # Use the keyspace
         self.session.execute(f"USE {CASSANDRA_KEYSPACE}")
-        
+
         # Get the distance metric
         distance_metric = self.DISTANCE_MAPPING.get(dataset.config.distance)
         vector_size = dataset.config.vector_size
-        
+
         # Create vector table
         # Using a simple schema that supports vector similarity search
         self.session.execute(
@@ -64,14 +70,14 @@ def recreate(self, dataset: Dataset, collection_params):
                 metadata map<text, text>
             )"""
         )
-        
+
         # Create vector index using the appropriate distance metric
         self.session.execute(
-            f"""CREATE CUSTOM INDEX IF NOT EXISTS vector_index ON {CASSANDRA_TABLE}(embedding) 
-            USING 'StorageAttachedIndex' 
+            f"""CREATE CUSTOM INDEX IF NOT EXISTS vector_index ON {CASSANDRA_TABLE}(embedding)
+            USING 'StorageAttachedIndex'
             WITH OPTIONS = {{ 'similarity_function': '{distance_metric}' }}"""
         )
-        
+
         # Add additional schema fields based on collection_params if needed
         for field_name, field_type in dataset.config.schema.items():
             if field_type in ["keyword", "text"]:
@@ -81,7 +87,7 @@ def recreate(self, dataset: Dataset, collection_params):
                 # For numeric fields that need separate indexing
                 # In a real implementation, we might alter the table to add these columns
                 pass
-        
+
         return collection_params
 
     def execution_params(self, distance, vector_size) -> dict:
@@ -90,7 +96,7 @@ def execution_params(self, distance, vector_size) -> dict:
 
     def delete_client(self):
         """Close the Cassandra connection"""
-        if hasattr(self, 'session') and self.session:
+        if hasattr(self, "session") and self.session:
             self.session.shutdown()
-        if hasattr(self, 'cluster') and self.cluster:
-            self.cluster.shutdown()
+        if hasattr(self, "cluster") and self.cluster:
+            self.cluster.shutdown()

engine/clients/cassandra/parser.py

@@ -11,23 +11,23 @@ def build_condition(
         Build a CQL condition expression that combines AND and OR subfilters
         """
         conditions = []
-        
+
         # Add AND conditions
         if and_subfilters and len(and_subfilters) > 0:
             and_conds = " AND ".join([f"({cond})" for cond in and_subfilters if cond])
             if and_conds:
                 conditions.append(f"({and_conds})")
-        
+
         # Add OR conditions
         if or_subfilters and len(or_subfilters) > 0:
             or_conds = " OR ".join([f"({cond})" for cond in or_subfilters if cond])
             if or_conds:
                 conditions.append(f"({or_conds})")
-        
+
         # Combine all conditions
         if not conditions:
             return None
-        
+
         return " AND ".join(conditions)
 
     def build_exact_match_filter(self, field_name: str, value: FieldValue) -> Any:
@@ -52,31 +52,31 @@ def build_range_filter(
         Build a CQL range filter condition
         """
         conditions = []
-        
+
         if lt is not None:
             if isinstance(lt, str):
                 conditions.append(f"metadata['{field_name}'] < '{lt}'")
             else:
                 conditions.append(f"metadata['{field_name}'] < '{str(lt)}'")
-                
+
         if gt is not None:
             if isinstance(gt, str):
                 conditions.append(f"metadata['{field_name}'] > '{gt}'")
             else:
                 conditions.append(f"metadata['{field_name}'] > '{str(gt)}'")
-                
+
         if lte is not None:
             if isinstance(lte, str):
                 conditions.append(f"metadata['{field_name}'] <= '{lte}'")
             else:
                 conditions.append(f"metadata['{field_name}'] <= '{str(lte)}'")
-                
+
         if gte is not None:
             if isinstance(gte, str):
                 conditions.append(f"metadata['{field_name}'] >= '{gte}'")
             else:
                 conditions.append(f"metadata['{field_name}'] >= '{str(gte)}'")
-        
+
         return " AND ".join(conditions)
 
     def build_geo_filter(
@@ -89,4 +89,4 @@ def build_geo_filter(
         """
         # In a real implementation with a geo extension, we'd implement proper geo filtering
         # For this benchmark, we'll return a placeholder condition that doesn't filter
-        return "1=1"  # Always true condition as a placeholder
+        return "1=1"  # Always true condition as a placeholder

engine/clients/cassandra/search.py

@@ -1,9 +1,13 @@
 import multiprocessing as mp
 from typing import List, Tuple
 
-from cassandra.cluster import Cluster, ExecutionProfile, EXEC_PROFILE_DEFAULT
-from cassandra.policies import DCAwareRoundRobinPolicy, TokenAwarePolicy, ExponentialReconnectionPolicy
 from cassandra import ConsistencyLevel, ProtocolVersion
+from cassandra.cluster import EXEC_PROFILE_DEFAULT, Cluster, ExecutionProfile
+from cassandra.policies import (
+    DCAwareRoundRobinPolicy,
+    ExponentialReconnectionPolicy,
+    TokenAwarePolicy,
+)
 
 from dataset_reader.base_reader import Query
 from engine.base_client.distances import Distance
@@ -24,20 +28,22 @@ def init_client(cls, host, distance, connection_params: dict, search_params: dic
         profile = ExecutionProfile(
             load_balancing_policy=TokenAwarePolicy(DCAwareRoundRobinPolicy()),
             consistency_level=ConsistencyLevel.LOCAL_ONE,  # Use LOCAL_ONE for faster reads
-            request_timeout=60
+            request_timeout=60,
         )
-        
+
         # Initialize Cassandra cluster connection
         cls.cluster = Cluster(
-            contact_points=[host], 
+            contact_points=[host],
             execution_profiles={EXEC_PROFILE_DEFAULT: profile},
-            reconnection_policy=ExponentialReconnectionPolicy(base_delay=1, max_delay=60),
+            reconnection_policy=ExponentialReconnectionPolicy(
+                base_delay=1, max_delay=60
+            ),
             protocol_version=ProtocolVersion.V4,
-            **connection_params
+            **connection_params,
         )
         cls.session = cls.cluster.connect(CASSANDRA_KEYSPACE)
         cls.search_params = search_params
-        
+
         # Update prepared statements with current search parameters
         cls.update_prepared_statements(distance)
 
@@ -50,7 +56,7 @@ def update_prepared_statements(cls, distance):
         """Create prepared statements for vector searches"""
         # Prepare a vector similarity search query
         limit = cls.search_params.get("top", 10)
-        
+
         if distance == Distance.COSINE:
             SIMILARITY_FUNC = "similarity_cosine"
         elif distance == Distance.L2:
@@ -61,48 +67,49 @@ def update_prepared_statements(cls, distance):
             raise ValueError(f"Unsupported distance metric: {distance}")
 
         cls.ann_search_stmt = cls.session.prepare(
-            f"""SELECT id, {SIMILARITY_FUNC}(embedding, ?) as distance 
-            FROM {CASSANDRA_TABLE} 
+            f"""SELECT id, {SIMILARITY_FUNC}(embedding, ?) as distance
+            FROM {CASSANDRA_TABLE}
             ORDER BY embedding ANN OF ?
             LIMIT {limit}"""
         )
-        
+
         # Prepare a statement for filtered vector search
-        cls.filtered_search_query_template = (
-            f"""SELECT id, {SIMILARITY_FUNC}(embedding, ?) as distance 
-            FROM {CASSANDRA_TABLE} 
+        cls.filtered_search_query_template = f"""SELECT id, {SIMILARITY_FUNC}(embedding, ?) as distance
+            FROM {CASSANDRA_TABLE}
             WHERE {{conditions}}
             ORDER BY embedding ANN OF ?
             LIMIT {limit}"""
-        )
 
     @classmethod
     def search_one(cls, query: Query, top: int) -> List[Tuple[int, float]]:
         """Execute a vector similarity search with optional filters"""
         # Convert query vector to a format Cassandra can use
-        query_vector = query.vector.tolist() if hasattr(query.vector, 'tolist') else query.vector
-        
+        query_vector = (
+            query.vector.tolist() if hasattr(query.vector, "tolist") else query.vector
+        )
+
         # Generate filter conditions if metadata conditions exist
         filter_conditions = cls.parser.parse(query.meta_conditions)
-        
+
         try:
             if filter_conditions:
                 # Use the filtered search query
-                query_with_conditions = cls.filtered_search_query_template.format(conditions=filter_conditions)
+                query_with_conditions = cls.filtered_search_query_template.format(
+                    conditions=filter_conditions
+                )
                 results = cls.session.execute(
                     cls.session.prepare(query_with_conditions),
-                    (query_vector, query_vector)
+                    (query_vector, query_vector),
                 )
             else:
                 # Use the basic ANN search query
                 results = cls.session.execute(
-                    cls.ann_search_stmt,
-                    (query_vector, query_vector)
+                    cls.ann_search_stmt, (query_vector, query_vector)
                 )
-            
+
             # Extract and return results
             return [(row.id, row.distance) for row in results]
-            
+
         except Exception as ex:
             print(f"Error during Cassandra vector search: {ex}")
             raise ex
@@ -113,4 +120,4 @@ def delete_client(cls):
         if cls.session:
             cls.session.shutdown()
         if cls.cluster:
-            cls.cluster.shutdown()
+            cls.cluster.shutdown()