adapta 3.2.6a482.dev7__tar.gz → 3.2.6a485.dev5__tar.gz

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: adapta
-Version: 3.2.6a482.dev7
+Version: 3.2.6a485.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
@@ -47,7 +47,7 @@ Requires-Dist: limits (>=3.7,<3.8)
 Requires-Dist: mlflow-skinny (>=2.4.1,<2.5.0) ; extra == "ml"
 Requires-Dist: pandas[performance] (>=2.0.0,<3.0)
 Requires-Dist: pandera (>=0.20.3,<1.0)
-Requires-Dist: polars (>=0.20,<2.0)
+Requires-Dist: polars (>=1.7,<2.0)
 Requires-Dist: pyarrow (>=7.0)
 Requires-Dist: pyodbc (>=4.0,<4.1) ; extra == "databases"
 Requires-Dist: redis[hiredis] (>=4.4.0,<4.5.0) ; extra == "caching"

adapta-3.2.6a485.dev5/adapta/_version.py

@@ -0,0 +1 @@
+__version__ = 'v3.2.6a485.dev5'

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/adapta/storage/delta_lake/v3/_functions.py

@@ -45,6 +45,7 @@ def load(  # pylint: disable=R0913
     columns: Optional[List[str]] = None,
     batch_size: Optional[int] = None,
     partition_filter_expressions: Optional[List[Tuple]] = None,
+    limit: Optional[int] = None,
 ) -> Union[MetaFrame, Iterator[MetaFrame]]:
     """
      Loads Delta Lake table from Azure or AWS storage and converts it to a pandas dataframe.
@@ -60,6 +61,7 @@ def load(  # pylint: disable=R0913
 
     :param columns: Optional list of columns to select when reading. Defaults to all columns of not provided.
     :param batch_size: Optional batch size when reading in batches. If not set, whole table will be loaded into memory.
+    :param limit: Optional limit on number of rows to read.
     :param partition_filter_expressions: Optional partitions filters. Examples:
 
        partition_filter_expressions = [("day", "=", "3")]
@@ -82,6 +84,9 @@ def load(  # pylint: disable=R0913
         filesystem=auth_client.get_pyarrow_filesystem(path),
     )
 
+    if limit:
+        pyarrow_ds = pyarrow_ds.head(limit)
+
     row_filter = (
         compile_expression(row_filter, ArrowFilterExpression) if isinstance(row_filter, Expression) else row_filter
     )

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/adapta/storage/distributed_object_store/v3/datastax_astra/_model_mappers.py

@@ -336,8 +336,6 @@ class PanderaPolarsMapper(CassandraModelMapper):
             polars.Datetime(time_unit="us"): (columns.DateTime,),
             polars.Datetime(time_unit="ns"): (columns.DateTime,),
             polars.Datetime(time_unit="ms"): (columns.DateTime,),
-            polars.Datetime(time_unit="ms", time_zone='UTC'): (columns.DateTime,),
-            polars.Datetime(time_unit="us", time_zone='UTC'): (columns.DateTime,),
         }
 
         column_type = mapping.get(type_to_map, None)

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/adapta/storage/distributed_object_store/v3/datastax_astra/astra_client.py

@@ -18,6 +18,7 @@
 #
 
 import base64
+import itertools
 import logging
 import math
 import os
@@ -62,7 +63,7 @@ from adapta import __version__
 from adapta.storage.distributed_object_store.v3.datastax_astra._models import SimilarityFunction, VectorSearchQuery
 from adapta.storage.models.filter_expression import Expression, AstraFilterExpression, compile_expression
 from adapta.utils import chunk_list, rate_limit
-from adapta.utils.metaframe import MetaFrame, concat, PolarsOptions
+from adapta.utils.metaframe import MetaFrame, concat
 from adapta.storage.distributed_object_store.v3.datastax_astra._model_mappers import get_mapper
 
 TModel = TypeVar("TModel")  # pylint: disable=C0103
@@ -238,6 +239,7 @@ class AstraClient:
         custom_indexes: Optional[List[str]] = None,
         deduplicate=False,
         num_threads: Optional[int] = None,
+        limit: Optional[int] = None,
     ) -> MetaFrame:
         """
         Run a filter query on the entity of type TModel backed by table `table_name`.
@@ -262,6 +264,7 @@ class AstraClient:
         :param: custom_indexes: An optional list of custom indexes, if it cannot be inferred, if it cannot be inferred from the data model.
         :param: deduplicate: Optionally deduplicate query result, for example when only the partition key part of a primary key is used to fetch results.
         :param: num_threads: Optionally run filtering using multiple threads. Setting this to -1 will cause this method to automatically evaluate number of threads based on filter expression size.
+        :param: limit: Optionally limit the number of results returned.
         """
 
         @on_exception(
@@ -274,12 +277,14 @@ class AstraClient:
             max_time=self._transient_error_max_wait_s,
             raise_on_giveup=True,
         )
-        def apply(model: Type[Model], key_column_filter: Dict[str, Any], columns_to_select: Optional[List[str]]):
-            model = model.filter(**key_column_filter).limit(None)
+        def apply(
+            model: Type[Model], key_column_filter: Dict[str, Any], columns_to_select: Optional[List[str]]
+        ) -> typing.Iterable[dict]:
+            model = model.filter(**key_column_filter).limit(limit)
             if columns_to_select:
-                return model.only(select_columns)
+                model = model.only(select_columns)
 
-            return model
+            return (dict(v.items()) for v in list(model))
 
         def normalize_column_name(column_name: str) -> str:
             filter_suffix = re.findall(self._filter_pattern, column_name)
@@ -288,15 +293,6 @@ class AstraClient:
 
             return column_name.replace(filter_suffix[0], "")
 
-        def to_frame(
-            model: Type[Model], key_column_filter: Dict[str, Any], columns_to_select: Optional[List[str]]
-        ) -> MetaFrame:
-            return MetaFrame(
-                [dict(v.items()) for v in list(apply(model, key_column_filter, columns_to_select))],
-                convert_to_polars=lambda x: polars.DataFrame(x, schema=select_columns),
-                convert_to_pandas=lambda x: pandas.DataFrame(x, columns=select_columns),
-            )
-
         assert (
             self._session is not None
         ), "Please instantiate an AstraClient using with AstraClient(...) before calling this method"
@@ -325,35 +321,33 @@ class AstraClient:
                 else num_threads
             )
             with ThreadPoolExecutor(max_workers=max_threads) as tpe:
-                result = concat(
-                    tpe.map(
-                        lambda args: to_frame(*args),
-                        [
-                            (cassandra_model, key_column_filter, select_columns)
-                            for key_column_filter in compiled_filter_values
-                        ],
-                        chunksize=max(int(len(compiled_filter_values) / num_threads), 1),
-                    ),
-                    options=[PolarsOptions(how="diagonal_relaxed")],
+                data = tpe.map(
+                    lambda args: apply(*args),
+                    [
+                        (cassandra_model, key_column_filter, select_columns)
+                        for key_column_filter in compiled_filter_values
+                    ],
+                    chunksize=max(int(len(compiled_filter_values) / num_threads), 1),
                 )
         else:
-            result = concat(
-                [
-                    MetaFrame(
-                        [dict(v.items()) for v in list(apply(cassandra_model, key_column_filter, select_columns))],
-                        convert_to_polars=(lambda x: polars.DataFrame(x, schema=select_columns))
-                        if not deduplicate
-                        else (lambda x: polars.DataFrame(x, schema=select_columns).unique()),
-                        convert_to_pandas=(lambda x: pandas.DataFrame(x, columns=select_columns))
-                        if not deduplicate
-                        else (lambda x: pandas.DataFrame(x, columns=select_columns).drop_duplicates()),
-                    )
-                    for key_column_filter in compiled_filter_values
-                ],
-                options=[PolarsOptions(how="diagonal_relaxed")],
+            data = (
+                apply(cassandra_model, key_column_filter, select_columns)
+                for key_column_filter in compiled_filter_values
             )
 
-        return result
+        data = itertools.chain.from_iterable(data)
+        if limit:
+            data = itertools.islice(data, limit)
+
+        return MetaFrame(
+            data,
+            convert_to_polars=(lambda x: polars.DataFrame(x, schema=select_columns))
+            if not deduplicate
+            else (lambda x: polars.DataFrame(x, schema=select_columns).unique()),
+            convert_to_pandas=(lambda x: pandas.DataFrame(x, columns=select_columns))
+            if not deduplicate
+            else (lambda x: pandas.DataFrame(x, columns=select_columns).drop_duplicates()),
+        )
 
     def get_entities_raw(self, query: str) -> MetaFrame:
         """

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/adapta/storage/models/format.py

@@ -193,6 +193,75 @@ class PolarsDataFrameJsonSerializationFormat(SerializationFormat[polars.DataFram
         return polars.read_json(io.BytesIO(data))
 
 
+class PolarsLazyFrameParquetSerializationFormat(SerializationFormat[polars.LazyFrame]):
+    """
+    Serializes lazyframes as parquet format.
+    """
+
+    def serialize(self, data: polars.LazyFrame) -> bytes:
+        """
+        Serializes lazyframe to bytes using parquet format.
+        :param data: Lazyframe to serialize.
+        :return: Parquet serialized lazyframe as byte array.
+        """
+        buffer = io.BytesIO()
+        data.collect().write_parquet(buffer)
+        return buffer.getvalue()
+
+    def deserialize(self, data: bytes) -> polars.LazyFrame:
+        """
+        Deserializes lazyframe from bytes using parquet format.
+        :param data: Lazyframe to deserialize in parquet format as bytes.
+        :return: Deserialized lazyframe.
+        """
+        return polars.scan_parquet(io.BytesIO(data))
+
+
+class PolarsLazyFrameCsvSerializationFormat(SerializationFormat[polars.LazyFrame]):
+    """
+    Serializes lazyframes as CSV format.
+    """
+
+    def serialize(self, data: polars.LazyFrame) -> bytes:
+        """
+        Serializes lazyframe to bytes using CSV format.
+        :param data: Lazyframe to serialize.
+        :return: CSV serialized Lazyframe as byte array.
+        """
+
+        return data.collect().write_csv().encode(encoding="utf-8")
+
+    def deserialize(self, data: bytes) -> polars.LazyFrame:
+        """
+        Deserializes lazyframe from bytes using CSV format.
+        :param data: LazyFrame to deserialize in CSV format as bytes.
+        :return: Deserialized lazyframe.
+        """
+        return polars.scan_csv(io.BytesIO(data))
+
+
+class PolarsLazyFrameJsonSerializationFormat(SerializationFormat[polars.LazyFrame]):
+    """
+    Serializes lazyframes as JSON format.
+    """
+
+    def serialize(self, data: polars.LazyFrame) -> bytes:
+        """
+        Serializes lazyframes to bytes using JSON format.
+        :param data: LazyFrame to serialize.
+        :return: JSON serialized lazyframe as byte array.
+        """
+        return data.collect().write_ndjson().encode(encoding="utf-8")
+
+    def deserialize(self, data: bytes) -> polars.LazyFrame:
+        """
+        Deserializes lazyframes from bytes using JSON format.
+        :param data: LazyFrame to deserialize in JSON format as bytes.
+        :return: Deserialized lazyframe.
+        """
+        return polars.scan_ndjson(io.BytesIO(data))
+
+
 class DictJsonSerializationFormat(SerializationFormat[dict]):
     """
     Serializes dictionaries as JSON format.

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/adapta/storage/query_enabled_store/_models.py

@@ -83,16 +83,16 @@ class QueryEnabledStore(Generic[TCredential, TSettings], ABC):
 
     @abstractmethod
     def _apply_filter(
-        self, path: DataPath, filter_expression: Expression, columns: list[str]
+        self, path: DataPath, filter_expression: Expression, columns: list[str], limit: Optional[int] = 10000
     ) -> Union[MetaFrame, Iterator[MetaFrame]]:
         """
-        Applies the provided filter expression to this Store and returns the result in a pandas DataFrame
+        Applies the provided filter expression to this Store and returns the result in a MetaFrame
         """
 
     @abstractmethod
     def _apply_query(self, query: str) -> Union[MetaFrame, Iterator[MetaFrame]]:
         """
-        Applies a plaintext query to this Store and returns the result in a pandas DataFrame
+        Applies a plaintext query to this Store and returns the result in a MetaFrame
         """
 
     @classmethod
@@ -139,6 +139,7 @@ class QueryConfigurationBuilder:
         self._path = path
         self._filter_expression: Optional[Expression] = None
         self._columns: list[str] = []
+        self._limit = 10000
 
     def filter(self, filter_expression: Expression) -> "QueryConfigurationBuilder":
         """
@@ -156,10 +157,20 @@ class QueryConfigurationBuilder:
         self._columns = list(columns)
         return self
 
+    def limit(self, limit: int) -> "QueryConfigurationBuilder":
+        """
+        Limit the number of results returned by the underlying store.
+        """
+        self._limit = limit
+        return self
+
     def read(self) -> Union[MetaFrame, Iterator[MetaFrame]]:
         """
         Execute the query on the underlying store.
         """
         return self._store._apply_filter(
-            path=self._path, filter_expression=self._filter_expression, columns=self._columns
+            path=self._path,
+            filter_expression=self._filter_expression,
+            columns=self._columns,
+            limit=self._limit,
         )

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/adapta/storage/query_enabled_store/_qes_astra.py

@@ -78,7 +78,7 @@ class AstraQueryEnabledStore(QueryEnabledStore[AstraCredential, AstraSettings]):
             self._astra_client.connect()
 
     def _apply_filter(
-        self, path: DataPath, filter_expression: Expression, columns: list[str]
+        self, path: DataPath, filter_expression: Expression, columns: list[str], limit: Optional[int] = 10000
     ) -> Union[MetaFrame, Iterator[MetaFrame]]:
         assert isinstance(path, AstraPath)
         astra_path: AstraPath = path
@@ -91,6 +91,7 @@ class AstraQueryEnabledStore(QueryEnabledStore[AstraCredential, AstraSettings]):
                     table_name=astra_path.table,
                     select_columns=columns,
                     num_threads=-1,  # auto-infer, see method documentation
+                    limit=limit,
                 )
 
         return self._astra_client.filter_entities(
@@ -100,6 +101,7 @@ class AstraQueryEnabledStore(QueryEnabledStore[AstraCredential, AstraSettings]):
             table_name=astra_path.table,
             select_columns=columns,
             num_threads=-1,  # auto-infer, see method documentation
+            limit=limit,
         )
 
     def _apply_query(self, query: str) -> Union[MetaFrame, Iterator[MetaFrame]]:

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/adapta/storage/query_enabled_store/_qes_delta.py

@@ -67,13 +67,18 @@ class DeltaQueryEnabledStore(QueryEnabledStore[DeltaCredential, DeltaSettings]):
         return cls(credentials=DeltaCredential.from_json(credentials), settings=DeltaSettings.from_json(settings))
 
     def _apply_filter(
-        self, path: DataPath, filter_expression: Expression, columns: list[str]
+        self,
+        path: DataPath,
+        filter_expression: Expression,
+        columns: list[str],
+        limit: Optional[int] = 10000,
     ) -> Union[MetaFrame, Iterator[MetaFrame]]:
         return load(
             auth_client=self.credentials.auth_client(credentials=self.credentials.auth_client_credentials()),
             path=path,
             row_filter=filter_expression,
             columns=columns,
+            limit=limit,
         )
 
     def _apply_query(self, query: str) -> Union[MetaFrame, Iterator[MetaFrame]]:

{adapta-3.2.6a482.dev7 → adapta-3.2.6a485.dev5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "adapta"
-version = "v3.2.6a482.dev7"
+version = "v3.2.6a485.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>']
@@ -17,7 +17,7 @@ pandas = { version = ">=2.0.0,<3.0", extras = ["performance"] }
 pyarrow = ">=7.0"
 dataclasses-json = "~0.6"
 limits = "~3.7"
-polars = ">=0.20,<2.0"
+polars = ">=1.7 <2.0"
 pandera = ">=0.20.3 <1.0"
 
 cassandra-driver = { version = "~3.29.1", optional = true }

adapta-3.2.6a482.dev7/adapta/_version.py

@@ -1 +0,0 @@
-__version__ = 'v3.2.6a482.dev7'