From 85b8dd6efdc30dc3b93280b143c6e6633beb32c4 Mon Sep 17 00:00:00 2001 From: JanusAsmussen Date: Thu, 26 Sep 2024 16:56:35 +0200 Subject: [PATCH] Set polars concatenation option to avoid issues with empty partitions --- .../v3/datastax_astra/astra_client.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/adapta/storage/distributed_object_store/v3/datastax_astra/astra_client.py b/adapta/storage/distributed_object_store/v3/datastax_astra/astra_client.py index 9ed3f438..d2d64f7a 100644 --- a/adapta/storage/distributed_object_store/v3/datastax_astra/astra_client.py +++ b/adapta/storage/distributed_object_store/v3/datastax_astra/astra_client.py @@ -62,7 +62,7 @@ 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 +from adapta.utils.metaframe import MetaFrame, concat, PolarsOptions from adapta.storage.distributed_object_store.v3.datastax_astra._model_mappers import get_mapper TModel = TypeVar("TModel") # pylint: disable=C0103 @@ -332,7 +332,8 @@ def to_frame( for key_column_filter in compiled_filter_values ], chunksize=max(int(len(compiled_filter_values) / num_threads), 1), - ) + ), + options=[PolarsOptions(how="diagonal_relaxed")], ) else: result = concat( @@ -347,7 +348,8 @@ def to_frame( else (lambda x: pandas.DataFrame(x, columns=select_columns).drop_duplicates()), ) for key_column_filter in compiled_filter_values - ] + ], + options=[PolarsOptions(how="diagonal_relaxed")], ) return result