From 898dba11ccc2a0c806fd146bd8a66e1a015faa3a Mon Sep 17 00:00:00 2001 From: frayle-ons <194791647+frayle-ons@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:25:38 +0100 Subject: [PATCH] updated result to json handling to consider duplicate dataframe ids --- src/classifai/indexers/main.py | 11 ++++++++--- src/classifai/servers/pydantic_models.py | 12 ++++++------ 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/classifai/indexers/main.py b/src/classifai/indexers/main.py index 44d22ae..e74035c 100644 --- a/src/classifai/indexers/main.py +++ b/src/classifai/indexers/main.py @@ -684,7 +684,7 @@ def reverse_search( # noqa: C901, PLR0912 # polars conversion paired_query = pl.DataFrame( {"id": query.id.astype(str).to_list(), "searched_doc_label": query.doc_label.astype(str).to_list()} - ) + ).with_row_index("_query_row_idx") # rename vectors dataframe for reverse search return column names and joining docs = self.vectors.rename({"label": "doc_label", "text": "doc_text"}).with_columns( @@ -701,9 +701,14 @@ def reverse_search( # noqa: C901, PLR0912 how="inner", ).rename({"doc_label_copy": "doc_label"}) - out = out.sort(by=["id", "searched_doc_label"], descending=[False, False]) if max_n_results != -1: - out = out.group_by("id").head(max_n_results) + out = out.sort(by=["_query_row_idx", "id", "searched_doc_label"], descending=[False, False, False]) + out = out.group_by(["_query_row_idx", "id", "searched_doc_label"], maintain_order=True).head( + max_n_results + ) + + # drop the helper column before final select + out = out.drop("_query_row_idx") # get formatted table final_table = out.select( diff --git a/src/classifai/servers/pydantic_models.py b/src/classifai/servers/pydantic_models.py index f908843..fdf566e 100644 --- a/src/classifai/servers/pydantic_models.py +++ b/src/classifai/servers/pydantic_models.py @@ -193,8 +193,8 @@ def convert_reverse_search_dataframe_to_pydantic_response( for original_query in original_input: input_id = original_query["id"] input_doc_label = original_query["doc_label"] - # Get the subset of the DataFrame corresponding to the current `id` - group_df = df[df["id"] == input_id] + # Get the subset of the DataFrame corresponding to the current `id` AND `searched_doc_label` + group_df = df[(df["id"] == input_id) & (df["searched_doc_label"] == input_doc_label)] if group_df.empty: # If there are no matches for this input_id, we still want to include it in the response, with empty results @@ -272,11 +272,11 @@ def convert_search_dataframe_to_pydantic_response(df: pd.DataFrame, meta_data: d ) ) - # Group rows by `query_id` - grouped = df.groupby("query_id") + # Group rows by `query_id` and 'query_text' in case the DataFrame contains multiple queries with the same 'qeuery_id'. + grouped = df.groupby(["query_id", "query_text"]) results_list = [] - for query_id, group_df in grouped: + for (query_id, query_text), group_df in grouped: # Convert group_df to a list of dictionaries rows_as_dicts = group_df.to_dict(orient="records") @@ -305,7 +305,7 @@ def convert_search_dataframe_to_pydantic_response(df: pd.DataFrame, meta_data: d results_list.append( SearchResponseSet( query_id=query_id, - query_text=group_df["query_text"].iloc[0], + query_text=query_text, entries=response_entries, ) )