adapta 3.2.3a475.dev3__tar.gz → 3.2.4a477.dev5__tar.gz

{adapta-3.2.3a475.dev3 → adapta-3.2.4a477.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: adapta
-Version: 3.2.3a475.dev3
+Version: 3.2.4a477.dev5
 Summary: Logging, data connectors, monitoring, secret handling and general lifehacks to make data people lives easier.
 Home-page: https://github.com/SneaksAndData/adapta
 License: Apache 2.0

adapta-3.2.4a477.dev5/adapta/_version.py

@@ -0,0 +1 @@
+__version__ = 'v3.2.4a477.dev5'

{adapta-3.2.3a475.dev3 → adapta-3.2.4a477.dev5}/adapta/storage/blob/local_storage_client.py

@@ -70,6 +70,8 @@ class LocalStorageClient(StorageClient):
         self, blob_path: DataPath, filter_predicate: Optional[Callable[[...], bool]] = None
     ) -> Iterator[DataPath]:
         for blob in os.listdir(cast_path(blob_path).path):
+            if filter_predicate is not None and not filter_predicate(blob):
+                continue
             yield LocalPath(path=blob)
 
     def read_blobs(
@@ -79,8 +81,14 @@ class LocalStorageClient(StorageClient):
         filter_predicate: Optional[Callable[[...], bool]] = None,
     ) -> Iterator[T]:
         dir_path = cast_path(blob_path).path
-        for blob in os.listdir(dir_path):
-            with open(os.path.join(dir_path, blob), "rb") as blob_file:
+        if os.path.isdir(dir_path):
+            for blob in os.listdir(dir_path):
+                if filter_predicate is not None and not filter_predicate(blob):
+                    continue
+                with open(os.path.join(dir_path, blob), "rb") as blob_file:
+                    yield serialization_format().deserialize(blob_file.read())
+        else:
+            with open(dir_path, "rb") as blob_file:
                 yield serialization_format().deserialize(blob_file.read())
 
     def download_blobs(

{adapta-3.2.3a475.dev3 → adapta-3.2.4a477.dev5}/adapta/storage/distributed_object_store/v3/datastax_astra/README.md

@@ -129,28 +129,33 @@ with AstraClient(
   ```
 
 ## Using the Vector Search
-Create a table in Astra and insert some rows:
+1. Create a table in Astra and insert some rows:
 ```cassandraql
 CREATE TABLE IF NOT EXISTS tmp.test_entity_with_embeddings (
     col_a TEXT PRIMARY KEY,
     col_b TEXT,
-    col_c VECTOR<FLOAT, 3>
+    col_c VECTOR<FLOAT, 3>,
+    col_d TEXT,
 );
 
 CREATE INDEX IF NOT EXISTS ann_index
   ON tmp.test_entity_with_embeddings(col_c)
   WITH OPTIONS = {'source_model': 'other'};
 
-INSERT INTO tmp.test_entity_with_embeddings (col_a, col_b, col_c)
-VALUES ('something1', 'different', [0.3, 0.4, 0.5]);
+CREATE INDEX IF NOT EXISTS col_b_index
+  ON tmp.test_entity_with_embeddings(col_b);
 
-INSERT INTO tmp.test_entity_with_embeddings (col_a, col_b, col_c)
-VALUES ('something2', 'different1', [0.1, 0.24, 0.25]);
-```
+INSERT INTO tmp.test_entity_with_embeddings (col_a, col_b, col_c, col_d)
+VALUES ('something1', 'different', [0.3, 0.4, 0.5], 'extra1');
 
+INSERT INTO tmp.test_entity_with_embeddings (col_a, col_b, col_c, col_d)
+VALUES ('something2', 'different1', [0.1, 0.24, 0.25], 'extra2');
+```
+2. Test out functionality in Python
 ```python
 from adapta.storage.distributed_object_store.v3.datastax_astra import AstraClient
 from adapta.storage.distributed_object_store.v3.datastax_astra import SimilarityFunction
+from adapta.storage.models.filter_expression import FilterField
 
 from dataclasses import dataclass, field
 
@@ -161,28 +166,61 @@ class TestEntityWithEmbeddings:
         "is_primary_key": True,
         "is_partition_key": True
     })
-    col_b: str = field(metadata={
-        "is_primary_key": True,
-        "is_partition_key": False
-    })
+    col_b: str
     col_c: list[float] = field(metadata={
         "is_vector_enabled": True
     })
+    col_d: str
+
+
+astra_client = AstraClient(
+    client_name='test',
+    keyspace='tmp',
+    secure_connect_bundle_bytes='base64string',
+    client_id='Astra Token client_id',
+    client_secret='Astra Token client_secret'
+)
+
+# Search in Astra
+with astra_client:
+    print(astra_client.ann_search(
+        entity_type=TestEntityWithEmbeddings,
+        vector_to_match=[0.1, 0.2, 0.3],
+        similarity_function=SimilarityFunction.DOT_PRODUCT
+        , num_results=2
+    ).to_pandas())
+
+    #         col_a       col_b   col_d  sim_value
+    # 0  something2  different1  extra2     0.5665
+    # 1  something1   different  extra1     0.6300
 
-
-# Apply the filters for Astra
-with AstraClient(
-        client_name='test',
-        keyspace='tmp',
-        secure_connect_bundle_bytes="base64 bundle string",
-        client_id='client id',
-        client_secret='client secret'
-) as ac:
-    # Filter expressions are compiled into specific target, in this case Astra filters, in filter_entities method
-    print(ac.ann_search(entity_type=TestEntityWithEmbeddings, vector_to_match=[0.1, 0.2, 0.3],
-                        similarity_function=SimilarityFunction.DOT_PRODUCT, num_results=2).to_pandas())
-
-    #         col_a       col_b  sim_value
-    # 0  something2  different1     0.5665
-    # 1  something1   different     0.6300
+    
+# Search with primary key filter in Astra (with dictionary)
+filter_expression = [{'col_a': 'something2', 'col_b': 'different1'}]
+with astra_client:
+    print(astra_client.ann_search(
+        entity_type=TestEntityWithEmbeddings,
+        vector_to_match=[0.1, 0.2, 0.3],
+        similarity_function=SimilarityFunction.DOT_PRODUCT,
+        num_results=2,
+        key_column_filter_values=filter_expression
+    ).to_pandas())
+    
+    #         col_a       col_b   col_d  sim_value
+    # 0  something2  different1  extra2     0.5665
+
+
+# Search with primary key filter in Astra (with Expression)
+filter_expression = (FilterField('col_a') == 'something2') & (FilterField('col_b').isin(['different1', 'doesnt_exist']))
+with astra_client:
+    print(astra_client.ann_search(
+        entity_type=TestEntityWithEmbeddings,
+        vector_to_match=[0.1, 0.2, 0.3],
+        similarity_function=SimilarityFunction.DOT_PRODUCT,
+        num_results=2,
+        key_column_filter_values=filter_expression,
+    ).to_pandas())
+
+    #         col_a       col_b   col_d  sim_value
+    # 0  something2  different1  extra2     0.5665
   ```

adapta-3.2.4a477.dev5/adapta/storage/distributed_object_store/v3/datastax_astra/_models.py

@@ -0,0 +1,134 @@
+"""
+ Models for Astra DB.
+"""
+#  Copyright (c) 2023-2024. ECCO Sneaks & Data
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+import re
+
+from enum import Enum
+from typing import Optional, Union, List, Dict, Any
+
+from adapta.storage.models.filter_expression import (
+    Expression,
+    compile_expression,
+    AstraFilterExpression,
+    FilterExpressionOperation,
+)
+
+
+class SimilarityFunction(Enum):
+    """
+    Supported similarity functions.
+    """
+
+    COSINE = "similarity_cosine"
+    DOT_PRODUCT = "similarity_dot_product"
+    EUCLIDEAN = "similarity_euclidean"
+
+
+class VectorSearchQuery:
+    """
+    Builder for a vector search query. Result would resemble the following:
+
+    select description, similarity_cosine(item_vector, [0.1, 0.15, 0.3, 0.12, 0.05])
+    from vsearch.products
+
+    order by item_vector ann of [0.1, 0.15, 0.3, 0.12, 0.05]
+    limit 1;
+
+    """
+
+    def __init__(
+        self,
+        table_fqn: str,
+        data_fields: list[str],
+        sim_func: SimilarityFunction,
+        vector: list[float],
+        field_name: str,
+        num_results=1,
+        key_column_filter_values: Optional[Union[Expression, List[Dict[str, Any]]]] = None,
+    ):
+        self._sim_func = sim_func
+        self._vector = vector
+        self._field_name = field_name
+        self._num_results = num_results
+        self._table_fqn = table_fqn
+        self._data_fields = data_fields
+        self._key_column_filter_values = key_column_filter_values
+
+    def _get_similarity_colum(self) -> str:
+        return f"{self._sim_func.value}({self._field_name}, {self._vector})"
+
+    def _get_order_by(self) -> str:
+        return f"order by {self._field_name} ann of {self._vector} limit {self._num_results};"
+
+    def _get_filter(self) -> str:
+        if self._key_column_filter_values is None:
+            return ""
+
+        def format_value_for_cql(val: Any) -> str:
+            if isinstance(val, (tuple, list, set)):
+                return f"({', '.join(format_value_for_cql(v) for v in val)})"
+            if isinstance(val, str):
+                return f"'{val}'"
+            return str(val)
+
+        def get_astra_operator(column_expression: str) -> str:
+            operator_map = {
+                FilterExpressionOperation.GT.value["astra"]: ">",
+                FilterExpressionOperation.GE.value["astra"]: ">=",
+                FilterExpressionOperation.LT.value["astra"]: "<",
+                FilterExpressionOperation.LE.value["astra"]: "<=",
+                FilterExpressionOperation.IN.value["astra"]: "IN",
+            }
+
+            for suffix, operator in operator_map.items():
+                if column_expression.endswith(suffix):
+                    return operator
+
+            return "="
+
+        def remove_operator_suffix(col: str) -> str:
+            """
+            Removes the double underscore and the subsequent operator suffix from a column expression.
+            Returns only the column name.
+            """
+            return re.sub(r"__\w+$", "", col)
+
+        compiled_filter_values = (
+            compile_expression(self._key_column_filter_values, AstraFilterExpression)
+            if isinstance(self._key_column_filter_values, Expression)
+            else self._key_column_filter_values
+        )
+        if len(compiled_filter_values) > 1:
+            raise ValueError("Restriction on key columns must not be nested under OR operator")
+
+        cql_filter_expressions = [
+            f"{remove_operator_suffix(col)} {get_astra_operator(col)} {format_value_for_cql(val)}"
+            for col, val in compiled_filter_values[0].items()
+        ]
+        return f"where {' and '.join(cql_filter_expressions)}"
+
+    def __str__(self):
+        return " ".join(
+            [
+                "select",
+                ", ".join(self._data_fields),
+                ", " f"{self._get_similarity_colum()} as sim_value",
+                f"from {self._table_fqn}",
+                self._get_filter(),
+                self._get_order_by(),
+            ]
+        )

{adapta-3.2.3a475.dev3 → adapta-3.2.4a477.dev5}/adapta/storage/distributed_object_store/v3/datastax_astra/astra_client.py

@@ -300,17 +300,16 @@ class AstraClient:
             self._session is not None
         ), "Please instantiate an AstraClient using with AstraClient(...) before calling this method"
 
-        mapper = get_mapper(
+        select_columns = list(map(normalize_column_name, select_columns)) if select_columns else None
+
+        cassandra_model = get_mapper(
             data_model=model_class,
             keyspace=keyspace,
             table_name=table_name,
             primary_keys=primary_keys,
             partition_keys=partition_keys,
             custom_indexes=custom_indexes,
-        )
-        cassandra_model = mapper.map()
-
-        select_columns = list(map(normalize_column_name, select_columns)) if select_columns else mapper.column_names
+        ).map()
 
         compiled_filter_values = (
             compile_expression(key_column_filter_values, AstraFilterExpression)
@@ -490,6 +489,7 @@ class AstraClient:
         entity_type: Type[TModel],
         vector_to_match: list[float],
         similarity_function: SimilarityFunction = SimilarityFunction.COSINE,
+        key_column_filter_values: Optional[Union[Expression, List[Dict[str, Any]]]] = None,
         table_name: Optional[str] = None,
         return_vector: bool = False,
         num_results=1,
@@ -509,6 +509,7 @@ class AstraClient:
         query = VectorSearchQuery(
             table_fqn=f"{self._keyspace}.{model_mapper.table_name}",
             data_fields=[f for f in model_mapper.column_names if f != model_mapper.vector_column or return_vector],
+            key_column_filter_values=key_column_filter_values,
             sim_func=similarity_function,
             vector=vector_to_match,
             field_name=model_mapper.vector_column,

{adapta-3.2.3a475.dev3 → adapta-3.2.4a477.dev5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "adapta"
-version = "v3.2.3a475.dev3"
+version = "v3.2.4a477.dev5"
 description = "Logging, data connectors, monitoring, secret handling and general lifehacks to make data people lives easier."
 authors = ["ECCO Sneaks & Data <esdsupport@ecco.com>"]
 maintainers = ['GZU <gzu@ecco.com>', 'JRB <ext-jrb@ecco.com>']

adapta-3.2.3a475.dev3/adapta/_version.py

@@ -1 +0,0 @@
-__version__ = 'v3.2.3a475.dev3'

adapta-3.2.3a475.dev3/adapta/storage/distributed_object_store/v3/datastax_astra/_models.py

@@ -1,75 +0,0 @@
-"""
- Models for Astra DB.
-"""
-#  Copyright (c) 2023-2024. ECCO Sneaks & Data
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  distributed under the License is distributed on an "AS IS" BASIS,
-#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#  See the License for the specific language governing permissions and
-#  limitations under the License.
-#
-
-from enum import Enum
-
-
-class SimilarityFunction(Enum):
-    """
-    Supported similarity functions.
-    """
-
-    COSINE = "similarity_cosine"
-    DOT_PRODUCT = "similarity_dot_product"
-    EUCLIDEAN = "similarity_euclidean"
-
-
-class VectorSearchQuery:
-    """
-    Builder for a vector search query. Result would resemble the following:
-
-    select description, similarity_cosine(item_vector, [0.1, 0.15, 0.3, 0.12, 0.05])
-    from vsearch.products
-
-    order by item_vector ann of [0.1, 0.15, 0.3, 0.12, 0.05]
-    limit 1;
-
-    """
-
-    def __init__(
-        self,
-        table_fqn: str,
-        data_fields: list[str],
-        sim_func: SimilarityFunction,
-        vector: list[float],
-        field_name: str,
-        num_results=1,
-    ):
-        self._sim_func = sim_func
-        self._vector = vector
-        self._field_name = field_name
-        self._num_results = num_results
-        self._table_fqn = table_fqn
-        self._data_fields = data_fields
-
-    def _get_similarity_colum(self) -> str:
-        return f"{self._sim_func.value}({self._field_name}, {self._vector})"
-
-    def _get_order_by(self) -> str:
-        return f"order by {self._field_name} ann of {self._vector} limit {self._num_results};"
-
-    def __str__(self):
-        return " ".join(
-            [
-                "select",
-                ", ".join(self._data_fields),
-                ", " f"{self._get_similarity_colum()} as sim_value",
-                f"from {self._table_fqn}",
-                self._get_order_by(),
-            ]
-        )