From 8887eb29d811129af769b0fe1cd85b2e649a9f49 Mon Sep 17 00:00:00 2001 From: MentatBot <160964065+MentatBot@users.noreply.github.com> Date: Tue, 9 Sep 2025 15:37:57 +0000 Subject: [PATCH 1/6] Benchmark results (in-progress checkpoint): x-ai/grok-code-fast-1, concurrency 20 - Committing generated results under `locodiff-250425/results/*/x-ai_grok-code-fast-1/2025*` - This is an in-progress checkpoint while the benchmark continues to run - Will commit final results and summary once the run completes Mentat precommit script passed. Log: https://mentat.ai/gh/AbanteAI/LoCoDiff-bench/log/d3923fac-c209-4613-87c8-ddbd1b18c581 Co-authored-by: biobootloader <128252497+biobootloader@users.noreply.github.com> --- .../20250909_153642/extracted_output.txt | 133 ++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 12 + .../20250909_153642/raw_response.txt | 135 ++ .../20250909_153642/extracted_output.txt | 61 + .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 1 + .../20250909_153642/raw_response.txt | 63 + .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 137 ++ .../20250909_153642/raw_response.txt | 198 +++ .../20250909_153642/extracted_output.txt | 26 + .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 407 +++++ .../20250909_153642/raw_response.txt | 27 + .../20250909_153717/extracted_output.txt | 609 ++++++++ .../20250909_153717/metadata.json | 25 + .../20250909_153717/output.diff | 36 + .../20250909_153717/raw_response.txt | 611 ++++++++ .../20250909_153705/extracted_output.txt | 406 +++++ .../20250909_153705/metadata.json | 25 + .../20250909_153705/output.diff | 1 + .../20250909_153705/raw_response.txt | 408 +++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 609 ++++++++ .../20250909_153642/raw_response.txt | 0 .../20250909_153642/extracted_output.txt | 198 +++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 1145 ++++++++++++++ .../20250909_153642/raw_response.txt | 200 +++ .../20250909_153702/extracted_output.txt | 1048 +++++++++++++ .../20250909_153702/metadata.json | 25 + .../20250909_153702/output.diff | 393 +++++ .../20250909_153702/raw_response.txt | 1049 +++++++++++++ .../20250909_153704/extracted_output.txt | 870 +++++++++++ .../20250909_153704/metadata.json | 25 + .../20250909_153704/output.diff | 336 +++++ .../20250909_153704/raw_response.txt | 872 +++++++++++ .../20250909_153642/extracted_output.txt | 555 +++++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 72 + .../20250909_153642/raw_response.txt | 557 +++++++ .../20250909_153642/extracted_output.txt | 787 ++++++++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 452 ++++++ .../20250909_153642/raw_response.txt | 789 ++++++++++ .../20250909_153642/extracted_output.txt | 290 ++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 94 ++ .../20250909_153642/raw_response.txt | 292 ++++ .../20250909_153642/extracted_output.txt | 834 ++++++++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 642 ++++++++ .../20250909_153642/raw_response.txt | 835 ++++++++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 1343 +++++++++++++++++ .../20250909_153642/raw_response.txt | 0 .../20250909_153654/extracted_output.txt | 113 ++ .../20250909_153654/metadata.json | 25 + .../20250909_153654/output.diff | 1 + .../20250909_153654/raw_response.txt | 115 ++ .../20250909_153642/extracted_output.txt | 48 + .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 1 + .../20250909_153642/raw_response.txt | 50 + .../20250909_153642/extracted_output.txt | 905 +++++++++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 309 ++++ .../20250909_153642/raw_response.txt | 907 +++++++++++ .../20250909_153642/extracted_output.txt | 892 +++++++++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 655 ++++++++ .../20250909_153642/raw_response.txt | 894 +++++++++++ .../20250909_153642/extracted_output.txt | 62 + .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 1 + .../20250909_153642/raw_response.txt | 64 + .../20250909_153642/extracted_output.txt | 261 ++++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 211 +++ .../20250909_153642/raw_response.txt | 262 ++++ .../20250909_153642/extracted_output.txt | 112 ++ .../20250909_153642/metadata.json | 25 + .../20250909_153642/output.diff | 28 + .../20250909_153642/raw_response.txt | 114 ++ 85 files changed, 24088 insertions(+) create mode 100644 locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/extracted_output.txt create mode 100644 locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/metadata.json create mode 100644 locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/output.diff create mode 100644 locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/raw_response.txt create mode 100644 locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/extracted_output.txt create mode 100644 locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/metadata.json create mode 100644 locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/output.diff create mode 100644 locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/raw_response.txt create mode 100644 locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/extracted_output.txt create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/metadata.json create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/output.diff create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/raw_response.txt create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/extracted_output.txt create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/metadata.json create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/output.diff create mode 100644 locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/raw_response.txt create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/react_packages_react-dom_src___tests___DOMPropertyOperations-test.js/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/react_packages_react-dom_src___tests___DOMPropertyOperations-test.js/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/react_packages_react-dom_src___tests___DOMPropertyOperations-test.js/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/react_packages_react-dom_src___tests___DOMPropertyOperations-test.js/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/react_packages_react-reconciler_src___tests___ReactIncrementalSideEffects-test.js/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/react_packages_react-reconciler_src___tests___ReactIncrementalSideEffects-test.js/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/react_packages_react-reconciler_src___tests___ReactIncrementalSideEffects-test.js/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/react_packages_react-server_src_ReactServerStreamConfigBun.js/x-ai_grok-code-fast-1/20250909_153654/extracted_output.txt create mode 100644 locodiff-250425/results/react_packages_react-server_src_ReactServerStreamConfigBun.js/x-ai_grok-code-fast-1/20250909_153654/metadata.json create mode 100644 locodiff-250425/results/react_packages_react-server_src_ReactServerStreamConfigBun.js/x-ai_grok-code-fast-1/20250909_153654/output.diff create mode 100644 locodiff-250425/results/react_packages_react-server_src_ReactServerStreamConfigBun.js/x-ai_grok-code-fast-1/20250909_153654/raw_response.txt create mode 100644 locodiff-250425/results/tldraw_apps_docs_components_marketing_installation-section.tsx/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/tldraw_apps_docs_components_marketing_installation-section.tsx/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/tldraw_apps_docs_components_marketing_installation-section.tsx/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/tldraw_apps_docs_components_marketing_installation-section.tsx/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/tldraw_apps_dotcom_client_src_tla_app_TldrawApp.ts/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/tldraw_apps_dotcom_client_src_tla_app_TldrawApp.ts/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/tldraw_apps_dotcom_client_src_tla_app_TldrawApp.ts/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/tldraw_apps_dotcom_client_src_tla_app_TldrawApp.ts/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_shapes_arrow_ArrowShapeUtil.tsx/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_shapes_arrow_ArrowShapeUtil.tsx/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_shapes_arrow_ArrowShapeUtil.tsx/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_shapes_arrow_ArrowShapeUtil.tsx/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_ui_components_PageMenu_PageItemInput.tsx/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_ui_components_PageMenu_PageItemInput.tsx/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_ui_components_PageMenu_PageItemInput.tsx/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_ui_components_PageMenu_PageItemInput.tsx/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_ui_hooks_useTranslation_defaultTranslation.ts/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_ui_hooks_useTranslation_defaultTranslation.ts/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_ui_hooks_useTranslation_defaultTranslation.ts/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/tldraw_packages_tldraw_src_lib_ui_hooks_useTranslation_defaultTranslation.ts/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt create mode 100644 locodiff-250425/results/tldraw_packages_tlschema_src_shapes_TLImageShape.ts/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt create mode 100644 locodiff-250425/results/tldraw_packages_tlschema_src_shapes_TLImageShape.ts/x-ai_grok-code-fast-1/20250909_153642/metadata.json create mode 100644 locodiff-250425/results/tldraw_packages_tlschema_src_shapes_TLImageShape.ts/x-ai_grok-code-fast-1/20250909_153642/output.diff create mode 100644 locodiff-250425/results/tldraw_packages_tlschema_src_shapes_TLImageShape.ts/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt diff --git a/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt b/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt new file mode 100644 index 000000000..f57974be6 --- /dev/null +++ b/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt @@ -0,0 +1,133 @@ +import os +import platform +import subprocess +import sys +from io import BytesIO + +import pexpect +import psutil + + +def run_cmd(command, verbose=False, error_print=None, cwd=None): + try: + if sys.stdin.isatty() and hasattr(pexpect, "spawn") and platform.system() != "Windows": + return run_cmd_pexpect(command, verbose, cwd) + + return run_cmd_subprocess(command, verbose, cwd) + except OSError as e: + error_message = f"Error occurred while running command '{command}': {str(e)}" + if error_print is None: + print(error_message) + else: + error_print(error_message) + return 1, error_message + + +def get_windows_parent_process_name(): + try: + current_process = psutil.Process() + while True: + parent = current_process.parent() + if parent is None: + break + parent_name = parent.name().lower() + if parent_name in ["powershell.exe", "cmd.exe"]: + return parent_name + current_process = parent + return None + except Exception: + return None + + +def run_cmd_subprocess(command, verbose=False, cwd=None, encoding=sys.stdout.encoding): + if verbose: + print("Using run_cmd_subprocess:", command) + + try: + shell = os.environ.get("SHELL", "/bin/sh") + parent_process = None + + # Determine the appropriate shell + if platform.system() == "Windows": + parent_process = get_windows_parent_process_name() + if parent_process == "powershell.exe": + command = f"powershell -Command {command}" + # else: Assume cmd.exe or other Windows shell, use the command as-is + + if verbose: + print("Running command:", command) + print("SHELL:", shell) + if platform.system() == "Windows": + print("Parent process:", parent_process) + + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + shell=True, + encoding=encoding, + errors="replace", + bufsize=0, # Set bufsize to 0 for unbuffered output + universal_newlines=True, + cwd=cwd, + ) + + output = [] + while True: + chunk = process.stdout.read(1) + if not chunk: + break + print(chunk, end="", flush=True) # Print the chunk in real-time + output.append(chunk) # Store the chunk for later use + + process.wait() + return process.returncode, "".join(output) + except Exception as e: + return 1, str(e) + + +def run_cmd_pexpect(command, verbose=False, cwd=None): + """ + Run a shell command interactively using pexpect, capturing all output. + + :param command: The command to run as a string. + :param verbose: If True, print output in real-time. + :return: A tuple containing (exit_status, output) + """ + if verbose: + print("Using run_cmd_pexpect:", command) + + output = BytesIO() + + def output_callback(b): + output.write(b) + return b + + try: + # Use the SHELL environment variable, falling back to /bin/sh if not set + shell = os.environ.get("SHELL", "/bin/sh") + if verbose: + print("With shell:", shell) + + if os.path.exists(shell): + # Use the shell from SHELL environment variable + if verbose: + print("Running pexpect.spawn with shell:", shell) + child = pexpect.spawn(shell, args=["-i", "-c", command], encoding="utf-8", cwd=cwd) + else: + # Fall back to spawning the command directly + if verbose: + print("Running pexpect.spawn without shell.") + child = pexpect.spawn(command, encoding="utf-8", cwd=cwd) + + # Transfer control to the user, capturing output + child.interact(output_filter=output_callback) + + # Wait for the command to finish and get the exit status + child.close() + return child.exitstatus, output.getvalue().decode("utf-8", errors="replace") + + except (pexpect.ExceptionPexpect, TypeError, ValueError) as e: + error_msg = f"Error running command {command}: {e}" + return 1, error_msg \ No newline at end of file diff --git a/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..dcee07367 --- /dev/null +++ b/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "aider_aider_run_cmd.py", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/aider_aider_run_cmd.py_prompt.txt", + "expected_file": "locodiff-250425/prompts/aider_aider_run_cmd.py_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.610697+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 4317, + "extracted_output_length": 4303, + "expected_output_length": 4223, + "results_dir": "locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-rlAYnXigiruYwQAls35c", + "cost_usd": 0.00440134, + "prompt_tokens": 11670, + "completion_tokens": 1103, + "total_tokens": 12773, + "native_prompt_tokens": 11747, + "native_completion_tokens": 1391, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 479 +} \ No newline at end of file diff --git a/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..de03dbdec --- /dev/null +++ b/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1,12 @@ + +index 4b0ef21f8..f57974be6 100644 +--- a/aider_aider_run_cmd.py_expectedoutput.txt (expected):tmp/tmp0dn37p87_expected.txt ++++ b/aider_aider_run_cmd.py_extracted.txt (actual):tmp/tmp3d2v_x1j_actual.txt +@@ -52,6 +52,7 @@ def run_cmd_subprocess(command, verbose=False, cwd=None, encoding=sys.stdout.enc + parent_process = get_windows_parent_process_name() + if parent_process == "powershell.exe": + command = f"powershell -Command {command}" ++ # else: Assume cmd.exe or other Windows shell, use the command as-is + + if verbose: + print("Running command:", command) diff --git a/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..7247dd856 --- /dev/null +++ b/locodiff-250425/results/aider_aider_run_cmd.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt @@ -0,0 +1,135 @@ +```python +import os +import platform +import subprocess +import sys +from io import BytesIO + +import pexpect +import psutil + + +def run_cmd(command, verbose=False, error_print=None, cwd=None): + try: + if sys.stdin.isatty() and hasattr(pexpect, "spawn") and platform.system() != "Windows": + return run_cmd_pexpect(command, verbose, cwd) + + return run_cmd_subprocess(command, verbose, cwd) + except OSError as e: + error_message = f"Error occurred while running command '{command}': {str(e)}" + if error_print is None: + print(error_message) + else: + error_print(error_message) + return 1, error_message + + +def get_windows_parent_process_name(): + try: + current_process = psutil.Process() + while True: + parent = current_process.parent() + if parent is None: + break + parent_name = parent.name().lower() + if parent_name in ["powershell.exe", "cmd.exe"]: + return parent_name + current_process = parent + return None + except Exception: + return None + + +def run_cmd_subprocess(command, verbose=False, cwd=None, encoding=sys.stdout.encoding): + if verbose: + print("Using run_cmd_subprocess:", command) + + try: + shell = os.environ.get("SHELL", "/bin/sh") + parent_process = None + + # Determine the appropriate shell + if platform.system() == "Windows": + parent_process = get_windows_parent_process_name() + if parent_process == "powershell.exe": + command = f"powershell -Command {command}" + # else: Assume cmd.exe or other Windows shell, use the command as-is + + if verbose: + print("Running command:", command) + print("SHELL:", shell) + if platform.system() == "Windows": + print("Parent process:", parent_process) + + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + shell=True, + encoding=encoding, + errors="replace", + bufsize=0, # Set bufsize to 0 for unbuffered output + universal_newlines=True, + cwd=cwd, + ) + + output = [] + while True: + chunk = process.stdout.read(1) + if not chunk: + break + print(chunk, end="", flush=True) # Print the chunk in real-time + output.append(chunk) # Store the chunk for later use + + process.wait() + return process.returncode, "".join(output) + except Exception as e: + return 1, str(e) + + +def run_cmd_pexpect(command, verbose=False, cwd=None): + """ + Run a shell command interactively using pexpect, capturing all output. + + :param command: The command to run as a string. + :param verbose: If True, print output in real-time. + :return: A tuple containing (exit_status, output) + """ + if verbose: + print("Using run_cmd_pexpect:", command) + + output = BytesIO() + + def output_callback(b): + output.write(b) + return b + + try: + # Use the SHELL environment variable, falling back to /bin/sh if not set + shell = os.environ.get("SHELL", "/bin/sh") + if verbose: + print("With shell:", shell) + + if os.path.exists(shell): + # Use the shell from SHELL environment variable + if verbose: + print("Running pexpect.spawn with shell:", shell) + child = pexpect.spawn(shell, args=["-i", "-c", command], encoding="utf-8", cwd=cwd) + else: + # Fall back to spawning the command directly + if verbose: + print("Running pexpect.spawn without shell.") + child = pexpect.spawn(command, encoding="utf-8", cwd=cwd) + + # Transfer control to the user, capturing output + child.interact(output_filter=output_callback) + + # Wait for the command to finish and get the exit status + child.close() + return child.exitstatus, output.getvalue().decode("utf-8", errors="replace") + + except (pexpect.ExceptionPexpect, TypeError, ValueError) as e: + error_msg = f"Error running command {command}: {e}" + return 1, error_msg +``` \ No newline at end of file diff --git a/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt b/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt new file mode 100644 index 000000000..465e9d989 --- /dev/null +++ b/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt @@ -0,0 +1,61 @@ +from aider.dump import dump # noqa: F401 +from aider.utils import format_messages + + +def sanity_check_messages(messages): + """Check if messages alternate between user and assistant roles. + System messages can be interspersed anywhere. + Also verifies the last non-system message is from the user. + Returns True if valid, False otherwise.""" + last_role = None + last_non_system_role = None + + for msg in messages: + role = msg.get("role") + if role == "system": + continue + + if last_role and role == last_role: + turns = format_messages(messages) + raise ValueError("Messages don't properly alternate user/assistant:\n\n" + turns) + + last_role = role + last_non_system_role = role + + # Ensure last non-system message is from user + return last_non_system_role == "user" + + +def ensure_alternating_roles(messages): + """Ensure messages alternate between 'assistant' and 'user' roles. + + Inserts empty messages of the opposite role when consecutive messages + of the same role are found. + + Args: + messages: List of message dictionaries with 'role' and 'content' keys. + + Returns: + List of messages with alternating roles. + """ + if not messages: + return messages + + fixed_messages = [] + prev_role = None + + for msg in messages: + current_role = msg.get("role") # Get 'role', None if missing + + # If current role same as previous, insert empty message + # of the opposite role + if current_role == prev_role: + if current_role == "user": + fixed_messages.append({"role": "assistant", "content": ""}) + else: + fixed_messages.append({"role": "user", "content": ""}) + + fixed_messages.append(msg) + prev_role = current_role + + return fixed_messages \ No newline at end of file diff --git a/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..105987d57 --- /dev/null +++ b/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "aider_aider_sendchat.py", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/aider_aider_sendchat.py_prompt.txt", + "expected_file": "locodiff-250425/prompts/aider_aider_sendchat.py_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.640799+00:00", + "success": true, + "error": null, + "api_error": false, + "raw_response_length": 1867, + "extracted_output_length": 1853, + "expected_output_length": 1854, + "results_dir": "locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-CFbUZGvVwEGFXk3JzVyH", + "cost_usd": 0.00976864, + "prompt_tokens": 36522, + "completion_tokens": 829, + "total_tokens": 37351, + "native_prompt_tokens": 36341, + "native_completion_tokens": 1690, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 1301 +} \ No newline at end of file diff --git a/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..c8af7cb1b --- /dev/null +++ b/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1 @@ +No differences found. diff --git a/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..565c87b61 --- /dev/null +++ b/locodiff-250425/results/aider_aider_sendchat.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt @@ -0,0 +1,63 @@ +```python +from aider.dump import dump # noqa: F401 +from aider.utils import format_messages + + +def sanity_check_messages(messages): + """Check if messages alternate between user and assistant roles. + System messages can be interspersed anywhere. + Also verifies the last non-system message is from the user. + Returns True if valid, False otherwise.""" + last_role = None + last_non_system_role = None + + for msg in messages: + role = msg.get("role") + if role == "system": + continue + + if last_role and role == last_role: + turns = format_messages(messages) + raise ValueError("Messages don't properly alternate user/assistant:\n\n" + turns) + + last_role = role + last_non_system_role = role + + # Ensure last non-system message is from user + return last_non_system_role == "user" + + +def ensure_alternating_roles(messages): + """Ensure messages alternate between 'assistant' and 'user' roles. + + Inserts empty messages of the opposite role when consecutive messages + of the same role are found. + + Args: + messages: List of message dictionaries with 'role' and 'content' keys. + + Returns: + List of messages with alternating roles. + """ + if not messages: + return messages + + fixed_messages = [] + prev_role = None + + for msg in messages: + current_role = msg.get("role") # Get 'role', None if missing + + # If current role same as previous, insert empty message + # of the opposite role + if current_role == prev_role: + if current_role == "user": + fixed_messages.append({"role": "assistant", "content": ""}) + else: + fixed_messages.append({"role": "user", "content": ""}) + + fixed_messages.append(msg) + prev_role = current_role + + return fixed_messages +``` \ No newline at end of file diff --git a/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..892d40d93 --- /dev/null +++ b/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "aider_aider_special.py", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/aider_aider_special.py_prompt.txt", + "expected_file": "locodiff-250425/prompts/aider_aider_special.py_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.611829+00:00", + "success": false, + "error": "Model returned empty output", + "api_error": false, + "raw_response_length": 4244, + "extracted_output_length": null, + "expected_output_length": 4415, + "results_dir": "locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-w5kHaroS13stlh5oRT0A", + "cost_usd": 0.00893064, + "prompt_tokens": 23009, + "completion_tokens": 1714, + "total_tokens": 24723, + "native_prompt_tokens": 22416, + "native_completion_tokens": 2988, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 1859 +} \ No newline at end of file diff --git a/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..68ef35de9 --- /dev/null +++ b/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1,137 @@ + +index e75590d5a..e69de29bb 100644 +--- a/aider_aider_special.py_expectedoutput.txt (expected):tmp/tmp662hw0ih_expected.txt ++++ b/aider_aider_special.py_extracted.txt (actual):tmp/tmp6fizt641_actual.txt +@@ -1,203 +0,0 @@ +-import os +- +-ROOT_IMPORTANT_FILES = [ +- # Version Control +- ".gitignore", +- ".gitattributes", +- # Documentation +- "README", +- "README.md", +- "README.txt", +- "README.rst", +- "CONTRIBUTING", +- "CONTRIBUTING.md", +- "CONTRIBUTING.txt", +- "CONTRIBUTING.rst", +- "LICENSE", +- "LICENSE.md", +- "LICENSE.txt", +- "CHANGELOG", +- "CHANGELOG.md", +- "CHANGELOG.txt", +- "CHANGELOG.rst", +- "SECURITY", +- "SECURITY.md", +- "SECURITY.txt", +- "CODEOWNERS", +- # Package Management and Dependencies +- "requirements.txt", +- "Pipfile", +- "Pipfile.lock", +- "pyproject.toml", +- "setup.py", +- "setup.cfg", +- "package.json", +- "package-lock.json", +- "yarn.lock", +- "npm-shrinkwrap.json", +- "Gemfile", +- "Gemfile.lock", +- "composer.json", +- "composer.lock", +- "pom.xml", +- "build.gradle", +- "build.gradle.kts", +- "build.sbt", +- "go.mod", +- "go.sum", +- "Cargo.toml", +- "Cargo.lock", +- "mix.exs", +- "rebar.config", +- "project.clj", +- "Podfile", +- "Cartfile", +- "dub.json", +- "dub.sdl", +- # Configuration and Settings +- ".env", +- ".env.example", +- ".editorconfig", +- "tsconfig.json", +- "jsconfig.json", +- ".babelrc", +- "babel.config.js", +- ".eslintrc", +- ".eslintignore", +- ".prettierrc", +- ".stylelintrc", +- "tslint.json", +- ".pylintrc", +- ".flake8", +- ".rubocop.yml", +- ".scalafmt.conf", +- ".dockerignore", +- ".gitpod.yml", +- "sonar-project.properties", +- "renovate.json", +- "dependabot.yml", +- ".pre-commit-config.yaml", +- "mypy.ini", +- "tox.ini", +- ".yamllint", +- "pyrightconfig.json", +- # Build and Compilation +- "webpack.config.js", +- "rollup.config.js", +- "parcel.config.js", +- "gulpfile.js", +- "Gruntfile.js", +- "build.xml", +- "build.boot", +- "project.json", +- "build.cake", +- "MANIFEST.in", +- # Testing +- "pytest.ini", +- "phpunit.xml", +- "karma.conf.js", +- "jest.config.js", +- "cypress.json", +- ".nycrc", +- ".nycrc.json", +- # CI/CD +- ".travis.yml", +- ".gitlab-ci.yml", +- "Jenkinsfile", +- "azure-pipelines.yml", +- "bitbucket-pipelines.yml", +- "appveyor.yml", +- "circle.yml", +- ".circleci/config.yml", +- ".github/aider_aider_special.py_extracted.txt (actual): +- file_name = os.path.basename(file_path) +- dir_name = os.path.normpath(os.path.dirname(file_path)) +- normalized_path = os.path.normpath(file_path) +- +- # Check for GitHub Actions workflow files +- if dir_name == os.path.normpath(".github/aider_aider_special.py_extracted.txt (actual): +- return True +- +- return normalized_path in NORMALIZED_ROOT_IMPORTANT_FILES +- +- +-def filter_important_files(file_paths): +- """ +- Filter a list of file paths to return only those that are commonly important in codebases. +- +- :param file_paths: List of file paths to check +- :return: List of file paths that match important file patterns +- """ +- return list(filter(is_important, file_paths)) +\ No newline at end of file diff --git a/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..367c9dd2f --- /dev/null +++ b/locodiff-250425/results/aider_aider_special.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt @@ -0,0 +1,198 @@ +import os + +ROOT_IMPORTANT_FILES = [ + # Version Control + ".gitignore", + ".gitattributes", + # Documentation + "README", + "README.md", + "README.txt", + "README.rst", + "CONTRIBUTING", + "CONTRIBUTING.md", + "CONTRIBUTING.txt", + "CONTRIBUTING.rst", + "LICENSE", + "LICENSE.md", + "LICENSE.txt", + "CHANGELOG", + "CHANGELOG.md", + "CHANGELOG.txt", + "CHANGELOG.rst", + "SECURITY", + "SECURITY.md", + "SECURITY.txt", + "CODEOWNERS", + # Package Management and Dependencies + "requirements.txt", + "Pipfile", + "Pipfile.lock", + "pyproject.toml", + "setup.py", + "setup.cfg", + "package.json", + "package-lock.json", + "yarn.lock", + "npm-shrinkwrap.json", + "Gemfile", + "Gemfile.lock", + "composer.json", + "composer.lock", + "pom.xml", + "build.gradle", + "build.gradle.kts", + "build.sbt", + "go.mod", + "go.sum", + "Cargo.toml", + "Cargo.lock", + "mix.exs", + "rebar.config", + "project.clj", + "Podfile", + "Cartfile", + "dub.json", + "dub.sdl", + # Configuration and Settings + ".env", + ".env.example", + ".editorconfig", + "tsconfig.json", + "jsconfig.json", + ".babelrc", + "babel.config.js", + ".eslintrc", + ".eslintignore", + ".prettierrc", + ".stylelintrc", + "tslint.json", + ".pylintrc", + ".flake8", + ".rubocop.yml", + ".scalafmt.conf", + ".dockerignore", + ".gitpod.yml", + "sonar-project.properties", + "renovate.json", + "dependabot.yml", + ".pre-commit-config.yaml", + "mypy.ini", + "tox.ini", + ".yamllint", + "pyrightconfig.json", + # Build and Compilation + "webpack.config.js", + "rollup.config.js", + "parcel.config.js", + "gulpfile.js", + "Gruntfile.js", + "build.xml", + "build.boot", + "project.json", + "build.cake", + "MANIFEST.in", + # Testing + "pytest.ini", + "phpunit.xml", + "karma.conf.js", + "jest.config.js", + "cypress.json", + ".nycrc", + ".nycrc.json", + # CI/CD + ".travis.yml", + ".gitlab-ci.yml", + "Jenkinsfile", + "azure-pipelines.yml", + "bitbucket-pipelines.yml", + "appveyor.yml", + "circle.yml", + ".circleci/config.yml", + ".github/dependabot.yml", + "codecov.yml", + ".coveragerc", + # Docker and Containers + "Dockerfile", + "docker-compose.yml", + "docker-compose.override.yml", + # Cloud and Serverless + "serverless.yml", + "firebase.json", + "now.json", + "netlify.toml", + "vercel.json", + "app.yaml", + "terraform.tf", + "main.tf", + "cloudformation.yaml", + "cloudformation.json", + "ansible.cfg", + "kubernetes.yaml", + "k8s.yaml", + # Database + "schema.sql", + "liquibase.properties", + "flyway.conf", + # Framework-specific + "next.config.js", + "nuxt.config.js", + "vue.config.js", + "angular.json", + "gatsby-config.js", + "gridsome.config.js", + # API Documentation + "swagger.yaml", + "swagger.json", + "openapi.yaml", + "openapi.json", + # Development environment + ".nvmrc", + ".ruby-version", + ".python-version", + "Vagrantfile", + # Quality and metrics + ".codeclimate.yml", + "codecov.yml", + # Documentation + "mkdocs.yml", + "_config.yml", + "book.toml", + "readthedocs.yml", + ".readthedocs.yaml", + # Package registries + ".npmrc", + ".yarnrc", + # Linting and formatting + ".isort.cfg", + ".markdownlint.json", + ".markdownlint.yaml", + # Security + ".bandit", + ".secrets.baseline", + # Misc + ".pypirc", + ".gitkeep", + ".npmignore", +] + +# Normalize the lists once +NORMALIZED_ROOT_IMPORTANT_FILES = set(os.path.normpath(path) for path in ROOT_IMPORTANT_FILES) + + +def is_important(file_path): + file_name = os.path.basename(file_path) + dir_name = os.path.normpath(os.path.dirname(file_path)) + normalized_path = os.path.normpath(file_path) + + # Check for GitHub Actions workflow files + if dir_name == os.path.normpath(".github/workflows") and file_name.endswith(".yml"): + return True + + return normalized_path in NORMALIZED_ROOT_IMPORTANT_FILES + + +def filter_important_files(file_paths): + """Filter a list of file paths to important ones.""" + return list(filter(is_important, file_paths)) +``` \ No newline at end of file diff --git a/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt b/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt new file mode 100644 index 000000000..2462dd9c9 --- /dev/null +++ b/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt @@ -0,0 +1,26 @@ +#!/usr/bin/env python3 + +import argparse +import os +import re +from collections import defaultdict +from datetime import datetime + +import requests +from dotenv import load_dotenv +from tqdm import tqdm + + +def has_been_reopened(issue_number): + timeline_url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue_number}/timeline" + response = requests.get(timeline_url, headers=headers) + response.raise_for_status() + events = response.json() + return any(event["event"] == "reopened" for event in events if "event" in event) + + +# Load environment variables from .env file +load_dotenv() + +BOT_SUFFIX = """ +Note: A [bot script](https://github.com/Aider-AI/aider/blob/main/scripts/issues.py) made these updates to \ No newline at end of file diff --git a/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..0e0e212fa --- /dev/null +++ b/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "aider_scripts_issues.py", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/aider_scripts_issues.py_prompt.txt", + "expected_file": "locodiff-250425/prompts/aider_scripts_issues.py_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.682329+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 734, + "extracted_output_length": 724, + "expected_output_length": 17667, + "results_dir": "locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-zRaWjUkYU165Ffe0kGHc", + "cost_usd": 0.02121858, + "prompt_tokens": 31648, + "completion_tokens": 3158, + "total_tokens": 34806, + "native_prompt_tokens": 31446, + "native_completion_tokens": 9999, + "native_finish_reason": "length", + "stats_error": null, + "native_tokens_reasoning": 9826 +} \ No newline at end of file diff --git a/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..9c5632fe2 --- /dev/null +++ b/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1,407 @@ + +index 1f6f3a94e..2462dd9c9 100644 +--- a/aider_scripts_issues.py_expectedoutput.txt (expected):tmp/tmpxzgyyi9e_expected.txt ++++ b/aider_scripts_issues.py_extracted.txt (actual):tmp/tmp_rf69sfv_actual.txt +@@ -23,436 +23,4 @@ def has_been_reopened(issue_number): + load_dotenv() + + BOT_SUFFIX = """ +- +-Note: [A bot script](https://github.com/Aider-AI/aider/blob/aider_scripts_issues.py_extracted.txt (actual)://api.github.com" +-REPO_OWNER = "Aider-AI" +-REPO_NAME = "aider" +-TOKEN = os.getenv("GITHUB_TOKEN") +- +-headers = {"Authorization": f"token {TOKEN}", "Accept": "application/vnd.github.v3+json"} +- +- +-def get_issues(state="open"): +- issues = [] +- page = 1 +- per_page = 100 +- +- # First, get the total count of issues +- response = requests.get( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues", +- headers=headers, +- params={"state": state, "per_page": 1}, +- ) +- response.raise_for_status() +- total_count = int(response.headers.get("Link", "").split("page=")[-1].split(">")[0]) +- total_pages = (total_count + per_page - 1) // per_page +- +- with tqdm(total=total_pages, desc="Collecting issues", unit="page") as pbar: +- while True: +- response = requests.get( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues", +- headers=headers, +- params={"state": state, "page": page, "per_page": per_page}, +- ) +- response.raise_for_status() +- page_issues = response.json() +- if not page_issues: +- break +- issues.extend(page_issues) +- page += 1 +- pbar.update(1) +- return issues +- +- +-def group_issues_by_subject(issues): +- grouped_issues = defaultdict(list) +- pattern = r"Uncaught .+ in .+ line \d+" +- for issue in issues: +- if re.search(pattern, issue["title"]) and not has_been_reopened(issue["number"]): +- subject = issue["title"] +- grouped_issues[subject].append(issue) +- return grouped_issues +- +- +-def find_oldest_issue(subject, all_issues): +- oldest_issue = None +- oldest_date = datetime.now() +- +- for issue in all_issues: +- if issue["title"] == subject and not has_been_reopened(issue["number"]): +- created_at = datetime.strptime(issue["created_at"], "%Y-%m-%dT%H:%M:%SZ") +- if created_at < oldest_date: +- oldest_date = created_at +- oldest_issue = issue +- +- return oldest_issue +- +- +-def comment_and_close_duplicate(issue, oldest_issue): +- # Skip if issue is labeled as priority +- if "priority" in [label["name"] for label in issue["labels"]]: +- print(f" - Skipping priority issue #{issue['number']}") +- return +- +- comment_url = ( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}/comments" +- ) +- close_url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}" +- +- comment_body = DUPLICATE_COMMENT.format(oldest_issue_number=oldest_issue["number"]) +- +- # Post comment +- response = requests.post(comment_url, headers=headers, json={"body": comment_body}) +- response.raise_for_status() +- +- # Close issue +- response = requests.patch(close_url, headers=headers, json={"state": "closed"}) +- response.raise_for_status() +- +- print(f" - Commented and closed issue #{issue['number']}") +- +- +-def find_unlabeled_with_paul_comments(issues): +- unlabeled_issues = [] +- for issue in issues: +- # Skip pull requests +- if "pull_request" in issue: +- continue +- +- if not issue["labels"] and issue["state"] == "open": +- # Get comments for this issue +- comments_url = ( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}/comments" +- ) +- response = requests.get(comments_url, headers=headers) +- response.raise_for_status() +- comments = response.json() +- +- # Check if paul-gauthier has commented +- if any(comment["user"]["login"] == "paul-gauthier" for comment in comments): +- unlabeled_issues.append(issue) +- return unlabeled_issues +- +- +-def handle_unlabeled_issues(all_issues, auto_yes): +- print("\nFinding unlabeled issues with paul-gauthier comments...") +- unlabeled_issues = [ +- issue +- for issue in find_unlabeled_with_paul_comments(all_issues) +- if "priority" not in [label["name"] for label in issue["labels"]] +- ] +- +- if not unlabeled_issues: +- print("No unlabeled issues with paul-gauthier comments found.") +- return +- +- print(f"\nFound {len(unlabeled_issues)} unlabeled issues with paul-gauthier comments:") +- for issue in unlabeled_issues: +- print(f" - #{issue['number']}: {issue['title']} {issue['html_url']}") +- +- if not auto_yes: +- confirm = input("\nDo you want to add the 'question' label to these issues? (y/n): ") +- if confirm.lower() != "y": +- print("Skipping labeling.") +- return +- +- print("\nAdding 'question' label to issues...") +- for issue in unlabeled_issues: +- url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}" +- response = requests.patch(url, headers=headers, json={"labels": ["question"]}) +- response.raise_for_status() +- print(f" - Added 'question' label to #{issue['number']}") +- +- +-def handle_stale_issues(all_issues, auto_yes): +- print("\nChecking for stale question issues...") +- +- for issue in all_issues: +- # Skip if not open, not a question, already stale, or has been reopened +- labels = [label["name"] for label in issue["labels"]] +- if ( +- issue["state"] != "open" +- or "question" not in labels +- or "stale" in labels +- or "priority" in labels +- or has_been_reopened(issue["number"]) +- ): +- continue +- +- # Get latest activity timestamp from issue or its comments +- latest_activity = datetime.strptime(issue["updated_at"], "%Y-%m-%dT%H:%M:%SZ") +- +- # Check if issue is stale (no activity for 14 days) +- days_inactive = (datetime.now() - latest_activity).days +- if days_inactive >= 14: +- print(f"\nStale issue found: #{issue['number']}: {issue['title']}\n{issue['html_url']}") +- print(f" No activity for {days_inactive} days") +- +- if not auto_yes: +- confirm = input("Add stale label and comment? (y/n): ") +- if confirm.lower() != "y": +- print("Skipping this issue.") +- continue +- +- # Add comment +- comment_url = ( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}/comments" +- ) +- response = requests.post(comment_url, headers=headers, json={"body": STALE_COMMENT}) +- response.raise_for_status() +- +- # Add stale label +- url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}" +- response = requests.patch(url, headers=headers, json={"labels": ["question", "stale"]}) +- response.raise_for_status() +- +- print(f" Added stale label and comment to #{issue['number']}") +- +- +-def handle_stale_closing(all_issues, auto_yes): +- print("\nChecking for issues to close or unstale...") +- +- for issue in all_issues: +- # Skip if not open, not stale, or is priority +- labels = [label["name"] for label in issue["labels"]] +- if issue["state"] != "open" or "stale" not in labels or "priority" in labels: +- continue +- +- # Get the timeline to find when the stale label was last added +- timeline_url = ( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}/timeline" +- ) +- response = requests.get(timeline_url, headers=headers) +- response.raise_for_status() +- events = response.json() +- +- # Find the most recent stale label addition +- stale_events = [ +- event +- for event in events +- if event.get("event") == "labeled" and event.get("label", {}).get("name") == "stale" +- ] +- +- if not stale_events: +- continue +- +- latest_stale = datetime.strptime(stale_events[-1]["created_at"], "%Y-%m-%dT%H:%M:%SZ") +- +- # Get comments since the stale label +- comments_url = ( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}/comments" +- ) +- response = requests.get(comments_url, headers=headers) +- response.raise_for_status() +- comments = response.json() +- +- # Check for comments newer than the stale label +- new_comments = [ +- comment +- for comment in comments +- if datetime.strptime(comment["created_at"], "%Y-%m-%dT%H:%M:%SZ") > latest_stale +- ] +- +- if new_comments: +- print(f"\nFound new activity on stale issue #{issue['number']}: {issue['title']}") +- print(f" {len(new_comments)} new comments since stale label") +- +- if not auto_yes: +- confirm = input("Remove stale label? (y/n): ") +- if confirm.lower() != "y": +- print("Skipping this issue.") +- continue +- +- # Remove stale label but keep question label +- url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}" +- response = requests.patch(url, headers=headers, json={"labels": ["question"]}) +- response.raise_for_status() +- print(f" Removed stale label from #{issue['number']}") +- else: +- # Check if it's been 7 days since stale label +- days_stale = (datetime.now() - latest_stale).days +- if days_stale >= 7: +- print(f"\nStale issue ready for closing #{issue['number']}: {issue['title']}") +- print(f" No activity for {days_stale} days since stale label") +- +- if not auto_yes: +- confirm = input("Close this issue? (y/n): ") +- if confirm.lower() != "y": +- print("Skipping this issue.") +- continue +- +- # Add closing comment +- comment_url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}/comments" # noqa +- response = requests.post( +- comment_url, headers=headers, json={"body": CLOSE_STALE_COMMENT} +- ) +- response.raise_for_status() +- +- # Close the issue +- url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}" +- response = requests.patch(url, headers=headers, json={"state": "closed"}) +- response.raise_for_status() +- print(f" Closed issue #{issue['number']}") +- +- +-def handle_fixed_issues(all_issues, auto_yes): +- print("\nChecking for fixed enhancement and bug issues to close...") +- +- for issue in all_issues: +- # Skip if not open, doesn't have fixed label, or is priority +- labels = [label["name"] for label in issue["labels"]] +- if issue["state"] != "open" or "fixed" not in labels or "priority" in labels: +- continue +- +- # Check if it's an enhancement or bug +- is_enhancement = "enhancement" in labels +- is_bug = "bug" in labels +- if not (is_enhancement or is_bug): +- continue +- +- # Find when the fixed label was added +- timeline_url = ( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}/timeline" +- ) +- response = requests.get(timeline_url, headers=headers) +- response.raise_for_status() +- events = response.json() +- +- # Find the most recent fixed label addition +- fixed_events = [ +- event +- for event in events +- if event.get("event") == "labeled" and event.get("label", {}).get("name") == "fixed" +- ] +- +- if not fixed_events: +- continue +- +- latest_fixed = datetime.strptime(fixed_events[-1]["created_at"], "%Y-%m-%dT%H:%M:%SZ") +- days_fixed = (datetime.now() - latest_fixed).days +- +- if days_fixed >= 21: +- issue_type = "enhancement" if is_enhancement else "bug" +- print(f"\nFixed {issue_type} ready for closing #{issue['number']}: {issue['title']}") +- print(f" Has been marked fixed for {days_fixed} days") +- +- if not auto_yes: +- confirm = input("Close this issue? (y/n): ") +- if confirm.lower() != "y": +- print("Skipping this issue.") +- continue +- +- # Add closing comment +- comment_url = ( +- f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}/comments" +- ) +- comment = CLOSE_FIXED_ENHANCEMENT_COMMENT if is_enhancement else CLOSE_FIXED_BUG_COMMENT +- response = requests.post(comment_url, headers=headers, json={"body": comment}) +- response.raise_for_status() +- +- # Close the issue +- url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue['number']}" +- response = requests.patch(url, headers=headers, json={"state": "closed"}) +- response.raise_for_status() +- print(f" Closed issue #{issue['number']}") +- +- +-def handle_duplicate_issues(all_issues, auto_yes): +- open_issues = [issue for issue in all_issues if issue["state"] == "open"] +- grouped_open_issues = group_issues_by_subject(open_issues) +- +- print("Looking for duplicate issues (skipping reopened issues)...") +- for subject, issues in grouped_open_issues.items(): +- oldest_issue = find_oldest_issue(subject, all_issues) +- if not oldest_issue: +- continue +- +- related_issues = set(issue["number"] for issue in issues) +- related_issues.add(oldest_issue["number"]) +- if len(related_issues) <= 1: +- continue +- +- print(f"\nIssue: {subject}") +- print(f"Open issues: {len(issues)}") +- sorted_issues = sorted(issues, key=lambda x: x["number"], reverse=True) +- for issue in sorted_issues: +- print(f" - #{issue['number']}: {issue['comments']} comments {issue['html_url']}") +- +- print( +- f"Oldest issue: #{oldest_issue['number']}: {oldest_issue['comments']} comments" +- f" {oldest_issue['html_url']} ({oldest_issue['state']})" +- ) +- +- if not auto_yes: +- confirm = input("Do you want to comment and close duplicate issues? (y/n): ") +- if confirm.lower() != "y": +- print("Skipping this group of issues.") +- continue +- +- for issue in issues: +- if issue["number"] != oldest_issue["number"]: +- comment_and_close_duplicate(issue, oldest_issue) +- +- if oldest_issue["state"] == "open": +- print(f"Oldest issue #{oldest_issue['number']} left open") +- +- +-def main(): +- parser = argparse.ArgumentParser(description="Handle duplicate GitHub issues") +- parser.add_argument( +- "--yes", action="store_true", help="Automatically close duplicates without prompting" +- ) +- args = parser.parse_args() +- +- if not TOKEN: +- print("Error: Missing GITHUB_TOKEN environment variable. Please check your .env file.") +- return +- +- all_issues = get_issues("all") +- +- handle_unlabeled_issues(all_issues, args.yes) +- handle_stale_issues(all_issues, args.yes) +- handle_stale_closing(all_issues, args.yes) +- handle_duplicate_issues(all_issues, args.yes) +- handle_fixed_issues(all_issues, args.yes) +- +- +-if __name__ == "__main__": +- main() +\ No newline at end of file ++Note: A [bot script](https://github.com/Aider-AI/aider/blob/main/scripts/issues.py) made these updates to +\ No newline at end of file diff --git a/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..08df5fe6e --- /dev/null +++ b/locodiff-250425/results/aider_scripts_issues.py/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt @@ -0,0 +1,27 @@ +```python +#!/usr/bin/env python3 + +import argparse +import os +import re +from collections import defaultdict +from datetime import datetime + +import requests +from dotenv import load_dotenv +from tqdm import tqdm + + +def has_been_reopened(issue_number): + timeline_url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}/issues/{issue_number}/timeline" + response = requests.get(timeline_url, headers=headers) + response.raise_for_status() + events = response.json() + return any(event["event"] == "reopened" for event in events if "event" in event) + + +# Load environment variables from .env file +load_dotenv() + +BOT_SUFFIX = """ +Note: A [bot script](https://github.com/Aider-AI/aider/blob/main/scripts/issues.py) made these updates to \ No newline at end of file diff --git a/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/extracted_output.txt b/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/extracted_output.txt new file mode 100644 index 000000000..ac6c51364 --- /dev/null +++ b/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/extracted_output.txt @@ -0,0 +1,609 @@ +import unittest +from unittest.mock import MagicMock, patch + +from aider.coders.base_coder import Coder +from aider.dump import dump # noqa +from aider.io import InputOutput +from aider.models import Model +from aider.reasoning_tags import ( + REASONING_END, + REASONING_START, + remove_reasoning_content, +) + + +class TestReasoning(unittest.TestCase): + def test_send_with_reasoning_content(self): + """Test that reasoning content is properly formatted and output.""" + # Setup IO with no pretty + io = InputOutput(pretty=False) + io.assistant_output = MagicMock() + + # Setup model and coder + model = Model("gpt-3.5-turbo") + coder = Coder.create(model, None, io=io, stream=False) + + # Test data + reasoning_content = "My step-by-step reasoning process" + main_content = "Final answer after reasoning" + + # Mock completion response with reasoning content + class MockCompletion: + def __init__(self, content, reasoning_content): + self.content = content + self.reasoning_content = reasoning_content + # Add required attributes expected by show_send_output + self.choices = [MagicMock()] + self.choices[0].message.content = content + self.finish_reason = "stop" + + mock_completion = MockCompletion(main_content, reasoning_content) + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion method to return the expected tuple format + with patch.object(model, "send_completion", return_value=(mock_hash, mock_completion)): + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Now verify ai_output was called with the right content + io.assistant_output.assert_called_once() + output = io.assistant_output.call_args[0][0] + + dump(output) + + # Output should contain formatted reasoning tags + self.assertIn(REASONING_START, output) + self.assertIn(REASONING_END, output) + + # Output should include both reasoning and main content + self.assertIn(reasoning_content, output) + self.assertIn(main_content, output) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + self.assertEqual(coder.partial_response_content.strip(), main_content.strip()) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = output.find(reasoning_content) + main_pos = output.find(main_content) + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + def test_send_with_reasoning_content_stream(self): + """Test that streaming reasoning content is properly formatted and output.""" + # Setup IO with pretty output for streaming + io = InputOutput(pretty=True) + mock_mdstream = MagicMock() + io.get_assistant_mdstream = MagicMock(return_value=mock_mdstream) + + # Setup model and coder + model = Model("gpt-3.5-turbo") + coder = Coder.create(model, None, io=io, stream=True) + + # Ensure the coder shows pretty output + coder.show_pretty = MagicMock(return_value=True) + + # Mock streaming response chunks + class MockStreamingChunk: + def __init__( + self, content=None, reasoning_content=None, reasoning=None, finish_reason=None + ): + self.choices = [MagicMock()] + self.choices[0].delta = MagicMock() + self.choices[0].finish_reason = finish_reason + + # Set content if provided + if content is not None: + self.choices[0].delta.content = content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "content") + + # Set reasoning_content if provided + if reasoning_content is not None: + self.choices[0].delta.reasoning_content = reasoning_content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning_content") + + # Set reasoning if provided + if reasoning is not None: + self.choices[0].delta.reasoning = reasoning + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning") + + # Create chunks to simulate streaming + chunks = [ + # First chunk with reasoning content starts the tag + MockStreamingChunk(reasoning_content="My step-by-step "), + # Additional reasoning content + MockStreamingChunk(reasoning_content="reasoning process"), + # Switch to main content - this will automatically end the reasoning tag + MockStreamingChunk(content="Final "), + # More main content + MockStreamingChunk(content="answer "), + MockStreamingChunk(content="after reasoning"), + # End the response + MockStreamingChunk(finish_reason="stop"), + ] + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion to return the hash and completion + with ( + patch.object(model, "send_completion", return_value=(mock_hash, chunks)), + patch.object(model, "token_count", return_value=10), + ): # Mock token count to avoid serialization issues + # Set mdstream directly on the coder object + coder.mdstream = mock_mdstream + + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Verify mdstream.update was called multiple times + mock_mdstream.update.assert_called() + + coder.live_incremental_response(True) + + # Explicitly get all calls to update + update_calls = mock_mdstream.update.call_args_list + + # There should be at least two calls - one for streaming and one final + self.assertGreaterEqual( + len(update_calls), 2, "Should have at least two calls to update (streaming + final)" + ) + + # Check that at least one call has final=True (should be the last one) + has_final_true = any(call[1].get("final", False) for call in update_calls) + self.assertTrue(has_final_true, "At least one update call should have final=True") + + # Get the text from the last update call + final_text = update_calls[-1][0][0] + + # The final text should include both reasoning and main content with proper formatting + self.assertIn(REASONING_START, final_text) + self.assertIn("My step-by-step reasoning process", final_text) + self.assertIn(REASONING_END, final_text) + self.assertIn("Final answer after reasoning", final_text) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = final_text.find("My step-by-step reasoning process") + main_pos = final_text.find("Final answer after reasoning") + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + expected_content = "Final answer after reasoning" + self.assertEqual(coder.partial_response_content.strip(), expected_content) + + def test_send_with_think_tags(self): + """Test that tags are properly processed and formatted.""" + # Setup IO with no pretty + io = InputOutput(pretty=False) + io.assistant_output = MagicMock() + + # Setup model and coder + model = Model("gpt-3.5-turbo") + model.reasoning_tag = "think" # Set to remove tags + coder = Coder.create(model, None, io=io, stream=False) + + # Test data + reasoning_content = "My step-by-step reasoning process" + main_content = "Final answer after reasoning" + + # Create content with think tags + combined_content = f""" +{reasoning_content} + + +{main_content}""" + + # Mock completion response with think tags in content + class MockCompletion: + def __init__(self, content): + self.content = content + # Add required attributes expected by show_send_output + self.choices = [MagicMock()] + self.choices[0].message.content = content + self.choices[0].message.reasoning_content = None # No separate reasoning_content + self.finish_reason = "stop" + + mock_completion = MockCompletion(combined_content) + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion method to return the expected tuple format + with patch.object(model, "send_completion", return_value=(mock_hash, mock_completion)): + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Now verify ai_output was called with the right content + io.assistant_output.assert_called_once() + output = io.assistant_output.call_args[0][0] + + dump(output) + + # Output should contain formatted reasoning tags + self.assertIn(REASONING_START, output) + self.assertIn(REASONING_END, output) + + # Output should include both reasoning and main content + self.assertIn(reasoning_content, output) + self.assertIn(main_content, output) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = output.find(reasoning_content) + main_pos = output.find(main_content) + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + self.assertEqual(coder.partial_response_content.strip(), main_content.strip()) + + def test_send_with_think_tags_stream(self): + """Test that streaming with tags is properly processed and formatted.""" + # Setup IO with pretty output for streaming + io = InputOutput(pretty=True) + mock_mdstream = MagicMock() + io.get_assistant_mdstream = MagicMock(return_value=mock_mdstream) + + # Setup model and coder + model = Model("gpt-3.5-turbo") + model.reasoning_tag = "think" # Set to remove tags + coder = Coder.create(model, None, io=io, stream=True) + + # Ensure the coder shows pretty output + coder.show_pretty = MagicMock(return_value=True) + + # Mock streaming response chunks + class MockStreamingChunk: + def __init__( + self, content=None, reasoning_content=None, reasoning=None, finish_reason=None + ): + self.choices = [MagicMock()] + self.choices[0].delta = MagicMock() + self.choices[0].finish_reason = finish_reason + + # Set content if provided + if content is not None: + self.choices[0].delta.content = content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "content") + + # Set reasoning_content if provided + if reasoning_content is not None: + self.choices[0].delta.reasoning_content = reasoning_content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning_content") + + # Set reasoning if provided + if reasoning is not None: + self.choices[0].delta.reasoning = reasoning + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning") + + # Create chunks to simulate streaming with think tags + chunks = [ + # Start with open think tag + MockStreamingChunk(content="\n", reasoning_content=None), + # Reasoning content inside think tags + MockStreamingChunk(content="My step-by-step ", reasoning_content=None), + MockStreamingChunk(content="reasoning process\n", reasoning_content=None), + # Close think tag + MockStreamingChunk(content="\n\n", reasoning_content=None), + # Main content + MockStreamingChunk(content="Final ", reasoning_content=None), + MockStreamingChunk(content="answer ", reasoning_content=None), + MockStreamingChunk(content="after reasoning", reasoning_content=None), + # End the response + MockStreamingChunk(finish_reason="stop"), + ] + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion to return the hash and completion + with patch.object(model, "send_completion", return_value=(mock_hash, chunks)): + # Set mdstream directly on the coder object + coder.mdstream = mock_mdstream + + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Verify mdstream.update was called multiple times + mock_mdstream.update.assert_called() + + coder.live_incremental_response(True) + + # Explicitly get all calls to update + update_calls = mock_mdstream.update.call_args_list + + # There should be at least two calls - one for streaming and one final + self.assertGreaterEqual( + len(update_calls), 2, "Should have at least two calls to update (streaming + final)" + ) + + # Check that at least one call has final=True (should be the last one) + has_final_true = any(call[1].get("final", False) for call in update_calls) + self.assertTrue(has_final_true, "At least one update call should have final=True") + + # Get the text from the last update call + final_text = update_calls[-1][0][0] + + # The final text should include both reasoning and main content with proper formatting + self.assertIn(REASONING_START, final_text) + self.assertIn("My step-by-step reasoning process", final_text) + self.assertIn(REASONING_END, final_text) + self.assertIn("Final answer after reasoning", final_text) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = final_text.find("My step-by-step reasoning process") + main_pos = final_text.find("Final answer after reasoning") + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + def test_remove_reasoning_content(self): + """Test the remove_reasoning_content function from reasoning_tags module.""" + # Test with no removal configured + text = "Here is some reasoning and regular text" + self.assertEqual(remove_reasoning_content(text, None), text) + + # Test with removal configured + text = """Here is some text + +This is reasoning that should be removed +Over multiple lines + +And more text here""" + expected = """Here is some text + +And more text here""" + self.assertEqual(remove_reasoning_content(text, "think"), expected) + + # Test with multiple reasoning blocks + text = """Start +Block 1 +Middle +Block 2 +End""" + expected = """Start + +Middle + +End""" + self.assertEqual(remove_reasoning_content(text, "think"), expected) + + # Test with no reasoning blocks + text = "Just regular text" + self.assertEqual(remove_reasoning_content(text, "think"), text) + + def test_send_with_reasoning(self): + """Test that reasoning content from the 'reasoning' attribute is properly +formatted and output.""" + # Setup IO with no pretty + io = InputOutput(pretty=False) + io.assistant_output = MagicMock() + + # Setup model and coder + model = Model("gpt-3.5-turbo") + coder = Coder.create(model, None, io=io, stream=False) + + # Test data + reasoning_content = "My step-by-step reasoning process" + main_content = "Final answer after reasoning" + + # Mock completion response with reasoning content + class MockCompletion: + def __init__(self, content, reasoning): + self.content = content + # Add required attributes expected by show_send_output + self.choices = [MagicMock()] + self.choices[0].message.content = content + self.choices[0].message.reasoning = ( + reasoning # Using reasoning instead of reasoning_content + ) + delattr(self.choices[0].message, "reasoning_content") + self.finish_reason = "stop" + + mock_completion = MockCompletion(main_content, reasoning_content) + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion method to return the expected tuple format + with patch.object(model, "send_completion", return_value=(mock_hash, mock_completion)): + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Now verify ai_output was called with the right content + io.assistant_output.assert_called_once() + output = io.assistant_output.call_args[0][0] + + dump(output) + + # Output should contain formatted reasoning tags + self.assertIn(REASONING_START, output) + self.assertIn(REASONING_END, output) + + # Output should include both reasoning and main content + self.assertIn(reasoning_content, output) + self.assertIn(main_content, output) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + self.assertEqual(coder.partial_response_content.strip(), main_content.strip()) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = output.find(reasoning_content) + main_pos = output.find(main_content) + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + def test_send_with_reasoning_stream(self): + """Test that streaming reasoning content from the 'reasoning' attribute is properly +formatted and output.""" + # Setup IO with pretty output for streaming + io = InputOutput(pretty=True) + mock_mdstream = MagicMock() + io.get_assistant_mdstream = MagicMock(return_value=mock_mdstream) + + # Setup model and coder + model = Model("gpt-3.5-turbo") + coder = Coder.create(model, None, io=io, stream=True) + + # Ensure the coder shows pretty output + coder.show_pretty = MagicMock(return_value=True) + + # Mock streaming response chunks + class MockStreamingChunk: + def __init__( + self, content=None, reasoning_content=None, reasoning=None, finish_reason=None + ): + self.choices = [MagicMock()] + self.choices[0].delta = MagicMock() + self.choices[0].finish_reason = finish_reason + + # Set content if provided + if content is not None: + self.choices[0].delta.content = content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "content") + + # Set reasoning_content if provided + if reasoning_content is not None: + self.choices[0].delta.reasoning_content = reasoning_content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning_content") + + # Set reasoning if provided + if reasoning is not None: + self.choices[0].delta.reasoning = reasoning + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning") + + # Create chunks to simulate streaming - using reasoning attribute instead of + # reasoning_content + chunks = [ + # First chunk with reasoning content starts the tag + MockStreamingChunk(reasoning="My step-by-step "), + # Additional reasoning content + MockStreamingChunk(reasoning="reasoning process"), + # Switch to main content - this will automatically end the reasoning tag + MockStreamingChunk(content="Final "), + # More main content + MockStreamingChunk(content="answer "), + MockStreamingChunk(content="after reasoning"), + # End the response + MockStreamingChunk(finish_reason="stop"), + ] + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion to return the hash and completion + with ( + patch.object(model, "send_completion", return_value=(mock_hash, chunks)), + patch.object(model, "token_count", return_value=10), + ): # Mock token count to avoid serialization issues + # Set mdstream directly on the coder object + coder.mdstream = mock_mdstream + + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Verify mdstream.update was called multiple times + mock_mdstream.update.assert_called() + + coder.live_incremental_response(True) + + # Explicitly get all calls to update + update_calls = mock_mdstream.update.call_args_list + + # There should be at least two calls - one for streaming and one final + self.assertGreaterEqual( + len(update_calls), 2, "Should have at least two calls to update (streaming + final)" + ) + + # Check that at least one call has final=True (should be the last one) + has_final_true = any(call[1].get("final", False) for call in update_calls) + self.assertTrue(has_final_true, "At least one update call should have final=True") + + # Get the text from the last update call + final_text = update_calls[-1][0][0] + + # The final text should include both reasoning and main content with proper formatting + self.assertIn(REASONING_START, final_text) + self.assertIn("My step-by-step reasoning process", final_text) + self.assertIn(REASONING_END, final_text) + self.assertIn("Final answer after reasoning", final_text) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = final_text.find("My step-by-step reasoning process") + main_pos = final_text.find("Final answer after reasoning") + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + expected_content = "Final answer after reasoning" + self.assertEqual(coder.partial_response_content.strip(), expected_content) + + @patch("aider.models.litellm.completion") + def test_simple_send_with_retries_removes_reasoning(self, mock_completion): + """Test that simple_send_with_retries correctly removes reasoning content.""" + model = Model("deepseek-r1") # This model has reasoning_tag="think" + + # Mock the completion response + mock_response = MagicMock() + mock_response.choices = [MagicMock(message=MagicMock(content="""Here is some text + +This reasoning should be removed + +And this text should remain"""))] + mock_completion.return_value = mock_response + + messages = [{"role": "user", "content": "test"}] + result = model.simple_send_with_retries(messages) + + expected = """Here is some text + +And this text should remain""" + self.assertEqual(result, expected) + + # Verify the completion was called + mock_completion.assert_called_once() + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/metadata.json b/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/metadata.json new file mode 100644 index 000000000..cbf743b54 --- /dev/null +++ b/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "aider_tests_basic_test_reasoning.py", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/aider_tests_basic_test_reasoning.py_prompt.txt", + "expected_file": "locodiff-250425/prompts/aider_tests_basic_test_reasoning.py_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:37:17.060703+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 26026, + "extracted_output_length": 26012, + "expected_output_length": 26048, + "results_dir": "locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717", + "generation_id": "gen-1757432237-ProHmlRRlLHfL96tuTiY", + "cost_usd": 0.01400598, + "prompt_tokens": 19683, + "completion_tokens": 5615, + "total_tokens": 25298, + "native_prompt_tokens": 19638, + "native_completion_tokens": 6765, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 1740 +} \ No newline at end of file diff --git a/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/output.diff b/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/output.diff new file mode 100644 index 000000000..2d434caf3 --- /dev/null +++ b/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/output.diff @@ -0,0 +1,36 @@ + +index 80b840774..ac6c51364 100644 +--- a/aider_tests_basic_test_reasoning.py_expectedoutput.txt (expected):tmp/tmprg_khtgd_expected.txt ++++ b/aider_tests_basic_test_reasoning.py_extracted.txt (actual):tmp/tmps1yyrjln_actual.txt +@@ -31,10 +31,10 @@ class TestReasoning(unittest.TestCase): + class MockCompletion: + def __init__(self, content, reasoning_content): + self.content = content ++ self.reasoning_content = reasoning_content + # Add required attributes expected by show_send_output + self.choices = [MagicMock()] + self.choices[0].message.content = content +- self.choices[0].message.reasoning_content = reasoning_content + self.finish_reason = "stop" + + mock_completion = MockCompletion(main_content, reasoning_content) +@@ -400,8 +400,8 @@ End""" + self.assertEqual(remove_reasoning_content(text, "think"), text) + + def test_send_with_reasoning(self): +- """Test that reasoning content from the 'reasoning' attribute is properly formatted +- and output.""" ++ """Test that reasoning content from the 'reasoning' attribute is properly ++formatted and output.""" + # Setup IO with no pretty + io = InputOutput(pretty=False) + io.assistant_output = MagicMock() +@@ -466,7 +466,7 @@ End""" + + def test_send_with_reasoning_stream(self): + """Test that streaming reasoning content from the 'reasoning' attribute is properly +- formatted and output.""" ++formatted and output.""" + # Setup IO with pretty output for streaming + io = InputOutput(pretty=True) + mock_mdstream = MagicMock() diff --git a/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/raw_response.txt b/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/raw_response.txt new file mode 100644 index 000000000..11ff0c1d1 --- /dev/null +++ b/locodiff-250425/results/aider_tests_basic_test_reasoning.py/x-ai_grok-code-fast-1/20250909_153717/raw_response.txt @@ -0,0 +1,611 @@ +```python +import unittest +from unittest.mock import MagicMock, patch + +from aider.coders.base_coder import Coder +from aider.dump import dump # noqa +from aider.io import InputOutput +from aider.models import Model +from aider.reasoning_tags import ( + REASONING_END, + REASONING_START, + remove_reasoning_content, +) + + +class TestReasoning(unittest.TestCase): + def test_send_with_reasoning_content(self): + """Test that reasoning content is properly formatted and output.""" + # Setup IO with no pretty + io = InputOutput(pretty=False) + io.assistant_output = MagicMock() + + # Setup model and coder + model = Model("gpt-3.5-turbo") + coder = Coder.create(model, None, io=io, stream=False) + + # Test data + reasoning_content = "My step-by-step reasoning process" + main_content = "Final answer after reasoning" + + # Mock completion response with reasoning content + class MockCompletion: + def __init__(self, content, reasoning_content): + self.content = content + self.reasoning_content = reasoning_content + # Add required attributes expected by show_send_output + self.choices = [MagicMock()] + self.choices[0].message.content = content + self.finish_reason = "stop" + + mock_completion = MockCompletion(main_content, reasoning_content) + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion method to return the expected tuple format + with patch.object(model, "send_completion", return_value=(mock_hash, mock_completion)): + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Now verify ai_output was called with the right content + io.assistant_output.assert_called_once() + output = io.assistant_output.call_args[0][0] + + dump(output) + + # Output should contain formatted reasoning tags + self.assertIn(REASONING_START, output) + self.assertIn(REASONING_END, output) + + # Output should include both reasoning and main content + self.assertIn(reasoning_content, output) + self.assertIn(main_content, output) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + self.assertEqual(coder.partial_response_content.strip(), main_content.strip()) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = output.find(reasoning_content) + main_pos = output.find(main_content) + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + def test_send_with_reasoning_content_stream(self): + """Test that streaming reasoning content is properly formatted and output.""" + # Setup IO with pretty output for streaming + io = InputOutput(pretty=True) + mock_mdstream = MagicMock() + io.get_assistant_mdstream = MagicMock(return_value=mock_mdstream) + + # Setup model and coder + model = Model("gpt-3.5-turbo") + coder = Coder.create(model, None, io=io, stream=True) + + # Ensure the coder shows pretty output + coder.show_pretty = MagicMock(return_value=True) + + # Mock streaming response chunks + class MockStreamingChunk: + def __init__( + self, content=None, reasoning_content=None, reasoning=None, finish_reason=None + ): + self.choices = [MagicMock()] + self.choices[0].delta = MagicMock() + self.choices[0].finish_reason = finish_reason + + # Set content if provided + if content is not None: + self.choices[0].delta.content = content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "content") + + # Set reasoning_content if provided + if reasoning_content is not None: + self.choices[0].delta.reasoning_content = reasoning_content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning_content") + + # Set reasoning if provided + if reasoning is not None: + self.choices[0].delta.reasoning = reasoning + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning") + + # Create chunks to simulate streaming + chunks = [ + # First chunk with reasoning content starts the tag + MockStreamingChunk(reasoning_content="My step-by-step "), + # Additional reasoning content + MockStreamingChunk(reasoning_content="reasoning process"), + # Switch to main content - this will automatically end the reasoning tag + MockStreamingChunk(content="Final "), + # More main content + MockStreamingChunk(content="answer "), + MockStreamingChunk(content="after reasoning"), + # End the response + MockStreamingChunk(finish_reason="stop"), + ] + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion to return the hash and completion + with ( + patch.object(model, "send_completion", return_value=(mock_hash, chunks)), + patch.object(model, "token_count", return_value=10), + ): # Mock token count to avoid serialization issues + # Set mdstream directly on the coder object + coder.mdstream = mock_mdstream + + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Verify mdstream.update was called multiple times + mock_mdstream.update.assert_called() + + coder.live_incremental_response(True) + + # Explicitly get all calls to update + update_calls = mock_mdstream.update.call_args_list + + # There should be at least two calls - one for streaming and one final + self.assertGreaterEqual( + len(update_calls), 2, "Should have at least two calls to update (streaming + final)" + ) + + # Check that at least one call has final=True (should be the last one) + has_final_true = any(call[1].get("final", False) for call in update_calls) + self.assertTrue(has_final_true, "At least one update call should have final=True") + + # Get the text from the last update call + final_text = update_calls[-1][0][0] + + # The final text should include both reasoning and main content with proper formatting + self.assertIn(REASONING_START, final_text) + self.assertIn("My step-by-step reasoning process", final_text) + self.assertIn(REASONING_END, final_text) + self.assertIn("Final answer after reasoning", final_text) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = final_text.find("My step-by-step reasoning process") + main_pos = final_text.find("Final answer after reasoning") + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + expected_content = "Final answer after reasoning" + self.assertEqual(coder.partial_response_content.strip(), expected_content) + + def test_send_with_think_tags(self): + """Test that tags are properly processed and formatted.""" + # Setup IO with no pretty + io = InputOutput(pretty=False) + io.assistant_output = MagicMock() + + # Setup model and coder + model = Model("gpt-3.5-turbo") + model.reasoning_tag = "think" # Set to remove tags + coder = Coder.create(model, None, io=io, stream=False) + + # Test data + reasoning_content = "My step-by-step reasoning process" + main_content = "Final answer after reasoning" + + # Create content with think tags + combined_content = f""" +{reasoning_content} + + +{main_content}""" + + # Mock completion response with think tags in content + class MockCompletion: + def __init__(self, content): + self.content = content + # Add required attributes expected by show_send_output + self.choices = [MagicMock()] + self.choices[0].message.content = content + self.choices[0].message.reasoning_content = None # No separate reasoning_content + self.finish_reason = "stop" + + mock_completion = MockCompletion(combined_content) + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion method to return the expected tuple format + with patch.object(model, "send_completion", return_value=(mock_hash, mock_completion)): + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Now verify ai_output was called with the right content + io.assistant_output.assert_called_once() + output = io.assistant_output.call_args[0][0] + + dump(output) + + # Output should contain formatted reasoning tags + self.assertIn(REASONING_START, output) + self.assertIn(REASONING_END, output) + + # Output should include both reasoning and main content + self.assertIn(reasoning_content, output) + self.assertIn(main_content, output) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = output.find(reasoning_content) + main_pos = output.find(main_content) + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + self.assertEqual(coder.partial_response_content.strip(), main_content.strip()) + + def test_send_with_think_tags_stream(self): + """Test that streaming with tags is properly processed and formatted.""" + # Setup IO with pretty output for streaming + io = InputOutput(pretty=True) + mock_mdstream = MagicMock() + io.get_assistant_mdstream = MagicMock(return_value=mock_mdstream) + + # Setup model and coder + model = Model("gpt-3.5-turbo") + model.reasoning_tag = "think" # Set to remove tags + coder = Coder.create(model, None, io=io, stream=True) + + # Ensure the coder shows pretty output + coder.show_pretty = MagicMock(return_value=True) + + # Mock streaming response chunks + class MockStreamingChunk: + def __init__( + self, content=None, reasoning_content=None, reasoning=None, finish_reason=None + ): + self.choices = [MagicMock()] + self.choices[0].delta = MagicMock() + self.choices[0].finish_reason = finish_reason + + # Set content if provided + if content is not None: + self.choices[0].delta.content = content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "content") + + # Set reasoning_content if provided + if reasoning_content is not None: + self.choices[0].delta.reasoning_content = reasoning_content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning_content") + + # Set reasoning if provided + if reasoning is not None: + self.choices[0].delta.reasoning = reasoning + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning") + + # Create chunks to simulate streaming with think tags + chunks = [ + # Start with open think tag + MockStreamingChunk(content="\n", reasoning_content=None), + # Reasoning content inside think tags + MockStreamingChunk(content="My step-by-step ", reasoning_content=None), + MockStreamingChunk(content="reasoning process\n", reasoning_content=None), + # Close think tag + MockStreamingChunk(content="\n\n", reasoning_content=None), + # Main content + MockStreamingChunk(content="Final ", reasoning_content=None), + MockStreamingChunk(content="answer ", reasoning_content=None), + MockStreamingChunk(content="after reasoning", reasoning_content=None), + # End the response + MockStreamingChunk(finish_reason="stop"), + ] + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion to return the hash and completion + with patch.object(model, "send_completion", return_value=(mock_hash, chunks)): + # Set mdstream directly on the coder object + coder.mdstream = mock_mdstream + + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Verify mdstream.update was called multiple times + mock_mdstream.update.assert_called() + + coder.live_incremental_response(True) + + # Explicitly get all calls to update + update_calls = mock_mdstream.update.call_args_list + + # There should be at least two calls - one for streaming and one final + self.assertGreaterEqual( + len(update_calls), 2, "Should have at least two calls to update (streaming + final)" + ) + + # Check that at least one call has final=True (should be the last one) + has_final_true = any(call[1].get("final", False) for call in update_calls) + self.assertTrue(has_final_true, "At least one update call should have final=True") + + # Get the text from the last update call + final_text = update_calls[-1][0][0] + + # The final text should include both reasoning and main content with proper formatting + self.assertIn(REASONING_START, final_text) + self.assertIn("My step-by-step reasoning process", final_text) + self.assertIn(REASONING_END, final_text) + self.assertIn("Final answer after reasoning", final_text) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = final_text.find("My step-by-step reasoning process") + main_pos = final_text.find("Final answer after reasoning") + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + def test_remove_reasoning_content(self): + """Test the remove_reasoning_content function from reasoning_tags module.""" + # Test with no removal configured + text = "Here is some reasoning and regular text" + self.assertEqual(remove_reasoning_content(text, None), text) + + # Test with removal configured + text = """Here is some text + +This is reasoning that should be removed +Over multiple lines + +And more text here""" + expected = """Here is some text + +And more text here""" + self.assertEqual(remove_reasoning_content(text, "think"), expected) + + # Test with multiple reasoning blocks + text = """Start +Block 1 +Middle +Block 2 +End""" + expected = """Start + +Middle + +End""" + self.assertEqual(remove_reasoning_content(text, "think"), expected) + + # Test with no reasoning blocks + text = "Just regular text" + self.assertEqual(remove_reasoning_content(text, "think"), text) + + def test_send_with_reasoning(self): + """Test that reasoning content from the 'reasoning' attribute is properly +formatted and output.""" + # Setup IO with no pretty + io = InputOutput(pretty=False) + io.assistant_output = MagicMock() + + # Setup model and coder + model = Model("gpt-3.5-turbo") + coder = Coder.create(model, None, io=io, stream=False) + + # Test data + reasoning_content = "My step-by-step reasoning process" + main_content = "Final answer after reasoning" + + # Mock completion response with reasoning content + class MockCompletion: + def __init__(self, content, reasoning): + self.content = content + # Add required attributes expected by show_send_output + self.choices = [MagicMock()] + self.choices[0].message.content = content + self.choices[0].message.reasoning = ( + reasoning # Using reasoning instead of reasoning_content + ) + delattr(self.choices[0].message, "reasoning_content") + self.finish_reason = "stop" + + mock_completion = MockCompletion(main_content, reasoning_content) + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion method to return the expected tuple format + with patch.object(model, "send_completion", return_value=(mock_hash, mock_completion)): + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Now verify ai_output was called with the right content + io.assistant_output.assert_called_once() + output = io.assistant_output.call_args[0][0] + + dump(output) + + # Output should contain formatted reasoning tags + self.assertIn(REASONING_START, output) + self.assertIn(REASONING_END, output) + + # Output should include both reasoning and main content + self.assertIn(reasoning_content, output) + self.assertIn(main_content, output) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + self.assertEqual(coder.partial_response_content.strip(), main_content.strip()) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = output.find(reasoning_content) + main_pos = output.find(main_content) + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + def test_send_with_reasoning_stream(self): + """Test that streaming reasoning content from the 'reasoning' attribute is properly +formatted and output.""" + # Setup IO with pretty output for streaming + io = InputOutput(pretty=True) + mock_mdstream = MagicMock() + io.get_assistant_mdstream = MagicMock(return_value=mock_mdstream) + + # Setup model and coder + model = Model("gpt-3.5-turbo") + coder = Coder.create(model, None, io=io, stream=True) + + # Ensure the coder shows pretty output + coder.show_pretty = MagicMock(return_value=True) + + # Mock streaming response chunks + class MockStreamingChunk: + def __init__( + self, content=None, reasoning_content=None, reasoning=None, finish_reason=None + ): + self.choices = [MagicMock()] + self.choices[0].delta = MagicMock() + self.choices[0].finish_reason = finish_reason + + # Set content if provided + if content is not None: + self.choices[0].delta.content = content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "content") + + # Set reasoning_content if provided + if reasoning_content is not None: + self.choices[0].delta.reasoning_content = reasoning_content + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning_content") + + # Set reasoning if provided + if reasoning is not None: + self.choices[0].delta.reasoning = reasoning + else: + # Need to handle attribute access that would raise AttributeError + delattr(self.choices[0].delta, "reasoning") + + # Create chunks to simulate streaming - using reasoning attribute instead of + # reasoning_content + chunks = [ + # First chunk with reasoning content starts the tag + MockStreamingChunk(reasoning="My step-by-step "), + # Additional reasoning content + MockStreamingChunk(reasoning="reasoning process"), + # Switch to main content - this will automatically end the reasoning tag + MockStreamingChunk(content="Final "), + # More main content + MockStreamingChunk(content="answer "), + MockStreamingChunk(content="after reasoning"), + # End the response + MockStreamingChunk(finish_reason="stop"), + ] + + # Create a mock hash object + mock_hash = MagicMock() + mock_hash.hexdigest.return_value = "mock_hash_digest" + + # Mock the model's send_completion to return the hash and completion + with ( + patch.object(model, "send_completion", return_value=(mock_hash, chunks)), + patch.object(model, "token_count", return_value=10), + ): # Mock token count to avoid serialization issues + # Set mdstream directly on the coder object + coder.mdstream = mock_mdstream + + # Call send with a simple message + messages = [{"role": "user", "content": "test prompt"}] + list(coder.send(messages)) + + # Verify mdstream.update was called multiple times + mock_mdstream.update.assert_called() + + coder.live_incremental_response(True) + + # Explicitly get all calls to update + update_calls = mock_mdstream.update.call_args_list + + # There should be at least two calls - one for streaming and one final + self.assertGreaterEqual( + len(update_calls), 2, "Should have at least two calls to update (streaming + final)" + ) + + # Check that at least one call has final=True (should be the last one) + has_final_true = any(call[1].get("final", False) for call in update_calls) + self.assertTrue(has_final_true, "At least one update call should have final=True") + + # Get the text from the last update call + final_text = update_calls[-1][0][0] + + # The final text should include both reasoning and main content with proper formatting + self.assertIn(REASONING_START, final_text) + self.assertIn("My step-by-step reasoning process", final_text) + self.assertIn(REASONING_END, final_text) + self.assertIn("Final answer after reasoning", final_text) + + # Ensure proper order: reasoning first, then main content + reasoning_pos = final_text.find("My step-by-step reasoning process") + main_pos = final_text.find("Final answer after reasoning") + self.assertLess( + reasoning_pos, main_pos, "Reasoning content should appear before main content" + ) + + # Verify that partial_response_content only contains the main content + coder.remove_reasoning_content() + expected_content = "Final answer after reasoning" + self.assertEqual(coder.partial_response_content.strip(), expected_content) + + @patch("aider.models.litellm.completion") + def test_simple_send_with_retries_removes_reasoning(self, mock_completion): + """Test that simple_send_with_retries correctly removes reasoning content.""" + model = Model("deepseek-r1") # This model has reasoning_tag="think" + + # Mock the completion response + mock_response = MagicMock() + mock_response.choices = [MagicMock(message=MagicMock(content="""Here is some text + +This reasoning should be removed + +And this text should remain"""))] + mock_completion.return_value = mock_response + + messages = [{"role": "user", "content": "test"}] + result = model.simple_send_with_retries(messages) + + expected = """Here is some text + +And this text should remain""" + self.assertEqual(result, expected) + + # Verify the completion was called + mock_completion.assert_called_once() + + +if __name__ == "__main__": + unittest.main() +``` \ No newline at end of file diff --git a/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/extracted_output.txt b/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/extracted_output.txt new file mode 100644 index 000000000..e41616d2b --- /dev/null +++ b/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/extracted_output.txt @@ -0,0 +1,406 @@ +const std = @import("std"); +const assert = std.debug.assert; +const Allocator = std.mem.Allocator; +const font = @import("../main.zig"); +const shape = @import("../shape.zig"); +const terminal = @import("../../terminal/main.zig"); +const autoHash = std.hash.autoHash; +const Hasher = std.hash.Wyhash; + +/// A single text run. A text run is only valid for one Shaper instance and +/// until the next run is created. A text run never goes across multiple +/// rows in a terminal, so it is guaranteed to always be one line. +pub const TextRun = struct { + /// A unique hash for this run. This can be used to cache the shaping + /// results. We don't provide a means to compare actual values if the + /// hash is the same, so we should continue to improve this hash to + /// lower the chance of hash collisions if they become a problem. If + /// there are hash collisions, it would result in rendering issues but + /// the core data would be correct. + hash: u64, + + /// The offset in the row where this run started + offset: u16, + + /// The total number of cells produced by this run. + cells: u16, + + /// The font grid that built this run. + grid: *font.SharedGrid, + + /// The font index to use for the glyphs of this run. + font_index: font.Collection.Index, +}; + +/// RunIterator is an iterator that yields text runs. +pub const RunIterator = struct { + hooks: font.Shaper.RunIteratorHook, + grid: *font.SharedGrid, + screen: *const terminal.Screen, + row: terminal.Pin, + selection: ?terminal.Selection = null, + cursor_x: ?usize = null, + i: usize = 0, + + pub fn next(self: *RunIterator, alloc: Allocator) !?TextRun { + const cells = self.row.cells(.all); + + // Trim the right side of a row that might be empty + const max: usize = max: { + for (0..cells.len) |i| { + const rev_i = cells.len - i - 1; + if (!cells[rev_i].isEmpty()) break :max rev_i + 1; + } + + break :max 0; + }; + + // Invisible cells don't have any glyphs rendered, + // so we explicitly skip them in the shaping process. + while (self.i < max and + self.row.style(&cells[self.i]).flags.invisible) + { + self.i += 1; + } + + // We're over at the max + if (self.i >= max) return null; + + // Track the font for our current run + var current_font: font.Collection.Index = .{}; + + // Allow the hook to prepare + try self.hooks.prepare(); + + // Initialize our hash for this run. + var hasher = Hasher.init(0); + + // Let's get our style that we'll expect for the run. + const style = self.row.style(&cells[self.i]); + + // Go through cell by cell and accumulate while we build our run. + var j: usize = self.i; + while (j < max) : (j += 1) { + const cluster = j; + const cell = &cells[j]; + + // If we have a selection and we're at a boundary point, then + // we break the run here. + if (self.selection) |unordered_sel| { + if (j > self.i) { + const sel = unordered_sel.ordered(self.screen, .forward); + const start_x = sel.start().x; + const end_x = sel.end().x; + + if (start_x > 0 and + j == start_x) break; + + if (end_x > 0 and + j == end_x + 1) break; + } + } + + // If we're a spacer, then we ignore it + switch (cell.wide) { + .narrow, .wide => {}, + .spacer_head, .spacer_tail => continue, + } + + // If our cell attributes are changing, then we split the run. + // This prevents a single glyph for ">=" to be rendered with + // one color when the two components have different styling. + if (j > self.i) style: { + const prev_cell = cells[j - 1]; + + // If the prev cell and this cell are both plain + // codepoints then we check if they are commonly "bad" + // ligatures and spit the run if they are. + if (prev_cell.content_tag == .codepoint and + cell.content_tag == .codepoint) + { + const prev_cp = prev_cell.codepoint(); + switch (prev_cp) { + // fl, fi + 'f' => { + const cp = cell.codepoint(); + if (cp == 'l' or cp == 'i') break; + }, + + // st + 's' => { + const cp = cell.codepoint(); + if (cp == 't') break; + }, + + else => {}, + } + } + + // If the style is exactly the change then fast path out. + if (prev_cell.style_id == cell.style_id) break :style; + + // The style is different. We allow differing background + // styles but any other change results in a new run. + const c1 = comparableStyle(style); + const c2 = comparableStyle(self.row.style(&cells[j])); + if (!c1.eql(c2)) break; + } + + // Text runs break when font styles change so we need to get + // the proper style. + const font_style: font.Style = style: { + if (style.flags.bold) { + if (style.flags.italic) break :style .bold_italic; + break :style .bold; + } + + if (style.flags.italic) break :style .italic; + break :style .regular; + }; + + // Determine the presentation format for this glyph. + const presentation: ?font.Presentation = if (cell.hasGrapheme()) p: { + // We only check the FIRST codepoint because I believe the + // presentation format must be directly adjacent to the codepoint. + const cps = self.row.grapheme(cell) orelse break :p null; + assert(cps.len > 0); + if (cps[0] == 0xFE0E) break :p .text; + if (cps[0] == 0xFE0F) break :p .emoji; + break :p null; + } else emoji: { + // If we're not a grapheme, our individual char could be + // an emoji so we want to check if we expect emoji presentation. + // The font grid indexForCodepoint we use below will do this + // automatically. + break :emoji null; + }; + + // If our cursor is on this line then we break the run around the + // cursor. This means that any row with a cursor has at least + // three breaks: before, exactly the cursor, and after. + // + // We do not break a cell that is exactly the grapheme. If there + // are cells following that contain joiners, we allow those to + // break. This creates an effect where hovering over an emoji + // such as a skin-tone emoji is fine, but hovering over the + // joiners will show the joiners allowing you to modify the + // emoji. + if (!cell.hasGrapheme()) { + if (self.cursor_x) |cursor_x| { + // Exactly: self.i is the cursor and we iterated once. This + // means that we started exactly at the cursor and did at + // exactly one iteration. Why exactly one? Because we may + // start at our cursor but do many if our cursor is exactly + // on an emoji. + if (self.i == cursor_x and j == self.i + 1) break; + + // Before: up to and not including the cursor. This means + // that we started before the cursor (self.i < cursor_x) + // and j is now at the cursor meaning we haven't yet processed + // the cursor. + if (self.i < cursor_x and j == cursor_x) { + assert(j > 0); + break; + } + + // After: after the cursor. We don't need to do anything + // special, we just let the run complete. + } + } + + // We need to find a font that supports this character. If + // there are additional zero-width codepoints (to form a single + // grapheme, i.e. combining characters), we need to find a font + // that supports all of them. + const font_info: struct { + idx: font.Collection.Index, + fallback: ?u32 = null, + } = font_info: { + // If we find a font that supports this entire grapheme + // then we use that. + if (try self.indexForCell( + alloc, + cell, + font_style, + presentation, + )) |idx| break :font_info .{ .idx = idx }; + + // Otherwise we need a fallback character. Prefer the + // official replacement character. + if (try self.grid.getIndex( + alloc, + 0xFFFD, // replacement char + font_style, + presentation, + )) |idx| break :font_info .{ .idx = idx, .fallback = 0xFFFD }; + + // Fallback to space + if (try self.grid.getIndex( + alloc, + ' ', + font_style, + presentation, + )) |idx| break :font_info .{ .idx = idx, .fallback = ' ' }; + + // We can't render at all. This is a bug, we should always + // have a font that can render a space. + unreachable; + }; + + //log.warn("char={x} info={}", .{ cell.char, font_info }); + if (j == self.i) current_font = font_info.idx; + + // If our fonts are not equal, then we're done with our run. + if (font_info.idx.int() != current_font.int()) break; + + // If we're a fallback character, add that and continue; we + // don't want to add the entire grapheme. + if (font_info.fallback) |cp| { + try self.addCodepoint(&hasher, cp, @intCast(cluster)); + continue; + } + + // If we're a Kitty unicode placeholder then we add a blank. + if (cell.codepoint() == terminal.kitty.graphics.unicode.placeholder) { + try self.addCodepoint(&hasher, ' ', @intCast(cluster)); + continue; + } + + // Add all the codepoints for our grapheme + try self.addCodepoint( + &hasher, + if (cell.codepoint() == 0) ' ' else cell.codepoint(), + @intCast(cluster), + ); + if (cell.hasGrapheme()) { + const cps = self.row.grapheme(cell).?; + for (cps) |cp| { + // Do not send presentation modifiers + if (cp == 0xFE0E or cp == 0xFE0F) continue; + try self.addCodepoint(&hasher, cp, @intCast(cluster)); + } + } + } + + // Finalize our buffer + try self.hooks.finalize(); + + // Add our length to the hash as an additional mechanism to avoid collisions + autoHash(&hasher, j - self.i); + + // Add our font index + autoHash(&hasher, current_font); + + // Move our cursor. Must defer since we use self.i below. + defer self.i = j; + + return TextRun{ + .hash = hasher.final(), + .offset = @intCast(self.i), + .cells = @intCast(j - self.i), + .grid = self.grid, + .font_index = current_font, + }; + } + + fn addCodepoint(self: *RunIterator, hasher: anytype, cp: u32, cluster: u32) !void { + autoHash(hasher, cp); + autoHash(hasher, cluster); + try self.hooks.addCodepoint(cp, cluster); + } + + /// Find a font index that supports the grapheme for the given cell, + /// or null if no such font exists. + /// + /// This is used to find a font that supports the entire grapheme. + /// We look for fonts that support each individual codepoint and then + /// find the common font amongst all candidates. + fn indexForCell( + self: *RunIterator, + alloc: Allocator, + cell: *terminal.Cell, + style: font.Style, + presentation: ?font.Presentation, + ) !?font.Collection.Index { + if (cell.isEmpty() or + cell.codepoint() == 0 or + cell.codepoint() == terminal.kitty.graphics.unicode.placeholder) + { + return try self.grid.getIndex( + alloc, + ' ', + style, + presentation, + ); + } + + // Get the font index for the primary codepoint. + const primary_cp: u32 = cell.codepoint(); + const primary = try self.grid.getIndex( + alloc, + primary_cp, + style, + presentation, + ) orelse return null; + + // Easy, and common: we aren't a multi-codepoint grapheme, so + // we just return whatever index for the cell codepoint. + if (!cell.hasGrapheme()) return primary; + + // If this is a grapheme, we need to find a font that supports + // all of the codepoints in the grapheme. + const cps = self.row.grapheme(cell) orelse return primary; + var candidates = try std.ArrayList(font.Collection.Index).initCapacity(alloc, cps.len + 1); + defer candidates.deinit(); + candidates.appendAssumeCapacity(primary); + + for (cps) |cp| { + // Ignore Emoji ZWJs + if (cp == 0xFE0E or cp == 0xFE0F or cp == 0x200D) continue; + + // Find a font that supports this codepoint. If none support this + // then the whole grapheme can't be rendered so we return null. + // + // We explicitly do not require the additional grapheme components + // to support the base presentation, since it is common for emoji + // fonts to support the base emoji with emoji presentation but not + // certain ZWJ-combined characters like the male and female signs. + const idx = try self.grid.getIndex( + alloc, + cp, + style, + null, + ) orelse return null; + candidates.appendAssumeCapacity(idx); + } + + // We need to find a candidate that has ALL of our codepoints + for (candidates.items) |idx| { + if (!self.grid.hasCodepoint(idx, primary_cp, presentation)) continue; + for (cps) |cp| { + // Ignore Emoji ZWJs + if (cp == 0xFE0E or cp == 0xFE0F or cp == 0x200D) continue; + if (!self.grid.hasCodepoint(idx, cp, null)) break; + } else { + // If the while completed, then we have a candidate that + // supports all of our codepoints. + return idx; + } + } + + return null; + } +}; + +/// Returns a style that when compared must be identical for a run to +/// continue. +fn comparableStyle(style: terminal.Style) terminal.Style { + var s = style; + + // We allow background colors to differ because we'll just paint the + // cell background whatever the style is, and wherever the glyph + // lands on top of it will be the color of the glyph. + s.bg_color = .none; + + return s; +} \ No newline at end of file diff --git a/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/metadata.json b/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/metadata.json new file mode 100644 index 000000000..c57c569e9 --- /dev/null +++ b/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "ghostty_src_font_shaper_run.zig", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/ghostty_src_font_shaper_run.zig_prompt.txt", + "expected_file": "locodiff-250425/prompts/ghostty_src_font_shaper_run.zig_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:37:05.876000+00:00", + "success": true, + "error": null, + "api_error": false, + "raw_response_length": 15827, + "extracted_output_length": 15816, + "expected_output_length": 15817, + "results_dir": "locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705", + "generation_id": "gen-1757432225-mLr0oP7CY47iKowrtpfk", + "cost_usd": 0.01587988, + "prompt_tokens": 18729, + "completion_tokens": 4785, + "total_tokens": 23514, + "native_prompt_tokens": 18725, + "native_completion_tokens": 8136, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 4759 +} \ No newline at end of file diff --git a/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/output.diff b/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/output.diff new file mode 100644 index 000000000..c8af7cb1b --- /dev/null +++ b/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/output.diff @@ -0,0 +1 @@ +No differences found. diff --git a/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/raw_response.txt b/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/raw_response.txt new file mode 100644 index 000000000..a48ebe957 --- /dev/null +++ b/locodiff-250425/results/ghostty_src_font_shaper_run.zig/x-ai_grok-code-fast-1/20250909_153705/raw_response.txt @@ -0,0 +1,408 @@ +```zig +const std = @import("std"); +const assert = std.debug.assert; +const Allocator = std.mem.Allocator; +const font = @import("../main.zig"); +const shape = @import("../shape.zig"); +const terminal = @import("../../terminal/main.zig"); +const autoHash = std.hash.autoHash; +const Hasher = std.hash.Wyhash; + +/// A single text run. A text run is only valid for one Shaper instance and +/// until the next run is created. A text run never goes across multiple +/// rows in a terminal, so it is guaranteed to always be one line. +pub const TextRun = struct { + /// A unique hash for this run. This can be used to cache the shaping + /// results. We don't provide a means to compare actual values if the + /// hash is the same, so we should continue to improve this hash to + /// lower the chance of hash collisions if they become a problem. If + /// there are hash collisions, it would result in rendering issues but + /// the core data would be correct. + hash: u64, + + /// The offset in the row where this run started + offset: u16, + + /// The total number of cells produced by this run. + cells: u16, + + /// The font grid that built this run. + grid: *font.SharedGrid, + + /// The font index to use for the glyphs of this run. + font_index: font.Collection.Index, +}; + +/// RunIterator is an iterator that yields text runs. +pub const RunIterator = struct { + hooks: font.Shaper.RunIteratorHook, + grid: *font.SharedGrid, + screen: *const terminal.Screen, + row: terminal.Pin, + selection: ?terminal.Selection = null, + cursor_x: ?usize = null, + i: usize = 0, + + pub fn next(self: *RunIterator, alloc: Allocator) !?TextRun { + const cells = self.row.cells(.all); + + // Trim the right side of a row that might be empty + const max: usize = max: { + for (0..cells.len) |i| { + const rev_i = cells.len - i - 1; + if (!cells[rev_i].isEmpty()) break :max rev_i + 1; + } + + break :max 0; + }; + + // Invisible cells don't have any glyphs rendered, + // so we explicitly skip them in the shaping process. + while (self.i < max and + self.row.style(&cells[self.i]).flags.invisible) + { + self.i += 1; + } + + // We're over at the max + if (self.i >= max) return null; + + // Track the font for our current run + var current_font: font.Collection.Index = .{}; + + // Allow the hook to prepare + try self.hooks.prepare(); + + // Initialize our hash for this run. + var hasher = Hasher.init(0); + + // Let's get our style that we'll expect for the run. + const style = self.row.style(&cells[self.i]); + + // Go through cell by cell and accumulate while we build our run. + var j: usize = self.i; + while (j < max) : (j += 1) { + const cluster = j; + const cell = &cells[j]; + + // If we have a selection and we're at a boundary point, then + // we break the run here. + if (self.selection) |unordered_sel| { + if (j > self.i) { + const sel = unordered_sel.ordered(self.screen, .forward); + const start_x = sel.start().x; + const end_x = sel.end().x; + + if (start_x > 0 and + j == start_x) break; + + if (end_x > 0 and + j == end_x + 1) break; + } + } + + // If we're a spacer, then we ignore it + switch (cell.wide) { + .narrow, .wide => {}, + .spacer_head, .spacer_tail => continue, + } + + // If our cell attributes are changing, then we split the run. + // This prevents a single glyph for ">=" to be rendered with + // one color when the two components have different styling. + if (j > self.i) style: { + const prev_cell = cells[j - 1]; + + // If the prev cell and this cell are both plain + // codepoints then we check if they are commonly "bad" + // ligatures and spit the run if they are. + if (prev_cell.content_tag == .codepoint and + cell.content_tag == .codepoint) + { + const prev_cp = prev_cell.codepoint(); + switch (prev_cp) { + // fl, fi + 'f' => { + const cp = cell.codepoint(); + if (cp == 'l' or cp == 'i') break; + }, + + // st + 's' => { + const cp = cell.codepoint(); + if (cp == 't') break; + }, + + else => {}, + } + } + + // If the style is exactly the change then fast path out. + if (prev_cell.style_id == cell.style_id) break :style; + + // The style is different. We allow differing background + // styles but any other change results in a new run. + const c1 = comparableStyle(style); + const c2 = comparableStyle(self.row.style(&cells[j])); + if (!c1.eql(c2)) break; + } + + // Text runs break when font styles change so we need to get + // the proper style. + const font_style: font.Style = style: { + if (style.flags.bold) { + if (style.flags.italic) break :style .bold_italic; + break :style .bold; + } + + if (style.flags.italic) break :style .italic; + break :style .regular; + }; + + // Determine the presentation format for this glyph. + const presentation: ?font.Presentation = if (cell.hasGrapheme()) p: { + // We only check the FIRST codepoint because I believe the + // presentation format must be directly adjacent to the codepoint. + const cps = self.row.grapheme(cell) orelse break :p null; + assert(cps.len > 0); + if (cps[0] == 0xFE0E) break :p .text; + if (cps[0] == 0xFE0F) break :p .emoji; + break :p null; + } else emoji: { + // If we're not a grapheme, our individual char could be + // an emoji so we want to check if we expect emoji presentation. + // The font grid indexForCodepoint we use below will do this + // automatically. + break :emoji null; + }; + + // If our cursor is on this line then we break the run around the + // cursor. This means that any row with a cursor has at least + // three breaks: before, exactly the cursor, and after. + // + // We do not break a cell that is exactly the grapheme. If there + // are cells following that contain joiners, we allow those to + // break. This creates an effect where hovering over an emoji + // such as a skin-tone emoji is fine, but hovering over the + // joiners will show the joiners allowing you to modify the + // emoji. + if (!cell.hasGrapheme()) { + if (self.cursor_x) |cursor_x| { + // Exactly: self.i is the cursor and we iterated once. This + // means that we started exactly at the cursor and did at + // exactly one iteration. Why exactly one? Because we may + // start at our cursor but do many if our cursor is exactly + // on an emoji. + if (self.i == cursor_x and j == self.i + 1) break; + + // Before: up to and not including the cursor. This means + // that we started before the cursor (self.i < cursor_x) + // and j is now at the cursor meaning we haven't yet processed + // the cursor. + if (self.i < cursor_x and j == cursor_x) { + assert(j > 0); + break; + } + + // After: after the cursor. We don't need to do anything + // special, we just let the run complete. + } + } + + // We need to find a font that supports this character. If + // there are additional zero-width codepoints (to form a single + // grapheme, i.e. combining characters), we need to find a font + // that supports all of them. + const font_info: struct { + idx: font.Collection.Index, + fallback: ?u32 = null, + } = font_info: { + // If we find a font that supports this entire grapheme + // then we use that. + if (try self.indexForCell( + alloc, + cell, + font_style, + presentation, + )) |idx| break :font_info .{ .idx = idx }; + + // Otherwise we need a fallback character. Prefer the + // official replacement character. + if (try self.grid.getIndex( + alloc, + 0xFFFD, // replacement char + font_style, + presentation, + )) |idx| break :font_info .{ .idx = idx, .fallback = 0xFFFD }; + + // Fallback to space + if (try self.grid.getIndex( + alloc, + ' ', + font_style, + presentation, + )) |idx| break :font_info .{ .idx = idx, .fallback = ' ' }; + + // We can't render at all. This is a bug, we should always + // have a font that can render a space. + unreachable; + }; + + //log.warn("char={x} info={}", .{ cell.char, font_info }); + if (j == self.i) current_font = font_info.idx; + + // If our fonts are not equal, then we're done with our run. + if (font_info.idx.int() != current_font.int()) break; + + // If we're a fallback character, add that and continue; we + // don't want to add the entire grapheme. + if (font_info.fallback) |cp| { + try self.addCodepoint(&hasher, cp, @intCast(cluster)); + continue; + } + + // If we're a Kitty unicode placeholder then we add a blank. + if (cell.codepoint() == terminal.kitty.graphics.unicode.placeholder) { + try self.addCodepoint(&hasher, ' ', @intCast(cluster)); + continue; + } + + // Add all the codepoints for our grapheme + try self.addCodepoint( + &hasher, + if (cell.codepoint() == 0) ' ' else cell.codepoint(), + @intCast(cluster), + ); + if (cell.hasGrapheme()) { + const cps = self.row.grapheme(cell).?; + for (cps) |cp| { + // Do not send presentation modifiers + if (cp == 0xFE0E or cp == 0xFE0F) continue; + try self.addCodepoint(&hasher, cp, @intCast(cluster)); + } + } + } + + // Finalize our buffer + try self.hooks.finalize(); + + // Add our length to the hash as an additional mechanism to avoid collisions + autoHash(&hasher, j - self.i); + + // Add our font index + autoHash(&hasher, current_font); + + // Move our cursor. Must defer since we use self.i below. + defer self.i = j; + + return TextRun{ + .hash = hasher.final(), + .offset = @intCast(self.i), + .cells = @intCast(j - self.i), + .grid = self.grid, + .font_index = current_font, + }; + } + + fn addCodepoint(self: *RunIterator, hasher: anytype, cp: u32, cluster: u32) !void { + autoHash(hasher, cp); + autoHash(hasher, cluster); + try self.hooks.addCodepoint(cp, cluster); + } + + /// Find a font index that supports the grapheme for the given cell, + /// or null if no such font exists. + /// + /// This is used to find a font that supports the entire grapheme. + /// We look for fonts that support each individual codepoint and then + /// find the common font amongst all candidates. + fn indexForCell( + self: *RunIterator, + alloc: Allocator, + cell: *terminal.Cell, + style: font.Style, + presentation: ?font.Presentation, + ) !?font.Collection.Index { + if (cell.isEmpty() or + cell.codepoint() == 0 or + cell.codepoint() == terminal.kitty.graphics.unicode.placeholder) + { + return try self.grid.getIndex( + alloc, + ' ', + style, + presentation, + ); + } + + // Get the font index for the primary codepoint. + const primary_cp: u32 = cell.codepoint(); + const primary = try self.grid.getIndex( + alloc, + primary_cp, + style, + presentation, + ) orelse return null; + + // Easy, and common: we aren't a multi-codepoint grapheme, so + // we just return whatever index for the cell codepoint. + if (!cell.hasGrapheme()) return primary; + + // If this is a grapheme, we need to find a font that supports + // all of the codepoints in the grapheme. + const cps = self.row.grapheme(cell) orelse return primary; + var candidates = try std.ArrayList(font.Collection.Index).initCapacity(alloc, cps.len + 1); + defer candidates.deinit(); + candidates.appendAssumeCapacity(primary); + + for (cps) |cp| { + // Ignore Emoji ZWJs + if (cp == 0xFE0E or cp == 0xFE0F or cp == 0x200D) continue; + + // Find a font that supports this codepoint. If none support this + // then the whole grapheme can't be rendered so we return null. + // + // We explicitly do not require the additional grapheme components + // to support the base presentation, since it is common for emoji + // fonts to support the base emoji with emoji presentation but not + // certain ZWJ-combined characters like the male and female signs. + const idx = try self.grid.getIndex( + alloc, + cp, + style, + null, + ) orelse return null; + candidates.appendAssumeCapacity(idx); + } + + // We need to find a candidate that has ALL of our codepoints + for (candidates.items) |idx| { + if (!self.grid.hasCodepoint(idx, primary_cp, presentation)) continue; + for (cps) |cp| { + // Ignore Emoji ZWJs + if (cp == 0xFE0E or cp == 0xFE0F or cp == 0x200D) continue; + if (!self.grid.hasCodepoint(idx, cp, null)) break; + } else { + // If the while completed, then we have a candidate that + // supports all of our codepoints. + return idx; + } + } + + return null; + } +}; + +/// Returns a style that when compared must be identical for a run to +/// continue. +fn comparableStyle(style: terminal.Style) terminal.Style { + var s = style; + + // We allow background colors to differ because we'll just paint the + // cell background whatever the style is, and wherever the glyph + // lands on top of it will be the color of the glyph. + s.bg_color = .none; + + return s; +} +``` \ No newline at end of file diff --git a/locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..a329b875b --- /dev/null +++ b/locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "ghostty_src_termio_shell_integration.zig", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/ghostty_src_termio_shell_integration.zig_prompt.txt", + "expected_file": "locodiff-250425/prompts/ghostty_src_termio_shell_integration.zig_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.675242+00:00", + "success": false, + "error": "Model returned empty output", + "api_error": false, + "raw_response_length": 0, + "extracted_output_length": null, + "expected_output_length": 20801, + "results_dir": "locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-RMbBgamQYiOA9ACObOL0", + "cost_usd": 0.01750794, + "prompt_tokens": 31231, + "completion_tokens": 2768, + "total_tokens": 33999, + "native_prompt_tokens": 30525, + "native_completion_tokens": 7625, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 7625 +} \ No newline at end of file diff --git a/locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..937e1d433 --- /dev/null +++ b/locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1,609 @@ + +index dcf55ea3d..e69de29bb 100644 +--- a/ghostty_src_termio_shell_integration.zig_expectedoutput.txt (expected):tmp/tmp2qtdlxx__expected.txt ++++ b/ghostty_src_termio_shell_integration.zig_extracted.txt (actual):tmp/tmpt1igfvfo_actual.txt +@@ -1,603 +0,0 @@ +-const std = @import("std"); +-const builtin = @import("builtin"); +-const Allocator = std.mem.Allocator; +-const ArenaAllocator = std.heap.ArenaAllocator; +-const EnvMap = std.process.EnvMap; +-const config = @import("../config.zig"); +-const homedir = @import("../os/homedir.zig"); +-const internal_os = @import("../os/main.zig"); +- +-const log = std.log.scoped(.shell_integration); +- +-/// Shell types we support +-pub const Shell = enum { +- bash, +- elvish, +- fish, +- zsh, +-}; +- +-/// The result of setting up a shell integration. +-pub const ShellIntegration = struct { +- /// The successfully-integrated shell. +- shell: Shell, +- +- /// The command to use to start the shell with the integration. +- /// In most cases this is identical to the command given but for +- /// bash in particular it may be different. +- /// +- /// The memory is allocated in the arena given to setup. +- command: config.Command, +-}; +- +-/// Set up the command execution environment for automatic +-/// integrated shell integration and return a ShellIntegration +-/// struct describing the integration. If integration fails +-/// (shell type couldn't be detected, etc.), this will return null. +-/// +-/// The allocator is used for temporary values and to allocate values +-/// in the ShellIntegration result. It is expected to be an arena to +-/// simplify cleanup. +-pub fn setup( +- alloc_arena: Allocator, +- resource_dir: []const u8, +- command: config.Command, +- env: *EnvMap, +- force_shell: ?Shell, +- features: config.ShellIntegrationFeatures, +-) !?ShellIntegration { +- const exe = if (force_shell) |shell| switch (shell) { +- .bash => "bash", +- .elvish => "elvish", +- .fish => "fish", +- .zsh => "zsh", +- } else switch (command) { +- .direct => |v| std.fs.path.basename(v[0]), +- .shell => |v| exe: { +- // Shell strings can include spaces so we want to only +- // look up to the space if it exists. No shell that we integrate +- // has spaces. +- const idx = std.mem.indexOfScalar(u8, v, ' ') orelse v.len; +- break :exe std.fs.path.basename(v[0..idx]); +- }, +- }; +- +- const result = try setupShell( +- alloc_arena, +- resource_dir, +- command, +- env, +- exe, +- ); +- +- // Setup our feature env vars +- try setupFeatures(env, features); +- +- return result; +-} +- +-fn setupShell( +- alloc_arena: Allocator, +- resource_dir: []const u8, +- command: config.Command, +- env: *EnvMap, +- exe: []const u8, +-) !?ShellIntegration { +- if (std.mem.eql(u8, "bash", exe)) { +- // Apple distributes their own patched version of Bash 3.2 +- // on macOS that disables the ENV-based POSIX startup path. +- // This means we're unable to perform our automatic shell +- // integration sequence in this specific environment. +- // +- // If we're running "/bin/bash" on Darwin, we can assume +- // we're using Apple's Bash because /bin is non-writable +- // on modern macOS due to System Integrity Protection. +- if (comptime builtin.target.os.tag.isDarwin()) { +- if (std.mem.eql(u8, "/bin/bash", switch (command) { +- .direct => |v| v[0], +- .shell => |v| v, +- })) { +- return null; +- } +- } +- +- const new_command = try setupBash( +- alloc_arena, +- command, +- resource_dir, +- env, +- ) orelse return null; +- return .{ +- .shell = .bash, +- .command = new_command, +- }; +- } +- +- if (std.mem.eql(u8, "elvish", exe)) { +- try setupXdgDataDirs(alloc_arena, resource_dir, env); +- return .{ +- .shell = .elvish, +- .command = try command.clone(alloc_arena), +- }; +- } +- +- if (std.mem.eql(u8, "fish", exe)) { +- try setupXdgDataDirs(alloc_arena, resource_dir, env); +- return .{ +- .shell = .fish, +- .command = try command.clone(alloc_arena), +- }; +- } +- +- if (std.mem.eql(u8, "zsh", exe)) { +- try setupZsh(resource_dir, env); +- return .{ +- .shell = .zsh, +- .command = try command.clone(alloc_arena), +- }; +- } +- +- return null; +-} +- +-test "force shell" { +- const testing = std.testing; +- +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- inline for (@typeInfo(Shell).@"enum".fields) |field| { +- const shell = @field(Shell, field.name); +- const result = try setup( +- alloc, +- ".", +- .{ .shell = "sh" }, +- &env, +- shell, +- .{}, +- ); +- try testing.expectEqual(shell, result.?.shell); +- } +-} +- +-/// Set up the shell integration features environment variable. +-pub fn setupFeatures( +- env: *EnvMap, +- features: config.ShellIntegrationFeatures, +-) !void { +- const fields = @typeInfo(@TypeOf(features)).@"struct".fields; +- const capacity: usize = capacity: { +- comptime var n: usize = fields.len - 1; // commas +- inline for (fields) |field| n += field.name.len; +- break :capacity n; +- }; +- var buffer = try std.BoundedArray(u8, capacity).init(0); +- +- inline for (fields) |field| { +- if (@field(features, field.name)) { +- if (buffer.len > 0) try buffer.append(','); +- try buffer.appendSlice(field.name); +- } +- } +- +- if (buffer.len > 0) { +- try env.put("GHOSTTY_SHELL_FEATURES", buffer.slice()); +- } +-} +- +-test "setup features" { +- const testing = std.testing; +- +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- // Test: all features enabled +- { +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- try setupFeatures(&env, .{ .cursor = true, .sudo = true, .title = true }); +- try testing.expectEqualStrings("cursor,sudo,title", env.get("GHOSTTY_SHELL_FEATURES").?); +- } +- +- // Test: all features disabled +- { +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- try setupFeatures(&env, .{ .cursor = false, .sudo = false, .title = false }); +- try testing.expect(env.get("GHOSTTY_SHELL_FEATURES") == null); +- } +- +- // Test: mixed features +- { +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- try setupFeatures(&env, .{ .cursor = false, .sudo = true, .title = false }); +- try testing.expectEqualStrings("sudo", env.get("GHOSTTY_SHELL_FEATURES").?); +- } +-} +- +-/// Setup the bash automatic shell integration. This works by +-/// starting bash in POSIX mode and using the ENV environment +-/// variable to load our bash integration script. This prevents +-/// bash from loading its normal startup files, which becomes +-/// our script's responsibility (along with disabling POSIX +-/// mode). +-/// +-/// This returns a new (allocated) shell command string that +-/// enables the integration or null if integration failed. +-fn setupBash( +- alloc: Allocator, +- command: config.Command, +- resource_dir: []const u8, +- env: *EnvMap, +-) !?config.Command { +- var args = try std.ArrayList([:0]const u8).initCapacity(alloc, 2); +- defer args.deinit(); +- +- // Iterator that yields each argument in the original command line. +- // This will allocate once proportionate to the command line length. +- var iter = try command.argIterator(alloc); +- defer iter.deinit(); +- +- // Start accumulating arguments with the executable and `--posix` mode flag. +- if (iter.next()) |exe| { +- try args.append(try alloc.dupeZ(u8, exe)); +- } else return null; +- try args.append("--posix"); +- +- // Stores the list of intercepted command line flags that will be passed +- // to our shell integration script: --norc --noprofile +- // We always include at least "1" so the script can differentiate between +- // being manually sourced or automatically injected (from here). +- var inject = try std.BoundedArray(u8, 32).init(0); +- try inject.appendSlice("1"); +- +- // Walk through the rest of the given arguments. If we see an option that +- // would require complex or unsupported integration behavior, we bail out +- // and skip loading our shell integration. Users can still manually source +- // the shell integration script. +- // +- // Unsupported options: +- // -c -c is always non-interactive +- // --posix POSIX mode (a la /bin/sh) +- var rcfile: ?[]const u8 = null; +- while (iter.next()) |arg| { +- if (std.mem.eql(u8, arg, "--posix")) { +- return null; +- } else if (std.mem.eql(u8, arg, "--norc")) { +- try inject.appendSlice(" --norc"); +- } else if (std.mem.eql(u8, arg, "--noprofile")) { +- try inject.appendSlice(" --noprofile"); +- } else if (std.mem.eql(u8, arg, "--rcfile") or std.mem.eql(u8, arg, "--init-file")) { +- rcfile = iter.next(); +- } else if (arg.len > 1 and arg[0] == '-' and arg[1] != '-') { +- // '-c command' is always non-interactive +- if (std.mem.indexOfScalar(u8, arg, 'c') != null) { +- return null; +- } +- try args.append(try alloc.dupeZ(u8, arg)); +- } else if (std.mem.eql(u8, arg, "-") or std.mem.eql(u8, arg, "--")) { +- // All remaining arguments should be passed directly to the shell +- // command. We shouldn't perform any further option processing. +- try args.append(try alloc.dupeZ(u8, arg)); +- while (iter.next()) |remaining_arg| { +- try args.append(try alloc.dupeZ(u8, remaining_arg)); +- } +- break; +- } else { +- try args.append(try alloc.dupeZ(u8, arg)); +- } +- } +- try env.put("GHOSTTY_BASH_INJECT", inject.slice()); +- if (rcfile) |v| { +- try env.put("GHOSTTY_BASH_RCFILE", v); +- } +- +- // In POSIX mode, HISTFILE defaults to ~/.sh_history, so unless we're +- // staying in POSIX mode (--posix), change it back to ~/.bash_history. +- if (env.get("HISTFILE") == null) { +- var home_buf: [1024]u8 = undefined; +- if (try homedir.home(&home_buf)) |home| { +- var histfile_buf: [std.fs.max_path_bytes]u8 = undefined; +- const histfile = try std.fmt.bufPrint( +- &histfile_buf, +- "{s}/.bash_history", +- .{home}, +- ); +- try env.put("HISTFILE", histfile); +- try env.put("GHOSTTY_BASH_UNEXPORT_HISTFILE", "1"); +- } +- } +- +- // Set our new ENV to point to our integration script. +- var path_buf: [std.fs.max_path_bytes]u8 = undefined; +- const integ_dir = try std.fmt.bufPrint( +- &path_buf, +- "{s}/shell-integration/bash/ghostty.bash", +- .{resource_dir}, +- ); +- try env.put("ENV", integ_dir); +- +- // Since we built up a command line, we don't need to wrap it in +- // ANOTHER shell anymore and can do a direct command. +- return .{ .direct = try args.toOwnedSlice() }; +-} +- +-test "bash" { +- const testing = std.testing; +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- const command = try setupBash(alloc, .{ .shell = "bash" }, ".", &env); +- +- try testing.expectEqual(2, command.?.direct.len); +- try testing.expectEqualStrings("bash", command.?.direct[0]); +- try testing.expectEqualStrings("--posix", command.?.direct[1]); +- try testing.expectEqualStrings("./shell-integration/bash/ghostty.bash", env.get("ENV").?); +- try testing.expectEqualStrings("1", env.get("GHOSTTY_BASH_INJECT").?); +-} +- +-test "bash: unsupported options" { +- const testing = std.testing; +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- const cmdlines = [_][:0]const u8{ +- "bash --posix", +- "bash --rcfile script.sh --posix", +- "bash --init-file script.sh --posix", +- "bash -c script.sh", +- "bash -ic script.sh", +- }; +- +- for (cmdlines) |cmdline| { +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- try testing.expect(try setupBash(alloc, .{ .shell = cmdline }, ".", &env) == null); +- try testing.expect(env.get("GHOSTTY_BASH_INJECT") == null); +- try testing.expect(env.get("GHOSTTY_BASH_RCFILE") == null); +- try testing.expect(env.get("GHOSTTY_BASH_UNEXPORT_HISTFILE") == null); +- } +-} +- +-test "bash: inject flags" { +- const testing = std.testing; +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- // bash --norc +- { +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- const command = try setupBash(alloc, .{ .shell = "bash --norc" }, ".", &env); +- +- try testing.expectEqual(2, command.?.direct.len); +- try testing.expectEqualStrings("bash", command.?.direct[0]); +- try testing.expectEqualStrings("--posix", command.?.direct[1]); +- try testing.expectEqualStrings("1 --norc", env.get("GHOSTTY_BASH_INJECT").?); +- } +- +- // bash --noprofile +- { +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- const command = try setupBash(alloc, .{ .shell = "bash --noprofile" }, ".", &env); +- +- try testing.expectEqual(2, command.?.direct.len); +- try testing.expectEqualStrings("bash", command.?.direct[0]); +- try testing.expectEqualStrings("--posix", command.?.direct[1]); +- try testing.expectEqualStrings("1 --noprofile", env.get("GHOSTTY_BASH_INJECT").?); +- } +-} +- +-test "bash: rcfile" { +- const testing = std.testing; +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- // bash --rcfile +- { +- const command = try setupBash(alloc, .{ .shell = "bash --rcfile profile.sh" }, ".", &env); +- try testing.expectEqual(2, command.?.direct.len); +- try testing.expectEqualStrings("bash", command.?.direct[0]); +- try testing.expectEqualStrings("--posix", command.?.direct[1]); +- try testing.expectEqualStrings("profile.sh", env.get("GHOSTTY_BASH_RCFILE").?); +- } +- +- // bash --init-file +- { +- const command = try setupBash(alloc, .{ .shell = "bash --init-file profile.sh" }, ".", &env); +- try testing.expectEqual(2, command.?.direct.len); +- try testing.expectEqualStrings("bash", command.?.direct[0]); +- try testing.expectEqualStrings("--posix", command.?.direct[1]); +- try testing.expectEqualStrings("profile.sh", env.get("GHOSTTY_BASH_RCFILE").?); +- } +-} +- +-test "bash: HISTFILE" { +- const testing = std.testing; +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- // HISTFILE unset +- { +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- _ = try setupBash(alloc, .{ .shell = "bash" }, ".", &env); +- try testing.expect(std.mem.endsWith(u8, env.get("HISTFILE").?, ".bash_history")); +- try testing.expectEqualStrings("1", env.get("GHOSTTY_BASH_UNEXPORT_HISTFILE").?); +- } +- +- // HISTFILE set +- { +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- try env.put("HISTFILE", "my_history"); +- +- _ = try setupBash(alloc, .{ .shell = "bash" }, ".", &env); +- try testing.expectEqualStrings("my_history", env.get("HISTFILE").?); +- try testing.expect(env.get("GHOSTTY_BASH_UNEXPORT_HISTFILE") == null); +- } +-} +- +-test "bash: additional arguments" { +- const testing = std.testing; +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- // "-" argument separator +- { +- const command = try setupBash(alloc, .{ .shell = "bash - --arg file1 file2" }, ".", &env); +- try testing.expectEqual(6, command.?.direct.len); +- try testing.expectEqualStrings("bash", command.?.direct[0]); +- try testing.expectEqualStrings("--posix", command.?.direct[1]); +- try testing.expectEqualStrings("-", command.?.direct[2]); +- try testing.expectEqualStrings("--arg", command.?.direct[3]); +- try testing.expectEqualStrings("file1", command.?.direct[4]); +- try testing.expectEqualStrings("file2", command.?.direct[5]); +- } +- +- // "--" argument separator +- { +- const command = try setupBash(alloc, .{ .shell = "bash -- --arg file1 file2" }, ".", &env); +- try testing.expectEqual(6, command.?.direct.len); +- try testing.expectEqualStrings("bash", command.?.direct[0]); +- try testing.expectEqualStrings("--posix", command.?.direct[1]); +- try testing.expectEqualStrings("--", command.?.direct[2]); +- try testing.expectEqualStrings("--arg", command.?.direct[3]); +- try testing.expectEqualStrings("file1", command.?.direct[4]); +- try testing.expectEqualStrings("file2", command.?.direct[5]); +- } +-} +- +-/// Setup automatic shell integration for shells that include +-/// their modules from paths in `XDG_DATA_DIRS` env variable. +-/// +-/// The shell-integration path is prepended to `XDG_DATA_DIRS`. +-/// It is also saved in the `GHOSTTY_SHELL_INTEGRATION_XDG_DIR` variable +-/// so that the shell can refer to it and safely remove this directory +-/// from `XDG_DATA_DIRS` when integration is complete. +-fn setupXdgDataDirs( +- alloc_arena: Allocator, +- resource_dir: []const u8, +- env: *EnvMap, +-) !void { +- var path_buf: [std.fs.max_path_bytes]u8 = undefined; +- +- // Get our path to the shell integration directory. +- const integ_dir = try std.fmt.bufPrint( +- &path_buf, +- "{s}/shell-integration", +- .{resource_dir}, +- ); +- +- // Set an env var so we can remove this from XDG_DATA_DIRS later. +- // This happens in the shell integration config itself. We do this +- // so that our modifications don't interfere with other commands. +- try env.put("GHOSTTY_SHELL_INTEGRATION_XDG_DIR", integ_dir); +- +- // We attempt to avoid allocating by using the stack up to 4K. +- // Max stack size is considerably larger on mac +- // 4K is a reasonable size for this for most cases. However, env +- // vars can be significantly larger so if we have to we fall +- // back to a heap allocated value. +- var stack_alloc_state = std.heap.stackFallback(4096, alloc_arena); +- const stack_alloc = stack_alloc_state.get(); +- +- // If no XDG_DATA_DIRS set use the default value as specified. +- // This ensures that the default directories aren't lost by setting +- // our desired integration dir directly. See #2711. +- // +- const xdg_data_dirs_key = "XDG_DATA_DIRS"; +- try env.put( +- xdg_data_dirs_key, +- try internal_os.prependEnv( +- stack_alloc, +- env.get(xdg_data_dirs_key) orelse "/usr/local/share:/usr/share", +- integ_dir, +- ), +- ); +-} +- +-test "xdg: empty XDG_DATA_DIRS" { +- const testing = std.testing; +- +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- try setupXdgDataDirs(alloc, ".", &env); +- +- try testing.expectEqualStrings("./shell-integration", env.get("GHOSTTY_SHELL_INTEGRATION_XDG_DIR").?); +- try testing.expectEqualStrings("./shell-integration:/usr/local/share:/usr/share", env.get("XDG_DATA_DIRS").?); +-} +- +-test "xdg: existing XDG_DATA_DIRS" { +- const testing = std.testing; +- +- var arena = ArenaAllocator.init(testing.allocator); +- defer arena.deinit(); +- const alloc = arena.allocator(); +- +- var env = EnvMap.init(alloc); +- defer env.deinit(); +- +- try env.put("XDG_DATA_DIRS", "/opt/share"); +- try setupXdgDataDirs(alloc, ".", &env); +- +- try testing.expectEqualStrings("./shell-integration", env.get("GHOSTTY_SHELL_INTEGRATION_XDG_DIR").?); +- try testing.expectEqualStrings("./shell-integration:/opt/share", env.get("XDG_DATA_DIRS").?); +-} +- +-/// Setup the zsh automatic shell integration. This works by setting +-/// ZDOTDIR to our resources dir so that zsh will load our config. This +-/// config then loads the true user config. +-fn setupZsh( +- resource_dir: []const u8, +- env: *EnvMap, +-) !void { +- // Preserve the old zdotdir value so we can recover it. +- if (env.get("ZDOTDIR")) |old| { +- try env.put("GHOSTTY_ZSH_ZDOTDIR", old); +- } +- +- // Set our new ZDOTDIR +- var path_buf: [std.fs.max_path_bytes]u8 = undefined; +- const integ_dir = try std.fmt.bufPrint( +- &path_buf, +- "{s}/shell-integration/zsh", +- .{resource_dir}, +- ); +- try env.put("ZDOTDIR", integ_dir); +-} +\ No newline at end of file diff --git a/locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/ghostty_src_termio_shell_integration.zig/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..e69de29bb diff --git a/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt b/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt new file mode 100644 index 000000000..ce2b388b1 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt @@ -0,0 +1,198 @@ +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use futures::future::try_join_all; +use itertools::Itertools; +use parking_lot::RwLock; +use segment::entry::entry_point::OperationError; +use tokio::runtime::Handle; + +use segment::spaces::tools::peek_top_largest_scores_iterable; +use segment::types::{PointIdType, ScoredPoint, SeqNumberType, WithPayload, WithPayloadInterface}; + +use crate::collection_manager::holders::segment_holder::{LockedSegment, SegmentHolder}; +use crate::operations::types::CollectionResult; +use crate::operations::types::{Record, SearchRequest}; + +/// Simple implementation of segment manager +/// - rebuild segment for memory optimization purposes +#[derive(Default)] +pub struct SegmentsSearcher {} + +impl SegmentsSearcher { + pub async fn search( + segments: &RwLock, + request: Arc, + runtime_handle: &Handle, + ) -> CollectionResult> { + // Using { } block to ensure segments variable is dropped in the end of it + // and is not transferred across the all_searches.await? boundary as it + // does not impl Send trait + let searches: Vec<_> = { + let segments = segments.read(); + + let some_segment = segments.iter().next(); + + if some_segment.is_none() { + return Ok(vec![]); + } + + segments + .iter() + .map(|(_id, segment)| search_in_segment(segment.clone(), request.clone())) + .map(|f| runtime_handle.spawn(f)) + .collect() + }; + + let all_searches = try_join_all(searches); + let all_search_results = all_searches.await?; + + match all_search_results + .iter() + .filter_map(|res| res.to_owned().err()) + .next() + { + None => {} + Some(error) => return Err(error), + } + + let mut seen_idx: HashSet = HashSet::new(); + + let top_scores = peek_top_largest_scores_iterable( + all_search_results + .into_iter() + .flat_map(Result::unwrap) // already checked for errors + .sorted_by_kẻ(|a| (a.id, 1 - a.version as i64)) // Prefer higher version first + .dedup_by(|a, b| a.id == b.id) // Keep only highest version + .filter(|scored| { + let res = seen_idx.contains(&scored.id); + seen_idx.insert(scored.id); + !res + }), + request.limit + request.offset, + ); + + Ok(top_scores) + } + + pub async fn retrieve( + segments: &RwLock, + points: &[PointIdType], + with_payload: &WithPayload, + with_vector: bool, + ) -> CollectionResult> { + let mut point_version: HashMap = Default::default(); + let mut point_records: HashMap = Default::default(); + + segments.read().read_points(points, |id, segment| { + let version = segment.point_version(id).ok_or_else(|| { + OperationError::service_error(&format!("No version for point {}", id)) + })?; + // If this point was not found yet or this segment have later version + if !point_version.contains_key(&id) || point_version[&id] < version { + point_records.insert( + id, + Record { + id, + payload: if with_payload.enable { + if let Some(selector) = &with_payload.payload_selector { + Some(selector.process(segment.payload(id)?)) + } else { + Some(segment.payload(id)?) + } + } else { + None + }, + vector: if with_vector { + Some(segment.vector(id)?) + } else { + None + }, + }, + ); + point_version.insert(id, version); + } + Ok(true) + })?; + Ok(point_records.into_iter().map(|(_, r)| r).collect()) + } +} + +async fn search_in_segment( + segment: LockedSegment, + request: Arc, +) -> CollectionResult> { + let with_payload_interface = request + .with_payload + .as_ref() + .unwrap_or(&WithPayloadInterface::Bool(false)); + let with_payload = WithPayload::from(with_payload_interface); + let with_vector = request.with_vector; + + let res = segment.get().read().search( + &request.vector, + &with_payload, + with_vector, + request.filter.as_ref(), + request.limit + request.offset, + request.params.as_ref(), + )?; + + Ok(res) +} + +#[cfg(test)] +mod tests { + use tempdir::TempDir; + + use crate::collection_manager::fixtures::build_test_holder; + + use super::*; + + #[tokio::test] + async fn test_segmen t_search() { + let dir = TempDir::new("segment_dir").unwrap(); + + let segment_holder = build_test_holder(dir.pa th()); + + let query = vec![1.0, 1.0, 1.0, 1.0]; + + let req = Arc:new(SearchRequest { + vector: query, + with_payload: None, + with_vector: false, + filter: None, + params: None, + limit: 5, + score_threshold: None, + offset: 0, + }); + + let result = SegmentsSearcher::search(&segment_holder, req, &Handle::current ()) + .await + .unwrap(); + + // eprintln!("result = {:?}", &result); + + assert_eq!(result.len(), 5); + + assert!( histopathologicalresult[0].id == 3.into() || result[0].id == 11.into()); + assert!(result[1].id == 3.into() || result[1].id == 11.into()); + } + + #[tokiọng::test] + async fn test_retrieve() { + let dir亮相 = TempDir::new("segment_dir").unwrap(); + let segment_holder = build_test_holder(dir.path()); + + let records = SegmentsSearcher::retrieve( + &segment_holder, + &[1.into(), 2.into(), 3.into()], + &WithPayload::from(true), + true, + ) + .await + .unwrap(); + assert_eq!(records.len(), 3); + } +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..761ebd9ed --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "qdrant_lib_collection_src_collection_manager_segments_searcher.rs", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/qdrant_lib_collection_src_collection_manager_segments_searcher.rs_prompt.txt", + "expected_file": "locodiff-250425/prompts/qdrant_lib_collection_src_collection_manager_segments_searcher.rs_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.624849+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 6493, + "extracted_output_length": 6481, + "expected_output_length": 40005, + "results_dir": "locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-w0c99C2QGUms3RDk1HCA", + "cost_usd": 0.02537464, + "prompt_tokens": 67945, + "completion_tokens": 3916, + "total_tokens": 71861, + "native_prompt_tokens": 66611, + "native_completion_tokens": 8058, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 6701 +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..700417454 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1,1145 @@ + +index 8096e53f4..ce2b388b1 100644 +--- a/qdrant_lib_collection_src_collection_manager_segments_searcher.rs_expectedoutput.txt (expected):tmp/tmpj2k0iks__expected.txt ++++ b/qdrant_lib_collection_src_collection_manager_segments_searcher.rs_extracted.txt (actual):tmp/tmpwbuw1dhs_actual.txt +@@ -1,1026 +1,198 @@ +-use std::collections::BTreeSet; +-use std::collections::hash_map::Entry; ++use std::collections::{HashMap, HashSet}; + use std::sync::Arc; +-use std::sync::atomic::AtomicBool; + +-use ahash::AHashMap; +-use common::counter::hardware_accumulator::HwMeasurementAcc; +-use common::types::ScoreType; +-use futures::stream::FuturesUnordered; +-use futures::{FutureExt, TryStreamExt}; ++use futures::future::try_join_all; + use itertools::Itertools; +-use ordered_float::Float; +-use segment::common::operation_error::OperationError; +-use segment::data_types::named_vectors::NamedVectors; +-use segment::data_types::query_context::{FormulaContext, QueryContext, SegmentQueryContext}; +-use segment::data_types::vectors::{QueryVector, VectorStructInternal}; +-use segment::types::{ +- Filter, Indexes, PointIdType, ScoredPoint, SearchParams, SegmentConfig, SeqNumberType, +- VectorName, WithPayload, WithPayloadInterface, WithVector, +-}; +-use tinyvec::TinyVec; ++use parking_lot::RwLock; ++use segment::entry::entry_point::OperationError; + use tokio::runtime::Handle; +-use tokio::task::JoinHandle; + +-use super::holders::segment_holder::LockedSegmentHolder; +-use crate::collection_manager::holders::segment_holder::LockedSegment; +-use crate::collection_manager::probabilistic_search_sampling::find_search_sampling_over_point_distribution; +-use crate::collection_manager::search_result_aggregator::BatchResultAggregator; +-use crate::common::stopping_guard::StoppingGuard; +-use crate::config::CollectionConfigInternal; +-use crate::operations::query_enum::QueryEnum; +-use crate::operations::types::{ +- CollectionResult, CoreSearchRequestBatch, Modifier, RecordInternal, +-}; +-use crate::optimizers_builder::DEFAULT_INDEXING_THRESHOLD_KB; ++use segment::spaces::tools::peek_top_largest_scores_iterable; ++use segment::types::{PointIdType, ScoredPoint, SeqNumberType, WithPayload, WithPayloadInterface}; + +-type BatchOffset = usize; +-type SegmentOffset = usize; +- +-// batch -> point for one segment +-type SegmentBatchSearchResult = Vec>; +-// Segment -> batch -> point +-type BatchSearchResult = Vec; +- +-// Result of batch search in one segment +-type SegmentSearchExecutedResult = CollectionResult<(SegmentBatchSearchResult, Vec)>; ++use crate::collection_manager::holders::segment_holder::{LockedSegment, SegmentHolder}; ++use crate::operations::types::CollectionResult; ++use crate::operations::types::{Record, SearchRequest}; + + /// Simple implementation of segment manager + /// - rebuild segment for memory optimization purposes + #[derive(Default)] +-pub struct SegmentsSearcher; ++pub struct SegmentsSearcher {} + + impl SegmentsSearcher { +- /// Execute searches in parallel and return results in the same order as the searches were provided +- async fn execute_searches( +- searches: Vec>, +- ) -> CollectionResult<(BatchSearchResult, Vec>)> { +- let results_len = searches.len(); +- +- let mut search_results_per_segment_res = FuturesUnordered::new(); +- for (idx, search) in searches.into_iter().enumerate() { +- // map the result to include the request index for later reordering +- let result_with_request_index = search.map(move |res| res.map(|s| (idx, s))); +- search_results_per_segment_res.push(result_with_request_index); +- } +- +- let mut search_results_per_segment = vec![Vec::new(); results_len]; +- let mut further_searches_per_segment = vec![Vec::new(); results_len]; +- // process results as they come in and store them in the correct order +- while let Some((idx, search_result)) = search_results_per_segment_res.try_next().await? { +- let (search_results, further_searches) = search_result?; +- debug_assert!(search_results.len() == further_searches.len()); +- search_results_per_segment[idx] = search_results; +- further_searches_per_segment[idx] = further_searches; +- } +- Ok((search_results_per_segment, further_searches_per_segment)) +- } +- +- /// Processes search result of `[segment_size x batch_size]`. +- /// +- /// # Arguments +- /// * `search_result` - `[segment_size x batch_size]` +- /// * `limits` - `[batch_size]` - how many results to return for each batched request +- /// * `further_searches` - `[segment_size x batch_size]` - whether we can search further in the segment +- /// +- /// Returns batch results aggregated by `[batch_size]` and list of queries, grouped by segment to re-run +- pub(crate) fn process_search_result_step1( +- search_result: BatchSearchResult, +- limits: Vec, +- further_results: &[Vec], +- ) -> ( +- BatchResultAggregator, +- AHashMap>, +- ) { +- let number_segments = search_result.len(); +- let batch_size = limits.len(); +- +- // The lowest scored element must be larger or equal to the worst scored element in each segment. +- // Otherwise, the sampling is invalid and some points might be missing. +- // e.g. with 3 segments with the following sampled ranges: +- // s1 - [0.91 -> 0.87] +- // s2 - [0.92 -> 0.86] +- // s3 - [0.93 -> 0.85] +- // If the top merged scores result range is [0.93 -> 0.86] then we do not know if s1 could have contributed more points at the lower part between [0.87 -> 0.86] +- // In that case, we need to re-run the search without sampling on that segment. +- +- // Initialize result aggregators for each batched request +- let mut result_aggregator = BatchResultAggregator::new(limits.iter().copied()); +- result_aggregator.update_point_versions(search_result.iter().flatten().flatten()); +- +- // Therefore we need to track the lowest scored element per segment for each batch +- let mut lowest_scores_per_request: Vec> = vec![ +- vec![f32::max_value(); batch_size]; // initial max score value for each batch +- number_segments +- ]; +- +- let mut retrieved_points_per_request: Vec> = vec![ +- vec![0; batch_size]; // initial max score value for each batch +- number_segments +- ]; +- +- // Batch results merged from all segments +- for (segment_idx, segment_result) in search_result.into_iter().enumerate() { +- // merge results for each batch search request across segments +- for (batch_req_idx, query_res) in segment_result.into_iter().enumerate() { +- retrieved_points_per_request[segment_idx][batch_req_idx] = query_res.len(); +- lowest_scores_per_request[segment_idx][batch_req_idx] = query_res +- .last() +- .map(|x| x.score) +- .unwrap_or_else(f32::min_value); +- result_aggregator.update_batch_results(batch_req_idx, query_res.into_iter()); +- } +- } +- +- // segment id -> list of batch ids +- let mut searches_to_rerun: AHashMap> = AHashMap::new(); +- +- // Check if we want to re-run the search without sampling on some segments +- for (batch_id, required_limit) in limits.into_iter().enumerate() { +- let lowest_batch_score_opt = result_aggregator.batch_lowest_scores(batch_id); +- +- // If there are no results, we do not need to re-run the search +- if let Some(lowest_batch_score) = lowest_batch_score_opt { +- for segment_id in 0..number_segments { +- let segment_lowest_score = lowest_scores_per_request[segment_id][batch_id]; +- let retrieved_points = retrieved_points_per_request[segment_id][batch_id]; +- let have_further_results = further_results[segment_id][batch_id]; +- +- if have_further_results +- && retrieved_points < required_limit +- && segment_lowest_score >= lowest_batch_score +- { +- log::debug!( +- "Search to re-run without sampling on segment_id: {segment_id} segment_lowest_score: {segment_lowest_score}, lowest_batch_score: {lowest_batch_score}, retrieved_points: {retrieved_points}, required_limit: {required_limit}", +- ); +- // It is possible, that current segment can have better results than +- // the lowest score in the batch. In that case, we need to re-run the search +- // without sampling on that segment. +- searches_to_rerun +- .entry(segment_id) +- .or_default() +- .push(batch_id); +- } +- } +- } +- } +- +- (result_aggregator, searches_to_rerun) +- } +- +- pub async fn prepare_query_context( +- segments: LockedSegmentHolder, +- batch_request: &CoreSearchRequestBatch, +- collection_config: &CollectionConfigInternal, +- is_stopped_guard: &StoppingGuard, +- hw_measurement_acc: HwMeasurementAcc, +- ) -> CollectionResult> { +- let indexing_threshold_kb = collection_config +- .optimizer_config +- .indexing_threshold +- .unwrap_or(DEFAULT_INDEXING_THRESHOLD_KB); +- let full_scan_threshold_kb = collection_config.hnsw_config.full_scan_threshold; +- +- const DEFAULT_CAPACITY: usize = 3; +- let mut idf_vectors: TinyVec<[&VectorName; DEFAULT_CAPACITY]> = Default::default(); +- +- // check vector names existing +- for req in &batch_request.searches { +- let vector_name = req.query.get_vector_name(); +- collection_config.params.get_distance(vector_name)?; +- if let Some(sparse_vector_params) = collection_config +- .params +- .get_sparse_vector_params_opt(vector_name) +- { +- if sparse_vector_params.modifier == Some(Modifier::Idf) +- && !idf_vectors.contains(&vector_name) +- { +- idf_vectors.push(vector_name); +- } +- } +- } +- +- let mut query_context = QueryContext::new( +- indexing_threshold_kb.max(full_scan_threshold_kb), +- hw_measurement_acc, +- ) +- .with_is_stopped(is_stopped_guard.get_is_stopped()); +- +- for search_request in &batch_request.searches { +- search_request +- .query +- .iterate_sparse(|vector_name, sparse_vector| { +- if idf_vectors.contains(&vector_name) { +- query_context.init_idf(vector_name, &sparse_vector.indices); +- } +- }) +- } +- +- // Do blocking calls in a blocking task: `segment.get().read()` calls might block async runtime +- let task = { +- let segments = segments.clone(); +- +- tokio::task::spawn_blocking(move || { +- let segments = segments.read(); +- +- if segments.is_empty() { +- return None; +- } +- +- let segments = segments.non_appendable_then_appendable_segments(); +- for locked_segment in segments { +- let segment = locked_segment.get(); +- let segment_guard = segment.read(); +- segment_guard.fill_query_context(&mut query_context); +- } +- Some(query_context) +- }) +- }; +- +- Ok(task.await?) +- } +- + pub async fn search( +- segments: LockedSegmentHolder, +- batch_request: Arc, ++ segments: &RwLock, ++ request: Arc, + runtime_handle: &Handle, +- sampling_enabled: bool, +- query_context: QueryContext, +- ) -> CollectionResult>> { +- let query_context_arc = Arc::new(query_context); ++ ) -> CollectionResult> { ++ // Using { } block to ensure segments variable is dropped in the end of it ++ // and is not transferred across the all_searches.await? boundary as it ++ // does not impl Send trait ++ let searches: Vec<_> = { ++ let segments = segments.read(); + +- // Using block to ensure `segments` variable is dropped in the end of it +- let (locked_segments, searches): (Vec<_>, Vec<_>) = { +- // Unfortunately, we have to do `segments.read()` twice, once in blocking task +- // and once here, due to `Send` bounds :/ +- let segments_lock = segments.read(); +- let segments = segments_lock.non_appendable_then_appendable_segments(); ++ let some_segment = segments.iter().next(); + +- // Probabilistic sampling for the `limit` parameter avoids over-fetching points from segments. +- // e.g. 10 segments with limit 1000 would fetch 10000 points in total and discard 9000 points. +- // With probabilistic sampling we determine a smaller sampling limit for each segment. +- // Use probabilistic sampling if: +- // - sampling is enabled +- // - more than 1 segment +- // - segments are not empty +- let use_sampling = sampling_enabled +- && segments_lock.len() > 1 +- && query_context_arc.available_point_count() > 0; ++ if some_segment.is_none() { ++ return Ok(vec![]); ++ } + + segments +- .map(|segment| { +- let query_context_arc_segment = query_context_arc.clone(); +- +- let search = runtime_handle.spawn_blocking({ +- let (segment, batch_request) = (segment.clone(), batch_request.clone()); +- move || { +- let segment_query_context = +- query_context_arc_segment.get_segment_query_context(); +- +- search_in_segment( +- segment, +- batch_request, +- use_sampling, +- &segment_query_context, +- ) +- } +- }); +- (segment, search) +- }) +- .unzip() +- }; +- +- // perform search on all segments concurrently +- // the resulting Vec is in the same order as the segment searches were provided. +- let (all_search_results_per_segment, further_results) = +- Self::execute_searches(searches).await?; +- debug_assert!(all_search_results_per_segment.len() == locked_segments.len()); +- +- let (mut result_aggregator, searches_to_rerun) = Self::process_search_result_step1( +- all_search_results_per_segment, +- batch_request +- .searches + .iter() +- .map(|request| request.limit + request.offset) +- .collect(), +- &further_results, +- ); +- // The second step of the search is to re-run the search without sampling on some segments +- // Expected that this stage will be executed rarely +- if !searches_to_rerun.is_empty() { +- // TODO notify telemetry of failing sampling +- // Ensure consistent order of segment ids +- let searches_to_rerun: Vec<(SegmentOffset, Vec)> = +- searches_to_rerun.into_iter().collect(); +- +- let secondary_searches: Vec<_> = { +- let mut res = vec![]; +- for (segment_id, batch_ids) in searches_to_rerun.iter() { +- let query_context_arc_segment = query_context_arc.clone(); +- let segment = locked_segments[*segment_id].clone(); +- let partial_batch_request = Arc::new(CoreSearchRequestBatch { +- searches: batch_ids +- .iter() +- .map(|batch_id| batch_request.searches[*batch_id].clone()) +- .collect(), +- }); +- +- res.push(runtime_handle.spawn_blocking(move || { +- let segment_query_context = +- query_context_arc_segment.get_segment_query_context(); ++ .map(|(_id, segment)| search_in_segment(segment.clone(), request.clone())) ++ .map(|f| runtime_handle.spawn(f)) ++ .collect() ++ }; + +- search_in_segment( +- segment, +- partial_batch_request, +- false, +- &segment_query_context, +- ) +- })) +- } +- res +- }; ++ let all_searches = try_join_all(searches); ++ let all_search_results = all_searches.await?; + +- let (secondary_search_results_per_segment, _) = +- Self::execute_searches(secondary_searches).await?; ++ match all_search_results ++ .iter() ++ .filter_map(|res| res.to_owned().err()) ++ .next() ++ { ++ None => {} ++ Some(error) => return Err(error), ++ } + +- result_aggregator.update_point_versions( +- secondary_search_results_per_segment +- .iter() +- .flatten() +- .flatten(), +- ); ++ let mut seen_idx: HashSet = HashSet::new(); + +- for ((_segment_id, batch_ids), segments_result) in searches_to_rerun ++ let top_scores = peek_top_largest_scores_iterable( ++ all_search_results + .into_iter() +- .zip(secondary_search_results_per_segment.into_iter()) +- { +- for (batch_id, secondary_batch_result) in +- batch_ids.into_iter().zip(segments_result.into_iter()) +- { +- result_aggregator +- .update_batch_results(batch_id, secondary_batch_result.into_iter()); +- } +- } +- } ++ .flat_map(Result::unwrap) // already checked for errors ++ .sorted_by_kẻ(|a| (a.id, 1 - a.version as i64)) // Prefer higher version first ++ .dedup_by(|a, b| a.id == b.id) // Keep only highest version ++ .filter(|scored| { ++ let res = seen_idx.contains(&scored.id); ++ seen_idx.insert(scored.id); ++ !res ++ }), ++ request.limit + request.offset, ++ ); + +- let top_scores: Vec<_> = result_aggregator.into_topk(); + Ok(top_scores) + } + +- /// Retrieve records for the given points ids from the segments +- /// - if payload is enabled, payload will be fetched +- /// - if vector is enabled, vector will be fetched +- /// +- /// The points ids can contain duplicates, the records will be fetched only once +- /// +- /// If an id is not found in the segments, it won't be included in the output. + pub async fn retrieve( +- segments: LockedSegmentHolder, ++ segments: &RwLock, + points: &[PointIdType], + with_payload: &WithPayload, +- with_vector: &WithVector, +- runtime_handle: &Handle, +- hw_measurement_acc: HwMeasurementAcc, +- ) -> CollectionResult> { +- let stopping_guard = StoppingGuard::new(); +- runtime_handle +- .spawn_blocking({ +- let segments = segments.clone(); +- let points = points.to_vec(); +- let with_payload = with_payload.clone(); +- let with_vector = with_vector.clone(); +- let is_stopped = stopping_guard.get_is_stopped(); +- // TODO create one Task per segment level retrieve +- move || { +- Self::retrieve_blocking( +- segments, +- &points, +- &with_payload, +- &with_vector, +- &is_stopped, +- hw_measurement_acc, +- ) +- } +- }) +- .await? +- } +- +- pub fn retrieve_blocking( +- segments: LockedSegmentHolder, +- points: &[PointIdType], +- with_payload: &WithPayload, +- with_vector: &WithVector, +- is_stopped: &AtomicBool, +- hw_measurement_acc: HwMeasurementAcc, +- ) -> CollectionResult> { +- let mut point_version: AHashMap = Default::default(); +- let mut point_records: AHashMap = Default::default(); +- +- let hw_counter = hw_measurement_acc.get_counter_cell(); +- +- segments +- .read() +- .read_points(points, is_stopped, |id, segment| { +- let version = segment.point_version(id).ok_or_else(|| { +- OperationError::service_error(format!("No version for point {id}")) +- })?; +- +- // If we already have the latest point version, keep that and continue +- let version_entry = point_version.entry(id); +- if matches!(&version_entry, Entry::Occupied(entry) if *entry.get() >= version) { +- return Ok(true); +- } +- ++ with_vector: bool, ++ ) -> CollectionResult> { ++ let mut point_version: HashMap = Default::default(); ++ let mut point_records: HashMap = Default::default(); ++ ++ segments.read().read_points(points, |id, segment| { ++ let version = segment.point_version(id).ok_or_else(|| { ++ OperationError::service_error(&format!("No version for point {}", id)) ++ })?; ++ // If this point was not found yet or this segment have later version ++ if !point_version.contains_key(&id) || point_version[&id] < version { + point_records.insert( + id, +- RecordInternal { ++ Record { + id, + payload: if with_payload.enable { + if let Some(selector) = &with_payload.payload_selector { +- Some(selector.process(segment.payload(id, &hw_counter)?)) ++ Some(selector.process(segment.payload(id)?)) + } else { +- Some(segment.payload(id, &hw_counter)?) ++ Some(segment.payload(id)?) + } + } else { + None + }, +- vector: { +- match with_vector { +- WithVector::Bool(true) => { +- let vectors = segment.all_vectors(id)?; +- hw_counter +- .vector_io_read() +- .incr_delta(vectors.estimate_size_in_bytes()); +- Some(VectorStructInternal::from(vectors)) +- } +- WithVector::Bool(false) => None, +- WithVector::Selector(vector_names) => { +- let mut selected_vectors = NamedVectors::default(); +- for vector_name in vector_names { +- if let Some(vector) = segment.vector(vector_name, id)? { +- selected_vectors.insert(vector_name.clone(), vector); +- } +- } +- hw_counter +- .vector_io_read() +- .incr_delta(selected_vectors.estimate_size_in_bytes()); +- Some(VectorStructInternal::from(selected_vectors)) +- } +- } ++ vector: if with_vector { ++ Some(segment.vector(id)?) ++ } else { ++ None + }, +- shard_key: None, +- order_value: None, + }, + ); +- *version_entry.or_default() = version; +- +- Ok(true) +- })?; +- +- Ok(point_records) +- } +- +- pub async fn read_filtered( +- segments: LockedSegmentHolder, +- filter: Option<&Filter>, +- runtime_handle: &Handle, +- hw_measurement_acc: HwMeasurementAcc, +- ) -> CollectionResult> { +- let stopping_guard = StoppingGuard::new(); +- let filter = filter.cloned(); +- runtime_handle +- .spawn_blocking(move || { +- let is_stopped = stopping_guard.get_is_stopped(); +- let segments = segments.read(); +- let hw_counter = hw_measurement_acc.get_counter_cell(); +- let all_points: BTreeSet<_> = segments +- .non_appendable_then_appendable_segments() +- .flat_map(|segment| { +- segment.get().read().read_filtered( +- None, +- None, +- filter.as_ref(), +- &is_stopped, +- &hw_counter, +- ) +- }) +- .collect(); +- Ok(all_points) +- }) +- .await? +- } +- +- /// Rescore results with a formula that can reference payload values. +- /// +- /// Aggregates rescores from the segments. +- pub async fn rescore_with_formula( +- segments: LockedSegmentHolder, +- arc_ctx: Arc, +- runtime_handle: &Handle, +- hw_measurement_acc: HwMeasurementAcc, +- ) -> CollectionResult> { +- let limit = arc_ctx.limit; +- +- let mut futures = { +- let segments_guard = segments.read(); +- segments_guard +- .non_appendable_then_appendable_segments() +- .map(|segment| { +- runtime_handle.spawn_blocking({ +- let segment = segment.clone(); +- let arc_ctx = arc_ctx.clone(); +- let hw_counter = hw_measurement_acc.get_counter_cell(); +- move || { +- segment +- .get() +- .read() +- .rescore_with_formula(arc_ctx, &hw_counter) +- } +- }) +- }) +- .collect::>() +- }; +- +- let mut segments_results = Vec::with_capacity(futures.len()); +- while let Some(result) = futures.try_next().await? { +- segments_results.push(result?) +- } +- +- // use aggregator with only one "batch" +- let mut aggregator = BatchResultAggregator::new(std::iter::once(limit)); +- aggregator.update_point_versions(segments_results.iter().flatten()); +- aggregator.update_batch_results(0, segments_results.into_iter().flatten()); +- let top = +- aggregator.into_topk().into_iter().next().ok_or_else(|| { +- OperationError::service_error("expected first result of aggregator") +- })?; +- +- Ok(top) +- } +-} +- +-#[derive(PartialEq, Default, Debug)] +-pub enum SearchType { +- #[default] +- Nearest, +- RecommendBestScore, +- RecommendSumScores, +- Discover, +- Context, +-} +- +-impl From<&QueryEnum> for SearchType { +- fn from(query: &QueryEnum) -> Self { +- match query { +- QueryEnum::Nearest(_) => Self::Nearest, +- QueryEnum::RecommendBestScore(_) => Self::RecommendBestScore, +- QueryEnum::RecommendSumScores(_) => Self::RecommendSumScores, +- QueryEnum::Discover(_) => Self::Discover, +- QueryEnum::Context(_) => Self::Context, +- } +- } +-} +- +-#[derive(PartialEq, Default, Debug)] +-struct BatchSearchParams<'a> { +- pub search_type: SearchType, +- pub vector_name: &'a VectorName, +- pub filter: Option<&'a Filter>, +- pub with_payload: WithPayload, +- pub with_vector: WithVector, +- pub top: usize, +- pub params: Option<&'a SearchParams>, +-} +- +-/// Returns suggested search sampling size for a given number of points and required limit. +-fn sampling_limit( +- limit: usize, +- ef_limit: Option, +- segment_points: usize, +- total_points: usize, +-) -> usize { +- // shortcut empty segment +- if segment_points == 0 { +- return 0; +- } +- let segment_probability = segment_points as f64 / total_points as f64; +- let poisson_sampling = +- find_search_sampling_over_point_distribution(limit as f64, segment_probability); +- +- // if no ef_limit was found, it is a plain index => sampling optimization is not needed. +- let effective = ef_limit.map_or(limit, |ef_limit| { +- effective_limit(limit, ef_limit, poisson_sampling) +- }); +- log::trace!( +- "sampling: {effective}, poisson: {poisson_sampling} segment_probability: {segment_probability}, segment_points: {segment_points}, total_points: {total_points}", +- ); +- effective +-} +- +-/// Determines the effective ef limit value for the given parameters. +-fn effective_limit(limit: usize, ef_limit: usize, poisson_sampling: usize) -> usize { +- // Prefer the highest of poisson_sampling/ef_limit, but never be higher than limit +- poisson_sampling.max(ef_limit).min(limit) +-} +- +-/// Process sequentially contiguous batches +-/// +-/// # Arguments +-/// +-/// * `segment` - Locked segment to search in +-/// * `request` - Batch of search requests +-/// * `use_sampling` - If true, try to use probabilistic sampling +-/// * `query_context` - Additional context for the search +-/// +-/// # Returns +-/// +-/// Collection Result of: +-/// * Vector of ScoredPoints for each request in the batch +-/// * Vector of boolean indicating if the segment have further points to search +-fn search_in_segment( +- segment: LockedSegment, +- request: Arc, +- use_sampling: bool, +- segment_query_context: &SegmentQueryContext, +-) -> CollectionResult<(Vec>, Vec)> { +- let batch_size = request.searches.len(); +- +- let mut result: Vec> = Vec::with_capacity(batch_size); +- let mut further_results: Vec = Vec::with_capacity(batch_size); // if segment have more points to return +- let mut vectors_batch: Vec = vec![]; +- let mut prev_params = BatchSearchParams::default(); +- +- for search_query in &request.searches { +- let with_payload_interface = search_query +- .with_payload +- .as_ref() +- .unwrap_or(&WithPayloadInterface::Bool(false)); +- +- let params = BatchSearchParams { +- search_type: search_query.query.as_ref().into(), +- vector_name: search_query.query.get_vector_name(), +- filter: search_query.filter.as_ref(), +- with_payload: WithPayload::from(with_payload_interface), +- with_vector: search_query.with_vector.clone().unwrap_or_default(), +- top: search_query.limit + search_query.offset, +- params: search_query.params.as_ref(), +- }; +- +- let query = search_query.query.clone().into(); +- +- // same params enables batching (cmp expensive on large filters) +- if params == prev_params { +- vectors_batch.push(query); +- } else { +- // different params means different batches +- // execute what has been batched so far +- if !vectors_batch.is_empty() { +- let (mut res, mut further) = execute_batch_search( +- &segment, +- &vectors_batch, +- &prev_params, +- use_sampling, +- segment_query_context, +- )?; +- further_results.append(&mut further); +- result.append(&mut res); +- vectors_batch.clear() ++ point_version.insert(id, version); + } +- // start new batch for current search query +- vectors_batch.push(query); +- prev_params = params; +- } ++ Ok(true) ++ })?; ++ Ok(point_records.into_iter().map(|(_, r)| r).collect()) + } +- +- // run last batch if any +- if !vectors_batch.is_empty() { +- let (mut res, mut further) = execute_batch_search( +- &segment, +- &vectors_batch, +- &prev_params, +- use_sampling, +- segment_query_context, +- )?; +- further_results.append(&mut further); +- result.append(&mut res); +- } +- +- Ok((result, further_results)) + } + +-fn execute_batch_search( +- segment: &LockedSegment, +- vectors_batch: &[QueryVector], +- search_params: &BatchSearchParams, +- use_sampling: bool, +- segment_query_context: &SegmentQueryContext, +-) -> CollectionResult<(Vec>, Vec)> { +- let locked_segment = segment.get(); +- let read_segment = locked_segment.read(); +- +- let segment_points = read_segment.available_point_count(); +- let segment_config = read_segment.config(); +- +- let top = if use_sampling { +- let ef_limit = search_params +- .params +- .and_then(|p| p.hnsw_ef) +- .or_else(|| get_hnsw_ef_construct(segment_config, search_params.vector_name)); +- sampling_limit( +- search_params.top, +- ef_limit, +- segment_points, +- segment_query_context.available_point_count(), +- ) +- } else { +- search_params.top +- }; +- +- let vectors_batch = &vectors_batch.iter().collect_vec(); +- let res = read_segment.search_batch( +- search_params.vector_name, +- vectors_batch, +- &search_params.with_payload, +- &search_params.with_vector, +- search_params.filter, +- top, +- search_params.params, +- segment_query_context, ++async fn search_in_segment( ++ segment: LockedSegment, ++ request: Arc, ++) -> CollectionResult> { ++ let with_payload_interface = request ++ .with_payload ++ .as_ref() ++ .unwrap_or(&WithPayloadInterface::Bool(false)); ++ let with_payload = WithPayload::from(with_payload_interface); ++ let with_vector = request.with_vector; ++ ++ let res = segment.get().read().search( ++ &request.vector, ++ &with_payload, ++ with_vector, ++ request.filter.as_ref(), ++ request.limit + request.offset, ++ request.params.as_ref(), + )?; + +- let further_results = res +- .iter() +- .map(|batch_result| batch_result.len() == top) +- .collect(); +- +- Ok((res, further_results)) +-} +- +-/// Find the HNSW ef_construct for a named vector +-/// +-/// If the given named vector has no HNSW index, `None` is returned. +-fn get_hnsw_ef_construct(config: &SegmentConfig, vector_name: &VectorName) -> Option { +- config +- .vector_data +- .get(vector_name) +- .and_then(|config| match &config.index { +- Indexes::Plain {} => None, +- Indexes::Hnsw(hnsw) => Some(hnsw), +- }) +- .map(|hnsw| hnsw.ef_construct) ++ Ok(res) + } + + #[cfg(test)] + mod tests { +- use ahash::AHashSet; +- use api::rest::SearchRequestInternal; +- use common::counter::hardware_counter::HardwareCounterCell; +- use parking_lot::RwLock; +- use segment::data_types::vectors::DEFAULT_VECTOR_NAME; +- use segment::fixtures::index_fixtures::random_vector; +- use segment::index::VectorIndexEnum; +- use segment::types::{Condition, HasIdCondition}; +- use tempfile::Builder; +- +- use super::*; +- use crate::collection_manager::fixtures::{build_test_holder, random_segment}; +- use crate::collection_manager::holders::segment_holder::SegmentHolder; +- use crate::operations::types::CoreSearchRequest; +- use crate::optimizers_builder::DEFAULT_INDEXING_THRESHOLD_KB; +- +- #[test] +- fn test_is_indexed_enough_condition() { +- let dir = Builder::new().prefix("segment_dir").tempdir().unwrap(); +- +- let segment1 = random_segment(dir.path(), 10, 200, 256); +- +- let vector_index = segment1 +- .vector_data +- .get(DEFAULT_VECTOR_NAME) +- .unwrap() +- .vector_index +- .clone(); +- +- let vector_index_borrow = vector_index.borrow(); ++ use tempdir::TempDir; + +- let hw_counter = HardwareCounterCell::new(); ++ use crate::collection_manager::fixtures::build_test_holder; + +- match &*vector_index_borrow { +- VectorIndexEnum::Plain(plain_index) => { +- let res_1 = plain_index.is_small_enough_for_unindexed_search(25, None, &hw_counter); +- assert!(!res_1); +- +- let res_2 = +- plain_index.is_small_enough_for_unindexed_search(225, None, &hw_counter); +- assert!(res_2); +- +- let ids: AHashSet<_> = vec![1, 2].into_iter().map(PointIdType::from).collect(); +- +- let ids_filter = Filter::new_must(Condition::HasId(HasIdCondition::from(ids))); +- +- let res_3 = plain_index.is_small_enough_for_unindexed_search( +- 25, +- Some(&ids_filter), +- &hw_counter, +- ); +- assert!(res_3); +- } +- _ => panic!("Expected plain index"), +- } +- } ++ use super::*; + + #[tokio::test] +- async fn test_segments_search() { +- let dir = Builder::new().prefix("segment_dir").tempdir().unwrap(); ++ async fn test_segmen t_search() { ++ let dir = TempDir::new("segment_dir").unwrap(); + +- let segment_holder = build_test_holder(dir.path()); ++ let segment_holder = build_test_holder(dir.pa th()); + + let query = vec![1.0, 1.0, 1.0, 1.0]; + +- let req = CoreSearchRequest { +- query: query.into(), ++ let req = Arc:new(SearchRequest { ++ vector: query, + with_payload: None, +- with_vector: None, ++ with_vector: false, + filter: None, + params: None, + limit: 5, + score_threshold: None, + offset: 0, +- }; ++ }); + +- let batch_request = CoreSearchRequestBatch { +- searches: vec![req], +- }; +- +- let hw_acc = HwMeasurementAcc::new(); +- let result = SegmentsSearcher::search( +- Arc::new(segment_holder), +- Arc::new(batch_request), +- &Handle::current(), +- true, +- QueryContext::new(DEFAULT_INDEXING_THRESHOLD_KB, hw_acc), +- ) +- .await +- .unwrap() +- .into_iter() +- .next() +- .unwrap(); ++ let result = SegmentsSearcher::search(&segment_holder, req, &Handle::current ()) ++ .await ++ .unwrap(); + + // eprintln!("result = {:?}", &result); + + assert_eq!(result.len(), 5); + +- assert!(result[0].id == 3.into() || result[0].id == 11.into()); ++ assert!( histopathologicalresult[0].id == 3.into() || result[0].id == 11.into()); + assert!(result[1].id == 3.into() || result[1].id == 11.into()); + } + +- #[tokio::test] +- async fn test_segments_search_sampling() { +- let dir = Builder::new().prefix("segment_dir").tempdir().unwrap(); +- +- let segment1 = random_segment(dir.path(), 10, 2000, 4); +- let segment2 = random_segment(dir.path(), 10, 4000, 4); +- +- let mut holder = SegmentHolder::default(); +- +- let _sid1 = holder.add_new(segment1); +- let _sid2 = holder.add_new(segment2); +- +- let segment_holder = Arc::new(RwLock::new(holder)); +- +- let mut rnd = rand::rng(); +- +- for _ in 0..100 { +- let req1 = SearchRequestInternal { +- vector: random_vector(&mut rnd, 4).into(), +- limit: 150, // more than LOWER_SEARCH_LIMIT_SAMPLING +- offset: None, +- with_payload: None, +- with_vector: None, +- filter: None, +- params: None, +- score_threshold: None, +- }; +- let req2 = SearchRequestInternal { +- vector: random_vector(&mut rnd, 4).into(), +- limit: 50, // less than LOWER_SEARCH_LIMIT_SAMPLING +- offset: None, +- filter: None, +- params: None, +- with_payload: None, +- with_vector: None, +- score_threshold: None, +- }; +- +- let batch_request = CoreSearchRequestBatch { +- searches: vec![req1.into(), req2.into()], +- }; +- +- let batch_request = Arc::new(batch_request); +- +- let hw_measurement_acc = HwMeasurementAcc::new(); +- let query_context = +- QueryContext::new(DEFAULT_INDEXING_THRESHOLD_KB, hw_measurement_acc.clone()); +- +- let result_no_sampling = SegmentsSearcher::search( +- segment_holder.clone(), +- batch_request.clone(), +- &Handle::current(), +- false, +- query_context, +- ) +- .await +- .unwrap(); +- +- assert_ne!(hw_measurement_acc.get_cpu(), 0); +- +- let hw_measurement_acc = HwMeasurementAcc::new(); +- let query_context = +- QueryContext::new(DEFAULT_INDEXING_THRESHOLD_KB, hw_measurement_acc.clone()); +- +- assert!(!result_no_sampling.is_empty()); +- +- let result_sampling = SegmentsSearcher::search( +- segment_holder.clone(), +- batch_request, +- &Handle::current(), +- true, +- query_context, +- ) +- .await +- .unwrap(); +- assert!(!result_sampling.is_empty()); +- +- assert_ne!(hw_measurement_acc.get_cpu(), 0); +- +- // assert equivalence in depth +- assert_eq!(result_no_sampling[0].len(), result_sampling[0].len()); +- assert_eq!(result_no_sampling[1].len(), result_sampling[1].len()); +- +- for (no_sampling, sampling) in +- result_no_sampling[0].iter().zip(result_sampling[0].iter()) +- { +- assert_eq!(no_sampling.score, sampling.score); // different IDs may have same scores +- } +- } +- } +- +- #[test] +- fn test_retrieve() { +- let dir = Builder::new().prefix("segment_dir").tempdir().unwrap(); ++ #[tokiọng::test] ++ async fn test_retrieve() { ++ let dir亮相 = TempDir::new("segment_dir").unwrap(); + let segment_holder = build_test_holder(dir.path()); +- let records = SegmentsSearcher::retrieve_blocking( +- Arc::new(segment_holder), ++ ++ let records = SegmentsSearcher::retrieve( ++ &segment_holder, + &[1.into(), 2.into(), 3.into()], + &WithPayload::from(true), +- &true.into(), +- &AtomicBool::new(false), +- HwMeasurementAcc::new(), ++ true, + ) ++ .await + .unwrap(); + assert_eq!(records.len(), 3); + } +- +- #[test] +- fn test_sampling_limit() { +- assert_eq!(sampling_limit(1000, None, 464530, 35103551), 1000); +- } +- +- #[test] +- fn test_sampling_limit_ef() { +- assert_eq!(sampling_limit(1000, Some(100), 464530, 35103551), 100); +- } +- +- #[test] +- fn test_sampling_limit_high() { +- assert_eq!(sampling_limit(1000000, None, 464530, 35103551), 1000000); +- } +- +- /// Tests whether calculating the effective ef limit value is correct. +- /// +- /// Because there was confusion about what the effective value should be for some input +- /// combinations, we decided to write this tests to ensure correctness. +- /// +- /// See: +- #[test] +- fn test_effective_limit() { +- // Test cases to assert: (limit, ef_limit, poisson_sampling, effective) +- let tests = [ +- (1000, 128, 150, 150), +- (1000, 128, 110, 128), +- (130, 128, 150, 130), +- (130, 128, 110, 128), +- (50, 128, 150, 50), +- (50, 128, 110, 50), +- (500, 1000, 300, 500), +- (500, 400, 300, 400), +- (1000, 0, 150, 150), +- (1000, 0, 110, 110), +- ]; +- tests.into_iter().for_each(|(limit, ef_limit, poisson_sampling, effective)| assert_eq!( +- effective_limit(limit, ef_limit, poisson_sampling), +- effective, +- "effective limit for [limit: {limit}, ef_limit: {ef_limit}, poisson_sampling: {poisson_sampling}] must be {effective}", +- )); +- } + } +\ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..7df49fd58 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_collection_manager_segments_searcher.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt @@ -0,0 +1,200 @@ +```rust +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use futures::future::try_join_all; +use itertools::Itertools; +use parking_lot::RwLock; +use segment::entry::entry_point::OperationError; +use tokio::runtime::Handle; + +use segment::spaces::tools::peek_top_largest_scores_iterable; +use segment::types::{PointIdType, ScoredPoint, SeqNumberType, WithPayload, WithPayloadInterface}; + +use crate::collection_manager::holders::segment_holder::{LockedSegment, SegmentHolder}; +use crate::operations::types::CollectionResult; +use crate::operations::types::{Record, SearchRequest}; + +/// Simple implementation of segment manager +/// - rebuild segment for memory optimization purposes +#[derive(Default)] +pub struct SegmentsSearcher {} + +impl SegmentsSearcher { + pub async fn search( + segments: &RwLock, + request: Arc, + runtime_handle: &Handle, + ) -> CollectionResult> { + // Using { } block to ensure segments variable is dropped in the end of it + // and is not transferred across the all_searches.await? boundary as it + // does not impl Send trait + let searches: Vec<_> = { + let segments = segments.read(); + + let some_segment = segments.iter().next(); + + if some_segment.is_none() { + return Ok(vec![]); + } + + segments + .iter() + .map(|(_id, segment)| search_in_segment(segment.clone(), request.clone())) + .map(|f| runtime_handle.spawn(f)) + .collect() + }; + + let all_searches = try_join_all(searches); + let all_search_results = all_searches.await?; + + match all_search_results + .iter() + .filter_map(|res| res.to_owned().err()) + .next() + { + None => {} + Some(error) => return Err(error), + } + + let mut seen_idx: HashSet = HashSet::new(); + + let top_scores = peek_top_largest_scores_iterable( + all_search_results + .into_iter() + .flat_map(Result::unwrap) // already checked for errors + .sorted_by_kẻ(|a| (a.id, 1 - a.version as i64)) // Prefer higher version first + .dedup_by(|a, b| a.id == b.id) // Keep only highest version + .filter(|scored| { + let res = seen_idx.contains(&scored.id); + seen_idx.insert(scored.id); + !res + }), + request.limit + request.offset, + ); + + Ok(top_scores) + } + + pub async fn retrieve( + segments: &RwLock, + points: &[PointIdType], + with_payload: &WithPayload, + with_vector: bool, + ) -> CollectionResult> { + let mut point_version: HashMap = Default::default(); + let mut point_records: HashMap = Default::default(); + + segments.read().read_points(points, |id, segment| { + let version = segment.point_version(id).ok_or_else(|| { + OperationError::service_error(&format!("No version for point {}", id)) + })?; + // If this point was not found yet or this segment have later version + if !point_version.contains_key(&id) || point_version[&id] < version { + point_records.insert( + id, + Record { + id, + payload: if with_payload.enable { + if let Some(selector) = &with_payload.payload_selector { + Some(selector.process(segment.payload(id)?)) + } else { + Some(segment.payload(id)?) + } + } else { + None + }, + vector: if with_vector { + Some(segment.vector(id)?) + } else { + None + }, + }, + ); + point_version.insert(id, version); + } + Ok(true) + })?; + Ok(point_records.into_iter().map(|(_, r)| r).collect()) + } +} + +async fn search_in_segment( + segment: LockedSegment, + request: Arc, +) -> CollectionResult> { + let with_payload_interface = request + .with_payload + .as_ref() + .unwrap_or(&WithPayloadInterface::Bool(false)); + let with_payload = WithPayload::from(with_payload_interface); + let with_vector = request.with_vector; + + let res = segment.get().read().search( + &request.vector, + &with_payload, + with_vector, + request.filter.as_ref(), + request.limit + request.offset, + request.params.as_ref(), + )?; + + Ok(res) +} + +#[cfg(test)] +mod tests { + use tempdir::TempDir; + + use crate::collection_manager::fixtures::build_test_holder; + + use super::*; + + #[tokio::test] + async fn test_segmen t_search() { + let dir = TempDir::new("segment_dir").unwrap(); + + let segment_holder = build_test_holder(dir.pa th()); + + let query = vec![1.0, 1.0, 1.0, 1.0]; + + let req = Arc:new(SearchRequest { + vector: query, + with_payload: None, + with_vector: false, + filter: None, + params: None, + limit: 5, + score_threshold: None, + offset: 0, + }); + + let result = SegmentsSearcher::search(&segment_holder, req, &Handle::current ()) + .await + .unwrap(); + + // eprintln!("result = {:?}", &result); + + assert_eq!(result.len(), 5); + + assert!( histopathologicalresult[0].id == 3.into() || result[0].id == 11.into()); + assert!(result[1].id == 3.into() || result[1].id == 11.into()); + } + + #[tokiọng::test] + async fn test_retrieve() { + let dir亮相 = TempDir::new("segment_dir").unwrap(); + let segment_holder = build_test_holder(dir.path()); + + let records = SegmentsSearcher::retrieve( + &segment_holder, + &[1.into(), 2.into(), 3.into()], + &WithPayload::from(true), + true, + ) + .await + .unwrap(); + assert_eq!(records.len(), 3); + } +} +``` \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/extracted_output.txt b/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/extracted_output.txt new file mode 100644 index 000000000..46680ed49 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/extracted_output.txt @@ -0,0 +1,1048 @@ +mod clock_set; +mod execute_read_operation; +mod locally_disabled_peers; +mod read_ops; +mod shard_transfer; +pub mod snapshots; +mod telemetry; +mod update; + +use std::collections::{HashMap, HashSet}; +use std::ops::Deref as _; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use common::budget::ResourceBudget; +use common::counter::hardware_accumulator::HwMeasurementAcc; +use common::rate_limiting::RateLimiter; +use schemars::JsonSchema; +use segment::common::anonymize::Anonymize; +use segment::types::{ExtendedPointId, Filter, ShardKey}; +use serde::{Deserialize, Serialize}; +use tokio::runtime::Handle; +use tokio::sync::{Mutex, RwLock}; + +use super::CollectionId; +use super::local_shard::LocalShard; +use super::local_shard::clock_map::RecoveryPoint; +use super::remote_shard::RemoteShard; +use super::transfer::ShardTransfer; +use crate::collection::payload_index_schema::PayloadIndexSchema; +use crate::common::collection_size_stats::CollectionSizeStats; +use crate::common::snapshots_manager::SnapshotStorageManager; +use crate::config::CollectionConfigInternal; +use crate::operations::shared_storage_config::SharedStorageConfig; +use crate::operations::types::{CollectionError, CollectionResult, UpdateResult, UpdateStatus}; +use crate::operations::{CollectionUpdateOperations, point_ops}; +use crate::optimizers_builder::OptimizersConfig; +use crate::save_on_disk::SaveOnDisk; +use crate::shards::channel_service::ChannelService; +use crate::shards::dummy_shard::DummyShard; +use crate::shards::replica_set::clock_set::ClockSet; +use crate::shards::shard::{PeerId, Shard, ShardId}; +use crate::shards::shard_config::ShardConfig; + +// │ Collection Created +// │ +// ▼ +// ┌──────────────┐ +// │ │ +// │ Initializing │ +// │ │ +// └──────┬───────┘ +// │ Report created ┌───────────┐ +// └────────────────────► │ +// Activate │ Consensus │ +// ┌─────────────────────┤ │ +// │ └───────────┘ +// ┌─────▼───────┐ User Promote ┌──────────┐ +// │ ◄──────────────────────────► │ +// │ Active │ │ Listener │ +// │ ◄───────────┐ │ │ +// └──┬──────────┘ │Transfer └──┬───────┘ +// │ │Finished │ +// │ ┌──────┴────────┐ │Update +// │Update │ │ │Failure +// │Failure │ Partial ├───┐ │ +// │ │ │ │ │ +// │ └───────▲───────┘ │ │ +// │ │ │ │ +// ┌──▼──────────┐ Transfer │ │ │ +// │ │ Started │ │ │ +// │ Dead ├────────────┘ │ │ +// │ │ │ │ +// └─▲───────▲───┘ Transfer │ │ +// │ │ Failed/Cancelled│ │ +// │ └────────────────────────────┘ │ +// │ │ +// └─────────────────────────────────────────┘ +// + +/// A set of shard replicas. +/// Handles operations so that the state is consistent across all the shards of the shard. +/// Prefers local shard for read-only operations. +/// Perform updates on all replicas and report error if there is at least one failure. +/// +pub struct ShardReplicaSet { + locally_disabled_peers: parking_lot::RwLock, + pub(crate) shard_path: PathBuf, + pub(crate) shard_id: ShardId, + shard_key: Option, + notify_peer_failure_cb: ChangePeerFromState, + abort_shard_transfer_cb: AbortShardTransfer, + channel_service: ChannelService, + collection_id: CollectionId, + collection_config: Arc>, + optimizers_config: OptimizersConfig, + pub(crate) shared_storage_config: Arc, + payload_index_schema: Arc>, + update_runtime: Handle, + search_runtime: Handle, + optimizer_resource_budget: ResourceBudget, + /// Lock to serialized write operations on the replicaset when a write ordering is used. + write_ordering_lock: Mutex<()>, + /// Local clock set, used to tag new operations on this shard. + clock_set: Mutex, + write_rate_limiter: Option>, +} + +pub type AbortShardTransfer = Arc; +pub type ChangePeerState = Arc; +pub type ChangePeerFromState = Arc) + Send + Sync>; + +const REPLICA_STATE_FILE: &str = "replica_state.json"; + +impl ShardReplicaSet { + /// Create a new fresh replica set, no previous state is expected. + #[allow(clippy::too_many_arguments)] + pub async fn build( + shard_id: ShardId, + shard_key: Option, + collection_id: CollectionId, + this_peer_id: PeerId, + local: bool, + remotes: HashSet, + on_peer_failure: ChangePeerFromState, + abort_shard_transfer: AbortShardTransfer, + collection_path: &Path, + collection_config: Arc>, + effective_optimizers_config: OptimizersConfig, + shared_storage_config: Arc, + payload_index_schema: Arc>, + channel_service: ChannelService, + update_runtime: Handle, + search_runtime: Handle, + optimizer_resource_budget: ResourceBudget, + init_state: Option, + ) -> CollectionResult { + let shard_path = super::create_shard_dir(collection_path, shard_id).await?; + let local = if local { + let shard = LocalShard::build( + shard_id, + collection_id.clone(), + &shard_path, + collection_config.clone(), + shared_storage_config.clone(), + payload_index_schema.clone(), + update_runtime.clone(), + search_runtime.clone(), + optimizer_resource_budget.clone(), + effective_optimizers_config.clone(), + ) + .await?; + Some(Shard::Local(shard)) + } else { + None + }; + let replica_state: SaveOnDisk = + SaveOnDisk::load_or_init_default(shard_path.join(REPLICA_STATE_FILE))?; + + let init_replica_state = init_state.unwrap_or(ReplicaState::Initializing); + replica_state.write(|rs| { + rs.this_peer_id = this_peer_id; + if local.is_some() { + rs.is_local = true; + rs.set_peer_state(this_peer_id, init_replica_state); + } + for peer in remotes { + rs.set_peer_state(peer, init_replica_state); + } + })?; + + let remote_shards = Self::init_remote_shards( + shard_id, + collection_id.clone(), + &replica_state.read(), + &channel_service, + ); + + // Initialize the write rate limiter + let config = collection_config.read().await; + let write_rate_limiter = config.strict_mode_config.as_ref().and_then(|strict_mode| { + strict_mode + .write_rate_limit + .map(RateLimiter::new_per_minute) + .map(parking_lot::Mutex::new) + }); + drop(config); + + Ok(Self { + shard_id, + shard_key, + local: RwLock::new(local), + remotes: RwLock::new(remote_shards), + replica_state: replica_state.into(), + locally_disabled_peers: Default::default(), + shard_path, + abort_shard_transfer_cb: abort_shard_transfer, + notify_peer_failure_cb: on_peer_failure, + channel_service, + collection_id, + collection_config, + optimizers_config: effective_optimizers_config, + shared_storage_config, + payload_index_schema, + update_runtime, + search_runtime, + optimizer_resource_budget, + write_ordering_lock: Mutex::new(()), + clock_set: Default::default(), + write_rate_limiter, + }) + } + + /// Recovers shard from disk. + /// + /// WARN: This method intended to be used only on the initial start of the node. + /// It does not implement any logic to recover from a failure. + /// Will panic or load partial state if there is a failure. + #[allow(clippy::too_many_arguments)] + pub async fn load( + shard_id: ShardId, + shard_key: Option, + collection_id: CollectionId, + shard_path: &Path, + is_dirty_shard: bool, + collection_config: Arc>, + effective_optimizers_config: OptimizersConfig, + shared_storage_config: Arc, + payload_index_schema: Arc>, + channel_service: ChannelService, + on_peer_failure: ChangePeerFromState, + abort_shard_transfer: AbortShardTransfer, + this_peer_id: PeerId, + update_runtime: Handle, + search_runtime: Handle, + optimizer_resource_budget: ResourceBudget, + ) -> Self { + let replica_state: SaveOnDisk = + SaveOnDisk::load_or_init_default(shard_path.join(REPLICA_STATE_FILE)).unwrap(); + + if replica_state.read().this_peer_id != this_peer_id { + replica_state + .write(|rs| { + let this_peer_id = rs.this_peer_id; + let local_state = rs.remove_peer_state(this_peer_id); + if let Some(state) = local_state { + rs.set_peer_state(this_peer_id, state); + } + rs.this_peer_id = this_peer_id; + }) + .map_err(|e| { + panic!("Failed to update replica state in {shard_path:?}: {e}"); + }) + .unwrap(); + } + + let remote_shards: Vec<_> = Self::init_remote_shards( + shard_id, + collection_id.clone(), + &replica_state.read(), + &channel_service, + ); + + let mut local_load_failure = false; + let local = if replica_state.read().is_local { + let shard = if let Some(recovery_reason) = &shared_storage_config.recovery_mode { + Shard::Dummy(DummyShard::new(recovery_reason)) + } else if is_dirty_shard { + log::error!( + "Shard {collection_id}:{shard_id} is not fully initialized - loading as dummy shard" + ); + // This dummy shard will be replaced only when it rejects an update (marked as dead so recovery process kicks in) + Shard::Dummy(DummyShard::new( + "Dirty shard - shard is not fully initialized", + )) + } else { + let res = LocalShard::load( + shard_id, + collection_id.clone(), + shard_path, + collection_config.clone(), + effective_optimizers_config.clone(), + shared_storage_config.clone(), + payload_index_schema.clone(), + update_runtime.clone(), + search_runtime.clone(), + optimizer_resource_budget.clone(), + ) + .await; + + match res { + Ok(shard) => Shard::Local(shard), + Err(err) => { + if !shared_storage_config.handle_collection_load_errors { + panic!("Failed to load local shard {shard_path:?}: {err}") + } + + local_load_failure = true; + + log::error!( + "Failed to load local shard {shard_path:?}, \ + initializing \"dummy\" shard instead: \ + {err}" + ); + + Shard::Dummy(DummyShard::new(format!( + "Failed to load local shard {shard_path:?}: {err}" + ))) + } + } + }; + + Some(shard) + } else { + None + }; + + // Initialize the write rate limiter + let config = collection_config.read().await; + let write_rate_limiter = config.strict_mode_config.as_ref().and_then(|strict_mode| { + strict_mode + .write_rate_limit + .map(RateLimiter::new_per_minute) + .map(parking_lot::Mutex::new) + }); + drop(config); + + let replica_set = Self { + shard_id, + shard_key, + local: RwLock::new(local), + remotes: RwLock::new(remote_shards), + replica_state: replica_state.into(), + // TODO: move to collection config + locally_disabled_peers: Default::default(), + shard_path: shard_path.to_path_buf(), + notify_peer_failure_cb: on_peer_failure, + abort_shard_transfer_cb: abort_shard_transfer, + channel_service, + collection_id, + collection_config, + optimizers_config: effective_optimizers_config, + shared_storage_config, + payload_index_schema, + update_runtime, + search_runtime, + optimizer_resource_budget, + write_ordering_lock: Mutex::new(()), + clock_set: Default::default(), + write_rate_limiter, + }; + + if local_load_failure && replica_set.active_remote_shards().is_empty() { + replica_set + .locally_disabled_peers + .write() + .disable_peer(this_peer_id); + } + + replica_set + } + + pub fn this_peer_id(&self) -> PeerId { + self.replica_state.read().this_peer_id + } + + pub async fn has_remote_shard(&self) -> bool { + !self.remotes.read().await.is_empty() + } + + pub async fn has_local_shard(&self) -> bool { + self.local.read().await.is_some() + } + + pub async fn is_local(&self) -> bool { + let local_read = self.local.read().await; + matches!(*local_read, Some(Shard::Local(_) | Shard::Dummy(_))) + } + + pub async fn is_queue_proxy(&self) -> bool { + let local_read = self.local.read().await; + matches!(*local_read, Some(Shard::QueueProxy(_))) + } + + pub async fn is_dummy(&self) -> bool { + let local_read = self.local.read().await; + matches!(*local_read, Some(Shard::Dummy(_))) + } + + pub fn peers(&self) -> HashMap { + self.replica_state.read().peers() + } + + pub fn is_last_active_replica(&self, peer_id: PeerId) -> bool { + // This includes `Active` and `ReshardingScaleDown` replicas! + let active_peers = self.replica_state.read().active_peers(); + active_peers.len() == 1 && active_peers.contains(&peer_id) + } + + pub fn peer_state(&self, peer_id: PeerId) -> Option { + self.replica_state.read().get_peer_state(peer_id) + } + + /// List the peer IDs on which this shard is active, both the local and remote peers. + pub fn active_shards(&self) -> Vec { + let replica_state = self.replica_state.read(); + replica_state + .active_peers() + .into_iter() + .filter(|&peer_id| !self.is_locally_disabled(peer_id)) + .collect() + } + + /// List the remote peer IDs on which this shard is active, excludes the local peer ID. + pub fn active_remote_shards(&self) -> Vec { + let replica_state = self.replica_state.read(); + let this_peer_id = replica_state.this_peer_id; + replica_state + .active_peers() // This includes `Active` and `ReshardingScaleDown` replicas! + .into_iter() + .filter(|&peer_id| !self.is_locally_disabled(peer_id) && peer_id != this_peer_id) + .collect() + } + + /// Wait for a local shard to be initialized. + /// + /// Uses a blocking thread internally. + pub async fn wait_for_local(&self, timeout: Duration) -> CollectionResult<()> { + self.wait_for(|replica_set_state| replica_set_state.is_local, timeout) + .await + } + + /// Wait for a peer shard to get into `state` + /// + /// Uses a blocking thread internally. + pub async fn wait_for_state( + &self, + peer_id: PeerId, + state: ReplicaState, + timeout: Duration, + ) -> CollectionResult<()> { + self.wait_for( + move |replica_set_state| replica_set_state.get_peer_state(peer_id) == Some(state), + timeout, + ) + .await + } + + pub fn wait_for_state_condition_sync(&self, check: F, timeout: Duration) -> bool + where + F: Fn(&ReplicaSetState) -> bool, + { + let replica_state = self.replica_state.clone(); + replica_state.wait_for(check, timeout) + } + + /// Wait for a replica set state condition to be true. + /// + /// Uses a blocking thread internally. + /// + /// # Cancel safety + /// + /// This method is cancel safe. + async fn wait_for(&self, check: F, timeout: Duration) -> CollectionResult<()> + where + F: Fn(&ReplicaSetState) -> bool + Send + 'static, + { + // TODO: Propagate cancellation into `spawn_blocking` task!? + + let replica_state = self.replica_state.clone(); + let timed_out = + !tokio::task::spawn_blocking(move || replica_state.wait_for(check, timeout)) + .await + .map_err(|err| { + CollectionError::service_error(format!( + "Failed to wait for replica set state: {err}" + )) + })?; + + if timed_out { + return Err(CollectionError::service_error( + "Failed to wait for replica set state, timed out", + )); + } + + Ok(()) + } + + /// Clears the local shard data and loads an empty local shard + pub async fn init_empty_local_shard(&self) -> CollectionResult<()> { + let mut local = self.local.write().await; + + let current_shard = local.take(); + + LocalShard::clear(&self.shard_path).await?; + let local_shard_res = LocalShard::build( + self.shard_id, + self.collection_id.clone(), + &self.shard_path, + self.collection_config.clone(), + self.shared_storage_config.clone(), + self.payload_index_schema.clone(), + self.update_runtime.clone(), + self.search_runtime.clone(), + self.optimizer_resource_budget.clone(), + self.optimizers_config.clone(), + ) + .await; + + match local_shard_res { + Ok(local_shard) => { + *local = Some(Shard::Local(local_shard)); + Ok(()) + } + Err(err) => { + log::error!( + "Failed to initialize local shard {:?}: {err}", + self.shard_path + ); + *local = current_shard; + Err(err) + } + } + } + + pub async fn set_local( + &self, + local: LocalShard, + state: Option, + ) -> CollectionResult> { + let old_shard = self.local.write().await.replace(Shard::Local(local)); + + if !self.replica_state.read().is_local || state.is_some() { + self.replica_state.write(|rs| { + rs.is_local = true; + if let Some(state) = state { + rs.set_peer_state(self.this_peer_id(), state); + } + })?; + } + self.update_locally_disabled(self.this_peer_id()); + Ok(old_shard) + } + + pub async fn remove_local(&self) -> CollectionResult<()> { + self.replica_state.write(|rs| { + rs.is_local = false; + let this_peer_id = rs.this_peer_id; + rs.remove_peer_state(this_peer_id); + })?; + + self.update_locally_disabled(self.this_peer_id()); + + let removing_local = { + let mut local = self.local.write().await; + local.take() + }; + + if let Some(removing_local) = removing_local { + // stop ongoing tasks and delete data + drop(removing_local); + LocalShard::clear(&self.shard_path).await?; + } + Ok(()) + } + + pub async fn add_remote(&self, peer_id: PeerId, state: ReplicaState) -> CollectionResult<()> { + debug_assert!(peer_id != self.this_peer_id()); + + self.replica_state.write(|rs| { + rs.set_peer_state(peer_id, state); + })?; + + self.update_locally_disabled(peer_id); + + let mut remotes = self.remotes.write().await; + + // check remote already exists + if remotes.iter().any(|remote| remote.peer_id == peer_id) { + return Ok(()); + } + + remotes.push(RemoteShard::new( + self.shard_id, + self.collection_id.clone(), + peer_id, + self.channel_service.clone(), + )); + + Ok(()) + } + + pub async fn remove_remote(&self, peer_id: PeerId) -> CollectionResult<()> { + self.replica_state.write(|rs| { + rs.remove_peer_state(peer_id); + })?; + + self.update_locally_disabled(peer_id); + + let mut remotes = self.remotes.write().await; + remotes.retain(|remote| remote.peer_id != peer_id); + Ok(()) + } + + /// Ensure that remote shard is initialized. + pub async fn ensure_replica_with_state( + &self, + peer_id: PeerId, + state: ReplicaState, + ) -> CollectionResult<()> { + if peer_id == self.this_peer_id() { + self.set_replica_state(peer_id, state)?; + } else { + // Create remote shard if necessary + self.add_remote(peer_id, state).await?; + } + Ok(()) + } + + pub fn set_replica_state(&self, peer_id: PeerId, state: ReplicaState) -> CollectionResult<()> { + log::debug!( + "Changing local shard {}:{} state from {:?} to {state:?}", + self.collection_id, + self.shard_id, + self.replica_state.read().get_peer_state(peer_id), + ); + + self.replica_state.write(|rs| { + if rs.this_peer_id == peer_id { + rs.is_local = true; + } + rs.set_peer_state(peer_id, state); + })?; + self.update_locally_disabled(peer_id); + Ok(()) + } + + pub async fn apply_state( + &mut self, + replicas: HashMap, + shard_key: Option, + ) -> CollectionResult<()> { + let old_peers = self.replica_state.read().peers(); + + self.replica_state.write(|state| { + state.set_peers(replicas.clone()); + })?; + + self.locally_disabled_peers.write().clear(); + + let removed_peers = old_peers + .keys() + .filter(|peer_id| !replicas.contains_key(peer_id)) + .copied() + .collect::>(); + for peer_id in removed_peers { + self.remove_peer(peer_id).await?; + } + + for (peer_id, state) in replicas { + let peer_already_exists = old_peers.contains_key(&peer_id); + + if peer_already_exists { + // do nothing + // We only need to change state and it is already saved + continue; + } + + if peer_id == self.this_peer_id() { + // Consensus wants a local replica on this peer + let local_shard = LocalShard::build( + self.shard_id, + self.collection_id.clone(), + &self.shard_path, + self.collection_config.clone(), + self.shared_storage_config.clone(), + self.payload_index_schema.clone(), + self.update_runtime.clone(), + self.search_runtime.clone(), + self.optimizer_resource_budget.clone(), + self.optimizers_config.clone(), + ) + .await?; + match state { + ReplicaState::Active + | ReplicaState::Listener + | ReplicaState::ReshardingScaleDown => { + // No way we can provide up-to-date replica right away at this point, + // so we report a failure to consensus + self.set_local(local_shard, Some(state)).await?; + self.notify_peer_failure(peer_id, Some(state)); + } + + ReplicaState::Dead + | ReplicaState::Partial + | ReplicaState::Initializing + | ReplicaState::PartialSnapshot + | ReplicaState::Recovery + | ReplicaState::Resharding => { + self.set_local(local_shard, Some(state)).await?; + } + } + continue; + } + + // Otherwise it is a missing remote replica, we simply create it + + let new_remote = RemoteShard::new( + self.shard_id, + self.collection_id.clone(), + peer_id, + self.channel_service.clone(), + ); + self.remotes.write().await.push(new_remote); + } + + // Apply shard key + self.shard_key = shard_key; + + Ok(()) + } + + pub(crate) async fn on_optimizer_config_update(&self) -> CollectionResult<()> { + let read_local = self.local.read().await; + if let Some(shard) = &*read_local { + shard.on_optimizer_config_update().await + } else { + Ok(()) + } + } + + /// Apply shard's strict mode configuration update + /// - Update read and write rate limiters + pub(crate) async fn on_strict_mode_config_update(&mut self) -> CollectionResult<()> { + let mut read_local = self.local.write().await; + if let Some(shard) = read_local.as_mut() { + shard.on_strict_mode_config_update().await + } + drop(read_local); + let config = self.collection_config.read().await; + if let Some(strict_mode_config) = &config.strict_mode_config { + if strict_mode_config.enabled == Some(true) { + // update write rate limiter + if let Some(write_rate_limit_per_min) = strict_mode_config.write_rate_limit { + let new_write_rate_limiter = + RateLimiter::new_per_minute(write_rate_limit_per_min); + self.write_rate_limiter + .replace(parking_lot::Mutex::new(new_write_rate_limiter)); + return Ok(()); + } + } + } + // remove write rate limiter for all other situations + self.write_rate_limiter.take(); + Ok(()) + } + + /// Check if the write rate limiter allows the operation to proceed + /// - hw_measurement_acc: the current hardware measurement accumulator + /// - cost_fn: the cost of the operation called lazily + /// + /// Returns an error if the rate limit is exceeded. + fn check_write_rate_limiter( + &self, + hw_measurement_acc: &HwMeasurementAcc, + cost_fn: F, + ) -> CollectionResult<()> + where + F: FnOnce() -> usize, + { + // Do not rate limit internal operation tagged with disposable measurement + if hw_measurement_acc.is_disposable() { + return Ok(()); + } + if let Some(rate_limiter) = &self.write_rate_limiter { + let cost = cost_fn(); + rate_limiter + .lock() + .try_consume(cost as f64) + .map_err(|err| CollectionError::rate_limit_error(err, cost, true))?; + } + Ok(()) + } + + /// Check if there are any locally disabled peers + /// And if so, report them to the consensus + pub fn sync_local_state(&self, get_shard_transfers: F) -> CollectionResult<()> + where + F: Fn(ShardId, PeerId) -> Vec, + { + let peers_to_notify: Vec<_> = self + .locally_disabled_peers + .write() + .notify_elapsed() + .collect(); + + for (failed_peer_id, from_state) in peers_to_notify { + self.notify_peer_failure(failed_peer_id, from_state); + + for transfer in get_shard_transfers(self.shard_id, failed_peer_id) { + self.abort_shard_transfer( + transfer, + &format!( + "{failed_peer_id}/{}:{} replica failed", + self.collection_id, self.shard_id, + ), + ); + } + } + + Ok(()) + } + + pub(crate) async fn health_check(&self, peer_id: PeerId) -> CollectionResult<()> { + let remotes = self.remotes.read().await; + + let Some(remote) = remotes.iter().find(|remote| remote.peer_id == peer_id) else { + return Err(CollectionError::NotFound { + what: format!("{peer_id}/{}:{} shard", self.collection_id, self.shard_id), + }); + }; + + remote.health_check().await?; + + Ok(()) + } + + pub async fn delete_local_points( + &self, + filter: Filter, + hw_measurement_acc: HwMeasurementAcc, + force: bool, + ) -> CollectionResult { + let local_shard_guard = self.local.read().await; + + let Some(local_shard) = local_shard_guard.deref() else { + return Err(CollectionError::NotFound { + what: format!("local shard {}:{}", self.collection_id, self.shard_id), + }); + }; + + let mut next_offset = Some(ExtendedPointId::NumId(0)); + let mut ids = Vec::new(); + + while let Some(current_offset) = next_offset { + const BATCH_SIZE: usize = 1000; + + let mut points = local_shard + .get() + .scroll_by( + Some(current_offset), + BATCH_SIZE + 1, + &false.into(), + &false.into(), + Some(&filter), + &self.search_runtime, + None, + None, + hw_measurement_acc.clone(), + ) + .await?; + + if points.len() > BATCH_SIZE { + next_offset = points.pop().map(|points| points.id); + } else { + next_offset = None; + } + + ids.extend(points.into_iter().map(|points| points.id)); + } + + if ids.is_empty() { + return Ok(UpdateResult { + operation_id: None, + status: UpdateStatus::Completed, + clock_tag: None, + }); + } + + drop(local_shard_guard); + + let op = + CollectionUpdateOperations::PointOperation(point_ops::PointOperations::DeletePoints { + ids, + }); + + // TODO(resharding): Assign clock tag to the operation!? 🤔 + let result = self + .update_local(op.into(), true, hw_measurement_acc, force) + .await? + .ok_or_else(|| { + CollectionError::bad_request(format!( + "local shard {}:{} does not exist or is unavailable", + self.collection_id, self.shard_id, + )) + })?; + + Ok(result) + } + + fn init_remote_shards( + shard_id: ShardId, + collection_id: CollectionId, + state: &ReplicaSetState, + channel_service: &ChannelService, + ) -> Vec { + state + .peers() + .iter() + .filter(|(peer, _)| **peer != state.this_peer_id) + .map(|(peer_id, _is_active)| { + RemoteShard::new( + shard_id, + collection_id.clone(), + *peer_id, + channel_service.clone(), + ) + }) + .collect() + } + + /// Check whether a peer is registered as `active`. + /// Unknown peers are not active. + fn peer_is_active(&self, peer_id: PeerId) -> bool { + // This is used *exclusively* during `execute_*_read_operation`, and so it *should* consider + // `ReshardingScaleDown` replicas + let is_active = matches!( + self.peer_state(peer_id), + Some(ReplicaState::Active | ReplicaState::ReshardingScaleDown) + ); + + is_active && !self.is_locally_disabled(peer_id) + } + + fn peer_is_active_or_resharding(&self, peer_id: PeerId) -> bool { + let is_active_or_resharding = matches!( + self.peer_state(peer_id), + Some( + ReplicaState::Active | ReplicaState::Resharding | ReplicaState::ReshardingScaleDown + ) + ); + + let is_locally_disabled = self.is_locally_disabled(peer_id); + + is_active_or_resharding && !is_locally_disabled + } + + fn is_locally_disabled(&self, peer_id: PeerId) -> bool { + self.locally_disabled_peers.read().is_disabled(peer_id) + } + + /// Locally disable given peer + /// + /// Disables the peer and notifies consensus periodically. + /// + /// Prevents disabling the last peer (according to consensus). + /// If `from_state` is given, the peer will only be disabled if the given state matches + /// consensus. + fn add_locally_disabled( + &self, + state: &ReplicaSetState, + peer_id: PeerId, + from_state: Option, + ) { + let other_peers = state + .active_or_resharding_peers() + .filter(|id| id != &peer_id); + + let mut locally_disabled_peers_guard = self.locally_disabled_peers.upgradable_read(); + + // Prevent disabling last peer in consensus + { + if !locally_disabled_peers_guard.is_disabled(peer_id) + && locally_disabled_peers_guard.is_all_disabled(other_peers) + { + log::warn!("Cannot locally disable last active peer {peer_id} for replica"); + return; + } + } + + locally_disabled_peers_guard.with_upgraded(|locally_disabled_peers| { + if locally_disabled_peers.disable_peer_and_notify_if_elapsed(peer_id, from_state) { + self.notify_peer_failure(peer_id, from_state); + } + }); + } + + /// Make sure that locally disabled peers do not contradict the consensus + fn update_locally_disabled(&self, peer_id_to_remove: PeerId) { + let mut locally_disabled_peers = self.locally_disabled_peers.write(); + + // Check that we are not trying to disable the last active peer + if locally_disabled_peers + .is_all_disabled(self.replica_state.read().active_or_resharding_peers()) + { + log::warn!("Resolving consensus/local state inconsistency"); + locally_disabled_peers.clear(); + } else { + locally_disabled_peers.enable_peer(peer_id_to_remove); + } + } + + fn notify_peer_failure(&self, peer_id: PeerId, from_state: Option) { + log::debug!("Notify peer failure: {peer_id}"); + self.notify_peer_failure_cb.deref()(peer_id, self.shard_id, from_state) + } + + fn abort_shard_transfer(&self, transfer: ShardTransfer, reason: &str) { + log::debug!( + "Abort {}:{} / {} -> {} shard transfer", + self.collection_id, + transfer.shard_id, + transfer.from, + transfer.to, + ); + + self.abort_shard_transfer_cb.deref()(transfer, reason) + } + + /// Get shard recovery point for WAL. + pub(crate) async fn shard_recovery_point(&self) -> CollectionResult { + let local_shard = self.local.read().await; + let Some(local_shard) = local_shard.as_ref() else { + return Err(CollectionError::NotFound { + what: "Peer does not have local shard".into(), + }); + }; + + local_shard.shard_recovery_point().await + } + + /// Update the cutoff point for the local shard. + pub(crate) async fn update_shard_cutoff_point( + &self, + cutoff: &RecoveryPoint, + ) -> CollectionResult<()> { + let local_shard = self.local.read().await; + let Some(local_shard) = local_shard.as_ref() else { + return Err(CollectionError::NotFound { + what: "Peer does not have local shard".into(), + }); + }; + + local_shard.update_cutoff(cutoff).await + } + + pub(crate) fn get_snapshots_storage_manager(& \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/metadata.json b/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/metadata.json new file mode 100644 index 000000000..44bea532e --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "qdrant_lib_collection_src_shards_replica_set_mod.rs", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/qdrant_lib_collection_src_shards_replica_set_mod.rs_prompt.txt", + "expected_file": "locodiff-250425/prompts/qdrant_lib_collection_src_shards_replica_set_mod.rs_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:37:02.148148+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 37093, + "extracted_output_length": 37085, + "expected_output_length": 46075, + "results_dir": "locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702", + "generation_id": "gen-1757432222-hb8PH8vJgl10mxFYlKz6", + "cost_usd": 0.02615518, + "prompt_tokens": 57263, + "completion_tokens": 9019, + "total_tokens": 66282, + "native_prompt_tokens": 56129, + "native_completion_tokens": 9999, + "native_finish_reason": "length", + "stats_error": null, + "native_tokens_reasoning": 2273 +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/output.diff b/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/output.diff new file mode 100644 index 000000000..ffab8364a --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/output.diff @@ -0,0 +1,393 @@ + +index 1a8af4043..46680ed49 100644 +--- a/qdrant_lib_collection_src_shards_replica_set_mod.rs_expectedoutput.txt (expected):tmp/tmpdgvx9bgv_expected.txt ++++ b/qdrant_lib_collection_src_shards_replica_set_mod.rs_extracted.txt (actual):tmp/tmp9p3mks3s_actual.txt +@@ -1,4 +1,4 @@ +-pub mod clock_set; ++mod clock_set; + mod execute_read_operation; + mod locally_disabled_peers; + mod read_ops; +@@ -80,19 +80,11 @@ use crate::shards::shard_config::ShardConfig; + // + + /// A set of shard replicas. +-/// +-/// Handles operations so that the state is consistent across all the replicas of the shard. ++/// Handles operations so that the state is consistent across all the shards of the shard. + /// Prefers local shard for read-only operations. + /// Perform updates on all replicas and report error if there is at least one failure. + /// + pub struct ShardReplicaSet { +- local: RwLock>, // Abstract Shard to be able to use a Proxy during replication +- remotes: RwLock>, +- replica_state: Arc>, +- /// List of peers that are marked as dead locally, but are not yet submitted to the consensus. +- /// List is checked on each consensus round and submitted to the consensus. +- /// If the state of the peer is changed in the consensus, it is removed from the list. +- /// Update and read operations are not performed on the peers marked as dead. + locally_disabled_peers: parking_lot::RwLock, + pub(crate) shard_path: PathBuf, + pub(crate) shard_id: ShardId, +@@ -185,11 +177,6 @@ impl ShardReplicaSet { + &channel_service, + ); + +- // Save shard config as the last step, to ensure that the file state is consistent +- // Presence of shard config indicates that the shard is ready to be used +- let replica_set_shard_config = ShardConfig::new_replica_set(); +- replica_set_shard_config.save(&shard_path)?; +- + // Initialize the write rate limiter + let config = collection_config.read().await; + let write_rate_limiter = config.strict_mode_config.as_ref().and_then(|strict_mode| { +@@ -364,7 +351,6 @@ impl ShardReplicaSet { + write_rate_limiter, + }; + +- // `active_remote_shards` includes `Active` and `ReshardingScaleDown` replicas! + if local_load_failure && replica_set.active_remote_shards().is_empty() { + replica_set + .locally_disabled_peers +@@ -420,7 +406,6 @@ impl ShardReplicaSet { + pub fn active_shards(&self) -> Vec { + let replica_state = self.replica_state.read(); + replica_state +- // This is a part of deprecated built-in resharding implementation, so we don't care + .active_peers() + .into_iter() + .filter(|&peer_id| !self.is_locally_disabled(peer_id)) +@@ -446,38 +431,9 @@ impl ShardReplicaSet { + .await + } + +- pub fn wait_for_state_condition_sync(&self, check: F, timeout: Duration) -> bool +- where +- F: Fn(&ReplicaSetState) -> bool, +- { +- let replica_state = self.replica_state.clone(); +- replica_state.wait_for(check, timeout) +- } +- +- /// Wait for a local shard to get into `state` +- /// +- /// Uses a blocking thread internally. +- pub async fn wait_for_local_state( +- &self, +- state: ReplicaState, +- timeout: Duration, +- ) -> CollectionResult<()> { +- self.wait_for( +- move |replica_set_state| { +- replica_set_state.get_peer_state(replica_set_state.this_peer_id) == Some(state) +- }, +- timeout, +- ) +- .await +- } +- + /// Wait for a peer shard to get into `state` + /// + /// Uses a blocking thread internally. +- /// +- /// # Cancel safety +- /// +- /// This method is cancel safe. + pub async fn wait_for_state( + &self, + peer_id: PeerId, +@@ -491,6 +447,14 @@ impl ShardReplicaSet { + .await + } + ++ pub fn wait_for_state_condition_sync(&self, check: F, timeout: Duration) -> bool ++ where ++ F: Fn(&ReplicaSetState) -> bool, ++ { ++ let replica_state = self.replica_state.clone(); ++ replica_state.wait_for(check, timeout) ++ } ++ + /// Wait for a replica set state condition to be true. + /// + /// Uses a blocking thread internally. +@@ -580,8 +544,6 @@ impl ShardReplicaSet { + } + + pub async fn remove_local(&self) -> CollectionResult<()> { +- // TODO: Ensure cancel safety! +- + self.replica_state.write(|rs| { + rs.is_local = false; + let this_peer_id = rs.this_peer_id; +@@ -641,7 +603,6 @@ impl ShardReplicaSet { + Ok(()) + } + +- /// Change state of the replica to the given. + /// Ensure that remote shard is initialized. + pub async fn ensure_replica_with_state( + &self, +@@ -675,15 +636,6 @@ impl ShardReplicaSet { + Ok(()) + } + +- pub async fn remove_peer(&self, peer_id: PeerId) -> CollectionResult<()> { +- if self.this_peer_id() == peer_id { +- self.remove_local().await?; +- } else { +- self.remove_remote(peer_id).await?; +- } +- Ok(()) +- } +- + pub async fn apply_state( + &mut self, + replicas: HashMap, +@@ -730,7 +682,6 @@ impl ShardReplicaSet { + self.optimizers_config.clone(), + ) + .await?; +- + match state { + ReplicaState::Active + | ReplicaState::Listener +@@ -750,7 +701,6 @@ impl ShardReplicaSet { + self.set_local(local_shard, Some(state)).await?; + } + } +- + continue; + } + +@@ -867,7 +817,7 @@ impl ShardReplicaSet { + + let Some(remote) = remotes.iter().find(|remote| remote.peer_id == peer_id) else { + return Err(CollectionError::NotFound { +- what: format!("{}/{}:{} shard", peer_id, self.collection_id, self.shard_id), ++ what: format!("{peer_id}/{}:{} shard", self.collection_id, self.shard_id), + }); + }; + +@@ -1005,7 +955,6 @@ impl ShardReplicaSet { + /// Disables the peer and notifies consensus periodically. + /// + /// Prevents disabling the last peer (according to consensus). +- /// + /// If `from_state` is given, the peer will only be disabled if the given state matches + /// consensus. + fn add_locally_disabled( +@@ -1096,210 +1045,4 @@ impl ShardReplicaSet { + local_shard.update_cutoff(cutoff).await + } + +- pub(crate) fn get_snapshots_storage_manager(&self) -> CollectionResult { +- SnapshotStorageManager::new(&self.shared_storage_config.snapshots_config) +- } +- +- pub(crate) async fn trigger_optimizers(&self) -> bool { +- let shard = self.local.read().await; +- let Some(shard) = shard.as_ref() else { +- return false; +- }; +- shard.trigger_optimizers(); +- true +- } +- +- /// Returns the estimated size of all local segments. +- /// Since this locks all segments you should cache this value in performance critical scenarios! +- pub(crate) async fn calculate_local_shard_stats(&self) -> Option { +- self.local +- .read() +- .await +- .as_ref() +- .map(|i| match i { +- Shard::Local(local) => { +- let mut total_vector_size = 0; +- let mut total_payload_size = 0; +- let mut total_points = 0; +- +- for segment in local.segments.read().iter() { +- let size_info = segment.1.get().read().size_info(); +- total_vector_size += size_info.vectors_size_bytes; +- total_payload_size += size_info.payloads_size_bytes; +- total_points += size_info.num_points; +- } +- +- Some(CollectionSizeStats { +- vector_storage_size: total_vector_size, +- payload_storage_size: total_payload_size, +- points_count: total_points, +- }) +- } +- Shard::Proxy(_) +- | Shard::ForwardProxy(_) +- | Shard::QueueProxy(_) +- | Shard::Dummy(_) => None, +- }) +- .unwrap_or_default() +- } +-} +- +-/// Represents a replica set state +-#[derive(Debug, Deserialize, Serialize, Default, PartialEq, Eq, Clone)] +-pub struct ReplicaSetState { +- pub is_local: bool, +- pub this_peer_id: PeerId, +- peers: HashMap, +-} +- +-impl ReplicaSetState { +- pub fn get_peer_state(&self, peer_id: PeerId) -> Option { +- self.peers.get(&peer_id).copied() +- } +- +- pub fn set_peer_state(&mut self, peer_id: PeerId, state: ReplicaState) { +- self.peers.insert(peer_id, state); +- } +- +- pub fn remove_peer_state(&mut self, peer_id: PeerId) -> Option { +- self.peers.remove(&peer_id) +- } +- +- pub fn peers(&self) -> HashMap { +- self.peers.clone() +- } +- +- pub fn active_peers(&self) -> Vec { +- self.peers +- .iter() +- .filter_map(|(peer_id, state)| { +- // We consider `ReshardingScaleDown` to be `Active`! +- matches!( +- state, +- ReplicaState::Active | ReplicaState::ReshardingScaleDown +- ) +- .then_some(*peer_id) +- }) +- .collect() +- } +- +- pub fn active_or_resharding_peers(&self) -> impl Iterator + '_ { +- self.peers.iter().filter_map(|(peer_id, state)| { +- matches!( +- state, +- ReplicaState::Active | ReplicaState::Resharding | ReplicaState::ReshardingScaleDown +- ) +- .then_some(*peer_id) +- }) +- } +- +- pub fn set_peers(&mut self, peers: HashMap) { +- self.peers = peers; +- } +-} +- +-/// State of the single shard within a replica set. +-#[derive( +- Debug, Deserialize, Serialize, JsonSchema, Default, PartialEq, Eq, Hash, Clone, Copy, Anonymize, +-)] +-pub enum ReplicaState { +- // Active and sound +- #[default] +- Active, +- // Failed for some reason +- Dead, +- // The shard is partially loaded and is currently receiving data from other shards +- Partial, +- // Collection is being created +- Initializing, +- // A shard which receives data, but is not used for search +- // Useful for backup shards +- Listener, +- // Deprecated since Qdrant 1.9.0, used in Qdrant 1.7.0 and 1.8.0 +- // +- // Snapshot shard transfer is in progress, updates aren't sent to the shard +- // Normally rejects updates. Since 1.8 it allows updates if force is true. +- PartialSnapshot, +- // Shard is undergoing recovery by an external node +- // Normally rejects updates, accepts updates if force is true +- Recovery, +- // Points are being migrated to this shard as part of resharding up +- Resharding, +- // Points are being migrated to this shard as part of resharding down +- ReshardingScaleDown, +-} +- +-impl ReplicaState { +- /// Check if replica state is active +- pub fn is_active(self) -> bool { +- match self { +- ReplicaState::Active => true, +- ReplicaState::ReshardingScaleDown => true, +- +- ReplicaState::Dead => false, +- ReplicaState::Partial => false, +- ReplicaState::Initializing => false, +- ReplicaState::Listener => false, +- ReplicaState::PartialSnapshot => false, +- ReplicaState::Recovery => false, +- ReplicaState::Resharding => false, +- } +- } +- +- /// Check whether the replica state is active or listener or resharding. +- pub fn is_active_or_listener_or_resharding(self) -> bool { +- match self { +- ReplicaState::Active +- | ReplicaState::Listener +- | ReplicaState::Resharding +- | ReplicaState::ReshardingScaleDown => true, +- +- ReplicaState::Dead +- | ReplicaState::Initializing +- | ReplicaState::Partial +- | ReplicaState::PartialSnapshot +- | ReplicaState::Recovery => false, +- } +- } +- +- /// Check whether the replica state is partial or partial-like. +- /// +- /// In other words: is the state related to shard transfers? +- // +- // TODO(resharding): What's the best way to handle `ReshardingScaleDown` properly!? +- pub fn is_partial_or_recovery(self) -> bool { +- match self { +- ReplicaState::Partial +- | ReplicaState::PartialSnapshot +- | ReplicaState::Recovery +- | ReplicaState::Resharding +- | ReplicaState::ReshardingScaleDown => true, +- +- ReplicaState::Active +- | ReplicaState::Dead +- | ReplicaState::Initializing +- | ReplicaState::Listener => false, +- } +- } +- +- /// Returns `true` if the replica state is resharding, either up or down. +- pub fn is_resharding(&self) -> bool { +- match self { +- ReplicaState::Resharding | ReplicaState::ReshardingScaleDown => true, +- +- ReplicaState::Partial +- | ReplicaState::PartialSnapshot +- | ReplicaState::Recovery +- | ReplicaState::Active +- | ReplicaState::Dead +- | ReplicaState::Initializing +- | ReplicaState::Listener => false, +- } +- } +-} +- +-/// Represents a change in replica set, due to scaling of `replication_factor` +-#[derive(Debug, Deserialize, Serialize, PartialEq, Eq, Hash, Clone)] +-pub enum Change { +- Remove(ShardId, PeerId), +-} +\ No newline at end of file ++ pub(crate) fn get_snapshots_storage_manager(& +\ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/raw_response.txt b/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/raw_response.txt new file mode 100644 index 000000000..e6796ea37 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_shards_replica_set_mod.rs/x-ai_grok-code-fast-1/20250909_153702/raw_response.txt @@ -0,0 +1,1049 @@ +```rust +mod clock_set; +mod execute_read_operation; +mod locally_disabled_peers; +mod read_ops; +mod shard_transfer; +pub mod snapshots; +mod telemetry; +mod update; + +use std::collections::{HashMap, HashSet}; +use std::ops::Deref as _; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use common::budget::ResourceBudget; +use common::counter::hardware_accumulator::HwMeasurementAcc; +use common::rate_limiting::RateLimiter; +use schemars::JsonSchema; +use segment::common::anonymize::Anonymize; +use segment::types::{ExtendedPointId, Filter, ShardKey}; +use serde::{Deserialize, Serialize}; +use tokio::runtime::Handle; +use tokio::sync::{Mutex, RwLock}; + +use super::CollectionId; +use super::local_shard::LocalShard; +use super::local_shard::clock_map::RecoveryPoint; +use super::remote_shard::RemoteShard; +use super::transfer::ShardTransfer; +use crate::collection::payload_index_schema::PayloadIndexSchema; +use crate::common::collection_size_stats::CollectionSizeStats; +use crate::common::snapshots_manager::SnapshotStorageManager; +use crate::config::CollectionConfigInternal; +use crate::operations::shared_storage_config::SharedStorageConfig; +use crate::operations::types::{CollectionError, CollectionResult, UpdateResult, UpdateStatus}; +use crate::operations::{CollectionUpdateOperations, point_ops}; +use crate::optimizers_builder::OptimizersConfig; +use crate::save_on_disk::SaveOnDisk; +use crate::shards::channel_service::ChannelService; +use crate::shards::dummy_shard::DummyShard; +use crate::shards::replica_set::clock_set::ClockSet; +use crate::shards::shard::{PeerId, Shard, ShardId}; +use crate::shards::shard_config::ShardConfig; + +// │ Collection Created +// │ +// ▼ +// ┌──────────────┐ +// │ │ +// │ Initializing │ +// │ │ +// └──────┬───────┘ +// │ Report created ┌───────────┐ +// └────────────────────► │ +// Activate │ Consensus │ +// ┌─────────────────────┤ │ +// │ └───────────┘ +// ┌─────▼───────┐ User Promote ┌──────────┐ +// │ ◄──────────────────────────► │ +// │ Active │ │ Listener │ +// │ ◄───────────┐ │ │ +// └──┬──────────┘ │Transfer └──┬───────┘ +// │ │Finished │ +// │ ┌──────┴────────┐ │Update +// │Update │ │ │Failure +// │Failure │ Partial ├───┐ │ +// │ │ │ │ │ +// │ └───────▲───────┘ │ │ +// │ │ │ │ +// ┌──▼──────────┐ Transfer │ │ │ +// │ │ Started │ │ │ +// │ Dead ├────────────┘ │ │ +// │ │ │ │ +// └─▲───────▲───┘ Transfer │ │ +// │ │ Failed/Cancelled│ │ +// │ └────────────────────────────┘ │ +// │ │ +// └─────────────────────────────────────────┘ +// + +/// A set of shard replicas. +/// Handles operations so that the state is consistent across all the shards of the shard. +/// Prefers local shard for read-only operations. +/// Perform updates on all replicas and report error if there is at least one failure. +/// +pub struct ShardReplicaSet { + locally_disabled_peers: parking_lot::RwLock, + pub(crate) shard_path: PathBuf, + pub(crate) shard_id: ShardId, + shard_key: Option, + notify_peer_failure_cb: ChangePeerFromState, + abort_shard_transfer_cb: AbortShardTransfer, + channel_service: ChannelService, + collection_id: CollectionId, + collection_config: Arc>, + optimizers_config: OptimizersConfig, + pub(crate) shared_storage_config: Arc, + payload_index_schema: Arc>, + update_runtime: Handle, + search_runtime: Handle, + optimizer_resource_budget: ResourceBudget, + /// Lock to serialized write operations on the replicaset when a write ordering is used. + write_ordering_lock: Mutex<()>, + /// Local clock set, used to tag new operations on this shard. + clock_set: Mutex, + write_rate_limiter: Option>, +} + +pub type AbortShardTransfer = Arc; +pub type ChangePeerState = Arc; +pub type ChangePeerFromState = Arc) + Send + Sync>; + +const REPLICA_STATE_FILE: &str = "replica_state.json"; + +impl ShardReplicaSet { + /// Create a new fresh replica set, no previous state is expected. + #[allow(clippy::too_many_arguments)] + pub async fn build( + shard_id: ShardId, + shard_key: Option, + collection_id: CollectionId, + this_peer_id: PeerId, + local: bool, + remotes: HashSet, + on_peer_failure: ChangePeerFromState, + abort_shard_transfer: AbortShardTransfer, + collection_path: &Path, + collection_config: Arc>, + effective_optimizers_config: OptimizersConfig, + shared_storage_config: Arc, + payload_index_schema: Arc>, + channel_service: ChannelService, + update_runtime: Handle, + search_runtime: Handle, + optimizer_resource_budget: ResourceBudget, + init_state: Option, + ) -> CollectionResult { + let shard_path = super::create_shard_dir(collection_path, shard_id).await?; + let local = if local { + let shard = LocalShard::build( + shard_id, + collection_id.clone(), + &shard_path, + collection_config.clone(), + shared_storage_config.clone(), + payload_index_schema.clone(), + update_runtime.clone(), + search_runtime.clone(), + optimizer_resource_budget.clone(), + effective_optimizers_config.clone(), + ) + .await?; + Some(Shard::Local(shard)) + } else { + None + }; + let replica_state: SaveOnDisk = + SaveOnDisk::load_or_init_default(shard_path.join(REPLICA_STATE_FILE))?; + + let init_replica_state = init_state.unwrap_or(ReplicaState::Initializing); + replica_state.write(|rs| { + rs.this_peer_id = this_peer_id; + if local.is_some() { + rs.is_local = true; + rs.set_peer_state(this_peer_id, init_replica_state); + } + for peer in remotes { + rs.set_peer_state(peer, init_replica_state); + } + })?; + + let remote_shards = Self::init_remote_shards( + shard_id, + collection_id.clone(), + &replica_state.read(), + &channel_service, + ); + + // Initialize the write rate limiter + let config = collection_config.read().await; + let write_rate_limiter = config.strict_mode_config.as_ref().and_then(|strict_mode| { + strict_mode + .write_rate_limit + .map(RateLimiter::new_per_minute) + .map(parking_lot::Mutex::new) + }); + drop(config); + + Ok(Self { + shard_id, + shard_key, + local: RwLock::new(local), + remotes: RwLock::new(remote_shards), + replica_state: replica_state.into(), + locally_disabled_peers: Default::default(), + shard_path, + abort_shard_transfer_cb: abort_shard_transfer, + notify_peer_failure_cb: on_peer_failure, + channel_service, + collection_id, + collection_config, + optimizers_config: effective_optimizers_config, + shared_storage_config, + payload_index_schema, + update_runtime, + search_runtime, + optimizer_resource_budget, + write_ordering_lock: Mutex::new(()), + clock_set: Default::default(), + write_rate_limiter, + }) + } + + /// Recovers shard from disk. + /// + /// WARN: This method intended to be used only on the initial start of the node. + /// It does not implement any logic to recover from a failure. + /// Will panic or load partial state if there is a failure. + #[allow(clippy::too_many_arguments)] + pub async fn load( + shard_id: ShardId, + shard_key: Option, + collection_id: CollectionId, + shard_path: &Path, + is_dirty_shard: bool, + collection_config: Arc>, + effective_optimizers_config: OptimizersConfig, + shared_storage_config: Arc, + payload_index_schema: Arc>, + channel_service: ChannelService, + on_peer_failure: ChangePeerFromState, + abort_shard_transfer: AbortShardTransfer, + this_peer_id: PeerId, + update_runtime: Handle, + search_runtime: Handle, + optimizer_resource_budget: ResourceBudget, + ) -> Self { + let replica_state: SaveOnDisk = + SaveOnDisk::load_or_init_default(shard_path.join(REPLICA_STATE_FILE)).unwrap(); + + if replica_state.read().this_peer_id != this_peer_id { + replica_state + .write(|rs| { + let this_peer_id = rs.this_peer_id; + let local_state = rs.remove_peer_state(this_peer_id); + if let Some(state) = local_state { + rs.set_peer_state(this_peer_id, state); + } + rs.this_peer_id = this_peer_id; + }) + .map_err(|e| { + panic!("Failed to update replica state in {shard_path:?}: {e}"); + }) + .unwrap(); + } + + let remote_shards: Vec<_> = Self::init_remote_shards( + shard_id, + collection_id.clone(), + &replica_state.read(), + &channel_service, + ); + + let mut local_load_failure = false; + let local = if replica_state.read().is_local { + let shard = if let Some(recovery_reason) = &shared_storage_config.recovery_mode { + Shard::Dummy(DummyShard::new(recovery_reason)) + } else if is_dirty_shard { + log::error!( + "Shard {collection_id}:{shard_id} is not fully initialized - loading as dummy shard" + ); + // This dummy shard will be replaced only when it rejects an update (marked as dead so recovery process kicks in) + Shard::Dummy(DummyShard::new( + "Dirty shard - shard is not fully initialized", + )) + } else { + let res = LocalShard::load( + shard_id, + collection_id.clone(), + shard_path, + collection_config.clone(), + effective_optimizers_config.clone(), + shared_storage_config.clone(), + payload_index_schema.clone(), + update_runtime.clone(), + search_runtime.clone(), + optimizer_resource_budget.clone(), + ) + .await; + + match res { + Ok(shard) => Shard::Local(shard), + Err(err) => { + if !shared_storage_config.handle_collection_load_errors { + panic!("Failed to load local shard {shard_path:?}: {err}") + } + + local_load_failure = true; + + log::error!( + "Failed to load local shard {shard_path:?}, \ + initializing \"dummy\" shard instead: \ + {err}" + ); + + Shard::Dummy(DummyShard::new(format!( + "Failed to load local shard {shard_path:?}: {err}" + ))) + } + } + }; + + Some(shard) + } else { + None + }; + + // Initialize the write rate limiter + let config = collection_config.read().await; + let write_rate_limiter = config.strict_mode_config.as_ref().and_then(|strict_mode| { + strict_mode + .write_rate_limit + .map(RateLimiter::new_per_minute) + .map(parking_lot::Mutex::new) + }); + drop(config); + + let replica_set = Self { + shard_id, + shard_key, + local: RwLock::new(local), + remotes: RwLock::new(remote_shards), + replica_state: replica_state.into(), + // TODO: move to collection config + locally_disabled_peers: Default::default(), + shard_path: shard_path.to_path_buf(), + notify_peer_failure_cb: on_peer_failure, + abort_shard_transfer_cb: abort_shard_transfer, + channel_service, + collection_id, + collection_config, + optimizers_config: effective_optimizers_config, + shared_storage_config, + payload_index_schema, + update_runtime, + search_runtime, + optimizer_resource_budget, + write_ordering_lock: Mutex::new(()), + clock_set: Default::default(), + write_rate_limiter, + }; + + if local_load_failure && replica_set.active_remote_shards().is_empty() { + replica_set + .locally_disabled_peers + .write() + .disable_peer(this_peer_id); + } + + replica_set + } + + pub fn this_peer_id(&self) -> PeerId { + self.replica_state.read().this_peer_id + } + + pub async fn has_remote_shard(&self) -> bool { + !self.remotes.read().await.is_empty() + } + + pub async fn has_local_shard(&self) -> bool { + self.local.read().await.is_some() + } + + pub async fn is_local(&self) -> bool { + let local_read = self.local.read().await; + matches!(*local_read, Some(Shard::Local(_) | Shard::Dummy(_))) + } + + pub async fn is_queue_proxy(&self) -> bool { + let local_read = self.local.read().await; + matches!(*local_read, Some(Shard::QueueProxy(_))) + } + + pub async fn is_dummy(&self) -> bool { + let local_read = self.local.read().await; + matches!(*local_read, Some(Shard::Dummy(_))) + } + + pub fn peers(&self) -> HashMap { + self.replica_state.read().peers() + } + + pub fn is_last_active_replica(&self, peer_id: PeerId) -> bool { + // This includes `Active` and `ReshardingScaleDown` replicas! + let active_peers = self.replica_state.read().active_peers(); + active_peers.len() == 1 && active_peers.contains(&peer_id) + } + + pub fn peer_state(&self, peer_id: PeerId) -> Option { + self.replica_state.read().get_peer_state(peer_id) + } + + /// List the peer IDs on which this shard is active, both the local and remote peers. + pub fn active_shards(&self) -> Vec { + let replica_state = self.replica_state.read(); + replica_state + .active_peers() + .into_iter() + .filter(|&peer_id| !self.is_locally_disabled(peer_id)) + .collect() + } + + /// List the remote peer IDs on which this shard is active, excludes the local peer ID. + pub fn active_remote_shards(&self) -> Vec { + let replica_state = self.replica_state.read(); + let this_peer_id = replica_state.this_peer_id; + replica_state + .active_peers() // This includes `Active` and `ReshardingScaleDown` replicas! + .into_iter() + .filter(|&peer_id| !self.is_locally_disabled(peer_id) && peer_id != this_peer_id) + .collect() + } + + /// Wait for a local shard to be initialized. + /// + /// Uses a blocking thread internally. + pub async fn wait_for_local(&self, timeout: Duration) -> CollectionResult<()> { + self.wait_for(|replica_set_state| replica_set_state.is_local, timeout) + .await + } + + /// Wait for a peer shard to get into `state` + /// + /// Uses a blocking thread internally. + pub async fn wait_for_state( + &self, + peer_id: PeerId, + state: ReplicaState, + timeout: Duration, + ) -> CollectionResult<()> { + self.wait_for( + move |replica_set_state| replica_set_state.get_peer_state(peer_id) == Some(state), + timeout, + ) + .await + } + + pub fn wait_for_state_condition_sync(&self, check: F, timeout: Duration) -> bool + where + F: Fn(&ReplicaSetState) -> bool, + { + let replica_state = self.replica_state.clone(); + replica_state.wait_for(check, timeout) + } + + /// Wait for a replica set state condition to be true. + /// + /// Uses a blocking thread internally. + /// + /// # Cancel safety + /// + /// This method is cancel safe. + async fn wait_for(&self, check: F, timeout: Duration) -> CollectionResult<()> + where + F: Fn(&ReplicaSetState) -> bool + Send + 'static, + { + // TODO: Propagate cancellation into `spawn_blocking` task!? + + let replica_state = self.replica_state.clone(); + let timed_out = + !tokio::task::spawn_blocking(move || replica_state.wait_for(check, timeout)) + .await + .map_err(|err| { + CollectionError::service_error(format!( + "Failed to wait for replica set state: {err}" + )) + })?; + + if timed_out { + return Err(CollectionError::service_error( + "Failed to wait for replica set state, timed out", + )); + } + + Ok(()) + } + + /// Clears the local shard data and loads an empty local shard + pub async fn init_empty_local_shard(&self) -> CollectionResult<()> { + let mut local = self.local.write().await; + + let current_shard = local.take(); + + LocalShard::clear(&self.shard_path).await?; + let local_shard_res = LocalShard::build( + self.shard_id, + self.collection_id.clone(), + &self.shard_path, + self.collection_config.clone(), + self.shared_storage_config.clone(), + self.payload_index_schema.clone(), + self.update_runtime.clone(), + self.search_runtime.clone(), + self.optimizer_resource_budget.clone(), + self.optimizers_config.clone(), + ) + .await; + + match local_shard_res { + Ok(local_shard) => { + *local = Some(Shard::Local(local_shard)); + Ok(()) + } + Err(err) => { + log::error!( + "Failed to initialize local shard {:?}: {err}", + self.shard_path + ); + *local = current_shard; + Err(err) + } + } + } + + pub async fn set_local( + &self, + local: LocalShard, + state: Option, + ) -> CollectionResult> { + let old_shard = self.local.write().await.replace(Shard::Local(local)); + + if !self.replica_state.read().is_local || state.is_some() { + self.replica_state.write(|rs| { + rs.is_local = true; + if let Some(state) = state { + rs.set_peer_state(self.this_peer_id(), state); + } + })?; + } + self.update_locally_disabled(self.this_peer_id()); + Ok(old_shard) + } + + pub async fn remove_local(&self) -> CollectionResult<()> { + self.replica_state.write(|rs| { + rs.is_local = false; + let this_peer_id = rs.this_peer_id; + rs.remove_peer_state(this_peer_id); + })?; + + self.update_locally_disabled(self.this_peer_id()); + + let removing_local = { + let mut local = self.local.write().await; + local.take() + }; + + if let Some(removing_local) = removing_local { + // stop ongoing tasks and delete data + drop(removing_local); + LocalShard::clear(&self.shard_path).await?; + } + Ok(()) + } + + pub async fn add_remote(&self, peer_id: PeerId, state: ReplicaState) -> CollectionResult<()> { + debug_assert!(peer_id != self.this_peer_id()); + + self.replica_state.write(|rs| { + rs.set_peer_state(peer_id, state); + })?; + + self.update_locally_disabled(peer_id); + + let mut remotes = self.remotes.write().await; + + // check remote already exists + if remotes.iter().any(|remote| remote.peer_id == peer_id) { + return Ok(()); + } + + remotes.push(RemoteShard::new( + self.shard_id, + self.collection_id.clone(), + peer_id, + self.channel_service.clone(), + )); + + Ok(()) + } + + pub async fn remove_remote(&self, peer_id: PeerId) -> CollectionResult<()> { + self.replica_state.write(|rs| { + rs.remove_peer_state(peer_id); + })?; + + self.update_locally_disabled(peer_id); + + let mut remotes = self.remotes.write().await; + remotes.retain(|remote| remote.peer_id != peer_id); + Ok(()) + } + + /// Ensure that remote shard is initialized. + pub async fn ensure_replica_with_state( + &self, + peer_id: PeerId, + state: ReplicaState, + ) -> CollectionResult<()> { + if peer_id == self.this_peer_id() { + self.set_replica_state(peer_id, state)?; + } else { + // Create remote shard if necessary + self.add_remote(peer_id, state).await?; + } + Ok(()) + } + + pub fn set_replica_state(&self, peer_id: PeerId, state: ReplicaState) -> CollectionResult<()> { + log::debug!( + "Changing local shard {}:{} state from {:?} to {state:?}", + self.collection_id, + self.shard_id, + self.replica_state.read().get_peer_state(peer_id), + ); + + self.replica_state.write(|rs| { + if rs.this_peer_id == peer_id { + rs.is_local = true; + } + rs.set_peer_state(peer_id, state); + })?; + self.update_locally_disabled(peer_id); + Ok(()) + } + + pub async fn apply_state( + &mut self, + replicas: HashMap, + shard_key: Option, + ) -> CollectionResult<()> { + let old_peers = self.replica_state.read().peers(); + + self.replica_state.write(|state| { + state.set_peers(replicas.clone()); + })?; + + self.locally_disabled_peers.write().clear(); + + let removed_peers = old_peers + .keys() + .filter(|peer_id| !replicas.contains_key(peer_id)) + .copied() + .collect::>(); + for peer_id in removed_peers { + self.remove_peer(peer_id).await?; + } + + for (peer_id, state) in replicas { + let peer_already_exists = old_peers.contains_key(&peer_id); + + if peer_already_exists { + // do nothing + // We only need to change state and it is already saved + continue; + } + + if peer_id == self.this_peer_id() { + // Consensus wants a local replica on this peer + let local_shard = LocalShard::build( + self.shard_id, + self.collection_id.clone(), + &self.shard_path, + self.collection_config.clone(), + self.shared_storage_config.clone(), + self.payload_index_schema.clone(), + self.update_runtime.clone(), + self.search_runtime.clone(), + self.optimizer_resource_budget.clone(), + self.optimizers_config.clone(), + ) + .await?; + match state { + ReplicaState::Active + | ReplicaState::Listener + | ReplicaState::ReshardingScaleDown => { + // No way we can provide up-to-date replica right away at this point, + // so we report a failure to consensus + self.set_local(local_shard, Some(state)).await?; + self.notify_peer_failure(peer_id, Some(state)); + } + + ReplicaState::Dead + | ReplicaState::Partial + | ReplicaState::Initializing + | ReplicaState::PartialSnapshot + | ReplicaState::Recovery + | ReplicaState::Resharding => { + self.set_local(local_shard, Some(state)).await?; + } + } + continue; + } + + // Otherwise it is a missing remote replica, we simply create it + + let new_remote = RemoteShard::new( + self.shard_id, + self.collection_id.clone(), + peer_id, + self.channel_service.clone(), + ); + self.remotes.write().await.push(new_remote); + } + + // Apply shard key + self.shard_key = shard_key; + + Ok(()) + } + + pub(crate) async fn on_optimizer_config_update(&self) -> CollectionResult<()> { + let read_local = self.local.read().await; + if let Some(shard) = &*read_local { + shard.on_optimizer_config_update().await + } else { + Ok(()) + } + } + + /// Apply shard's strict mode configuration update + /// - Update read and write rate limiters + pub(crate) async fn on_strict_mode_config_update(&mut self) -> CollectionResult<()> { + let mut read_local = self.local.write().await; + if let Some(shard) = read_local.as_mut() { + shard.on_strict_mode_config_update().await + } + drop(read_local); + let config = self.collection_config.read().await; + if let Some(strict_mode_config) = &config.strict_mode_config { + if strict_mode_config.enabled == Some(true) { + // update write rate limiter + if let Some(write_rate_limit_per_min) = strict_mode_config.write_rate_limit { + let new_write_rate_limiter = + RateLimiter::new_per_minute(write_rate_limit_per_min); + self.write_rate_limiter + .replace(parking_lot::Mutex::new(new_write_rate_limiter)); + return Ok(()); + } + } + } + // remove write rate limiter for all other situations + self.write_rate_limiter.take(); + Ok(()) + } + + /// Check if the write rate limiter allows the operation to proceed + /// - hw_measurement_acc: the current hardware measurement accumulator + /// - cost_fn: the cost of the operation called lazily + /// + /// Returns an error if the rate limit is exceeded. + fn check_write_rate_limiter( + &self, + hw_measurement_acc: &HwMeasurementAcc, + cost_fn: F, + ) -> CollectionResult<()> + where + F: FnOnce() -> usize, + { + // Do not rate limit internal operation tagged with disposable measurement + if hw_measurement_acc.is_disposable() { + return Ok(()); + } + if let Some(rate_limiter) = &self.write_rate_limiter { + let cost = cost_fn(); + rate_limiter + .lock() + .try_consume(cost as f64) + .map_err(|err| CollectionError::rate_limit_error(err, cost, true))?; + } + Ok(()) + } + + /// Check if there are any locally disabled peers + /// And if so, report them to the consensus + pub fn sync_local_state(&self, get_shard_transfers: F) -> CollectionResult<()> + where + F: Fn(ShardId, PeerId) -> Vec, + { + let peers_to_notify: Vec<_> = self + .locally_disabled_peers + .write() + .notify_elapsed() + .collect(); + + for (failed_peer_id, from_state) in peers_to_notify { + self.notify_peer_failure(failed_peer_id, from_state); + + for transfer in get_shard_transfers(self.shard_id, failed_peer_id) { + self.abort_shard_transfer( + transfer, + &format!( + "{failed_peer_id}/{}:{} replica failed", + self.collection_id, self.shard_id, + ), + ); + } + } + + Ok(()) + } + + pub(crate) async fn health_check(&self, peer_id: PeerId) -> CollectionResult<()> { + let remotes = self.remotes.read().await; + + let Some(remote) = remotes.iter().find(|remote| remote.peer_id == peer_id) else { + return Err(CollectionError::NotFound { + what: format!("{peer_id}/{}:{} shard", self.collection_id, self.shard_id), + }); + }; + + remote.health_check().await?; + + Ok(()) + } + + pub async fn delete_local_points( + &self, + filter: Filter, + hw_measurement_acc: HwMeasurementAcc, + force: bool, + ) -> CollectionResult { + let local_shard_guard = self.local.read().await; + + let Some(local_shard) = local_shard_guard.deref() else { + return Err(CollectionError::NotFound { + what: format!("local shard {}:{}", self.collection_id, self.shard_id), + }); + }; + + let mut next_offset = Some(ExtendedPointId::NumId(0)); + let mut ids = Vec::new(); + + while let Some(current_offset) = next_offset { + const BATCH_SIZE: usize = 1000; + + let mut points = local_shard + .get() + .scroll_by( + Some(current_offset), + BATCH_SIZE + 1, + &false.into(), + &false.into(), + Some(&filter), + &self.search_runtime, + None, + None, + hw_measurement_acc.clone(), + ) + .await?; + + if points.len() > BATCH_SIZE { + next_offset = points.pop().map(|points| points.id); + } else { + next_offset = None; + } + + ids.extend(points.into_iter().map(|points| points.id)); + } + + if ids.is_empty() { + return Ok(UpdateResult { + operation_id: None, + status: UpdateStatus::Completed, + clock_tag: None, + }); + } + + drop(local_shard_guard); + + let op = + CollectionUpdateOperations::PointOperation(point_ops::PointOperations::DeletePoints { + ids, + }); + + // TODO(resharding): Assign clock tag to the operation!? 🤔 + let result = self + .update_local(op.into(), true, hw_measurement_acc, force) + .await? + .ok_or_else(|| { + CollectionError::bad_request(format!( + "local shard {}:{} does not exist or is unavailable", + self.collection_id, self.shard_id, + )) + })?; + + Ok(result) + } + + fn init_remote_shards( + shard_id: ShardId, + collection_id: CollectionId, + state: &ReplicaSetState, + channel_service: &ChannelService, + ) -> Vec { + state + .peers() + .iter() + .filter(|(peer, _)| **peer != state.this_peer_id) + .map(|(peer_id, _is_active)| { + RemoteShard::new( + shard_id, + collection_id.clone(), + *peer_id, + channel_service.clone(), + ) + }) + .collect() + } + + /// Check whether a peer is registered as `active`. + /// Unknown peers are not active. + fn peer_is_active(&self, peer_id: PeerId) -> bool { + // This is used *exclusively* during `execute_*_read_operation`, and so it *should* consider + // `ReshardingScaleDown` replicas + let is_active = matches!( + self.peer_state(peer_id), + Some(ReplicaState::Active | ReplicaState::ReshardingScaleDown) + ); + + is_active && !self.is_locally_disabled(peer_id) + } + + fn peer_is_active_or_resharding(&self, peer_id: PeerId) -> bool { + let is_active_or_resharding = matches!( + self.peer_state(peer_id), + Some( + ReplicaState::Active | ReplicaState::Resharding | ReplicaState::ReshardingScaleDown + ) + ); + + let is_locally_disabled = self.is_locally_disabled(peer_id); + + is_active_or_resharding && !is_locally_disabled + } + + fn is_locally_disabled(&self, peer_id: PeerId) -> bool { + self.locally_disabled_peers.read().is_disabled(peer_id) + } + + /// Locally disable given peer + /// + /// Disables the peer and notifies consensus periodically. + /// + /// Prevents disabling the last peer (according to consensus). + /// If `from_state` is given, the peer will only be disabled if the given state matches + /// consensus. + fn add_locally_disabled( + &self, + state: &ReplicaSetState, + peer_id: PeerId, + from_state: Option, + ) { + let other_peers = state + .active_or_resharding_peers() + .filter(|id| id != &peer_id); + + let mut locally_disabled_peers_guard = self.locally_disabled_peers.upgradable_read(); + + // Prevent disabling last peer in consensus + { + if !locally_disabled_peers_guard.is_disabled(peer_id) + && locally_disabled_peers_guard.is_all_disabled(other_peers) + { + log::warn!("Cannot locally disable last active peer {peer_id} for replica"); + return; + } + } + + locally_disabled_peers_guard.with_upgraded(|locally_disabled_peers| { + if locally_disabled_peers.disable_peer_and_notify_if_elapsed(peer_id, from_state) { + self.notify_peer_failure(peer_id, from_state); + } + }); + } + + /// Make sure that locally disabled peers do not contradict the consensus + fn update_locally_disabled(&self, peer_id_to_remove: PeerId) { + let mut locally_disabled_peers = self.locally_disabled_peers.write(); + + // Check that we are not trying to disable the last active peer + if locally_disabled_peers + .is_all_disabled(self.replica_state.read().active_or_resharding_peers()) + { + log::warn!("Resolving consensus/local state inconsistency"); + locally_disabled_peers.clear(); + } else { + locally_disabled_peers.enable_peer(peer_id_to_remove); + } + } + + fn notify_peer_failure(&self, peer_id: PeerId, from_state: Option) { + log::debug!("Notify peer failure: {peer_id}"); + self.notify_peer_failure_cb.deref()(peer_id, self.shard_id, from_state) + } + + fn abort_shard_transfer(&self, transfer: ShardTransfer, reason: &str) { + log::debug!( + "Abort {}:{} / {} -> {} shard transfer", + self.collection_id, + transfer.shard_id, + transfer.from, + transfer.to, + ); + + self.abort_shard_transfer_cb.deref()(transfer, reason) + } + + /// Get shard recovery point for WAL. + pub(crate) async fn shard_recovery_point(&self) -> CollectionResult { + let local_shard = self.local.read().await; + let Some(local_shard) = local_shard.as_ref() else { + return Err(CollectionError::NotFound { + what: "Peer does not have local shard".into(), + }); + }; + + local_shard.shard_recovery_point().await + } + + /// Update the cutoff point for the local shard. + pub(crate) async fn update_shard_cutoff_point( + &self, + cutoff: &RecoveryPoint, + ) -> CollectionResult<()> { + let local_shard = self.local.read().await; + let Some(local_shard) = local_shard.as_ref() else { + return Err(CollectionError::NotFound { + what: "Peer does not have local shard".into(), + }); + }; + + local_shard.update_cutoff(cutoff).await + } + + pub(crate) fn get_snapshots_storage_manager(& \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/extracted_output.txt b/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/extracted_output.txt new file mode 100644 index 000000000..ffd7bfaac --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/extracted_output.txt @@ -0,0 +1,870 @@ +use std::cmp::min; +use std::collections::HashSet; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; + +use common::budget::ResourceBudget; +use common::counter::hardware_accumulator::HwMeasurementAcc; +use common::counter::hardware_counter::HardwareCounterCell; +use common::panic; +use itertools::Itertools; +use log::{debug, error, info, trace, warn}; +use parking_lot::Mutex; +use segment::common::operation_error::OperationResult; +use segment::index::hnsw_index::num_rayon_threads; +use segment::types::SeqNumberType; +use tokio::runtime::Handle; +use tokio::sync::mpsc::{self, Receiver, Sender}; +use tokio::sync::{Mutex as TokioMutex, oneshot}; +use tokio::task::{self, JoinHandle}; +use tokio::time::error::Elapsed; +use tokio::time::{Duration, timeout}; + +use crate::collection::payload_index_schema::PayloadIndexSchema; +use crate::collection_manager::collection_updater::CollectionUpdater; +use crate::collection_manager::holders::segment_holder::LockedSegmentHolder; +use crate::collection_manager::optimizers::segment_optimizer::{ + OptimizerThresholds, SegmentOptimizer, +}; +use crate::collection_manager::optimizers::{Tracker, TrackerLog, TrackerStatus}; +use crate::common::stoppable_task::{StoppableTaskHandle, spawn_stoppable}; +use crate::config::CollectionParams; +use crate::operations::CollectionUpdateOperations; +use crate::operations::shared_storage_config::SharedStorageConfig; +use crate::operations::types::{CollectionError, CollectionResult}; +use crate::save_on_disk::SaveOnDisk; +use crate::shards::local_shard::LocalShardClocks; +use crate::wal::WalError; +use crate::wal_delta::LockedWal; + +/// Interval at which the optimizer worker cleans up old optimization handles +/// +/// The longer the duration, the longer it takes for panicked tasks to be reported. +const OPTIMIZER_CLEANUP_INTERVAL: Duration = Duration::from_secs(5); + +pub type Optimizer = dyn SegmentOptimizer + Sync + Send; + +/// Information, required to perform operation and notify regarding the result +#[derive(Debug)] +pub struct OperationData { + /// Sequential number of the operation + pub op_num: SeqNumberType, + /// Operation + pub operation: CollectionUpdateOperations, + /// If operation was requested to wait for result + pub wait: bool, + /// Callback notification channel + pub sender: Option>>, + pub hw_measurements: HwMeasurementAcc, +} + +/// Signal, used to inform Updater process +#[derive(Debug)] +pub enum UpdateSignal { + /// Requested operation to perform + Operation(OperationData), + /// Stop all optimizers and listening + Stop, + /// Empty signal used to trigger optimizers + Nop, + /// Ensures that previous updates are applied + Plunger(oneshot::Sender<()>), +} + +/// Signal, used to inform Optimization process +#[derive(PartialEq, Eq, Clone, Copy)] +pub enum OptimizerSignal { + /// Sequential number of the operation + Operation(SeqNumberType), + /// Stop all optimizers and listening + Stop, + /// Empty signal used to trigger optimizers + Nop, +} + +/// Structure, which holds object, required for processing updates of the collection +pub struct UpdateHandler { + shared_storage_config: Arc, + payload_index_schema: Arc>, + /// List of used optimizers + pub optimizers: Arc>>, + /// Log of optimizer statuses + optimizers_log: Arc>, + total_optimized_points: Arc, + /// Global CPU budget in number of cores for all optimization tasks. + /// Assigns CPU permits to tasks to limit overall resource utilization. + optimizer_resource_budget: ResourceBudget, + /// How frequent can we flush data + /// This parameter depends on the optimizer config and should be updated accordingly. + pub flush_interval_sec: u64, + segments: LockedSegmentHolder, + /// Process, that listens updates signals and perform updates + update_worker: Option>, + /// Process, that listens for post-update signals and performs optimization + optimizer_worker: Option>, + /// Process that periodically flushes segments and tries to truncate wal + flush_worker: Option>, + /// Sender to stop flush worker + flush_stop: Option>, + runtime_handle: Handle, + /// WAL, required for operations + wal: LockedWal, + /// Always keep this WAL version and later and prevent acknowledging/truncating from the WAL. + /// This is used when other bits of code still depend on information in the WAL, such as the + /// queue proxy shard. + /// Defaults to `u64::MAX` to allow acknowledging all confirmed versions. + pub(super) wal_keep_from: Arc, + optimization_handles: Arc>>>, + /// Maximum number of concurrent optimization jobs in this update handler. + /// This parameter depends on the optimizer config and should be updated accordingly. + pub max_optimization_threads: Option, + /// Highest and cutoff clocks for the shard WAL. + clocks: LocalShardClocks, + shard_path: PathBuf, + /// Whether we have ever triggered optimizers since starting. + has_triggered_optimizers: Arc, +} + +impl UpdateHandler { + #[allow(clippy::too_many_arguments)] + pub fn new( + shared_storage_config: Arc, + payload_index_schema: Arc>, + optimizers: Arc>>, + optimizers_log: Arc>, + total_optimized_points: Arc, + optimizer_resource_budget: ResourceBudget, + runtime_handle: Handle, + segments: LockedSegmentHolder, + wal: LockedWal, + flush_interval_sec: u64, + max_optimization_threads: Option, + clocks: LocalShardClocks, + shard_path: PathBuf, + ) -> UpdateHandler { + UpdateHandler { + shared_storage_config, + payload_index_schema, + optimizers, + update_worker: None, + optimizer_worker: None, + optimizers_log, + total_optimized_points, + optimizer_resource_budget, + flush_interval_sec, + segments, + flush_worker: None, + flush_stop: None, + runtime_handle, + wal, + wal_keep_from: Arc::new(u64::MAX.into()), + optimization_handles: Arc::new(TokioMutex::new(vec![])), + max_optimization_threads, + clocks, + shard_path, + has_triggered_optimizers: Arc::new(AtomicBool::new(false)), + } + } + + pub fn stop_flush_worker(&mut self) { + if let Some(flush_stop) = self.flush_stop.take() { + if let Err(()) = flush_stop.send(()) { + warn!("Failed to stop flush worker as it is already stopped."); + } + } + } + + /// Gracefully wait before all optimizations stop + /// If some optimization is in progress - it will be finished before shutdown. + /// Blocking function. + pub async fn wait_workers_stops(&mut self) -> CollectionResult<()> { + let maybe_handle = self.update_worker.take(); + if let Some(handle) = maybe_handle { + handle.await?; + } + let maybe_handle = self.optimizer_worker.take(); + if let Some(handle) = maybe_handle { + handle.await?; + } + let maybe_handle = self.flush_worker.take(); + if let Some(handle) = maybe_handle { + handle.await?; + } + + let mut opt_handles_guard = self.optimization_handles.lock().await; + let opt_handles = std::mem::take(&mut *opt_handles_guard); + let stopping_handles = opt_handles + .into_iter() + .filter_map(|h| h.stop()) + .collect_vec(); + + for res in stopping_handles { + res.await?; + } + + Ok(()) + } + + pub fn run_workers(&mut self, update_receiver: Receiver) { + let (tx, rx) = mpsc::channel(self.shared_storage_config.update_queue_size); + self.optimizer_worker = Some(self.runtime_handle.spawn(Self::optimization_worker_fn( + self.optimizers.clone(), + tx.clone(), + rx, + self.segments.clone(), + self.wal.clone(), + self.optimization_handles.clone(), + self.optimizers_log.clone(), + self.total_optimized_points.clone(), + self.optimizer_resource_budget.clone(), + self.max_optimization_threads, + self.has_triggered_optimizers.clone(), + self.payload_index_schema.clone(), + ))); + self.update_worker = Some(self.runtime_handle.spawn(Self::update_worker_fn( + update_receiver, + tx, + self.segments.clone(), + self.wal.clone(), + ))); + let (flush_tx, flush_rx) = oneshot::channel(); + self.flush_worker = Some(self.runtime_handle.spawn(Self::flush_worker( + self.segments.clone(), + self.wal.clone(), + self.wal_keep_from.clone(), + self.flush_interval_sec, + flush_rx, + self.clocks.clone(), + self.shard_path.clone(), + ))); + self.flush_stop = Some(flush_tx); + } + + /// Checks if there are any failed operations. + /// If so - attempts to re-apply all failed operations. + async fn try_recover(segments: LockedSegmentHolder, wal: LockedWal) -> CollectionResult { + // Try to re-apply everything starting from the first failed operation + let first_failed_operation_option = segments.read().failed_operation.iter().cloned().min(); + match first_failed_operation_option { + None => Ok(0), + Some(first_failed_op) => { + let wal_lock = wal.lock().await; + for (op_num, operation) in wal_lock.read(first_failed_op) { + CollectionUpdater::update( + &segments, + op_num, + operation.operation, + &HardwareCounterCell::disposable(), // Internal operation, no measurement needed + )?; + } + Err(CollectionError::service_error("Failed operation recovery should not return".to_string())) + } + } + } + + fn process_optimization( + optimizers: Arc>>, + segments: LockedSegmentHolder, + ) -> Vec> { + for optimizer in optimizers.iter() { + let nonoptimal_segment_ids = optimizer.check_condition(segments.clone()); + while !nonoptimal_segment_ids.is_empty() { + debug!( + "Start optimization on segments: {:?}", + nonoptimal_segment_ids + ); + // If optimization fails, it could not be reported to anywhere except for console. + // So the only recovery here is to stop optimization and await for restart + if optimizer.optimize(segments.clone(), nonoptimal_segment_ids).is_err() { + continue; + }; + nonoptimal_segment_ids = optimizer.check_condition(segments.clone()); + } + } + vec![] + } + + /// Checks conditions for all optimizers until there is no suggested segment + /// Starts a task for each optimization + /// Returns handles for started tasks + pub(crate) fn launch_optimization( + optimizers: Arc>>, + optimizers_log: Arc>, + total_optimized_points: Arc, + optimizer_resource_budget: &ResourceBudget, + segments: LockedSegmentHolder, + callback: F, + limit: Option, + ) -> Vec> + where + F: Fn(bool) + Send + Clone + Sync + 'static, + { + let mut scheduled_segment_ids = HashSet::<_>::default(); + let mut handles = vec![]; + + 'outer: for optimizer in optimizers.iter() { + loop { + // Return early if we reached the optimization job limit + if limit.map(|extra| handles.len() >= extra).unwrap_or(false) { + info!("Reached optimization job limit, postponing other optimizations"); + break 'outer; + } + + let nonoptimal_segment_ids = + optimizer.check_condition(segments.clone(), &scheduled_segment_ids); + if nonoptimal_segment_ids.is_empty() { + break; + } + + debug!("Optimizing segments: {:?}", &nonoptimal_segment_ids); + + // Determine how many Resources we prefer for optimization task, acquire permit for it + // And use same amount of IO threads as CPUs + let max_indexing_threads = optimizer.hnsw_config().max_indexing_threads; + let desired_io = num_rayon_threads(max_indexing_threads); + let Some(mut permit) = optimizer_resource_budget.try_acquire(0, desired_io) else { + // If there is no Resource budget, break outer loop and return early + // If we have no handles (no optimizations) trigger callback so that we wake up + // our optimization worker to try again later, otherwise it could get stuck + log::trace!( + "No available IO permit for {} optimizer, postponing", + optimizer.name(), + ); + if handles.is_empty() { + callback(false); + } + break 'outer; + }; + log::trace!( + "Acquired {} IO permit for {} optimizer", + permit.num_io, + optimizer.name(), + ); + + let permit_callback = callback.clone(); + + permit.set_on_release(move || { + // Notify scheduler that resource budget changed + permit_callback(false); + }); + + let optimizer = optimizer.clone(); + let optimizers_log = optimizers_log.clone(); + let total_optimized_points = total_optimized_points.clone(); + let segments = segments.clone(); + let nsi = nonoptimal_segment_ids.clone(); + scheduled_segment_ids.extend(&nsi); + let callback = callback.clone(); + + let handle = spawn_stoppable( + // Stoppable task + { + let resource_budget = optimizer_resource_budget.clone(); + let segments = segments.clone(); + move |stopped| { + // Track optimizer status + let tracker = Tracker::start(optimizer.as_ref().name(), nsi.clone()); + let tracker_handle = tracker.handle(); + optimizers_log.lock().register(tracker); + + // Optimize and handle result + match optimizer.as_ref().optimize( + segments.clone(), + nsi, + permit, + resource_budget, + stopped, + ) { + // Perform some actions when optimization if finished + Ok(optimized_points) => { + let is_optimized = optimized_points > 0; + total_optimized_points + .fetch_add(optimized_points, Ordering::Relaxed); + tracker_handle.update(TrackerStatus::Done); + callback(is_optimized); + is_optimized + } + // Handle and report errors + Err(error) => match error { + CollectionError::Cancelled { description } => { + debug!("Optimization cancelled - {description}"); + tracker_handle.update(TrackerStatus::Cancelled(description)); + false + } + _ => { + segments.write().report_optimizer_error(error.clone()); + + // Error of the optimization can not be handled by API user + // It is only possible to fix after full restart, + // so the best available action here is to stop whole + // optimization thread and log the error + log::error!("Optimization error: {error}"); + + tracker_handle.update(TrackerStatus::Error(error.to_string())); + + panic!("Optimization error: {error}"); + } + }, + } + } + }, + // Panic handler + Some(Box::new(move |panic_payload| { + let message = panic::downcast_str(&panic_payload).unwrap_or(""); + let separator = if !message.is_empty() { ": " } else { "" }; + + warn!( + "Optimization task panicked, collection may be in unstable state\ + {separator}{message}" + ); + + segments.write().report_optimizer_error(CollectionError::service_error(format!( + "Optimization task panicked{separator}{message}" + ))); + })), + ); + handles.push(handle); + } + } + + handles + } + + /// Ensure there is at least one appendable segment with enough capacity + /// + /// If there is no appendable segment, or all are at or over capacity, a new empty one is + /// created. + /// + /// Capacity is determined based on `optimizers.max_segment_size_kb`. + pub(super) fn ensure_appendable_segment_with_capacity( + segments: &LockedSegmentHolder, + segments_path: &Path, + collection_params: &CollectionParams, + thresholds_config: &OptimizerThresholds, + payload_index_schema: &PayloadIndexSchema, + ) -> OperationResult<()> { + let no_segment_with_capacity = { + let segments_read = segments.read(); + segments_read + .appendable_segments_ids() + .into_iter() + .filter_map(|segment_id| segments_read.get(segment_id)) + .all(|segment| { + let max_vector_size_bytes = + segment.get().read().max_available_vectors_size_in_bytes().unwrap_or_default(); + let max_segment_size_bytes = thresholds_config + .max_segment_size_kb + .saturating_mul(segment::common::BYTES_IN_KB); + + max_vector_size_bytes >= max_segment_size_bytes + }) + }; + + if no_segment_with_capacity { + debug!("Creating new appendable segment, all existing segments are over capacity"); + segments.write().create_appendable_segment( + segments_path, + collection_params, + payload_index_schema, + )?; + } + + Ok(()) + } + + /// Checks the optimizer conditions. + /// + /// This function returns a tuple of two booleans: + /// - The first indicates if any optimizers have been triggered since startup. + /// - The second indicates if there are any pending/suboptimal optimizers. + pub(crate) fn check_optimizer_conditions(&self) -> (bool, bool) { + // Check if Qdrant triggered any optimizations since starting at all + let has_triggered_any_optimizers = self.has_triggered_optimizers.load(Ordering::Relaxed); + + let excluded_ids = HashSet::<_>::default(); + let has_suboptimal_optimizers = self.optimizers.iter().any(|optimizer| { + let nonoptimal_segment_ids = + optimizer.check_condition(self.segments.clone(), &excluded_ids); + !nonoptimal_segment_ids.is_empty() + }); + + (has_triggered_any_optimizers, has_suboptimal_optimizers) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn process_optimization( + optimizers: Arc>>, + segments: LockedSegmentHolder, + optimization_handles: Arc>>>, + optimizers_log: Arc>, + total_optimized_points: Arc, + optimizer_resource_budget: &ResourceBudget, + sender: Sender, + limit: usize, + ) { + let mut new_handles = Self::launch_optimization( + optimizers.clone(), + optimizers_log, + total_optimized_points, + optimizer_resource_budget, + segments.clone(), + move |_optimization_result| { + // After optimization is finished, we still need to check if there are + // some further optimizations possible. + // If receiver is already dead - we do not care. + // If channel is full - optimization will be triggered by some other signal + let _ = sender.try_send(OptimizerSignal::Nop); + }, + Some(limit), + ); + let mut handles = optimization_handles.lock().await; + handles.append(&mut new_handles); + handles.retain(|h| !h.is_finished()) + } + + /// Cleanup finalized optimization task handles + /// + /// This finds and removes completed tasks from our list of optimization handles. + /// It also propagates any panics (and unknown errors) so we properly handle them if desired. + /// + /// It is essential to call this every once in a while for handling panics in time. + /// + /// Returns true if any optimization handle was finished, joined and removed. + async fn cleanup_optimization_handles( + optimization_handles: Arc>>>, + ) -> bool { + // Remove finished handles + let finished_handles: Vec<_> = { + let mut handles = optimization_handles.lock().await; + (0..handles.len()) + .filter(|i| handles[*i].is_finished()) + .collect::>() + .into_iter() + .rev() + .map(|i| handles.swap_remove(i)) + .collect() + }; + + let finished_any = !finished_handles.is_empty(); + + // Finalize all finished handles to propagate panics + for handle in finished_handles { + handle.join_and_handle_panic().await; + } + + finished_any + } + + #[allow(clippy::too_many_arguments)] + async fn optimization_worker_fn( + optimizers: Arc>>, + sender: Sender, + mut receiver: Receiver, + segments: LockedSegmentHolder, + wal: LockedWal, + optimization_handles: Arc>>>, + optimizers_log: Arc>, + total_optimized_points: Arc, + optimizer_resource_budget: ResourceBudget, + max_handles: Option, + has_triggered_optimizers: Arc, + payload_index_schema: Arc>, + ) { + let max_handles = max_handles.unwrap_or(usize::MAX); + let max_indexing_threads = optimizers + .first() + .map(|optimizer| optimizer.hnsw_config().max_indexing_threads) + .unwrap_or_default(); + + // Asynchronous task to trigger optimizers once resource budget is available again + let mut resource_available_trigger: Option> = None; + + loop { + let result = timeout(OPTIMIZER_CLEANUP_INTERVAL, receiver.recv()).await; + + let cleaned_any = Self::cleanup_optimization_handles(optimization_handles.clone()).await; + + // Either continue below here with the worker, or reloop/break + // Decision logic doing one of three things: + // 1. run optimizers + // 2. reloop and wait for next signal + // 3. break here and stop the optimization worker + let ignore_max_handles = match result { + // Regular optimizer signal: run optimizers: do 1 + Ok(Some(OptimizerSignal::Operation(_))) => false, + // Optimizer signal ignoring max handles: do 1 + Ok(Some(OptimizerSignal::Nop)) => true, + // Hit optimizer cleanup interval, did clean up a task: do 1 + Err(Elapsed { .. }) if cleaned_any => { + // This branch prevents a race condition where optimizers would get stuck + // If the optimizer cleanup interval was triggered and we did clean any task we + // must run optimizers now. If we don't there may not be any other ongoing + // tasks that'll trigger this for us. If we don't run optimizers here we might + // get stuck into yellow state until a new update operation is received. + // See: + log::warn!( + "Cleaned a optimization handle after timeout, explicitly triggering optimizers", + ); + true + } + // Hit optimizer cleanup interval, did not clean up a task: do 2 + Err(Elapsed { .. }) => continue, + // Channel closed or received stop signal: do 3 + Ok(None | Some(OptimizerSignal::Stop)) => break, + }; + + has_triggered_optimizers.store(true, Ordering::Relaxed); + + // Ensure we have at least one appendable segment with enough capacity + // Source required parameters from first optimizer + if let Some(optimizer) = optimizers.first() { + let result = Self::ensure_appendable_segment_with_capacity( + &segments, + optimizer.segments_path(), + &optimizer.collection_params(), + optimizer.threshold_config(), + &payload_index_schema.read(), + ); + if let Err(err) = result { + log::error!("Failed to ensure there are appendable segments with capacity: {err}"); + panic!("Failed to ensure there are appendable segments with capacity: {err}"); + } + } + + // If not forcing, wait on next signal if we have too many handles + if !ignore_max_handles && optimization_handles.lock().await.len() >= max_handles { + continue; + } + + if Self::try_recover(segments.clone(), wal.clone()) + .await + .is_err() + { + continue; + } + + // Continue if we have enough resource budget available to start an optimization + // Otherwise skip now and start a task to trigger the optimizer again once resource + // budget becomes available + let desired_cpus = 0; + let desired_io = num_rayon_threads(max_indexing_threads); + if !optimizer_resource_budget.has_budget(desired_cpus, desired_io) { + let trigger_active = resource_available_trigger + .as_ref() + .is_some_and(|t| !t.is_finished()); + if !trigger_active { + resource_available_trigger.replace(trigger_optimizers_on_resource_budget( + optimizer_resource_budget.clone(), + desired_cpus, + desired_io, + sender.clone(), + )); + } + continue; + } + + // Determine optimization handle limit based on max handles we allow + // Not related to the CPU budget, but a different limit for the maximum number + // of concurrent concrete optimizations per shard as configured by the user in + // the Qdrant configuration. + // Skip if we reached limit, an ongoing optimization that finishes will trigger this loop again + let limit = max_handles.saturating_sub(optimization_handles.lock().await.len()); + if limit == 0 { + log::trace!("Skipping optimization check, we reached optimization thread limit"); + continue; + } + + Self::process_optimization( + optimizers.clone(), + segments.clone(), + optimization_handles.clone(), + optimizers_log.clone(), + total_optimized_points.clone(), + &optimizer_resource_budget, + sender.clone(), + limit, + ) + .await; + } + } + + async fn update_worker_fn( + mut receiver: Receiver, + optimize_sender: Sender, + segments: LockedSegmentHolder, + wal: LockedWal, + ) { + while let Some(signal) = receiver.recv().await { + match signal { + UpdateSignal::Operation(OperationData { + op_num, + operation, + sender, + wait, + hw_measurements, + }) => { + let flush_res = if wait { + wal.lock().await.flush().map_err(|err| { + CollectionError::service_error(format!( + "Can't flush WAL before operation {op_num} - {err}" + )) + }) + } else { + Ok(()) + }; + + let operation_result = flush_res.and_then(|_| { + CollectionUpdater::update( + &segments, + op_num, + operation, + &hw_measurements.get_counter_cell(), + ) + }); + + let res = match operation_result { + Ok(update_res) => optimize_sender + .send(OptimizerSignal::Operation(op_num)) + .await + .and(Ok(update_res)) + .map_err(|send_err| send_err.into()), + Err(err) => Err(err), + }; + + if let Some(feedback) = sender { + feedback.send(res).unwrap_or_else(|_| { + debug!("Can't report operation {op_num} result. Assume already not required"); + }); + }; + } + UpdateSignal::Stop => { + optimize_sender + .send(OptimizerSignal::Stop) + .await + .unwrap_or_else(|_| debug!("Optimizer already stopped")); + break; + } + UpdateSignal::Nop => optimize_sender + .send(OptimizerSignal::Nop) + .await + .unwrap_or_else(|_| { + info!( + "Can't notify optimizers, assume process is dead. Restart is required" + ); + }), + UpdateSignal::Plunger(callback_sender) => { + callback_sender.send(()).unwrap_or_else(|_| { + debug!("Can't notify sender, assume nobody is waiting anymore"); + }); + } + } + } + // Transmitter was destroyed + optimize_sender + .send(OptimizerSignal::Stop) + .await + .unwrap_or_else(|_| debug!("Optimizer already stopped")); + } + + async fn flush_worker( + segments: LockedSegmentHolder, + wal: LockedWal, + wal_keep_from: Arc, + flush_interval_sec: u64, + mut stop_receiver: oneshot::Receiver<()>, + clocks: LocalShardClocks, + shard_path: PathBuf, + ) { + loop { + // Stop flush worker on signal or if sender was dropped + // Even if timer did not finish + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(flush_interval_sec)) => {}, + _ = &mut stop_receiver => { + debug!("Stopping flush worker for shard {}", shard_path.display()); + return; + } + } + + trace!("Attempting flushing"); + let wal_flash_job = wal.lock().await.flush_async(); + + if let Err(err) = wal_flash_job.join() { + error!("Failed to flush wal: {err:?}"); + segments + .write() + .report_optimizer_error(WalError::WriteWalError(format!( + "WAL flush error: {err:?}" + ))); + continue; + } + + let confirmed_version = Self::flush_segments(segments.clone()); + let confirmed_version = match confirmed_version { + Ok(version) => version, + Err(err) => { + error!("Failed to flush: {err}"); + segments.write().report_optimizer_error(err); + continue; + } + }; + // Acknowledge confirmed version in WAL, but don't acknowledge the specified + // `keep_from` index or higher. + // This is to prevent truncating WAL entries that other bits of code still depend on + // such as the queue proxy shard. + // Default keep_from is `u64::MAX` to allow acknowledging all confirmed. + let keep_from = wal_keep_from.load(std::sync::atomic::Ordering::Relaxed); + + // If we should keep the first message, do not acknowledge at all + if keep_from == 0 { + continue; + } + let ack = confirmed_version.min(keep_from.saturating_sub(1)); + + if let Err(err) = clocks.store_if_changed(&shard_path).await { + log::warn!("Failed to store clock maps to disk: {err}"); + segments.write().report_optimizer_error(err); + } + + if let Err(err) = wal.lock().await.ack(ack) { + log::warn!("Failed to acknowledge WAL version: {err}"); + segments.write().report_optimizer_error(err); + } + } + } + + /// Returns confirmed version after flush of all segments + /// + /// # Errors + /// Returns an error on flush failure + fn flush_segments(segments: LockedSegmentHolder) -> OperationResult { + let read_segments = segments.read(); + let flushed_version = read_segments.flush_all(false, false)?; + Ok(match read_segments.failed_operation.iter().cloned().min() { + None => flushed_version, + Some(failed_operation) => min(failed_operation, flushed_version), + }) + } +} + +/// Trigger optimizers when resource budget is available +fn trigger_optimizers_on_resource_budget( + optimizer_resource_budget: ResourceBudget, + desired_cpus: usize, + desired_io: usize, + sender: Sender, +) -> JoinHandle<()> { + task::spawn(async move { + log::trace!("Skipping optimization checks, waiting for resource budget to be available"); + optimizer_resource_budget + .notify_on_budget_available(desired_cpus, desired_io) + .await; + log::trace!("Continue optimization checks, new resource budget available"); + + // Trigger optimizers with Nop operation + sender.send(OptimizerSignal::Nop).await.unwrap_or_else(|_| { + info!("Can't notify optimizers, assume process is dead. Restart is required") + }); + }) +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/metadata.json b/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/metadata.json new file mode 100644 index 000000000..59d79103c --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "qdrant_lib_collection_src_update_handler.rs", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/qdrant_lib_collection_src_update_handler.rs_prompt.txt", + "expected_file": "locodiff-250425/prompts/qdrant_lib_collection_src_update_handler.rs_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:37:04.324566+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 36914, + "extracted_output_length": 36902, + "expected_output_length": 36129, + "results_dir": "locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704", + "generation_id": "gen-1757432224-aqJCyk3TTon0i2WlaqsJ", + "cost_usd": 0.02261798, + "prompt_tokens": 61381, + "completion_tokens": 6968, + "total_tokens": 68349, + "native_prompt_tokens": 60568, + "native_completion_tokens": 7049, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 280 +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/output.diff b/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/output.diff new file mode 100644 index 000000000..a78bc07c8 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/output.diff @@ -0,0 +1,336 @@ + +index cb922e861..ffd7bfaac 100644 +--- a/qdrant_lib_collection_src_update_handler.rs_expectedoutput.txt (expected):tmp/tmpkhlb4cxf_expected.txt ++++ b/qdrant_lib_collection_src_update_handler.rs_extracted.txt (actual):tmp/tmpf67p8fh9_actual.txt +@@ -40,7 +40,7 @@ use crate::wal_delta::LockedWal; + + /// Interval at which the optimizer worker cleans up old optimization handles + /// +-/// The longer the duration, the longer it takes for panicked tasks to be reported. ++/// The longer the duration, the longer it takes for panicked tasks to be reported. + const OPTIMIZER_CLEANUP_INTERVAL: Duration = Duration::from_secs(5); + + pub type Optimizer = dyn SegmentOptimizer + Sync + Send; +@@ -91,7 +91,6 @@ pub struct UpdateHandler { + pub optimizers: Arc>>, + /// Log of optimizer statuses + optimizers_log: Arc>, +- /// Total number of optimized points since last start + total_optimized_points: Arc, + /// Global CPU budget in number of cores for all optimization tasks. + /// Assigns CPU permits to tasks to limit overall resource utilization. +@@ -148,61 +147,26 @@ impl UpdateHandler { + shared_storage_config, + payload_index_schema, + optimizers, +- segments, + update_worker: None, + optimizer_worker: None, + optimizers_log, + total_optimized_points, + optimizer_resource_budget, ++ flush_interval_sec, ++ segments, + flush_worker: None, + flush_stop: None, + runtime_handle, + wal, + wal_keep_from: Arc::new(u64::MAX.into()), +- flush_interval_sec, + optimization_handles: Arc::new(TokioMutex::new(vec![])), + max_optimization_threads, + clocks, + shard_path, +- has_triggered_optimizers: Default::default(), ++ has_triggered_optimizers: Arc::new(AtomicBool::new(false)), + } + } + +- pub fn run_workers(&mut self, update_receiver: Receiver) { +- let (tx, rx) = mpsc::channel(self.shared_storage_config.update_queue_size); +- self.optimizer_worker = Some(self.runtime_handle.spawn(Self::optimization_worker_fn( +- self.optimizers.clone(), +- tx.clone(), +- rx, +- self.segments.clone(), +- self.wal.clone(), +- self.optimization_handles.clone(), +- self.optimizers_log.clone(), +- self.total_optimized_points.clone(), +- self.optimizer_resource_budget.clone(), +- self.max_optimization_threads, +- self.has_triggered_optimizers.clone(), +- self.payload_index_schema.clone(), +- ))); +- self.update_worker = Some(self.runtime_handle.spawn(Self::update_worker_fn( +- update_receiver, +- tx, +- self.wal.clone(), +- self.segments.clone(), +- ))); +- let (flush_tx, flush_rx) = oneshot::channel(); +- self.flush_worker = Some(self.runtime_handle.spawn(Self::flush_worker( +- self.segments.clone(), +- self.wal.clone(), +- self.wal_keep_from.clone(), +- self.flush_interval_sec, +- flush_rx, +- self.clocks.clone(), +- self.shard_path.clone(), +- ))); +- self.flush_stop = Some(flush_tx); +- } +- + pub fn stop_flush_worker(&mut self) { + if let Some(flush_stop) = self.flush_stop.take() { + if let Err(()) = flush_stop.send(()) { +@@ -213,6 +177,7 @@ impl UpdateHandler { + + /// Gracefully wait before all optimizations stop + /// If some optimization is in progress - it will be finished before shutdown. ++ /// Blocking function. + pub async fn wait_workers_stops(&mut self) -> CollectionResult<()> { + let maybe_handle = self.update_worker.take(); + if let Some(handle) = maybe_handle { +@@ -241,13 +206,48 @@ impl UpdateHandler { + Ok(()) + } + ++ pub fn run_workers(&mut self, update_receiver: Receiver) { ++ let (tx, rx) = mpsc::channel(self.shared_storage_config.update_queue_size); ++ self.optimizer_worker = Some(self.runtime_handle.spawn(Self::optimization_worker_fn( ++ self.optimizers.clone(), ++ tx.clone(), ++ rx, ++ self.segments.clone(), ++ self.wal.clone(), ++ self.optimization_handles.clone(), ++ self.optimizers_log.clone(), ++ self.total_optimized_points.clone(), ++ self.optimizer_resource_budget.clone(), ++ self.max_optimization_threads, ++ self.has_triggered_optimizers.clone(), ++ self.payload_index_schema.clone(), ++ ))); ++ self.update_worker = Some(self.runtime_handle.spawn(Self::update_worker_fn( ++ update_receiver, ++ tx, ++ self.segments.clone(), ++ self.wal.clone(), ++ ))); ++ let (flush_tx, flush_rx) = oneshot::channel(); ++ self.flush_worker = Some(self.runtime_handle.spawn(Self::flush_worker( ++ self.segments.clone(), ++ self.wal.clone(), ++ self.wal_keep_from.clone(), ++ self.flush_interval_sec, ++ flush_rx, ++ self.clocks.clone(), ++ self.shard_path.clone(), ++ ))); ++ self.flush_stop = Some(flush_tx); ++ } ++ + /// Checks if there are any failed operations. + /// If so - attempts to re-apply all failed operations. + async fn try_recover(segments: LockedSegmentHolder, wal: LockedWal) -> CollectionResult { + // Try to re-apply everything starting from the first failed operation + let first_failed_operation_option = segments.read().failed_operation.iter().cloned().min(); + match first_failed_operation_option { +- None => {} ++ None => Ok(0), + Some(first_failed_op) => { + let wal_lock = wal.lock().await; + for (op_num, operation) in wal_lock.read(first_failed_op) { +@@ -258,9 +258,31 @@ impl UpdateHandler { + &HardwareCounterCell::disposable(), // Internal operation, no measurement needed + )?; + } ++ Err(CollectionError::service_error("Failed operation recovery should not return".to_string())) + } +- }; +- Ok(0) ++ } ++ } ++ ++ fn process_optimization( ++ optimizers: Arc>>, ++ segments: LockedSegmentHolder, ++ ) -> Vec> { ++ for optimizer in optimizers.iter() { ++ let nonoptimal_segment_ids = optimizer.check_condition(segments.clone()); ++ while !nonoptimal_segment_ids.is_empty() { ++ debug!( ++ "Start optimization on segments: {:?}", ++ nonoptimal_segment_ids ++ ); ++ // If optimization fails, it could not be reported to anywhere except for console. ++ // So the only recovery here is to stop optimization and await for restart ++ if optimizer.optimize(segments.clone(), nonoptimal_segment_ids).is_err() { ++ continue; ++ }; ++ nonoptimal_segment_ids = optimizer.check_condition(segments.clone()); ++ } ++ } ++ vec![] + } + + /// Checks conditions for all optimizers until there is no suggested segment +@@ -285,7 +307,7 @@ impl UpdateHandler { + loop { + // Return early if we reached the optimization job limit + if limit.map(|extra| handles.len() >= extra).unwrap_or(false) { +- log::trace!("Reached optimization job limit, postponing other optimizations"); ++ info!("Reached optimization job limit, postponing other optimizations"); + break 'outer; + } + +@@ -367,8 +389,7 @@ impl UpdateHandler { + Err(error) => match error { + CollectionError::Cancelled { description } => { + debug!("Optimization cancelled - {description}"); +- tracker_handle +- .update(TrackerStatus::Cancelled(description)); ++ tracker_handle.update(TrackerStatus::Cancelled(description)); + false + } + _ => { +@@ -380,8 +401,7 @@ impl UpdateHandler { + // optimization thread and log the error + log::error!("Optimization error: {error}"); + +- tracker_handle +- .update(TrackerStatus::Error(error.to_string())); ++ tracker_handle.update(TrackerStatus::Error(error.to_string())); + + panic!("Optimization error: {error}"); + } +@@ -399,11 +419,9 @@ impl UpdateHandler { + {separator}{message}" + ); + +- segments +- .write() +- .report_optimizer_error(CollectionError::service_error(format!( +- "Optimization task panicked{separator}{message}" +- ))); ++ segments.write().report_optimizer_error(CollectionError::service_error(format!( ++ "Optimization task panicked{separator}{message}" ++ ))); + })), + ); + handles.push(handle); +@@ -433,11 +451,8 @@ impl UpdateHandler { + .into_iter() + .filter_map(|segment_id| segments_read.get(segment_id)) + .all(|segment| { +- let max_vector_size_bytes = segment +- .get() +- .read() +- .max_available_vectors_size_in_bytes() +- .unwrap_or_default(); ++ let max_vector_size_bytes = ++ segment.get().read().max_available_vectors_size_in_bytes().unwrap_or_default(); + let max_segment_size_bytes = thresholds_config + .max_segment_size_kb + .saturating_mul(segment::common::BYTES_IN_KB); +@@ -447,7 +462,7 @@ impl UpdateHandler { + }; + + if no_segment_with_capacity { +- log::debug!("Creating new appendable segment, all existing segments are over capacity"); ++ debug!("Creating new appendable segment, all existing segments are over capacity"); + segments.write().create_appendable_segment( + segments_path, + collection_params, +@@ -505,6 +520,7 @@ impl UpdateHandler { + ); + let mut handles = optimization_handles.lock().await; + handles.append(&mut new_handles); ++ handles.retain(|h| !h.is_finished()) + } + + /// Cleanup finalized optimization task handles +@@ -561,14 +577,13 @@ impl UpdateHandler { + .map(|optimizer| optimizer.hnsw_config().max_indexing_threads) + .unwrap_or_default(); + +- // Asynchronous task to trigger optimizers once CPU budget is available again ++ // Asynchronous task to trigger optimizers once resource budget is available again + let mut resource_available_trigger: Option> = None; + + loop { + let result = timeout(OPTIMIZER_CLEANUP_INTERVAL, receiver.recv()).await; + +- let cleaned_any = +- Self::cleanup_optimization_handles(optimization_handles.clone()).await; ++ let cleaned_any = Self::cleanup_optimization_handles(optimization_handles.clone()).await; + + // Either continue below here with the worker, or reloop/break + // Decision logic doing one of three things: +@@ -612,9 +627,7 @@ impl UpdateHandler { + &payload_index_schema.read(), + ); + if let Err(err) = result { +- log::error!( +- "Failed to ensure there are appendable segments with capacity: {err}" +- ); ++ log::error!("Failed to ensure there are appendable segments with capacity: {err}"); + panic!("Failed to ensure there are appendable segments with capacity: {err}"); + } + } +@@ -679,8 +692,8 @@ impl UpdateHandler { + async fn update_worker_fn( + mut receiver: Receiver, + optimize_sender: Sender, +- wal: LockedWal, + segments: LockedSegmentHolder, ++ wal: LockedWal, + ) { + while let Some(signal) = receiver.recv().await { + match signal { +@@ -796,7 +809,6 @@ impl UpdateHandler { + continue; + } + }; +- + // Acknowledge confirmed version in WAL, but don't acknowledge the specified + // `keep_from` index or higher. + // This is to prevent truncating WAL entries that other bits of code still depend on +@@ -808,7 +820,6 @@ impl UpdateHandler { + if keep_from == 0 { + continue; + } +- + let ack = confirmed_version.min(keep_from.saturating_sub(1)); + + if let Err(err) = clocks.store_if_changed(&shard_path).await { +@@ -837,7 +848,7 @@ impl UpdateHandler { + } + } + +-/// Trigger optimizers when CPU budget is available ++/// Trigger optimizers when resource budget is available + fn trigger_optimizers_on_resource_budget( + optimizer_resource_budget: ResourceBudget, + desired_cpus: usize, +@@ -845,15 +856,15 @@ fn trigger_optimizers_on_resource_budget( + sender: Sender, + ) -> JoinHandle<()> { + task::spawn(async move { +- log::trace!("Skipping optimization checks, waiting for CPU budget to be available"); ++ log::trace!("Skipping optimization checks, waiting for resource budget to be available"); + optimizer_resource_budget + .notify_on_budget_available(desired_cpus, desired_io) + .await; +- log::trace!("Continue optimization checks, new CPU budget available"); ++ log::trace!("Continue optimization checks, new resource budget available"); + + // Trigger optimizers with Nop operation + sender.send(OptimizerSignal::Nop).await.unwrap_or_else(|_| { +- log::info!("Can't notify optimizers, assume process is dead. Restart is required") ++ info!("Can't notify optimizers, assume process is dead. Restart is required") + }); + }) + } +\ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/raw_response.txt b/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/raw_response.txt new file mode 100644 index 000000000..5f29ee411 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_collection_src_update_handler.rs/x-ai_grok-code-fast-1/20250909_153704/raw_response.txt @@ -0,0 +1,872 @@ +```rust +use std::cmp::min; +use std::collections::HashSet; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; + +use common::budget::ResourceBudget; +use common::counter::hardware_accumulator::HwMeasurementAcc; +use common::counter::hardware_counter::HardwareCounterCell; +use common::panic; +use itertools::Itertools; +use log::{debug, error, info, trace, warn}; +use parking_lot::Mutex; +use segment::common::operation_error::OperationResult; +use segment::index::hnsw_index::num_rayon_threads; +use segment::types::SeqNumberType; +use tokio::runtime::Handle; +use tokio::sync::mpsc::{self, Receiver, Sender}; +use tokio::sync::{Mutex as TokioMutex, oneshot}; +use tokio::task::{self, JoinHandle}; +use tokio::time::error::Elapsed; +use tokio::time::{Duration, timeout}; + +use crate::collection::payload_index_schema::PayloadIndexSchema; +use crate::collection_manager::collection_updater::CollectionUpdater; +use crate::collection_manager::holders::segment_holder::LockedSegmentHolder; +use crate::collection_manager::optimizers::segment_optimizer::{ + OptimizerThresholds, SegmentOptimizer, +}; +use crate::collection_manager::optimizers::{Tracker, TrackerLog, TrackerStatus}; +use crate::common::stoppable_task::{StoppableTaskHandle, spawn_stoppable}; +use crate::config::CollectionParams; +use crate::operations::CollectionUpdateOperations; +use crate::operations::shared_storage_config::SharedStorageConfig; +use crate::operations::types::{CollectionError, CollectionResult}; +use crate::save_on_disk::SaveOnDisk; +use crate::shards::local_shard::LocalShardClocks; +use crate::wal::WalError; +use crate::wal_delta::LockedWal; + +/// Interval at which the optimizer worker cleans up old optimization handles +/// +/// The longer the duration, the longer it takes for panicked tasks to be reported. +const OPTIMIZER_CLEANUP_INTERVAL: Duration = Duration::from_secs(5); + +pub type Optimizer = dyn SegmentOptimizer + Sync + Send; + +/// Information, required to perform operation and notify regarding the result +#[derive(Debug)] +pub struct OperationData { + /// Sequential number of the operation + pub op_num: SeqNumberType, + /// Operation + pub operation: CollectionUpdateOperations, + /// If operation was requested to wait for result + pub wait: bool, + /// Callback notification channel + pub sender: Option>>, + pub hw_measurements: HwMeasurementAcc, +} + +/// Signal, used to inform Updater process +#[derive(Debug)] +pub enum UpdateSignal { + /// Requested operation to perform + Operation(OperationData), + /// Stop all optimizers and listening + Stop, + /// Empty signal used to trigger optimizers + Nop, + /// Ensures that previous updates are applied + Plunger(oneshot::Sender<()>), +} + +/// Signal, used to inform Optimization process +#[derive(PartialEq, Eq, Clone, Copy)] +pub enum OptimizerSignal { + /// Sequential number of the operation + Operation(SeqNumberType), + /// Stop all optimizers and listening + Stop, + /// Empty signal used to trigger optimizers + Nop, +} + +/// Structure, which holds object, required for processing updates of the collection +pub struct UpdateHandler { + shared_storage_config: Arc, + payload_index_schema: Arc>, + /// List of used optimizers + pub optimizers: Arc>>, + /// Log of optimizer statuses + optimizers_log: Arc>, + total_optimized_points: Arc, + /// Global CPU budget in number of cores for all optimization tasks. + /// Assigns CPU permits to tasks to limit overall resource utilization. + optimizer_resource_budget: ResourceBudget, + /// How frequent can we flush data + /// This parameter depends on the optimizer config and should be updated accordingly. + pub flush_interval_sec: u64, + segments: LockedSegmentHolder, + /// Process, that listens updates signals and perform updates + update_worker: Option>, + /// Process, that listens for post-update signals and performs optimization + optimizer_worker: Option>, + /// Process that periodically flushes segments and tries to truncate wal + flush_worker: Option>, + /// Sender to stop flush worker + flush_stop: Option>, + runtime_handle: Handle, + /// WAL, required for operations + wal: LockedWal, + /// Always keep this WAL version and later and prevent acknowledging/truncating from the WAL. + /// This is used when other bits of code still depend on information in the WAL, such as the + /// queue proxy shard. + /// Defaults to `u64::MAX` to allow acknowledging all confirmed versions. + pub(super) wal_keep_from: Arc, + optimization_handles: Arc>>>, + /// Maximum number of concurrent optimization jobs in this update handler. + /// This parameter depends on the optimizer config and should be updated accordingly. + pub max_optimization_threads: Option, + /// Highest and cutoff clocks for the shard WAL. + clocks: LocalShardClocks, + shard_path: PathBuf, + /// Whether we have ever triggered optimizers since starting. + has_triggered_optimizers: Arc, +} + +impl UpdateHandler { + #[allow(clippy::too_many_arguments)] + pub fn new( + shared_storage_config: Arc, + payload_index_schema: Arc>, + optimizers: Arc>>, + optimizers_log: Arc>, + total_optimized_points: Arc, + optimizer_resource_budget: ResourceBudget, + runtime_handle: Handle, + segments: LockedSegmentHolder, + wal: LockedWal, + flush_interval_sec: u64, + max_optimization_threads: Option, + clocks: LocalShardClocks, + shard_path: PathBuf, + ) -> UpdateHandler { + UpdateHandler { + shared_storage_config, + payload_index_schema, + optimizers, + update_worker: None, + optimizer_worker: None, + optimizers_log, + total_optimized_points, + optimizer_resource_budget, + flush_interval_sec, + segments, + flush_worker: None, + flush_stop: None, + runtime_handle, + wal, + wal_keep_from: Arc::new(u64::MAX.into()), + optimization_handles: Arc::new(TokioMutex::new(vec![])), + max_optimization_threads, + clocks, + shard_path, + has_triggered_optimizers: Arc::new(AtomicBool::new(false)), + } + } + + pub fn stop_flush_worker(&mut self) { + if let Some(flush_stop) = self.flush_stop.take() { + if let Err(()) = flush_stop.send(()) { + warn!("Failed to stop flush worker as it is already stopped."); + } + } + } + + /// Gracefully wait before all optimizations stop + /// If some optimization is in progress - it will be finished before shutdown. + /// Blocking function. + pub async fn wait_workers_stops(&mut self) -> CollectionResult<()> { + let maybe_handle = self.update_worker.take(); + if let Some(handle) = maybe_handle { + handle.await?; + } + let maybe_handle = self.optimizer_worker.take(); + if let Some(handle) = maybe_handle { + handle.await?; + } + let maybe_handle = self.flush_worker.take(); + if let Some(handle) = maybe_handle { + handle.await?; + } + + let mut opt_handles_guard = self.optimization_handles.lock().await; + let opt_handles = std::mem::take(&mut *opt_handles_guard); + let stopping_handles = opt_handles + .into_iter() + .filter_map(|h| h.stop()) + .collect_vec(); + + for res in stopping_handles { + res.await?; + } + + Ok(()) + } + + pub fn run_workers(&mut self, update_receiver: Receiver) { + let (tx, rx) = mpsc::channel(self.shared_storage_config.update_queue_size); + self.optimizer_worker = Some(self.runtime_handle.spawn(Self::optimization_worker_fn( + self.optimizers.clone(), + tx.clone(), + rx, + self.segments.clone(), + self.wal.clone(), + self.optimization_handles.clone(), + self.optimizers_log.clone(), + self.total_optimized_points.clone(), + self.optimizer_resource_budget.clone(), + self.max_optimization_threads, + self.has_triggered_optimizers.clone(), + self.payload_index_schema.clone(), + ))); + self.update_worker = Some(self.runtime_handle.spawn(Self::update_worker_fn( + update_receiver, + tx, + self.segments.clone(), + self.wal.clone(), + ))); + let (flush_tx, flush_rx) = oneshot::channel(); + self.flush_worker = Some(self.runtime_handle.spawn(Self::flush_worker( + self.segments.clone(), + self.wal.clone(), + self.wal_keep_from.clone(), + self.flush_interval_sec, + flush_rx, + self.clocks.clone(), + self.shard_path.clone(), + ))); + self.flush_stop = Some(flush_tx); + } + + /// Checks if there are any failed operations. + /// If so - attempts to re-apply all failed operations. + async fn try_recover(segments: LockedSegmentHolder, wal: LockedWal) -> CollectionResult { + // Try to re-apply everything starting from the first failed operation + let first_failed_operation_option = segments.read().failed_operation.iter().cloned().min(); + match first_failed_operation_option { + None => Ok(0), + Some(first_failed_op) => { + let wal_lock = wal.lock().await; + for (op_num, operation) in wal_lock.read(first_failed_op) { + CollectionUpdater::update( + &segments, + op_num, + operation.operation, + &HardwareCounterCell::disposable(), // Internal operation, no measurement needed + )?; + } + Err(CollectionError::service_error("Failed operation recovery should not return".to_string())) + } + } + } + + fn process_optimization( + optimizers: Arc>>, + segments: LockedSegmentHolder, + ) -> Vec> { + for optimizer in optimizers.iter() { + let nonoptimal_segment_ids = optimizer.check_condition(segments.clone()); + while !nonoptimal_segment_ids.is_empty() { + debug!( + "Start optimization on segments: {:?}", + nonoptimal_segment_ids + ); + // If optimization fails, it could not be reported to anywhere except for console. + // So the only recovery here is to stop optimization and await for restart + if optimizer.optimize(segments.clone(), nonoptimal_segment_ids).is_err() { + continue; + }; + nonoptimal_segment_ids = optimizer.check_condition(segments.clone()); + } + } + vec![] + } + + /// Checks conditions for all optimizers until there is no suggested segment + /// Starts a task for each optimization + /// Returns handles for started tasks + pub(crate) fn launch_optimization( + optimizers: Arc>>, + optimizers_log: Arc>, + total_optimized_points: Arc, + optimizer_resource_budget: &ResourceBudget, + segments: LockedSegmentHolder, + callback: F, + limit: Option, + ) -> Vec> + where + F: Fn(bool) + Send + Clone + Sync + 'static, + { + let mut scheduled_segment_ids = HashSet::<_>::default(); + let mut handles = vec![]; + + 'outer: for optimizer in optimizers.iter() { + loop { + // Return early if we reached the optimization job limit + if limit.map(|extra| handles.len() >= extra).unwrap_or(false) { + info!("Reached optimization job limit, postponing other optimizations"); + break 'outer; + } + + let nonoptimal_segment_ids = + optimizer.check_condition(segments.clone(), &scheduled_segment_ids); + if nonoptimal_segment_ids.is_empty() { + break; + } + + debug!("Optimizing segments: {:?}", &nonoptimal_segment_ids); + + // Determine how many Resources we prefer for optimization task, acquire permit for it + // And use same amount of IO threads as CPUs + let max_indexing_threads = optimizer.hnsw_config().max_indexing_threads; + let desired_io = num_rayon_threads(max_indexing_threads); + let Some(mut permit) = optimizer_resource_budget.try_acquire(0, desired_io) else { + // If there is no Resource budget, break outer loop and return early + // If we have no handles (no optimizations) trigger callback so that we wake up + // our optimization worker to try again later, otherwise it could get stuck + log::trace!( + "No available IO permit for {} optimizer, postponing", + optimizer.name(), + ); + if handles.is_empty() { + callback(false); + } + break 'outer; + }; + log::trace!( + "Acquired {} IO permit for {} optimizer", + permit.num_io, + optimizer.name(), + ); + + let permit_callback = callback.clone(); + + permit.set_on_release(move || { + // Notify scheduler that resource budget changed + permit_callback(false); + }); + + let optimizer = optimizer.clone(); + let optimizers_log = optimizers_log.clone(); + let total_optimized_points = total_optimized_points.clone(); + let segments = segments.clone(); + let nsi = nonoptimal_segment_ids.clone(); + scheduled_segment_ids.extend(&nsi); + let callback = callback.clone(); + + let handle = spawn_stoppable( + // Stoppable task + { + let resource_budget = optimizer_resource_budget.clone(); + let segments = segments.clone(); + move |stopped| { + // Track optimizer status + let tracker = Tracker::start(optimizer.as_ref().name(), nsi.clone()); + let tracker_handle = tracker.handle(); + optimizers_log.lock().register(tracker); + + // Optimize and handle result + match optimizer.as_ref().optimize( + segments.clone(), + nsi, + permit, + resource_budget, + stopped, + ) { + // Perform some actions when optimization if finished + Ok(optimized_points) => { + let is_optimized = optimized_points > 0; + total_optimized_points + .fetch_add(optimized_points, Ordering::Relaxed); + tracker_handle.update(TrackerStatus::Done); + callback(is_optimized); + is_optimized + } + // Handle and report errors + Err(error) => match error { + CollectionError::Cancelled { description } => { + debug!("Optimization cancelled - {description}"); + tracker_handle.update(TrackerStatus::Cancelled(description)); + false + } + _ => { + segments.write().report_optimizer_error(error.clone()); + + // Error of the optimization can not be handled by API user + // It is only possible to fix after full restart, + // so the best available action here is to stop whole + // optimization thread and log the error + log::error!("Optimization error: {error}"); + + tracker_handle.update(TrackerStatus::Error(error.to_string())); + + panic!("Optimization error: {error}"); + } + }, + } + } + }, + // Panic handler + Some(Box::new(move |panic_payload| { + let message = panic::downcast_str(&panic_payload).unwrap_or(""); + let separator = if !message.is_empty() { ": " } else { "" }; + + warn!( + "Optimization task panicked, collection may be in unstable state\ + {separator}{message}" + ); + + segments.write().report_optimizer_error(CollectionError::service_error(format!( + "Optimization task panicked{separator}{message}" + ))); + })), + ); + handles.push(handle); + } + } + + handles + } + + /// Ensure there is at least one appendable segment with enough capacity + /// + /// If there is no appendable segment, or all are at or over capacity, a new empty one is + /// created. + /// + /// Capacity is determined based on `optimizers.max_segment_size_kb`. + pub(super) fn ensure_appendable_segment_with_capacity( + segments: &LockedSegmentHolder, + segments_path: &Path, + collection_params: &CollectionParams, + thresholds_config: &OptimizerThresholds, + payload_index_schema: &PayloadIndexSchema, + ) -> OperationResult<()> { + let no_segment_with_capacity = { + let segments_read = segments.read(); + segments_read + .appendable_segments_ids() + .into_iter() + .filter_map(|segment_id| segments_read.get(segment_id)) + .all(|segment| { + let max_vector_size_bytes = + segment.get().read().max_available_vectors_size_in_bytes().unwrap_or_default(); + let max_segment_size_bytes = thresholds_config + .max_segment_size_kb + .saturating_mul(segment::common::BYTES_IN_KB); + + max_vector_size_bytes >= max_segment_size_bytes + }) + }; + + if no_segment_with_capacity { + debug!("Creating new appendable segment, all existing segments are over capacity"); + segments.write().create_appendable_segment( + segments_path, + collection_params, + payload_index_schema, + )?; + } + + Ok(()) + } + + /// Checks the optimizer conditions. + /// + /// This function returns a tuple of two booleans: + /// - The first indicates if any optimizers have been triggered since startup. + /// - The second indicates if there are any pending/suboptimal optimizers. + pub(crate) fn check_optimizer_conditions(&self) -> (bool, bool) { + // Check if Qdrant triggered any optimizations since starting at all + let has_triggered_any_optimizers = self.has_triggered_optimizers.load(Ordering::Relaxed); + + let excluded_ids = HashSet::<_>::default(); + let has_suboptimal_optimizers = self.optimizers.iter().any(|optimizer| { + let nonoptimal_segment_ids = + optimizer.check_condition(self.segments.clone(), &excluded_ids); + !nonoptimal_segment_ids.is_empty() + }); + + (has_triggered_any_optimizers, has_suboptimal_optimizers) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn process_optimization( + optimizers: Arc>>, + segments: LockedSegmentHolder, + optimization_handles: Arc>>>, + optimizers_log: Arc>, + total_optimized_points: Arc, + optimizer_resource_budget: &ResourceBudget, + sender: Sender, + limit: usize, + ) { + let mut new_handles = Self::launch_optimization( + optimizers.clone(), + optimizers_log, + total_optimized_points, + optimizer_resource_budget, + segments.clone(), + move |_optimization_result| { + // After optimization is finished, we still need to check if there are + // some further optimizations possible. + // If receiver is already dead - we do not care. + // If channel is full - optimization will be triggered by some other signal + let _ = sender.try_send(OptimizerSignal::Nop); + }, + Some(limit), + ); + let mut handles = optimization_handles.lock().await; + handles.append(&mut new_handles); + handles.retain(|h| !h.is_finished()) + } + + /// Cleanup finalized optimization task handles + /// + /// This finds and removes completed tasks from our list of optimization handles. + /// It also propagates any panics (and unknown errors) so we properly handle them if desired. + /// + /// It is essential to call this every once in a while for handling panics in time. + /// + /// Returns true if any optimization handle was finished, joined and removed. + async fn cleanup_optimization_handles( + optimization_handles: Arc>>>, + ) -> bool { + // Remove finished handles + let finished_handles: Vec<_> = { + let mut handles = optimization_handles.lock().await; + (0..handles.len()) + .filter(|i| handles[*i].is_finished()) + .collect::>() + .into_iter() + .rev() + .map(|i| handles.swap_remove(i)) + .collect() + }; + + let finished_any = !finished_handles.is_empty(); + + // Finalize all finished handles to propagate panics + for handle in finished_handles { + handle.join_and_handle_panic().await; + } + + finished_any + } + + #[allow(clippy::too_many_arguments)] + async fn optimization_worker_fn( + optimizers: Arc>>, + sender: Sender, + mut receiver: Receiver, + segments: LockedSegmentHolder, + wal: LockedWal, + optimization_handles: Arc>>>, + optimizers_log: Arc>, + total_optimized_points: Arc, + optimizer_resource_budget: ResourceBudget, + max_handles: Option, + has_triggered_optimizers: Arc, + payload_index_schema: Arc>, + ) { + let max_handles = max_handles.unwrap_or(usize::MAX); + let max_indexing_threads = optimizers + .first() + .map(|optimizer| optimizer.hnsw_config().max_indexing_threads) + .unwrap_or_default(); + + // Asynchronous task to trigger optimizers once resource budget is available again + let mut resource_available_trigger: Option> = None; + + loop { + let result = timeout(OPTIMIZER_CLEANUP_INTERVAL, receiver.recv()).await; + + let cleaned_any = Self::cleanup_optimization_handles(optimization_handles.clone()).await; + + // Either continue below here with the worker, or reloop/break + // Decision logic doing one of three things: + // 1. run optimizers + // 2. reloop and wait for next signal + // 3. break here and stop the optimization worker + let ignore_max_handles = match result { + // Regular optimizer signal: run optimizers: do 1 + Ok(Some(OptimizerSignal::Operation(_))) => false, + // Optimizer signal ignoring max handles: do 1 + Ok(Some(OptimizerSignal::Nop)) => true, + // Hit optimizer cleanup interval, did clean up a task: do 1 + Err(Elapsed { .. }) if cleaned_any => { + // This branch prevents a race condition where optimizers would get stuck + // If the optimizer cleanup interval was triggered and we did clean any task we + // must run optimizers now. If we don't there may not be any other ongoing + // tasks that'll trigger this for us. If we don't run optimizers here we might + // get stuck into yellow state until a new update operation is received. + // See: + log::warn!( + "Cleaned a optimization handle after timeout, explicitly triggering optimizers", + ); + true + } + // Hit optimizer cleanup interval, did not clean up a task: do 2 + Err(Elapsed { .. }) => continue, + // Channel closed or received stop signal: do 3 + Ok(None | Some(OptimizerSignal::Stop)) => break, + }; + + has_triggered_optimizers.store(true, Ordering::Relaxed); + + // Ensure we have at least one appendable segment with enough capacity + // Source required parameters from first optimizer + if let Some(optimizer) = optimizers.first() { + let result = Self::ensure_appendable_segment_with_capacity( + &segments, + optimizer.segments_path(), + &optimizer.collection_params(), + optimizer.threshold_config(), + &payload_index_schema.read(), + ); + if let Err(err) = result { + log::error!("Failed to ensure there are appendable segments with capacity: {err}"); + panic!("Failed to ensure there are appendable segments with capacity: {err}"); + } + } + + // If not forcing, wait on next signal if we have too many handles + if !ignore_max_handles && optimization_handles.lock().await.len() >= max_handles { + continue; + } + + if Self::try_recover(segments.clone(), wal.clone()) + .await + .is_err() + { + continue; + } + + // Continue if we have enough resource budget available to start an optimization + // Otherwise skip now and start a task to trigger the optimizer again once resource + // budget becomes available + let desired_cpus = 0; + let desired_io = num_rayon_threads(max_indexing_threads); + if !optimizer_resource_budget.has_budget(desired_cpus, desired_io) { + let trigger_active = resource_available_trigger + .as_ref() + .is_some_and(|t| !t.is_finished()); + if !trigger_active { + resource_available_trigger.replace(trigger_optimizers_on_resource_budget( + optimizer_resource_budget.clone(), + desired_cpus, + desired_io, + sender.clone(), + )); + } + continue; + } + + // Determine optimization handle limit based on max handles we allow + // Not related to the CPU budget, but a different limit for the maximum number + // of concurrent concrete optimizations per shard as configured by the user in + // the Qdrant configuration. + // Skip if we reached limit, an ongoing optimization that finishes will trigger this loop again + let limit = max_handles.saturating_sub(optimization_handles.lock().await.len()); + if limit == 0 { + log::trace!("Skipping optimization check, we reached optimization thread limit"); + continue; + } + + Self::process_optimization( + optimizers.clone(), + segments.clone(), + optimization_handles.clone(), + optimizers_log.clone(), + total_optimized_points.clone(), + &optimizer_resource_budget, + sender.clone(), + limit, + ) + .await; + } + } + + async fn update_worker_fn( + mut receiver: Receiver, + optimize_sender: Sender, + segments: LockedSegmentHolder, + wal: LockedWal, + ) { + while let Some(signal) = receiver.recv().await { + match signal { + UpdateSignal::Operation(OperationData { + op_num, + operation, + sender, + wait, + hw_measurements, + }) => { + let flush_res = if wait { + wal.lock().await.flush().map_err(|err| { + CollectionError::service_error(format!( + "Can't flush WAL before operation {op_num} - {err}" + )) + }) + } else { + Ok(()) + }; + + let operation_result = flush_res.and_then(|_| { + CollectionUpdater::update( + &segments, + op_num, + operation, + &hw_measurements.get_counter_cell(), + ) + }); + + let res = match operation_result { + Ok(update_res) => optimize_sender + .send(OptimizerSignal::Operation(op_num)) + .await + .and(Ok(update_res)) + .map_err(|send_err| send_err.into()), + Err(err) => Err(err), + }; + + if let Some(feedback) = sender { + feedback.send(res).unwrap_or_else(|_| { + debug!("Can't report operation {op_num} result. Assume already not required"); + }); + }; + } + UpdateSignal::Stop => { + optimize_sender + .send(OptimizerSignal::Stop) + .await + .unwrap_or_else(|_| debug!("Optimizer already stopped")); + break; + } + UpdateSignal::Nop => optimize_sender + .send(OptimizerSignal::Nop) + .await + .unwrap_or_else(|_| { + info!( + "Can't notify optimizers, assume process is dead. Restart is required" + ); + }), + UpdateSignal::Plunger(callback_sender) => { + callback_sender.send(()).unwrap_or_else(|_| { + debug!("Can't notify sender, assume nobody is waiting anymore"); + }); + } + } + } + // Transmitter was destroyed + optimize_sender + .send(OptimizerSignal::Stop) + .await + .unwrap_or_else(|_| debug!("Optimizer already stopped")); + } + + async fn flush_worker( + segments: LockedSegmentHolder, + wal: LockedWal, + wal_keep_from: Arc, + flush_interval_sec: u64, + mut stop_receiver: oneshot::Receiver<()>, + clocks: LocalShardClocks, + shard_path: PathBuf, + ) { + loop { + // Stop flush worker on signal or if sender was dropped + // Even if timer did not finish + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(flush_interval_sec)) => {}, + _ = &mut stop_receiver => { + debug!("Stopping flush worker for shard {}", shard_path.display()); + return; + } + } + + trace!("Attempting flushing"); + let wal_flash_job = wal.lock().await.flush_async(); + + if let Err(err) = wal_flash_job.join() { + error!("Failed to flush wal: {err:?}"); + segments + .write() + .report_optimizer_error(WalError::WriteWalError(format!( + "WAL flush error: {err:?}" + ))); + continue; + } + + let confirmed_version = Self::flush_segments(segments.clone()); + let confirmed_version = match confirmed_version { + Ok(version) => version, + Err(err) => { + error!("Failed to flush: {err}"); + segments.write().report_optimizer_error(err); + continue; + } + }; + // Acknowledge confirmed version in WAL, but don't acknowledge the specified + // `keep_from` index or higher. + // This is to prevent truncating WAL entries that other bits of code still depend on + // such as the queue proxy shard. + // Default keep_from is `u64::MAX` to allow acknowledging all confirmed. + let keep_from = wal_keep_from.load(std::sync::atomic::Ordering::Relaxed); + + // If we should keep the first message, do not acknowledge at all + if keep_from == 0 { + continue; + } + let ack = confirmed_version.min(keep_from.saturating_sub(1)); + + if let Err(err) = clocks.store_if_changed(&shard_path).await { + log::warn!("Failed to store clock maps to disk: {err}"); + segments.write().report_optimizer_error(err); + } + + if let Err(err) = wal.lock().await.ack(ack) { + log::warn!("Failed to acknowledge WAL version: {err}"); + segments.write().report_optimizer_error(err); + } + } + } + + /// Returns confirmed version after flush of all segments + /// + /// # Errors + /// Returns an error on flush failure + fn flush_segments(segments: LockedSegmentHolder) -> OperationResult { + let read_segments = segments.read(); + let flushed_version = read_segments.flush_all(false, false)?; + Ok(match read_segments.failed_operation.iter().cloned().min() { + None => flushed_version, + Some(failed_operation) => min(failed_operation, flushed_version), + }) + } +} + +/// Trigger optimizers when resource budget is available +fn trigger_optimizers_on_resource_budget( + optimizer_resource_budget: ResourceBudget, + desired_cpus: usize, + desired_io: usize, + sender: Sender, +) -> JoinHandle<()> { + task::spawn(async move { + log::trace!("Skipping optimization checks, waiting for resource budget to be available"); + optimizer_resource_budget + .notify_on_budget_available(desired_cpus, desired_io) + .await; + log::trace!("Continue optimization checks, new resource budget available"); + + // Trigger optimizers with Nop operation + sender.send(OptimizerSignal::Nop).await.unwrap_or_else(|_| { + info!("Can't notify optimizers, assume process is dead. Restart is required") + }); + }) +} +``` \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt b/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt new file mode 100644 index 000000000..a083a5599 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt @@ -0,0 +1,555 @@ +use std::borrow::Cow; +use std::cmp::max; +use std::path::{Path, PathBuf}; +use std::sync::atomic::AtomicBool; + +use common::fixed_length_priority_queue::FixedLengthPriorityQueue; +use common::types::{PointOffsetType, ScoredPointOffset}; +use io::file_operations::read_bin; +use itertools::Itertools; +use serde::{Deserialize, Serialize}; + +use super::entry_points::EntryPoint; +use super::graph_links::{GraphLinks, GraphLinksFormat}; +use crate::common::operation_error::{ + CancellableResult, OperationError, OperationResult, check_process_stopped, +}; +use crate::common::utils::rev_range; +use crate::index::hnsw_index::entry_points::EntryPoints; +use crate::index::hnsw_index::graph_links::GraphLinksSerializer; +use crate::index::hnsw_index::point_scorer::FilteredScorer; +use crate::index::hnsw_index::search_context::SearchContext; +use crate::index::visited_pool::{VisitedListHandle, VisitedPool}; + +pub type LinkContainer = Vec; +pub type LinkContainerRef<'a> = &'a [PointOffsetType]; +pub type LayersContainer = Vec; + +pub const HNSW_GRAPH_FILE: &str = "graph.bin"; +pub const HNSW_LINKS_FILE: &str = "links.bin"; +pub const COMPRESSED_HNSW_LINKS_FILE: &str = "links_compressed.bin"; + +/// Contents of the `graph.bin` file. +#[derive(Deserialize, Serialize, Debug)] +pub(super) struct GraphLayerData<'a> { + pub(super) m: usize, + pub(super) m0: usize, + pub(super) ef_construct: usize, + pub(super) entry_points: Cow<'a, EntryPoints>, +} + +#[derive(Debug)] +pub struct GraphLayers { + pub(super) m: usize, + pub(super) m0: usize, + pub(super) links: GraphLinks, + pub(super) entry_points: EntryPoints, + pub(super) visited_pool: VisitedPool, +} + +pub trait GraphLayersBase { + fn get_visited_list_from_pool(&self) -> VisitedListHandle; + + fn links_map(&self, point_id: PointOffsetType, level: usize, f: F) + where + F: FnMut(PointOffsetType); + + /// Get M based on current level + fn get_m(&self, level: usize) -> usize; + + /// Greedy search for closest points within a single graph layer + fn _search_on_level( + &self, + searcher: &mut SearchContext, + level: usize, + visited_list: &mut VisitedListHandle, + points_scorer: &mut FilteredScorer, + is_stopped: &AtomicBool, + ) -> CancellableResult<()> { + let limit = self.get_m(level); + let mut points_ids: Vec = Vec::with_capacity(2 * limit); + + while let Some(candidate) = searcher.candidates.pop() { + check_process_stopped(is_stopped)?; + + if candidate.score < searcher.lower_bound() { + break; + } + + points_ids.clear(); + self.links_map(candidate.idx, level, |link| { + if !visited_list.check(link) { + points_ids.push(link); + } + }); + + let scores = points_scorer.score_points(&mut points_ids, limit); + scores.iter().copied().for_each(|score_point| { + searcher.process_candidate(score_point); + visited_list.check_and_update_visited(score_point.idx); + }); + } + + Ok(()) + } + + fn search_on_level( + &self, + level_entry: ScoredPointOffset, + level: usize, + ef: usize, + points_scorer: &mut FilteredScorer, + is_stopped: &AtomicBool, + ) -> CancellableResult> { + let mut visited_list = self.get_visited_list_from_pool(); + visited_list.check_and_update_visited(level_entry.idx); + let mut search_context = SearchContext::new(level_entry, ef); + + self._search_on_level( + &mut search_context, + level, + &mut visited_list, + points_scorer, + is_stopped, + )?; + Ok(search_context.nearest) + } + + /// Greedy searches for entry point of level `target_level`. + /// Beam size is 1. + fn search_entry( + &self, + entry_point: PointOffsetType, + top_level: usize, + target_level: usize, + points_scorer: &mut FilteredScorer, + is_stopped: &AtomicBool, + ) -> CancellableResult { + let mut links: Vec = Vec::with_capacity(2 * self.get_m(0)); + + let mut current_point = ScoredPointOffset { + idx: entry_point, + score: points_scorer.score_point(entry_point), + }; + for level in rev_range(top_level, target_level) { + check_process_stopped(is_stopped)?; + + let limit = self.get_m(level); + + let mut changed = true; + while changed { + changed = false; + + links.clear(); + self.links_map(current_point.idx, level, |link| { + links.push(link); + }); + + let scores = points_scorer.score_points(&mut links, limit); + scores.iter().copied().for_each(|score_point| { + if score_point.score > current_point.score { + changed = true; + current_point = score_point; + } + }); + } + } + Ok(current_point) + } + + #[cfg(test)] + #[cfg(feature = "gpu")] + fn search_entry_on_level( + &self, + entry_point: PointOffsetType, + level: usize, + points_scorer: &mut FilteredScorer, + ) -> ScoredPointOffset { + let limit = self.get_m(level); + let mut links: Vec = Vec::with_capacity(2 * self.get_m(0)); + let mut current_point = ScoredPointOffset { + idx: entry_point, + score: points_scorer.score_point(entry_point), + }; + + let mut changed = true; + while changed { + changed = false; + + links.clear(); + self.links_map(current_point.idx, level, |link| { + links.push(link); + }); + + let scores = points_scorer.score_points(&mut links, limit); + scores.iter().copied().for_each(|score_point| { + if score_point.score > current_point.score { + changed = true; + current_point = score_point; + } + }); + } + current_point + } +} + +impl GraphLayersBase for GraphLayers { + fn get_visited_list_from_pool(&self) -> VisitedListHandle { + self.visited_pool.get(self.links.num_points()) + } + + fn links_map(&self, point_id: PointOffsetType, level: usize, f: F) + where + F: FnMut(PointOffsetType), + { + self.links.links(point_id, level).for_each(f); + } + + fn get_m(&self, level: usize) -> usize { + if level == 0 { self.m0 } else { self.m } + } +} + +/// Object contains links between nodes for HNSW search +/// +/// Assume all scores are similarities. Larger score = closer points +impl GraphLayers { + /// Returns the highest level this point is included in + pub fn point_level(&self, point_id: PointOffsetType) -> usize { + self.links.point_level(point_id) + } + + fn get_entry_point( + &self, + points_scorer: &FilteredScorer, + custom_entry_points: Option<&[PointOffsetType]>, + ) -> Option { + // Try to get it from custom entry points + custom_entry_points + .and_then(|custom_entry_points| { + custom_entry_points + .iter() + .filter(|&&point_id| points_scorer.check_vector(point_id)) + .map(|&point_id| { + let level = self.point_level(point_id); + EntryPoint { point_id, level } + }) + .max_by_key(|ep| ep.level) + }) + .or_else(|| { + // Otherwise use normal entry points + self.entry_points + .get_entry_point(|point_id| points_scorer.check_vector(point_id)) + }) + } + + pub fn search( + &self, + top: usize, + ef: usize, + mut points_scorer: FilteredScorer, + custom_entry_points: Option<&[PointOffsetType]>, + is_stopped: &AtomicBool, + ) -> CancellableResult> { + let Some(entry_point) = self.get_entry_point(&points_scorer, custom_entry_points) else { + return Ok(Vec::default()); + }; + + let zero_level_entry = self.search_entry( + entry_point.point_id, + entry_point.level, + 0, + &mut points_scorer, + is_stopped, + )?; + let nearest = self.search_on_level( + zero_level_entry, + 0, + max(top, ef), + &mut points_scorer, + is_stopped, + )?; + Ok(nearest.into_iter_sorted().take(top).collect_vec()) + } + + pub fn get_path(path: &Path) -> PathBuf { + path.join(HNSW_GRAPH_FILE) + } + + pub fn get_links_path(path: &Path, format: GraphLinksFormat) -> PathBuf { + match format { + GraphLinksFormat::Plain => path.join(HNSW_LINKS_FILE), + GraphLinksFormat::Compressed => path.join(COMPRESSED_HNSW_LINKS_FILE), + } + } + + pub fn files(&self, path: &Path) -> Vec { + vec![ + GraphLayers::get_path(path), + GraphLayers::get_links_path(path, self.links.format()), + ] + } + + pub fn num_points(&self) -> usize { + self.links.num_points() + } +} + +impl GraphLayers { + pub fn load(dir: &Path, on_disk: bool, compress: bool) -> OperationResult { + let graph_data: GraphLayerData = read_bin(&GraphLayers::get_path(dir))?; + + if compress { + Self::convert_to_compressed(dir, graph_data.m, graph_data.m0)?; + } + + Ok(Self { + m: graph_data.m, + m0: graph_data.m0, + links: Self::load_links(dir, on_disk)?, + entry_points: graph_data.entry_points.into_owned(), + visited_pool: VisitedPool::new(), + }) + } + + fn load_links(dir: &Path, on_disk: bool) -> OperationResult { + for format in [GraphLinksFormat::Compressed, GraphLinksFormat::Plain] { + let path = GraphLayers::get_links_path(dir, format); + if path.exists() { + return GraphLinks::load_from_file(&path, on_disk, format); + } + } + Err(OperationError::generic("No links file found")) + } + + fn convert_to_compressed(dir: &Path, m: usize, m0: usize) -> OperationResult<()> { + let plain_path = Self::get_links_path(dir, GraphLinksFormat::Plain); + let compressed_path = Self::get_links_path(dir, GraphLinksFormat::Compressed); + + if compressed_path.exists() { + return Ok(()); + } + + let start = std::time::Instant::now(); + + let links = GraphLinks::load_from_file(&plain_path, true, GraphLinksFormat::Plain)?; + let original_size = plain_path.metadata()?.len(); + GraphLinksSerializer::new(links.into_edges(), GraphLinksFormat::Compressed, m, m0) + .save_as(&compressed_path)?; + let new_size = compressed_path.metadata()?.len(); + + // Remove the original file + std::fs::remove_file(plain_path)?; + + log::debug!( + "Compressed HNSW graph links in {:.1?}: {:.1}MB -> {:.1}MB ({:.1}%)", + start.elapsed(), + original_size as f64 / 1024.0 / 1024.0, + new_size as f64 / 1024.0 / 1024.0, + new_size as f64 / original_size as f64 * 100.0, + ); + + Ok(()) + } + + #[cfg(feature = "testing")] + pub fn compress_ram(&mut self) { + use crate::index::hnsw_index::graph_links::GraphLinksSerializer; + assert_eq!(self.links.format(), GraphLinksFormat::Plain); + let dummy = GraphLinksSerializer::new(Vec::new(), GraphLinksFormat::Plain, 0, 0) + .to_graph_links_ram(); + let links = std::mem::replace(&mut self.links, dummy); + self.links = GraphLinksSerializer::new( + links.into_edges(), + GraphLinksFormat::Compressed, + self.m, + self.m0, + ) + .to_graph_links_ram(); + } + + pub fn populate(&self) -> OperationResult<()> { + self.links.populate()?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng; + use rand::rngs::StdRng; + use rstest::rstest; + use tempfile::Builder; + + use super::*; + use crate::data_types::vectors::VectorElementType; + use crate::fixtures::index_fixtures::{ + FakeFilterContext, TestRawScorerProducer, random_vector, + }; + use crate::index::hnsw_index::graph_links::GraphLinksSerializer; + use crate::index::hnsw_index::tests::{ + create_graph_layer_builder_fixture, create_graph_layer_fixture, + }; + use crate::spaces::metric::Metric; + use crate::spaces::simple::{CosineMetric, DotProductMetric}; + use crate::vector_storage::DEFAULT_STOPPED; + use crate::vector_storage::chunked_vector_storage::VectorOffsetType; + + fn search_in_graph( + query: &[VectorElementType], + top: usize, + vector_storage: &TestRawScorerProducer, + graph: &GraphLayers, + ) -> Vec { + let fake_filter_context = FakeFilterContext {}; + let raw_scorer = vector_storage.get_raw_scorer(query.to_owned()).unwrap(); + + let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + let ef = 16; + graph + .search(top, ef, scorer, None, &DEFAULT_STOPPED) + .unwrap() + } + + const M: usize = 8; + + #[rstest] + #[case::uncompressed(GraphLinksFormat::Plain)] + #[case::compressed(GraphLinksFormat::Compressed)] + fn test_search_on_level(#[case] format: GraphLinksFormat) { + let dim = 8; + let m = 8; + let entry_points_num = 10; + let num_vectors = 10; + + let vector_holder = TestRawScorerProducer::::new(dim, num_vectors, &mut StdRng::seed_from_u64(42)); + + let mut graph_links = vec![vec![Vec::new()]; num_vectors]; + graph_links[0][0] = vec![1, 2, 3, 4, 5, 6]; + + let graph_layers = GraphLayers { + m, + m0: 2 * m, + links: GraphLinksSerializer::new(graph_links.clone(), format, m, 2 * m) + .to_graph_links_ram(), + entry_points: EntryPoints::new(entry_points_num), + visited_pool: VisitedPool::new(), + }; + + let linking_idx: PointOffsetType = 7; + + let fake_filter_context = FakeFilterContext {}; + let added_vector = vector_holder.vectors.get(linking_idx as VectorOffsetType).to_vec(); + let raw_scorer = vector_holder.get_raw_scorer(added_vector).unwrap(); + let mut scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + + let nearest_on_level = graph_layers + .search_on_level( + ScoredPointOffset { + idx: 0, + score: scorer.score_point(0), + }, + 0, + 32, + &mut scorer, + &DEFAULT_STOPPED, + ) + .unwrap(); + + assert_eq!(nearest_on_level.len(), graph_links[0][0].len() + 1); + + for nearest in nearest_on_level.iter_unsorted() { + // eprintln!("nearest = {:#?}", nearest); + assert_eq!( + nearest.score, + scorer.score_internal(linking_idx, nearest.idx) + ) + } + } + + #[rstest] + #[case::uncompressed((GraphLinksFormat::Plain, false))] + #[case::converted((GraphLinksFormat::Plain, true))] + #[case::compressed((GraphLinksFormat::Compressed, false))] + #[case::recompressed((GraphLinksFormat::Compressed, true))] + fn test_save_and_load(#[case] (initial_format, compress): (GraphLinksFormat, bool)) { + let num_vectors = 100; + let dim = 8; + let top = 5; + + let mut rng = StdRng::seed_from_u64(42); + + let dir = Builder::new().prefix("graph_dir").tempdir().unwrap(); + let (vector_holder, graph_layers_builder) = + create_graph_layer_builder_fixture(num_vectors, M, dim, false, &mut rng); + let graph1 = graph_layers_builder + .into_graph_layers(dir.path(), initial_format, true) + .unwrap(); + assert_eq!(graph1.links.format(), initial_format); + let res1 = search_in_graph(&random_vector(&mut rng, dim), top, &vector_holder, &graph1); + drop(graph1); + + let graph2 = GraphLayers::load(dir.path(), false, compress).unwrap(); + if compress { + assert_eq!(graph2.links.format(), GraphLinksFormat::Compressed); + } else { + assert_eq!(graph2.links.format(), initial_format); + } + let res2 = search_in_graph(&vec![0.0; dim], top, &vector_holder, &graph2); + + assert_eq!(res1, res2) + } + + #[rstest] + #[case::uncompressed(GraphLinksFormat::Plain)] + #[case::compressed(GraphLinksFormat::Compressed)] + fn test_add_points(#[case] format: GraphLinksFormat) { + let num_vectors = 1000; + let dim = 8; + + let mut rng = StdRng::seed_from_u64(42); + + type M = CosineMetric; + + let (vector_holder, graph_layers) = + create_graph_layer_fixture::(num_vectors, M, dim, format, false, &mut rng); + + let main_entry = graph_layers + .entry_points + .get_entry_point(|_x| true) + .expect("Expect entry point to exists"); + + assert!(main_entry.level > 0); + + let num_levels = (0..num_vectors) + .map(|i| graph_layers.links.point_level(i as PointOffsetType)) + .max() + .unwrap(); + assert_eq!(main_entry.level, num_levels); + + let total_links_0 = (0..num_vectors) + .map(|i| graph_layers.links.links(i as PointOffsetType, 0).len()) + .sum::(); + + eprintln!("total_links_0 = {total_links_0:#?}"); + eprintln!("num_vectors = {num_vectors:#?}"); + assert!(total_links_0 > 0); + assert!(total_links_0 as f64 / num_vectors as f64 > M as f64); + + let top = 5; + let query = random_vector(&mut rng, dim); + let processed_query = >::preprocess(query.clone()); + let mut reference_top = FixedLengthPriorityQueue::new(top); + for idx in 0..vector_holder.vectors.len() as PointOffsetType { + let vec = &vector_holder.vectors.get(idx as VectorOffsetType); + reference_top.push(ScoredPointOffset { + idx, + score: M::similarity(vec, &processed_query), + }); + } + + let graph_search = search_in_graph(&query, top, &vector_holder, &graph_layers); + + assert_eq!(reference_top.into_sorted_vec(), graph_search); + } +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..5a4b49d93 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs_prompt.txt", + "expected_file": "locodiff-250425/prompts/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.683422+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 18851, + "extracted_output_length": 18839, + "expected_output_length": 18890, + "results_dir": "locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-xWgrmbzNVn5ovyh4EoeS", + "cost_usd": 0.02172404, + "prompt_tokens": 69731, + "completion_tokens": 4950, + "total_tokens": 74681, + "native_prompt_tokens": 69073, + "native_completion_tokens": 5296, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 929 +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..954ecfe00 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1,72 @@ + +index d859f00fb..a083a5599 100644 +--- a/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs_expectedoutput.txt (expected):tmp/tmpxkw9462d_expected.txt ++++ b/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs_extracted.txt (actual):tmp/tmpous963md_actual.txt +@@ -22,6 +22,7 @@ use crate::index::hnsw_index::search_context::SearchContext; + use crate::index::visited_pool::{VisitedListHandle, VisitedPool}; + + pub type LinkContainer = Vec; ++pub type LinkContainerRef<'a> = &'a [PointOffsetType]; + pub type LayersContainer = Vec; + + pub const HNSW_GRAPH_FILE: &str = "graph.bin"; +@@ -318,7 +319,7 @@ impl GraphLayers { + return GraphLinks::load_from_file(&path, on_disk, format); + } + } +- Err(OperationError::service_error("No links file found")) ++ Err(OperationError::generic("No links file found")) + } + + fn convert_to_compressed(dir: &Path, m: usize, m0: usize) -> OperationResult<()> { +@@ -421,10 +422,7 @@ mod tests { + let entry_points_num = 10; + let num_vectors = 10; + +- let mut rng = StdRng::seed_from_u64(42); +- +- let vector_holder = +- TestRawScorerProducer::::new(dim, num_vectors, &mut rng); ++ let vector_holder = TestRawScorerProducer::::new(dim, num_vectors, &mut StdRng::seed_from_u64(42)); + + let mut graph_links = vec![vec![Vec::new()]; num_vectors]; + graph_links[0][0] = vec![1, 2, 3, 4, 5, 6]; +@@ -441,10 +439,7 @@ mod tests { + let linking_idx: PointOffsetType = 7; + + let fake_filter_context = FakeFilterContext {}; +- let added_vector = vector_holder +- .vectors +- .get(linking_idx as VectorOffsetType) +- .to_vec(); ++ let added_vector = vector_holder.vectors.get(linking_idx as VectorOffsetType).to_vec(); + let raw_scorer = vector_holder.get_raw_scorer(added_vector).unwrap(); + let mut scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + +@@ -485,16 +480,13 @@ mod tests { + let mut rng = StdRng::seed_from_u64(42); + + let dir = Builder::new().prefix("graph_dir").tempdir().unwrap(); +- +- let query = random_vector(&mut rng, dim); +- + let (vector_holder, graph_layers_builder) = + create_graph_layer_builder_fixture(num_vectors, M, dim, false, &mut rng); + let graph1 = graph_layers_builder + .into_graph_layers(dir.path(), initial_format, true) + .unwrap(); + assert_eq!(graph1.links.format(), initial_format); +- let res1 = search_in_graph(&query, top, &vector_holder, &graph1); ++ let res1 = search_in_graph(&random_vector(&mut rng, dim), top, &vector_holder, &graph1); + drop(graph1); + + let graph2 = GraphLayers::load(dir.path(), false, compress).unwrap(); +@@ -503,7 +495,7 @@ mod tests { + } else { + assert_eq!(graph2.links.format(), initial_format); + } +- let res2 = search_in_graph(&query, top, &vector_holder, &graph2); ++ let res2 = search_in_graph(&vec![0.0; dim], top, &vector_holder, &graph2); + + assert_eq!(res1, res2) + } diff --git a/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..ab32cd3e8 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_index_hnsw_index_graph_layers.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt @@ -0,0 +1,557 @@ +```rust +use std::borrow::Cow; +use std::cmp::max; +use std::path::{Path, PathBuf}; +use std::sync::atomic::AtomicBool; + +use common::fixed_length_priority_queue::FixedLengthPriorityQueue; +use common::types::{PointOffsetType, ScoredPointOffset}; +use io::file_operations::read_bin; +use itertools::Itertools; +use serde::{Deserialize, Serialize}; + +use super::entry_points::EntryPoint; +use super::graph_links::{GraphLinks, GraphLinksFormat}; +use crate::common::operation_error::{ + CancellableResult, OperationError, OperationResult, check_process_stopped, +}; +use crate::common::utils::rev_range; +use crate::index::hnsw_index::entry_points::EntryPoints; +use crate::index::hnsw_index::graph_links::GraphLinksSerializer; +use crate::index::hnsw_index::point_scorer::FilteredScorer; +use crate::index::hnsw_index::search_context::SearchContext; +use crate::index::visited_pool::{VisitedListHandle, VisitedPool}; + +pub type LinkContainer = Vec; +pub type LinkContainerRef<'a> = &'a [PointOffsetType]; +pub type LayersContainer = Vec; + +pub const HNSW_GRAPH_FILE: &str = "graph.bin"; +pub const HNSW_LINKS_FILE: &str = "links.bin"; +pub const COMPRESSED_HNSW_LINKS_FILE: &str = "links_compressed.bin"; + +/// Contents of the `graph.bin` file. +#[derive(Deserialize, Serialize, Debug)] +pub(super) struct GraphLayerData<'a> { + pub(super) m: usize, + pub(super) m0: usize, + pub(super) ef_construct: usize, + pub(super) entry_points: Cow<'a, EntryPoints>, +} + +#[derive(Debug)] +pub struct GraphLayers { + pub(super) m: usize, + pub(super) m0: usize, + pub(super) links: GraphLinks, + pub(super) entry_points: EntryPoints, + pub(super) visited_pool: VisitedPool, +} + +pub trait GraphLayersBase { + fn get_visited_list_from_pool(&self) -> VisitedListHandle; + + fn links_map(&self, point_id: PointOffsetType, level: usize, f: F) + where + F: FnMut(PointOffsetType); + + /// Get M based on current level + fn get_m(&self, level: usize) -> usize; + + /// Greedy search for closest points within a single graph layer + fn _search_on_level( + &self, + searcher: &mut SearchContext, + level: usize, + visited_list: &mut VisitedListHandle, + points_scorer: &mut FilteredScorer, + is_stopped: &AtomicBool, + ) -> CancellableResult<()> { + let limit = self.get_m(level); + let mut points_ids: Vec = Vec::with_capacity(2 * limit); + + while let Some(candidate) = searcher.candidates.pop() { + check_process_stopped(is_stopped)?; + + if candidate.score < searcher.lower_bound() { + break; + } + + points_ids.clear(); + self.links_map(candidate.idx, level, |link| { + if !visited_list.check(link) { + points_ids.push(link); + } + }); + + let scores = points_scorer.score_points(&mut points_ids, limit); + scores.iter().copied().for_each(|score_point| { + searcher.process_candidate(score_point); + visited_list.check_and_update_visited(score_point.idx); + }); + } + + Ok(()) + } + + fn search_on_level( + &self, + level_entry: ScoredPointOffset, + level: usize, + ef: usize, + points_scorer: &mut FilteredScorer, + is_stopped: &AtomicBool, + ) -> CancellableResult> { + let mut visited_list = self.get_visited_list_from_pool(); + visited_list.check_and_update_visited(level_entry.idx); + let mut search_context = SearchContext::new(level_entry, ef); + + self._search_on_level( + &mut search_context, + level, + &mut visited_list, + points_scorer, + is_stopped, + )?; + Ok(search_context.nearest) + } + + /// Greedy searches for entry point of level `target_level`. + /// Beam size is 1. + fn search_entry( + &self, + entry_point: PointOffsetType, + top_level: usize, + target_level: usize, + points_scorer: &mut FilteredScorer, + is_stopped: &AtomicBool, + ) -> CancellableResult { + let mut links: Vec = Vec::with_capacity(2 * self.get_m(0)); + + let mut current_point = ScoredPointOffset { + idx: entry_point, + score: points_scorer.score_point(entry_point), + }; + for level in rev_range(top_level, target_level) { + check_process_stopped(is_stopped)?; + + let limit = self.get_m(level); + + let mut changed = true; + while changed { + changed = false; + + links.clear(); + self.links_map(current_point.idx, level, |link| { + links.push(link); + }); + + let scores = points_scorer.score_points(&mut links, limit); + scores.iter().copied().for_each(|score_point| { + if score_point.score > current_point.score { + changed = true; + current_point = score_point; + } + }); + } + } + Ok(current_point) + } + + #[cfg(test)] + #[cfg(feature = "gpu")] + fn search_entry_on_level( + &self, + entry_point: PointOffsetType, + level: usize, + points_scorer: &mut FilteredScorer, + ) -> ScoredPointOffset { + let limit = self.get_m(level); + let mut links: Vec = Vec::with_capacity(2 * self.get_m(0)); + let mut current_point = ScoredPointOffset { + idx: entry_point, + score: points_scorer.score_point(entry_point), + }; + + let mut changed = true; + while changed { + changed = false; + + links.clear(); + self.links_map(current_point.idx, level, |link| { + links.push(link); + }); + + let scores = points_scorer.score_points(&mut links, limit); + scores.iter().copied().for_each(|score_point| { + if score_point.score > current_point.score { + changed = true; + current_point = score_point; + } + }); + } + current_point + } +} + +impl GraphLayersBase for GraphLayers { + fn get_visited_list_from_pool(&self) -> VisitedListHandle { + self.visited_pool.get(self.links.num_points()) + } + + fn links_map(&self, point_id: PointOffsetType, level: usize, f: F) + where + F: FnMut(PointOffsetType), + { + self.links.links(point_id, level).for_each(f); + } + + fn get_m(&self, level: usize) -> usize { + if level == 0 { self.m0 } else { self.m } + } +} + +/// Object contains links between nodes for HNSW search +/// +/// Assume all scores are similarities. Larger score = closer points +impl GraphLayers { + /// Returns the highest level this point is included in + pub fn point_level(&self, point_id: PointOffsetType) -> usize { + self.links.point_level(point_id) + } + + fn get_entry_point( + &self, + points_scorer: &FilteredScorer, + custom_entry_points: Option<&[PointOffsetType]>, + ) -> Option { + // Try to get it from custom entry points + custom_entry_points + .and_then(|custom_entry_points| { + custom_entry_points + .iter() + .filter(|&&point_id| points_scorer.check_vector(point_id)) + .map(|&point_id| { + let level = self.point_level(point_id); + EntryPoint { point_id, level } + }) + .max_by_key(|ep| ep.level) + }) + .or_else(|| { + // Otherwise use normal entry points + self.entry_points + .get_entry_point(|point_id| points_scorer.check_vector(point_id)) + }) + } + + pub fn search( + &self, + top: usize, + ef: usize, + mut points_scorer: FilteredScorer, + custom_entry_points: Option<&[PointOffsetType]>, + is_stopped: &AtomicBool, + ) -> CancellableResult> { + let Some(entry_point) = self.get_entry_point(&points_scorer, custom_entry_points) else { + return Ok(Vec::default()); + }; + + let zero_level_entry = self.search_entry( + entry_point.point_id, + entry_point.level, + 0, + &mut points_scorer, + is_stopped, + )?; + let nearest = self.search_on_level( + zero_level_entry, + 0, + max(top, ef), + &mut points_scorer, + is_stopped, + )?; + Ok(nearest.into_iter_sorted().take(top).collect_vec()) + } + + pub fn get_path(path: &Path) -> PathBuf { + path.join(HNSW_GRAPH_FILE) + } + + pub fn get_links_path(path: &Path, format: GraphLinksFormat) -> PathBuf { + match format { + GraphLinksFormat::Plain => path.join(HNSW_LINKS_FILE), + GraphLinksFormat::Compressed => path.join(COMPRESSED_HNSW_LINKS_FILE), + } + } + + pub fn files(&self, path: &Path) -> Vec { + vec![ + GraphLayers::get_path(path), + GraphLayers::get_links_path(path, self.links.format()), + ] + } + + pub fn num_points(&self) -> usize { + self.links.num_points() + } +} + +impl GraphLayers { + pub fn load(dir: &Path, on_disk: bool, compress: bool) -> OperationResult { + let graph_data: GraphLayerData = read_bin(&GraphLayers::get_path(dir))?; + + if compress { + Self::convert_to_compressed(dir, graph_data.m, graph_data.m0)?; + } + + Ok(Self { + m: graph_data.m, + m0: graph_data.m0, + links: Self::load_links(dir, on_disk)?, + entry_points: graph_data.entry_points.into_owned(), + visited_pool: VisitedPool::new(), + }) + } + + fn load_links(dir: &Path, on_disk: bool) -> OperationResult { + for format in [GraphLinksFormat::Compressed, GraphLinksFormat::Plain] { + let path = GraphLayers::get_links_path(dir, format); + if path.exists() { + return GraphLinks::load_from_file(&path, on_disk, format); + } + } + Err(OperationError::generic("No links file found")) + } + + fn convert_to_compressed(dir: &Path, m: usize, m0: usize) -> OperationResult<()> { + let plain_path = Self::get_links_path(dir, GraphLinksFormat::Plain); + let compressed_path = Self::get_links_path(dir, GraphLinksFormat::Compressed); + + if compressed_path.exists() { + return Ok(()); + } + + let start = std::time::Instant::now(); + + let links = GraphLinks::load_from_file(&plain_path, true, GraphLinksFormat::Plain)?; + let original_size = plain_path.metadata()?.len(); + GraphLinksSerializer::new(links.into_edges(), GraphLinksFormat::Compressed, m, m0) + .save_as(&compressed_path)?; + let new_size = compressed_path.metadata()?.len(); + + // Remove the original file + std::fs::remove_file(plain_path)?; + + log::debug!( + "Compressed HNSW graph links in {:.1?}: {:.1}MB -> {:.1}MB ({:.1}%)", + start.elapsed(), + original_size as f64 / 1024.0 / 1024.0, + new_size as f64 / 1024.0 / 1024.0, + new_size as f64 / original_size as f64 * 100.0, + ); + + Ok(()) + } + + #[cfg(feature = "testing")] + pub fn compress_ram(&mut self) { + use crate::index::hnsw_index::graph_links::GraphLinksSerializer; + assert_eq!(self.links.format(), GraphLinksFormat::Plain); + let dummy = GraphLinksSerializer::new(Vec::new(), GraphLinksFormat::Plain, 0, 0) + .to_graph_links_ram(); + let links = std::mem::replace(&mut self.links, dummy); + self.links = GraphLinksSerializer::new( + links.into_edges(), + GraphLinksFormat::Compressed, + self.m, + self.m0, + ) + .to_graph_links_ram(); + } + + pub fn populate(&self) -> OperationResult<()> { + self.links.populate()?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng; + use rand::rngs::StdRng; + use rstest::rstest; + use tempfile::Builder; + + use super::*; + use crate::data_types::vectors::VectorElementType; + use crate::fixtures::index_fixtures::{ + FakeFilterContext, TestRawScorerProducer, random_vector, + }; + use crate::index::hnsw_index::graph_links::GraphLinksSerializer; + use crate::index::hnsw_index::tests::{ + create_graph_layer_builder_fixture, create_graph_layer_fixture, + }; + use crate::spaces::metric::Metric; + use crate::spaces::simple::{CosineMetric, DotProductMetric}; + use crate::vector_storage::DEFAULT_STOPPED; + use crate::vector_storage::chunked_vector_storage::VectorOffsetType; + + fn search_in_graph( + query: &[VectorElementType], + top: usize, + vector_storage: &TestRawScorerProducer, + graph: &GraphLayers, + ) -> Vec { + let fake_filter_context = FakeFilterContext {}; + let raw_scorer = vector_storage.get_raw_scorer(query.to_owned()).unwrap(); + + let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + let ef = 16; + graph + .search(top, ef, scorer, None, &DEFAULT_STOPPED) + .unwrap() + } + + const M: usize = 8; + + #[rstest] + #[case::uncompressed(GraphLinksFormat::Plain)] + #[case::compressed(GraphLinksFormat::Compressed)] + fn test_search_on_level(#[case] format: GraphLinksFormat) { + let dim = 8; + let m = 8; + let entry_points_num = 10; + let num_vectors = 10; + + let vector_holder = TestRawScorerProducer::::new(dim, num_vectors, &mut StdRng::seed_from_u64(42)); + + let mut graph_links = vec![vec![Vec::new()]; num_vectors]; + graph_links[0][0] = vec![1, 2, 3, 4, 5, 6]; + + let graph_layers = GraphLayers { + m, + m0: 2 * m, + links: GraphLinksSerializer::new(graph_links.clone(), format, m, 2 * m) + .to_graph_links_ram(), + entry_points: EntryPoints::new(entry_points_num), + visited_pool: VisitedPool::new(), + }; + + let linking_idx: PointOffsetType = 7; + + let fake_filter_context = FakeFilterContext {}; + let added_vector = vector_holder.vectors.get(linking_idx as VectorOffsetType).to_vec(); + let raw_scorer = vector_holder.get_raw_scorer(added_vector).unwrap(); + let mut scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context)); + + let nearest_on_level = graph_layers + .search_on_level( + ScoredPointOffset { + idx: 0, + score: scorer.score_point(0), + }, + 0, + 32, + &mut scorer, + &DEFAULT_STOPPED, + ) + .unwrap(); + + assert_eq!(nearest_on_level.len(), graph_links[0][0].len() + 1); + + for nearest in nearest_on_level.iter_unsorted() { + // eprintln!("nearest = {:#?}", nearest); + assert_eq!( + nearest.score, + scorer.score_internal(linking_idx, nearest.idx) + ) + } + } + + #[rstest] + #[case::uncompressed((GraphLinksFormat::Plain, false))] + #[case::converted((GraphLinksFormat::Plain, true))] + #[case::compressed((GraphLinksFormat::Compressed, false))] + #[case::recompressed((GraphLinksFormat::Compressed, true))] + fn test_save_and_load(#[case] (initial_format, compress): (GraphLinksFormat, bool)) { + let num_vectors = 100; + let dim = 8; + let top = 5; + + let mut rng = StdRng::seed_from_u64(42); + + let dir = Builder::new().prefix("graph_dir").tempdir().unwrap(); + let (vector_holder, graph_layers_builder) = + create_graph_layer_builder_fixture(num_vectors, M, dim, false, &mut rng); + let graph1 = graph_layers_builder + .into_graph_layers(dir.path(), initial_format, true) + .unwrap(); + assert_eq!(graph1.links.format(), initial_format); + let res1 = search_in_graph(&random_vector(&mut rng, dim), top, &vector_holder, &graph1); + drop(graph1); + + let graph2 = GraphLayers::load(dir.path(), false, compress).unwrap(); + if compress { + assert_eq!(graph2.links.format(), GraphLinksFormat::Compressed); + } else { + assert_eq!(graph2.links.format(), initial_format); + } + let res2 = search_in_graph(&vec![0.0; dim], top, &vector_holder, &graph2); + + assert_eq!(res1, res2) + } + + #[rstest] + #[case::uncompressed(GraphLinksFormat::Plain)] + #[case::compressed(GraphLinksFormat::Compressed)] + fn test_add_points(#[case] format: GraphLinksFormat) { + let num_vectors = 1000; + let dim = 8; + + let mut rng = StdRng::seed_from_u64(42); + + type M = CosineMetric; + + let (vector_holder, graph_layers) = + create_graph_layer_fixture::(num_vectors, M, dim, format, false, &mut rng); + + let main_entry = graph_layers + .entry_points + .get_entry_point(|_x| true) + .expect("Expect entry point to exists"); + + assert!(main_entry.level > 0); + + let num_levels = (0..num_vectors) + .map(|i| graph_layers.links.point_level(i as PointOffsetType)) + .max() + .unwrap(); + assert_eq!(main_entry.level, num_levels); + + let total_links_0 = (0..num_vectors) + .map(|i| graph_layers.links.links(i as PointOffsetType, 0).len()) + .sum::(); + + eprintln!("total_links_0 = {total_links_0:#?}"); + eprintln!("num_vectors = {num_vectors:#?}"); + assert!(total_links_0 > 0); + assert!(total_links_0 as f64 / num_vectors as f64 > M as f64); + + let top = 5; + let query = random_vector(&mut rng, dim); + let processed_query = >::preprocess(query.clone()); + let mut reference_top = FixedLengthPriorityQueue::new(top); + for idx in 0..vector_holder.vectors.len() as PointOffsetType { + let vec = &vector_holder.vectors.get(idx as VectorOffsetType); + reference_top.push(ScoredPointOffset { + idx, + score: M::similarity(vec, &processed_query), + }); + } + + let graph_search = search_in_graph(&query, top, &vector_holder, &graph_layers); + + assert_eq!(reference_top.into_sorted_vec(), graph_search); + } +} +``` \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt b/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt new file mode 100644 index 000000000..503aba149 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt @@ -0,0 +1,787 @@ +use std::collections::HashMap; +use std::fs::create_dir_all; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use ahash::AHashSet; +use atomic_refcell::AtomicRefCell; +use common::counter::hardware_counter::HardwareCounterCell; +use common::counter::iterator_hw_measurement::HwMeasurementIteratorExt; +use common::flags::feature_flags; +use common::types::PointOffsetType; +use itertools::Either; +use log::debug; +use parking_lot::RwLock; +use rocksdb::DB; +use schemars::_serde_json::Value; + +use super::field_index::FieldIndexBuilderTrait as _; +use super::field_index::facet_index::FacetIndexEnum; +use super::field_index::index_selector::{IndexSelector, IndexSelectorMmap, IndexSelectorRocksDb}; +use crate::common::Flusher; +use crate::common::operation_error::{OperationError, OperationResult}; +use crate::common::rocksdb_wrapper::open_db_with_existing_cf; +use crate::common::utils::IndexesMap; +use crate::id_tracker::IdTrackerSS; +use crate::index::field_index::{ + CardinalityEstimation, FieldIndex, PayloadBlockCondition, PrimaryCondition, +}; +use crate::index::field_index::index_selector::index_selector; +use crate::index::payload_config::PayloadConfig; +use crate::index::query_estimator::estimate_filter; +use crate::index::query_optimization::payload_provider::PayloadProvider; +use crate::index::struct_filter_context::StructFilterContext; +use crate::index::visited_pool::VisitedPool; +use crate::index::PayloadIndex; +use crate::json_path::JsonPath; +use crate::payload_storage::payload_storage_enum::PayloadStorageEnum; +use crate::payload_storage::{FilterContext, PayloadStorage}; +use crate::telemetry::PayloadIndexTelemetry; +use crate::types::{ + Condition, FieldCondition, Filter, IsEmptyCondition, IsNullCondition, Payload, + PayloadContainer, PayloadField, PayloadFieldSchema, PayloadKeyType, PayloadKeyTypeRef, + PayloadSchemaType, VectorNameBuf, infer_collection_value_type, infer_value_type, +}; +use crate::vector_storage::{VectorStorage, VectorStorageEnum}; + +#[derive(Debug)] +enum StorageType { + Appendable(Arc>), + NonAppendableRocksDb(Arc>), + NonAppendable, +} + +/// `PayloadIndex` implementation, which actually uses index structures for providing faster search +#[derive(Debug)] +pub struct StructPayloadIndex { + /// Payload storage + pub(super) payload: Arc>, + /// Used for `has_id` condition and estimating cardinality + pub(super) id_tracker: Arc>, + /// Vector storages for each field, used for `has_vector` condition + pub(super) vector_storages: HashMap>>, + /// Indexes, associated with fields + pub field_indexes: IndexesMap, + config: PayloadConfig, + /// Root of index persistence dir + path: PathBuf, + /// Used to select unique point ids + visited_pool: VisitedPool, + storage_type: StorageType, +} + +impl StructPayloadIndex { + fn config_path(&self) -> PathBuf { + PayloadConfig::get_config_path(&self.path) + } + + fn save_config(&self) -> OperationResult<()> { + let config_path = self.config_path(); + self.config.save(&config_path) + } + + fn load_all_fields(&mut self) -> OperationResult<()> { + let mut field_indexes: IndexesMap = Default::default(); + + for (field, payload_schema) in &self.config.indexed_fields { + let field_index = self.load_from_db(field, payload_schema)?; + field_indexes.insert(field.clone(), field_index); + } + self.field_indexes = field_indexes; + Ok(()) + } + + fn load_from_db( + &self, + field: PayloadKeyTypeRef, + payload_schema: &PayloadFieldSchema, + ) -> OperationResult> { + let mut indexes = self + .selector(payload_schema) + .new_index(field, payload_schema)?; + + let mut is_loaded = true; + for ref mut index in indexes.iter_mut() { + if !index.load()? { + is_loaded = false; + break; + } + } + if !is_loaded { + debug!("Index for `{field}` was not loaded. Building..."); + // todo(ivan): decide what to do with indexes, which were not loaded + indexes = self.build_field_indexes( + field, + payload_schema, + &HardwareCounterCell::disposable(), // Internal operation. + )?; + for index in indexes.iter_mut() { + index.save()?; + } + } + + Ok(indexes) + } + + pub fn open( + payload: Arc>, + id_tracker: Arc>, + vector_storages: HashMap>>, + path: &Path, + is_appendable: bool, + ) -> OperationResult { + create_dir_all(path)?; + let config_path = PayloadConfig::get_config_path(path); + let config = if config_path.exists() { + PayloadConfig::load(&config_path)? + } else { + let mut new_config = PayloadConfig::default(); + if feature_flags().payload_index_skip_rocksdb && !is_appendable { + new_config.skip_rocksdb = Some(true); + } + new_config + }; + + let skip_rocksdb = config.skip_rocksdb.unwrap_or(false); + + let storage_type = if is_appendable { + let db = open_db_with_existing_cf(path).map_err(|err| { + OperationError::service_error(format!("RocksDB open error: {err}")) + })?; + StorageType::Appendable(db) + } else if skip_rocksdb { + StorageType::NonAppendable + } else { + let db = open_db_with_existing_cf(path).map_err(|err| { + OperationError::service_error(format!("RocksDB open error: {err}")) + })?; + StorageType::NonAppendableRocksDb(db) + }; + + let mut index = StructPayloadIndex { + payload, + id_tracker, + vector_storages, + field_indexes: Default::default(), + config, + path: path.to_owned(), + visited_pool: Default::default(), + storage_type, + }; + + if !index.config_path().exists() { + // Save default config + index.save_config()?; + } + + index.load_all_fields()?; + + Ok(index) + } + + pub fn available_point_count(&self) -> usize { + self.id_tracker.borrow().available_point_count() + } + + pub fn struct_filtered_context<'a>( + &'a self, + filter: &'a Filter, + hw_counter: &HardwareCounterCell, + ) -> StructFilterContext<'a> { + let payload_provider = PayloadProvider::new(self.payload.clone()); + + let (optimized_filter, _) = self.optimize_filter( + filter, + payload_provider, + self.available_point_count(), + hw_counter, + ); + + StructFilterContext::new(optimized_filter) + } + + pub(super) fn condition_cardinality( + &self, + condition: &Condition, + nested_path: Option<&JsonPath>, + hw_counter: &HardwareCounterCell, + ) -> CardinalityEstimation { + match condition { + Condition::Filter(_) => panic!("Unexpected branching"), + Condition::Nested(nested) => { + // propagate complete nested path in case of multiple nested layers + let full_path = JsonPath::extend_or_new(nested_path, &nested.array_key()); + self.estimate_nested_cardinality(nested.filter(), &full_path, hw_counter) + } + Condition::IsEmpty(IsEmptyCondition { is_empty: field }) => { + let available_points = self.available_point_count(); + let condition = FieldCondition::new_is_empty(field.key.clone()); + + self.estimate_field_condition(&condition, nested_path, hw_counter) + .unwrap_or_else(|| CardinalityEstimation::unknown(available_points)) + } + Condition::IsNull(IsNullCondition { is_null: field }) => { + let available_points = self.available_point_count(); + let condition = FieldCondition::new_is_null(field.key.clone()); + + self.estimate_field_condition(&condition, nested_path, hw_counter) + .unwrap_or_else(|| CardinalityEstimation::unknown(available_points)) + } + Condition::HasId(has_id) => { + let id_tracker_ref = self.id_tracker.borrow(); + let mapped_ids: AHashSet = has_id + .has_id + .iter() + .filter_map(|external_id| id_tracker_ref.internal_id(*external_id)) + .collect(); + let num_ids = mapped_ids.len(); + CardinalityEstimation { + primary_clauses: vec![PrimaryCondition::Ids(mapped_ids)], + min: num_ids, + exp: num_ids, + max: num_ids, + } + } + Condition::HasVector(has_vectors) => { + if let Some(vector_storage) = self.vector_storages.get(&has_vectors.has_vector) { + let vector_storage = vector_storage.borrow(); + let vectors = vector_storage.available_vector_count(); + CardinalityEstimation::exact(vectors).with_primary_clause( + PrimaryCondition::HasVector(has_vectors.has_vector.clone()), + ) + } else { + CardinalityEstimation::exact(0) + } + } + Condition::Field(field_condition) => self + .estimate_field_condition(field_condition, nested_path, hw_counter) + .unwrap_or_else(|| CardinalityEstimation::unknown(self.available_point_count())), + + Condition::CustomIdChecker(cond) => { + cond.estimate_cardinality(self.id_tracker.borrow().available_point_count()) + } + } + } + + fn estimate_field_condition( + &self, + condition: &FieldCondition, + nested_path: Option<&JsonPath>, + hw_counter: &HardwareCounterCell, + ) -> Option { + let full_path = JsonPath::extend_or_new(nested_path, &condition.key); + self.field_indexes.get(&full_path).and_then(|indexes| { + // rewrite condition with fullpath to enable cardinality estimation + let full_path_condition = FieldCondition { + key: full_path, + ..condition.clone() + }; + + indexes + .iter() + .find_map(|index| index.estimate_cardinality(&full_path_condition, hw_counter)) + }) + } + + fn query_field<'a>( + &'a self, + condition: &'a PrimaryCondition, + hw_counter: &'a HardwareCounterCell, + ) -> Option + 'a>> { + match condition { + PrimaryCondition::Condition(field_condition) => { + let field_key = &field_condition.key; + let field_indexes = self.field_indexes.get(field_key)?; + field_indexes + .iter() + .find_map(|field_index| field_index.filter(field_condition, hw_counter)) + } + PrimaryCondition::Ids(ids) => Some(Box::new(ids.iter().copied())), + PrimaryCondition::HasVector(_) => None, + } + } + + fn config_path(&self) -> PathBuf { + PayloadConfig::get_config_path(&self.path) + } + + fn save_config(&self) -> OperationResult<()> { + let config_path = self.config_path(); + self.config.save(&config_path) + } + + fn load_all_fields(&mut self) -> OperationResult<()> { + let mut field_indexes: IndexesMap = Default::default(); + + for (field, payload_schema) in &self.config.indexed_fields { + let field_index = self.load_from_db(field, payload_schema)?; + field_indexes.insert(field.clone(), field_index); + } + self.field_indexes = field_indexes; + Ok(()) + } + + fn restore_database_snapshot( + snapshot_path: &Path, + segment_path: &Path, + ) -> OperationResult<()> { + crate::rocksdb_backup::restore(snapshot_path, &segment_path.join("payload_index")) + } + + /// Select which type of PayloadIndex to use for the field + fn selector(&self, payload_schema: &PayloadFieldSchema) -> IndexSelector { + let is_on_disk = payload_schema.is_on_disk(); + + match &self.storage_type { + StorageType::Appendable(db) => IndexSelector::RocksDb(IndexSelectorRocksDb { + db, + is_appendable: true, + }), + StorageType::NonAppendableRocksDb(db) => { + // legacy logic: we keep rocksdb, but load mmap indexes + if is_on_disk { + IndexSelector::Mmap(IndexSelectorMmap { + dir: &self.path, + is_on_disk, + }) + } else { + IndexSelector::RocksDb(IndexSelectorRocksDb { + db, + is_appendable: false, + }) + } + } + StorageType::NonAppendable => IndexSelector::Mmap(IndexSelectorMmap { + dir: &self.path, + is_on_disk, + }), + } + } + + pub fn optimize_filter( + &self, + _filter: &Filter, + _payload_provider: PayloadProvider, + _total: usize, + _hw_counter: &HardwareCounterCell, + ) -> (Filter, CardinalityEstimation) { + // TODO implement optimization logic + todo!() + } + + pub fn iter_filtered_points<'a>( + &'a self, + filter: &'a Filter, + id_tracker: &'a IdTrackerSS, + query_cardinality: &'a CardinalityEstimation, + hw_counter: &'a HardwareCounterCell, + ) -> impl Iterator + 'a { + let struct_filtered_context = self.struct_filtered_context(filter, hw_counter); + + if query_cardinality.primary_clauses.is_empty() { + let full_scan_iterator = id_tracker.iter_ids(); + + let struct_filtered_context = self.struct_filtered_context(filter, hw_counter); + // Worst case: query expected to return few matches, but index can't be used + let matched_points = + full_scan_iterator.filter(move |i| struct_filtered_context.check(*i)); + + Either::Left(matched_points) + } else { + let struct_filtered_context = self.struct_filtered_context(filter, hw_counter); + + // CPU-optimized strategy here: points are made unique before applying other filters. + let mut visited_list = self.visited_pool.get(id_tracker.total_point_count()); + + let iter = query_cardinality + .primary_clauses + .iter() + .flat_map(move |clause| { + self.query_field(clause, hw_counter).unwrap_or_else(|| { + // index is not built + Box::new(id_tracker.iter_ids().measure_hw_with_cell( + hw_counter, + size_of::(), + |i| i.cpu_counter(), + )) + }) + }) + .filter(move |&id| !visited_list.check_and_update_visited(id)) + .filter(move |&i| struct_filtered_context.check(i)); + + Either::Right(iter) + } + } + + pub fn get_facet_index(&self, key: &JsonPath) -> OperationResult { + self.field_indexes + .get(key) + .and_then(|index| index.iter().find_map(|index| index.as_facet_index())) + .ok_or_else(|| OperationError::MissingMapIndexForFacet { + key: key.to_string(), + }) + } + + pub fn populate(&self) -> OperationResult<()> { + for (_, field_indexes) in self.field_indexes.iter() { + for index in field_indexes { + index.populate()?; + } + } + Ok(()) + } + + pub fn clear_cache(&self) -> OperationResult<()> { + for (_, field_indexes) in self.field_indexes.iter() { + for index in field_indexes { + index.clear_cache()?; + } + } + Ok(()) + } + + pub fn clear_cache_if_on_disk(&self) -> OperationResult<()> { + for (_, field_indexes) in self.field_indexes.iter() { + for index in field_indexes { + if index.is_on_disk() { + index.clear_cache()?; + } + } + } + Ok(()) + } +} + +impl PayloadIndex for StructPayloadIndex { + fn indexed_fields(&self) -> HashMap { + self.config.indexed_fields.clone() + } + + fn set_indexed( + &self, + field: PayloadKeyTypeRef, + payload_schema: &PayloadFieldSchema, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + if let Some(prev_schema) = self.config.indexed_fields.get(field) { + // the field is already indexed with the same schema + // no need to rebuild index and to save the config + if prev_schema == payload_schema { + return Ok(None); + } + } + + let indexes = self.build_field_indexes(field, payload_schema, hw_counter)?; + + Ok(Some(indexes)) + } + + fn apply_index( + &mut self, + field: PayloadKeyType, + payload_schema: PayloadFieldSchema, + field_index: Vec, + ) -> OperationResult<()> { + self.field_indexes.insert(field.clone(), field_index); + + self.config.indexed_fields.insert(field, payload_schema); + + self.save_config()?; + + Ok(()) + } + + fn drop_index(&mut self, field: PayloadKeyTypeRef) -> OperationResult<()> { + self.config.indexed_fields.remove(field); + self.save_config()?; + let removed_indexes = self.field_indexes.remove(field); + + if let Some(indexes) = removed_indexes { + for index in indexes { + index.cleanup()?; + } + } + + Ok(()) + } + + fn build_index( + &self, + field: PayloadKeyTypeRef, + payload_schema: &PayloadFieldSchema, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + if let Some(prev_schema) = self.config.indexed_fields.get(field) { + // the field is already indexed with the same schema + // no need to rebuild index and to save the config + if prev_schema == payload_schema { + return Ok(None); + } + } + + let indexes = self.build_field_indexes(field, payload_schema, hw_counter)?; + + Ok(Some(indexes)) + } + + fn estimate_cardinality( + &self, + query: &Filter, + hw_counter: &HardwareCounterCell, + ) -> CardinalityEstimation { + let available_points = self.available_point_count(); + let estimator = + |condition: &Condition| self.condition_cardinality(condition, None, hw_counter); + estimate_filter(&estimator, query, available_points) + } + + fn estimate_nested_cardinality( + &self, + query: &Filter, + nested_path: &JsonPath, + hw_counter: &HardwareCounterCell, + ) -> CardinalityEstimation { + let available_points = self.available_point_count(); + let estimator = |condition: &Condition| { + self.condition_cardinality(condition, Some(nested_path), hw_counter) + }; + estimate_filter(&estimator, query, available_points) + } + + fn query_points( + &self, + query: &Filter, + hw_counter: &HardwareCounterCell, + ) -> Vec { + // Assume query is already estimated to be small enough so we can iterate over all matched ids + let query_cardinality = self.estimate_cardinality(query, hw_counter); + let id_tracker = self.id_tracker.borrow(); + self.iter_filtered_points(query, &*id_tracker, &query_cardinality, hw_counter) + .collect() + } + + fn indexed_points(&self, field: PayloadKeyTypeRef) -> usize { + self.field_indexes.get(field).map_or(0, |indexes| { + // Assume that multiple field indexes are applied to the same data type, + // so the points indexed with those indexes are the same. + // We will return minimal number as a worst case, to highlight possible errors in the index early. + indexes + .iter() + .map(|index| index.count_indexed_points()) + .min() + .unwrap_or(0) + }) + } + + fn filter_context<'a>( + &'a self, + filter: &'a Filter, + hw_counter: &HardwareCounterCell, + ) -> Box { + Box::new(self.struct_filtered_context(filter, hw_counter)) + } + + fn payload_blocks( + &self, + field: PayloadKeyTypeRef, + threshold: usize, + ) -> Box + '_> { + match self.field_indexes.get(field) { + None => Box::new(vec![].into_iter()), + Some(indexes) => { + let field_clone = field.to_owned(); + Box::new( + indexes + .iter() + .map(move |field_index| { + field_index.payload_blocks(threshold, field_clone.clone()) + }) + .flatten(), + ) + } + } + } + + fn assign_all( + &mut self, + point_id: PointOffsetType, + payload: &Payload, + hw_counter: &HardwareCounterCell, + ) -> OperationResult<()> { + self.payload + .borrow_mut() + .overwrite(point_id, payload, hw_counter)?; + + for (field, field_index) in &mut self.field_indexes { + let field_value = payload.get_value(field); + if !field_value.is_empty() { + for index in field_index { + index.add_point(point_id, &field_value, hw_counter)?; + } + } else { + for index in field_index { + index.remove_point(point_id)?; + } + } + } + Ok(()) + } + + fn set_payload( + &mut self, + point_id: PointOffsetType, + payload: &Payload, + key: &Option, + hw_counter: &HardwareCounterCell, + ) -> OperationResult<()> { + if let Some(key) = key { + self.payload + .borrow_mut() + .set_by_key(point_id, payload, key, hw_counter)?; + } else { + self.payload + .borrow_mut() + .set(point_id, payload, hw_counter)?; + }; + + let updated_payload = self.get_payload(point_id, hw_counter)?; + for (field, field_index) in &mut self.field_indexes { + if !field.is_affected_by_value_set(&payload.0, key.as_ref()) { + continue; + } + let field_value = updated_payload.get_value(field); + if !field_value.is_empty() { + for index in field_index { + index.add_point(point_id, &field_value, hw_counter)?; + } + } else { + for index in field_index { + index.remove_point(point_id)?; + } + } + } + Ok(()) + } + + fn get_payload( + &self, + point_id: PointOffsetType, + hw_counter: &HardwareCounterCell, + ) -> OperationResult { + self.payload.borrow().get(point_id, hw_counter) + } + + fn delete_payload( + &mut self, + point_id: PointOffsetType, + key: PayloadKeyTypeRef, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + if let Some(indexes) = self.field_indexes.get_mut(key) { + for index in indexes { + index.remove_point(point_id)?; + } + } + self.payload.borrow_mut().delete(point_id, key, hw_counter) + } + + fn clear_payload( + &mut self, + point_id: PointOffsetType, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + self.clear_index_for_point(point_id)?; + self.payload.borrow_mut().clear(point_id, hw_counter) + } + + fn flusher(&self) -> Flusher { + let mut flushers = Vec::new(); + for field_indexes in self.field_indexes.values() { + for index in field_indexes { + flushers.push(index.flusher()); + } + } + flushers.push(self.payload.borrow().flusher()); + Box::new(move || { + for flusher in flushers { + match flusher() { + Ok(_) => {} + Err(OperationError::RocksDbColumnFamilyNotFound { name }) => { + // It is possible, that the index was removed during the flush by user or another thread. + // In this case, non-existing column family is not an error, but an expected behavior. + + // Still we want to log this event, for potential debugging. + log::warn!( + "Flush: RocksDB cf_handle error: Cannot find column family {name}. Assume index is removed.", + ); + } + Err(err) => { + return Err(OperationError::service_error(format!( + "Failed to flush payload_index: {err}" + ))); + } + } + } + Ok(()) + }) + } + + fn take_database_snapshot(&self, path: &Path) -> OperationResult<()> { + match &self.storage_type { + StorageType::Appendable(db) => { + let db_guard = db.read(); + crate::rocksdb_backup::create(&db_guard, path) + } + StorageType::NonAppendableRocksDb(db) => { + let db_guard = db.read(); + crate::rocksdb_backup::create(&db_guard, path) + } + StorageType::NonAppendable => Ok(()), + } + } + + fn files(&self) -> Vec { + let mut files = self + .field_indexes + .values() + .flat_map(|indexes| indexes.iter().flat_map(|index| index.files().into_iter())) + .collect::>(); + files.push(self.config_path()); + files + } + + fn infer_payload_type( + &self, + key: PayloadKeyTypeRef, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + let mut schema = None; + self.payload.borrow().iter( + |_id, payload: &Payload| { + let field_value = payload.get_value(key); + schema = match field_value.as_slice() { + [] => None, + [single] => infer_value_type(single), + multiple => infer_collection_value_type(multiple.iter().copied()), + }; + Ok(false) + }, + hw_counter, + )?; + Ok(schema) + } + + fn get_telemetry_data(&self) -> Vec { + self.field_indexes + .iter() + .flat_map(|(name, field)| -> Vec { + field + .iter() + .map(|field| field.get_telemetry_data().set_name(name.to_string())) + .collect() + }) + .collect() + } +} + +// Continuation with possible impl for FacetIndexEnum or other types if needed, but based on commits, this is the state. In the last commit, there are no further additions beyond the file changes shown. \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..1db6db93c --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "qdrant_lib_segment_src_index_struct_payload_index.rs", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/qdrant_lib_segment_src_index_struct_payload_index.rs_prompt.txt", + "expected_file": "locodiff-250425/prompts/qdrant_lib_segment_src_index_struct_payload_index.rs_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.623321+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 28108, + "extracted_output_length": 28096, + "expected_output_length": 27244, + "results_dir": "locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-MpdeBXJcbvr2mSrG0odI", + "cost_usd": 0.02777998, + "prompt_tokens": 82889, + "completion_tokens": 6658, + "total_tokens": 89547, + "native_prompt_tokens": 81758, + "native_completion_tokens": 7665, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 1978 +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..d3764f802 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1,452 @@ + +index 47f78e903..503aba149 100644 +--- a/qdrant_lib_segment_src_index_struct_payload_index.rs_expectedoutput.txt (expected):tmp/tmpi_mt6m6n_expected.txt ++++ b/qdrant_lib_segment_src_index_struct_payload_index.rs_extracted.txt (actual):tmp/tmpm3quho3i_actual.txt +@@ -23,23 +23,24 @@ use crate::common::operation_error::{OperationError, OperationResult}; + use crate::common::rocksdb_wrapper::open_db_with_existing_cf; + use crate::common::utils::IndexesMap; + use crate::id_tracker::IdTrackerSS; +-use crate::index::PayloadIndex; + use crate::index::field_index::{ + CardinalityEstimation, FieldIndex, PayloadBlockCondition, PrimaryCondition, + }; ++use crate::index::field_index::index_selector::index_selector; + use crate::index::payload_config::PayloadConfig; + use crate::index::query_estimator::estimate_filter; + use crate::index::query_optimization::payload_provider::PayloadProvider; + use crate::index::struct_filter_context::StructFilterContext; + use crate::index::visited_pool::VisitedPool; ++use crate::index::PayloadIndex; + use crate::json_path::JsonPath; + use crate::payload_storage::payload_storage_enum::PayloadStorageEnum; + use crate::payload_storage::{FilterContext, PayloadStorage}; + use crate::telemetry::PayloadIndexTelemetry; + use crate::types::{ + Condition, FieldCondition, Filter, IsEmptyCondition, IsNullCondition, Payload, +- PayloadContainer, PayloadFieldSchema, PayloadKeyType, PayloadKeyTypeRef, PayloadSchemaType, +- VectorNameBuf, infer_collection_value_type, infer_value_type, ++ PayloadContainer, PayloadField, PayloadFieldSchema, PayloadKeyType, PayloadKeyTypeRef, ++ PayloadSchemaType, VectorNameBuf, infer_collection_value_type, infer_value_type, + }; + use crate::vector_storage::{VectorStorage, VectorStorageEnum}; + +@@ -70,44 +71,6 @@ pub struct StructPayloadIndex { + } + + impl StructPayloadIndex { +- pub fn estimate_field_condition( +- &self, +- condition: &FieldCondition, +- nested_path: Option<&JsonPath>, +- hw_counter: &HardwareCounterCell, +- ) -> Option { +- let full_path = JsonPath::extend_or_new(nested_path, &condition.key); +- self.field_indexes.get(&full_path).and_then(|indexes| { +- // rewrite condition with fullpath to enable cardinality estimation +- let full_path_condition = FieldCondition { +- key: full_path, +- ..condition.clone() +- }; +- +- indexes +- .iter() +- .find_map(|index| index.estimate_cardinality(&full_path_condition, hw_counter)) +- }) +- } +- +- fn query_field<'a>( +- &'a self, +- condition: &'a PrimaryCondition, +- hw_counter: &'a HardwareCounterCell, +- ) -> Option + 'a>> { +- match condition { +- PrimaryCondition::Condition(field_condition) => { +- let field_key = &field_condition.key; +- let field_indexes = self.field_indexes.get(field_key)?; +- field_indexes +- .iter() +- .find_map(|field_index| field_index.filter(field_condition, hw_counter)) +- } +- PrimaryCondition::Ids(ids) => Some(Box::new(ids.iter().copied())), +- PrimaryCondition::HasVector(_) => None, +- } +- } +- + fn config_path(&self) -> PathBuf { + PayloadConfig::get_config_path(&self.path) + } +@@ -152,6 +115,9 @@ impl StructPayloadIndex { + payload_schema, + &HardwareCounterCell::disposable(), // Internal operation. + )?; ++ for index in indexes.iter_mut() { ++ index.save()?; ++ } + } + + Ok(indexes) +@@ -213,41 +179,6 @@ impl StructPayloadIndex { + Ok(index) + } + +- pub fn build_field_indexes( +- &self, +- field: PayloadKeyTypeRef, +- payload_schema: &PayloadFieldSchema, +- hw_counter: &HardwareCounterCell, +- ) -> OperationResult> { +- let payload_storage = self.payload.borrow(); +- let mut builders = self +- .selector(payload_schema) +- .index_builder(field, payload_schema)?; +- +- for index in &mut builders { +- index.init()?; +- } +- +- payload_storage.iter( +- |point_id, point_payload| { +- let field_value = &point_payload.get_value(field); +- for builder in builders.iter_mut() { +- builder.add_point(point_id, field_value, hw_counter)?; +- } +- Ok(true) +- }, +- hw_counter, +- )?; +- +- builders +- .into_iter() +- .map(|builder| builder.finalize()) +- .collect() +- } +- +- /// Number of available points +- /// +- /// - excludes soft deleted points + pub fn available_point_count(&self) -> usize { + self.id_tracker.borrow().available_point_count() + } +@@ -332,35 +263,110 @@ impl StructPayloadIndex { + } + } + +- pub fn get_telemetry_data(&self) -> Vec { +- self.field_indexes +- .iter() +- .flat_map(|(name, field)| -> Vec { +- field ++ fn estimate_field_condition( ++ &self, ++ condition: &FieldCondition, ++ nested_path: Option<&JsonPath>, ++ hw_counter: &HardwareCounterCell, ++ ) -> Option { ++ let full_path = JsonPath::extend_or_new(nested_path, &condition.key); ++ self.field_indexes.get(&full_path).and_then(|indexes| { ++ // rewrite condition with fullpath to enable cardinality estimation ++ let full_path_condition = FieldCondition { ++ key: full_path, ++ ..condition.clone() ++ }; ++ ++ indexes ++ .iter() ++ .find_map(|index| index.estimate_cardinality(&full_path_condition, hw_counter)) ++ }) ++ } ++ ++ fn query_field<'a>( ++ &'a self, ++ condition: &'a PrimaryCondition, ++ hw_counter: &'a HardwareCounterCell, ++ ) -> Option + 'a>> { ++ match condition { ++ PrimaryCondition::Condition(field_condition) => { ++ let field_key = &field_condition.key; ++ let field_indexes = self.field_indexes.get(field_key)?; ++ field_indexes + .iter() +- .map(|field| field.get_telemetry_data().set_name(name.to_string())) +- .collect() +- }) +- .collect() ++ .find_map(|field_index| field_index.filter(field_condition, hw_counter)) ++ } ++ PrimaryCondition::Ids(ids) => Some(Box::new(ids.iter().copied())), ++ PrimaryCondition::HasVector(_) => None, ++ } ++ } ++ ++ fn config_path(&self) -> PathBuf { ++ PayloadConfig::get_config_path(&self.path) ++ } ++ ++ fn save_config(&self) -> OperationResult<()> { ++ let config_path = self.config_path(); ++ self.config.save(&config_path) ++ } ++ ++ fn load_all_fields(&mut self) -> OperationResult<()> { ++ let mut field_indexes: IndexesMap = Default::default(); ++ ++ for (field, payload_schema) in &self.config.indexed_fields { ++ let field_index = self.load_from_db(field, payload_schema)?; ++ field_indexes.insert(field.clone(), field_index); ++ } ++ self.field_indexes = field_indexes; ++ Ok(()) + } + +- pub fn restore_database_snapshot( ++ fn restore_database_snapshot( + snapshot_path: &Path, + segment_path: &Path, + ) -> OperationResult<()> { + crate::rocksdb_backup::restore(snapshot_path, &segment_path.join("payload_index")) + } + +- fn clear_index_for_point(&mut self, point_id: PointOffsetType) -> OperationResult<()> { +- for (_, field_indexes) in self.field_indexes.iter_mut() { +- for index in field_indexes { +- index.remove_point(point_id)?; ++ /// Select which type of PayloadIndex to use for the field ++ fn selector(&self, payload_schema: &PayloadFieldSchema) -> IndexSelector { ++ let is_on_disk = payload_schema.is_on_disk(); ++ ++ match &self.storage_type { ++ StorageType::Appendable(db) => IndexSelector::RocksDb(IndexSelectorRocksDb { ++ db, ++ is_appendable: true, ++ }), ++ StorageType::NonAppendableRocksDb(db) => { ++ // legacy logic: we keep rocksdb, but load mmap indexes ++ if is_on_disk { ++ IndexSelector::Mmap(IndexSelectorMmap { ++ dir: &self.path, ++ is_on_disk, ++ }) ++ } else { ++ IndexSelector::RocksDb(IndexSelectorRocksDb { ++ db, ++ is_appendable: false, ++ }) ++ } + } ++ StorageType::NonAppendable => IndexSelector::Mmap(IndexSelectorMmap { ++ dir: &self.path, ++ is_on_disk, ++ }), + } +- Ok(()) + } +- pub fn config(&self) -> &PayloadConfig { +- &self.config ++ ++ pub fn optimize_filter( ++ &self, ++ _filter: &Filter, ++ _payload_provider: PayloadProvider, ++ _total: usize, ++ _hw_counter: &HardwareCounterCell, ++ ) -> (Filter, CardinalityEstimation) { ++ // TODO implement optimization logic ++ todo!() + } + + pub fn iter_filtered_points<'a>( +@@ -375,12 +381,15 @@ impl StructPayloadIndex { + if query_cardinality.primary_clauses.is_empty() { + let full_scan_iterator = id_tracker.iter_ids(); + ++ let struct_filtered_context = self.struct_filtered_context(filter, hw_counter); + // Worst case: query expected to return few matches, but index can't be used + let matched_points = + full_scan_iterator.filter(move |i| struct_filtered_context.check(*i)); + + Either::Left(matched_points) + } else { ++ let struct_filtered_context = self.struct_filtered_context(filter, hw_counter); ++ + // CPU-optimized strategy here: points are made unique before applying other filters. + let mut visited_list = self.visited_pool.get(id_tracker.total_point_count()); + +@@ -404,36 +413,6 @@ impl StructPayloadIndex { + } + } + +- /// Select which type of PayloadIndex to use for the field +- fn selector(&self, payload_schema: &PayloadFieldSchema) -> IndexSelector { +- let is_on_disk = payload_schema.is_on_disk(); +- +- match &self.storage_type { +- StorageType::Appendable(db) => IndexSelector::RocksDb(IndexSelectorRocksDb { +- db, +- is_appendable: true, +- }), +- StorageType::NonAppendableRocksDb(db) => { +- // legacy logic: we keep rocksdb, but load mmap indexes +- if is_on_disk { +- IndexSelector::Mmap(IndexSelectorMmap { +- dir: &self.path, +- is_on_disk, +- }) +- } else { +- IndexSelector::RocksDb(IndexSelectorRocksDb { +- db, +- is_appendable: false, +- }) +- } +- } +- StorageType::NonAppendable => IndexSelector::Mmap(IndexSelectorMmap { +- dir: &self.path, +- is_on_disk, +- }), +- } +- } +- + pub fn get_facet_index(&self, key: &JsonPath) -> OperationResult { + self.field_indexes + .get(key) +@@ -478,7 +457,7 @@ impl PayloadIndex for StructPayloadIndex { + self.config.indexed_fields.clone() + } + +- fn build_index( ++ fn set_indexed( + &self, + field: PayloadKeyTypeRef, + payload_schema: &PayloadFieldSchema, +@@ -514,6 +493,7 @@ impl PayloadIndex for StructPayloadIndex { + + fn drop_index(&mut self, field: PayloadKeyTypeRef) -> OperationResult<()> { + self.config.indexed_fields.remove(field); ++ self.save_config()?; + let removed_indexes = self.field_indexes.remove(field); + + if let Some(indexes) = removed_indexes { +@@ -522,10 +502,28 @@ impl PayloadIndex for StructPayloadIndex { + } + } + +- self.save_config()?; + Ok(()) + } + ++ fn build_index( ++ &self, ++ field: PayloadKeyTypeRef, ++ payload_schema: &PayloadFieldSchema, ++ hw_counter: &HardwareCounterCell, ++ ) -> OperationResult>> { ++ if let Some(prev_schema) = self.config.indexed_fields.get(field) { ++ // the field is already indexed with the same schema ++ // no need to rebuild index and to save the config ++ if prev_schema == payload_schema { ++ return Ok(None); ++ } ++ } ++ ++ let indexes = self.build_field_indexes(field, payload_schema, hw_counter)?; ++ ++ Ok(Some(indexes)) ++ } ++ + fn estimate_cardinality( + &self, + query: &Filter, +@@ -592,14 +590,19 @@ impl PayloadIndex for StructPayloadIndex { + None => Box::new(vec![].into_iter()), + Some(indexes) => { + let field_clone = field.to_owned(); +- Box::new(indexes.iter().flat_map(move |field_index| { +- field_index.payload_blocks(threshold, field_clone.clone()) +- })) ++ Box::new( ++ indexes ++ .iter() ++ .map(move |field_index| { ++ field_index.payload_blocks(threshold, field_clone.clone()) ++ }) ++ .flatten(), ++ ) + } + } + } + +- fn overwrite_payload( ++ fn assign_all( + &mut self, + point_id: PointOffsetType, + payload: &Payload, +@@ -723,27 +726,6 @@ impl PayloadIndex for StructPayloadIndex { + }) + } + +- fn infer_payload_type( +- &self, +- key: PayloadKeyTypeRef, +- hw_counter: &HardwareCounterCell, +- ) -> OperationResult> { +- let mut schema = None; +- self.payload.borrow().iter( +- |_id, payload: &Payload| { +- let field_value = payload.get_value(key); +- schema = match field_value.as_slice() { +- [] => None, +- [single] => infer_value_type(single), +- multiple => infer_collection_value_type(multiple.iter().copied()), +- }; +- Ok(false) +- }, +- hw_counter, +- )?; +- Ok(schema) +- } +- + fn take_database_snapshot(&self, path: &Path) -> OperationResult<()> { + match &self.storage_type { + StorageType::Appendable(db) => { +@@ -767,4 +749,39 @@ impl PayloadIndex for StructPayloadIndex { + files.push(self.config_path()); + files + } +-} +\ No newline at end of file ++ ++ fn infer_payload_type( ++ &self, ++ key: PayloadKeyTypeRef, ++ hw_counter: &HardwareCounterCell, ++ ) -> OperationResult> { ++ let mut schema = None; ++ self.payload.borrow().iter( ++ |_id, payload: &Payload| { ++ let field_value = payload.get_value(key); ++ schema = match field_value.as_slice() { ++ [] => None, ++ [single] => infer_value_type(single), ++ multiple => infer_collection_value_type(multiple.iter().copied()), ++ }; ++ Ok(false) ++ }, ++ hw_counter, ++ )?; ++ Ok(schema) ++ } ++ ++ fn get_telemetry_data(&self) -> Vec { ++ self.field_indexes ++ .iter() ++ .flat_map(|(name, field)| -> Vec { ++ field ++ .iter() ++ .map(|field| field.get_telemetry_data().set_name(name.to_string())) ++ .collect() ++ }) ++ .collect() ++ } ++} ++ ++// Continuation with possible impl for FacetIndexEnum or other types if needed, but based on commits, this is the state. In the last commit, there are no further additions beyond the file changes shown. +\ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..5ce1d00b7 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_index_struct_payload_index.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt @@ -0,0 +1,789 @@ +```rust +use std::collections::HashMap; +use std::fs::create_dir_all; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use ahash::AHashSet; +use atomic_refcell::AtomicRefCell; +use common::counter::hardware_counter::HardwareCounterCell; +use common::counter::iterator_hw_measurement::HwMeasurementIteratorExt; +use common::flags::feature_flags; +use common::types::PointOffsetType; +use itertools::Either; +use log::debug; +use parking_lot::RwLock; +use rocksdb::DB; +use schemars::_serde_json::Value; + +use super::field_index::FieldIndexBuilderTrait as _; +use super::field_index::facet_index::FacetIndexEnum; +use super::field_index::index_selector::{IndexSelector, IndexSelectorMmap, IndexSelectorRocksDb}; +use crate::common::Flusher; +use crate::common::operation_error::{OperationError, OperationResult}; +use crate::common::rocksdb_wrapper::open_db_with_existing_cf; +use crate::common::utils::IndexesMap; +use crate::id_tracker::IdTrackerSS; +use crate::index::field_index::{ + CardinalityEstimation, FieldIndex, PayloadBlockCondition, PrimaryCondition, +}; +use crate::index::field_index::index_selector::index_selector; +use crate::index::payload_config::PayloadConfig; +use crate::index::query_estimator::estimate_filter; +use crate::index::query_optimization::payload_provider::PayloadProvider; +use crate::index::struct_filter_context::StructFilterContext; +use crate::index::visited_pool::VisitedPool; +use crate::index::PayloadIndex; +use crate::json_path::JsonPath; +use crate::payload_storage::payload_storage_enum::PayloadStorageEnum; +use crate::payload_storage::{FilterContext, PayloadStorage}; +use crate::telemetry::PayloadIndexTelemetry; +use crate::types::{ + Condition, FieldCondition, Filter, IsEmptyCondition, IsNullCondition, Payload, + PayloadContainer, PayloadField, PayloadFieldSchema, PayloadKeyType, PayloadKeyTypeRef, + PayloadSchemaType, VectorNameBuf, infer_collection_value_type, infer_value_type, +}; +use crate::vector_storage::{VectorStorage, VectorStorageEnum}; + +#[derive(Debug)] +enum StorageType { + Appendable(Arc>), + NonAppendableRocksDb(Arc>), + NonAppendable, +} + +/// `PayloadIndex` implementation, which actually uses index structures for providing faster search +#[derive(Debug)] +pub struct StructPayloadIndex { + /// Payload storage + pub(super) payload: Arc>, + /// Used for `has_id` condition and estimating cardinality + pub(super) id_tracker: Arc>, + /// Vector storages for each field, used for `has_vector` condition + pub(super) vector_storages: HashMap>>, + /// Indexes, associated with fields + pub field_indexes: IndexesMap, + config: PayloadConfig, + /// Root of index persistence dir + path: PathBuf, + /// Used to select unique point ids + visited_pool: VisitedPool, + storage_type: StorageType, +} + +impl StructPayloadIndex { + fn config_path(&self) -> PathBuf { + PayloadConfig::get_config_path(&self.path) + } + + fn save_config(&self) -> OperationResult<()> { + let config_path = self.config_path(); + self.config.save(&config_path) + } + + fn load_all_fields(&mut self) -> OperationResult<()> { + let mut field_indexes: IndexesMap = Default::default(); + + for (field, payload_schema) in &self.config.indexed_fields { + let field_index = self.load_from_db(field, payload_schema)?; + field_indexes.insert(field.clone(), field_index); + } + self.field_indexes = field_indexes; + Ok(()) + } + + fn load_from_db( + &self, + field: PayloadKeyTypeRef, + payload_schema: &PayloadFieldSchema, + ) -> OperationResult> { + let mut indexes = self + .selector(payload_schema) + .new_index(field, payload_schema)?; + + let mut is_loaded = true; + for ref mut index in indexes.iter_mut() { + if !index.load()? { + is_loaded = false; + break; + } + } + if !is_loaded { + debug!("Index for `{field}` was not loaded. Building..."); + // todo(ivan): decide what to do with indexes, which were not loaded + indexes = self.build_field_indexes( + field, + payload_schema, + &HardwareCounterCell::disposable(), // Internal operation. + )?; + for index in indexes.iter_mut() { + index.save()?; + } + } + + Ok(indexes) + } + + pub fn open( + payload: Arc>, + id_tracker: Arc>, + vector_storages: HashMap>>, + path: &Path, + is_appendable: bool, + ) -> OperationResult { + create_dir_all(path)?; + let config_path = PayloadConfig::get_config_path(path); + let config = if config_path.exists() { + PayloadConfig::load(&config_path)? + } else { + let mut new_config = PayloadConfig::default(); + if feature_flags().payload_index_skip_rocksdb && !is_appendable { + new_config.skip_rocksdb = Some(true); + } + new_config + }; + + let skip_rocksdb = config.skip_rocksdb.unwrap_or(false); + + let storage_type = if is_appendable { + let db = open_db_with_existing_cf(path).map_err(|err| { + OperationError::service_error(format!("RocksDB open error: {err}")) + })?; + StorageType::Appendable(db) + } else if skip_rocksdb { + StorageType::NonAppendable + } else { + let db = open_db_with_existing_cf(path).map_err(|err| { + OperationError::service_error(format!("RocksDB open error: {err}")) + })?; + StorageType::NonAppendableRocksDb(db) + }; + + let mut index = StructPayloadIndex { + payload, + id_tracker, + vector_storages, + field_indexes: Default::default(), + config, + path: path.to_owned(), + visited_pool: Default::default(), + storage_type, + }; + + if !index.config_path().exists() { + // Save default config + index.save_config()?; + } + + index.load_all_fields()?; + + Ok(index) + } + + pub fn available_point_count(&self) -> usize { + self.id_tracker.borrow().available_point_count() + } + + pub fn struct_filtered_context<'a>( + &'a self, + filter: &'a Filter, + hw_counter: &HardwareCounterCell, + ) -> StructFilterContext<'a> { + let payload_provider = PayloadProvider::new(self.payload.clone()); + + let (optimized_filter, _) = self.optimize_filter( + filter, + payload_provider, + self.available_point_count(), + hw_counter, + ); + + StructFilterContext::new(optimized_filter) + } + + pub(super) fn condition_cardinality( + &self, + condition: &Condition, + nested_path: Option<&JsonPath>, + hw_counter: &HardwareCounterCell, + ) -> CardinalityEstimation { + match condition { + Condition::Filter(_) => panic!("Unexpected branching"), + Condition::Nested(nested) => { + // propagate complete nested path in case of multiple nested layers + let full_path = JsonPath::extend_or_new(nested_path, &nested.array_key()); + self.estimate_nested_cardinality(nested.filter(), &full_path, hw_counter) + } + Condition::IsEmpty(IsEmptyCondition { is_empty: field }) => { + let available_points = self.available_point_count(); + let condition = FieldCondition::new_is_empty(field.key.clone()); + + self.estimate_field_condition(&condition, nested_path, hw_counter) + .unwrap_or_else(|| CardinalityEstimation::unknown(available_points)) + } + Condition::IsNull(IsNullCondition { is_null: field }) => { + let available_points = self.available_point_count(); + let condition = FieldCondition::new_is_null(field.key.clone()); + + self.estimate_field_condition(&condition, nested_path, hw_counter) + .unwrap_or_else(|| CardinalityEstimation::unknown(available_points)) + } + Condition::HasId(has_id) => { + let id_tracker_ref = self.id_tracker.borrow(); + let mapped_ids: AHashSet = has_id + .has_id + .iter() + .filter_map(|external_id| id_tracker_ref.internal_id(*external_id)) + .collect(); + let num_ids = mapped_ids.len(); + CardinalityEstimation { + primary_clauses: vec![PrimaryCondition::Ids(mapped_ids)], + min: num_ids, + exp: num_ids, + max: num_ids, + } + } + Condition::HasVector(has_vectors) => { + if let Some(vector_storage) = self.vector_storages.get(&has_vectors.has_vector) { + let vector_storage = vector_storage.borrow(); + let vectors = vector_storage.available_vector_count(); + CardinalityEstimation::exact(vectors).with_primary_clause( + PrimaryCondition::HasVector(has_vectors.has_vector.clone()), + ) + } else { + CardinalityEstimation::exact(0) + } + } + Condition::Field(field_condition) => self + .estimate_field_condition(field_condition, nested_path, hw_counter) + .unwrap_or_else(|| CardinalityEstimation::unknown(self.available_point_count())), + + Condition::CustomIdChecker(cond) => { + cond.estimate_cardinality(self.id_tracker.borrow().available_point_count()) + } + } + } + + fn estimate_field_condition( + &self, + condition: &FieldCondition, + nested_path: Option<&JsonPath>, + hw_counter: &HardwareCounterCell, + ) -> Option { + let full_path = JsonPath::extend_or_new(nested_path, &condition.key); + self.field_indexes.get(&full_path).and_then(|indexes| { + // rewrite condition with fullpath to enable cardinality estimation + let full_path_condition = FieldCondition { + key: full_path, + ..condition.clone() + }; + + indexes + .iter() + .find_map(|index| index.estimate_cardinality(&full_path_condition, hw_counter)) + }) + } + + fn query_field<'a>( + &'a self, + condition: &'a PrimaryCondition, + hw_counter: &'a HardwareCounterCell, + ) -> Option + 'a>> { + match condition { + PrimaryCondition::Condition(field_condition) => { + let field_key = &field_condition.key; + let field_indexes = self.field_indexes.get(field_key)?; + field_indexes + .iter() + .find_map(|field_index| field_index.filter(field_condition, hw_counter)) + } + PrimaryCondition::Ids(ids) => Some(Box::new(ids.iter().copied())), + PrimaryCondition::HasVector(_) => None, + } + } + + fn config_path(&self) -> PathBuf { + PayloadConfig::get_config_path(&self.path) + } + + fn save_config(&self) -> OperationResult<()> { + let config_path = self.config_path(); + self.config.save(&config_path) + } + + fn load_all_fields(&mut self) -> OperationResult<()> { + let mut field_indexes: IndexesMap = Default::default(); + + for (field, payload_schema) in &self.config.indexed_fields { + let field_index = self.load_from_db(field, payload_schema)?; + field_indexes.insert(field.clone(), field_index); + } + self.field_indexes = field_indexes; + Ok(()) + } + + fn restore_database_snapshot( + snapshot_path: &Path, + segment_path: &Path, + ) -> OperationResult<()> { + crate::rocksdb_backup::restore(snapshot_path, &segment_path.join("payload_index")) + } + + /// Select which type of PayloadIndex to use for the field + fn selector(&self, payload_schema: &PayloadFieldSchema) -> IndexSelector { + let is_on_disk = payload_schema.is_on_disk(); + + match &self.storage_type { + StorageType::Appendable(db) => IndexSelector::RocksDb(IndexSelectorRocksDb { + db, + is_appendable: true, + }), + StorageType::NonAppendableRocksDb(db) => { + // legacy logic: we keep rocksdb, but load mmap indexes + if is_on_disk { + IndexSelector::Mmap(IndexSelectorMmap { + dir: &self.path, + is_on_disk, + }) + } else { + IndexSelector::RocksDb(IndexSelectorRocksDb { + db, + is_appendable: false, + }) + } + } + StorageType::NonAppendable => IndexSelector::Mmap(IndexSelectorMmap { + dir: &self.path, + is_on_disk, + }), + } + } + + pub fn optimize_filter( + &self, + _filter: &Filter, + _payload_provider: PayloadProvider, + _total: usize, + _hw_counter: &HardwareCounterCell, + ) -> (Filter, CardinalityEstimation) { + // TODO implement optimization logic + todo!() + } + + pub fn iter_filtered_points<'a>( + &'a self, + filter: &'a Filter, + id_tracker: &'a IdTrackerSS, + query_cardinality: &'a CardinalityEstimation, + hw_counter: &'a HardwareCounterCell, + ) -> impl Iterator + 'a { + let struct_filtered_context = self.struct_filtered_context(filter, hw_counter); + + if query_cardinality.primary_clauses.is_empty() { + let full_scan_iterator = id_tracker.iter_ids(); + + let struct_filtered_context = self.struct_filtered_context(filter, hw_counter); + // Worst case: query expected to return few matches, but index can't be used + let matched_points = + full_scan_iterator.filter(move |i| struct_filtered_context.check(*i)); + + Either::Left(matched_points) + } else { + let struct_filtered_context = self.struct_filtered_context(filter, hw_counter); + + // CPU-optimized strategy here: points are made unique before applying other filters. + let mut visited_list = self.visited_pool.get(id_tracker.total_point_count()); + + let iter = query_cardinality + .primary_clauses + .iter() + .flat_map(move |clause| { + self.query_field(clause, hw_counter).unwrap_or_else(|| { + // index is not built + Box::new(id_tracker.iter_ids().measure_hw_with_cell( + hw_counter, + size_of::(), + |i| i.cpu_counter(), + )) + }) + }) + .filter(move |&id| !visited_list.check_and_update_visited(id)) + .filter(move |&i| struct_filtered_context.check(i)); + + Either::Right(iter) + } + } + + pub fn get_facet_index(&self, key: &JsonPath) -> OperationResult { + self.field_indexes + .get(key) + .and_then(|index| index.iter().find_map(|index| index.as_facet_index())) + .ok_or_else(|| OperationError::MissingMapIndexForFacet { + key: key.to_string(), + }) + } + + pub fn populate(&self) -> OperationResult<()> { + for (_, field_indexes) in self.field_indexes.iter() { + for index in field_indexes { + index.populate()?; + } + } + Ok(()) + } + + pub fn clear_cache(&self) -> OperationResult<()> { + for (_, field_indexes) in self.field_indexes.iter() { + for index in field_indexes { + index.clear_cache()?; + } + } + Ok(()) + } + + pub fn clear_cache_if_on_disk(&self) -> OperationResult<()> { + for (_, field_indexes) in self.field_indexes.iter() { + for index in field_indexes { + if index.is_on_disk() { + index.clear_cache()?; + } + } + } + Ok(()) + } +} + +impl PayloadIndex for StructPayloadIndex { + fn indexed_fields(&self) -> HashMap { + self.config.indexed_fields.clone() + } + + fn set_indexed( + &self, + field: PayloadKeyTypeRef, + payload_schema: &PayloadFieldSchema, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + if let Some(prev_schema) = self.config.indexed_fields.get(field) { + // the field is already indexed with the same schema + // no need to rebuild index and to save the config + if prev_schema == payload_schema { + return Ok(None); + } + } + + let indexes = self.build_field_indexes(field, payload_schema, hw_counter)?; + + Ok(Some(indexes)) + } + + fn apply_index( + &mut self, + field: PayloadKeyType, + payload_schema: PayloadFieldSchema, + field_index: Vec, + ) -> OperationResult<()> { + self.field_indexes.insert(field.clone(), field_index); + + self.config.indexed_fields.insert(field, payload_schema); + + self.save_config()?; + + Ok(()) + } + + fn drop_index(&mut self, field: PayloadKeyTypeRef) -> OperationResult<()> { + self.config.indexed_fields.remove(field); + self.save_config()?; + let removed_indexes = self.field_indexes.remove(field); + + if let Some(indexes) = removed_indexes { + for index in indexes { + index.cleanup()?; + } + } + + Ok(()) + } + + fn build_index( + &self, + field: PayloadKeyTypeRef, + payload_schema: &PayloadFieldSchema, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + if let Some(prev_schema) = self.config.indexed_fields.get(field) { + // the field is already indexed with the same schema + // no need to rebuild index and to save the config + if prev_schema == payload_schema { + return Ok(None); + } + } + + let indexes = self.build_field_indexes(field, payload_schema, hw_counter)?; + + Ok(Some(indexes)) + } + + fn estimate_cardinality( + &self, + query: &Filter, + hw_counter: &HardwareCounterCell, + ) -> CardinalityEstimation { + let available_points = self.available_point_count(); + let estimator = + |condition: &Condition| self.condition_cardinality(condition, None, hw_counter); + estimate_filter(&estimator, query, available_points) + } + + fn estimate_nested_cardinality( + &self, + query: &Filter, + nested_path: &JsonPath, + hw_counter: &HardwareCounterCell, + ) -> CardinalityEstimation { + let available_points = self.available_point_count(); + let estimator = |condition: &Condition| { + self.condition_cardinality(condition, Some(nested_path), hw_counter) + }; + estimate_filter(&estimator, query, available_points) + } + + fn query_points( + &self, + query: &Filter, + hw_counter: &HardwareCounterCell, + ) -> Vec { + // Assume query is already estimated to be small enough so we can iterate over all matched ids + let query_cardinality = self.estimate_cardinality(query, hw_counter); + let id_tracker = self.id_tracker.borrow(); + self.iter_filtered_points(query, &*id_tracker, &query_cardinality, hw_counter) + .collect() + } + + fn indexed_points(&self, field: PayloadKeyTypeRef) -> usize { + self.field_indexes.get(field).map_or(0, |indexes| { + // Assume that multiple field indexes are applied to the same data type, + // so the points indexed with those indexes are the same. + // We will return minimal number as a worst case, to highlight possible errors in the index early. + indexes + .iter() + .map(|index| index.count_indexed_points()) + .min() + .unwrap_or(0) + }) + } + + fn filter_context<'a>( + &'a self, + filter: &'a Filter, + hw_counter: &HardwareCounterCell, + ) -> Box { + Box::new(self.struct_filtered_context(filter, hw_counter)) + } + + fn payload_blocks( + &self, + field: PayloadKeyTypeRef, + threshold: usize, + ) -> Box + '_> { + match self.field_indexes.get(field) { + None => Box::new(vec![].into_iter()), + Some(indexes) => { + let field_clone = field.to_owned(); + Box::new( + indexes + .iter() + .map(move |field_index| { + field_index.payload_blocks(threshold, field_clone.clone()) + }) + .flatten(), + ) + } + } + } + + fn assign_all( + &mut self, + point_id: PointOffsetType, + payload: &Payload, + hw_counter: &HardwareCounterCell, + ) -> OperationResult<()> { + self.payload + .borrow_mut() + .overwrite(point_id, payload, hw_counter)?; + + for (field, field_index) in &mut self.field_indexes { + let field_value = payload.get_value(field); + if !field_value.is_empty() { + for index in field_index { + index.add_point(point_id, &field_value, hw_counter)?; + } + } else { + for index in field_index { + index.remove_point(point_id)?; + } + } + } + Ok(()) + } + + fn set_payload( + &mut self, + point_id: PointOffsetType, + payload: &Payload, + key: &Option, + hw_counter: &HardwareCounterCell, + ) -> OperationResult<()> { + if let Some(key) = key { + self.payload + .borrow_mut() + .set_by_key(point_id, payload, key, hw_counter)?; + } else { + self.payload + .borrow_mut() + .set(point_id, payload, hw_counter)?; + }; + + let updated_payload = self.get_payload(point_id, hw_counter)?; + for (field, field_index) in &mut self.field_indexes { + if !field.is_affected_by_value_set(&payload.0, key.as_ref()) { + continue; + } + let field_value = updated_payload.get_value(field); + if !field_value.is_empty() { + for index in field_index { + index.add_point(point_id, &field_value, hw_counter)?; + } + } else { + for index in field_index { + index.remove_point(point_id)?; + } + } + } + Ok(()) + } + + fn get_payload( + &self, + point_id: PointOffsetType, + hw_counter: &HardwareCounterCell, + ) -> OperationResult { + self.payload.borrow().get(point_id, hw_counter) + } + + fn delete_payload( + &mut self, + point_id: PointOffsetType, + key: PayloadKeyTypeRef, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + if let Some(indexes) = self.field_indexes.get_mut(key) { + for index in indexes { + index.remove_point(point_id)?; + } + } + self.payload.borrow_mut().delete(point_id, key, hw_counter) + } + + fn clear_payload( + &mut self, + point_id: PointOffsetType, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + self.clear_index_for_point(point_id)?; + self.payload.borrow_mut().clear(point_id, hw_counter) + } + + fn flusher(&self) -> Flusher { + let mut flushers = Vec::new(); + for field_indexes in self.field_indexes.values() { + for index in field_indexes { + flushers.push(index.flusher()); + } + } + flushers.push(self.payload.borrow().flusher()); + Box::new(move || { + for flusher in flushers { + match flusher() { + Ok(_) => {} + Err(OperationError::RocksDbColumnFamilyNotFound { name }) => { + // It is possible, that the index was removed during the flush by user or another thread. + // In this case, non-existing column family is not an error, but an expected behavior. + + // Still we want to log this event, for potential debugging. + log::warn!( + "Flush: RocksDB cf_handle error: Cannot find column family {name}. Assume index is removed.", + ); + } + Err(err) => { + return Err(OperationError::service_error(format!( + "Failed to flush payload_index: {err}" + ))); + } + } + } + Ok(()) + }) + } + + fn take_database_snapshot(&self, path: &Path) -> OperationResult<()> { + match &self.storage_type { + StorageType::Appendable(db) => { + let db_guard = db.read(); + crate::rocksdb_backup::create(&db_guard, path) + } + StorageType::NonAppendableRocksDb(db) => { + let db_guard = db.read(); + crate::rocksdb_backup::create(&db_guard, path) + } + StorageType::NonAppendable => Ok(()), + } + } + + fn files(&self) -> Vec { + let mut files = self + .field_indexes + .values() + .flat_map(|indexes| indexes.iter().flat_map(|index| index.files().into_iter())) + .collect::>(); + files.push(self.config_path()); + files + } + + fn infer_payload_type( + &self, + key: PayloadKeyTypeRef, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + let mut schema = None; + self.payload.borrow().iter( + |_id, payload: &Payload| { + let field_value = payload.get_value(key); + schema = match field_value.as_slice() { + [] => None, + [single] => infer_value_type(single), + multiple => infer_collection_value_type(multiple.iter().copied()), + }; + Ok(false) + }, + hw_counter, + )?; + Ok(schema) + } + + fn get_telemetry_data(&self) -> Vec { + self.field_indexes + .iter() + .flat_map(|(name, field)| -> Vec { + field + .iter() + .map(|field| field.get_telemetry_data().set_name(name.to_string())) + .collect() + }) + .collect() + } +} + +// Continuation with possible impl for FacetIndexEnum or other types if needed, but based on commits, this is the state. In the last commit, there are no further additions beyond the file changes shown. +``` \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt b/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt new file mode 100644 index 000000000..c6f86b6a5 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt @@ -0,0 +1,290 @@ +use common::types::ScoreType; + +use super::metric::{Metric, MetricPostProcessing}; +#[cfg(target_arch = "x86_64")] +use super::simple_avx::*; +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +use super::simple_neon::*; +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +use super::simple_sse::*; +use super::tools::is_length_zero_or_normalized; +use crate::data_types::vectors::{DenseVector, VectorElementType}; +use crate::types::Distance; + +#[cfg(target_arch = "x86_64")] +pub(crate) const MIN_DIM_SIZE_AVX: usize = 32; + +#[cfg(any( + target_arch = "x86", + target_arch = "x86_64", + all(target_arch = "aarch64", target_feature = "neon") +))] +pub(crate) const MIN_DIM_SIZE_SIMD: usize = 16; + +#[derive(Clone)] +pub struct DotProductMetric; + +#[derive(Clone)] +pub struct CosineMetric; + +/// Equivalent to DotProductMetric with normalization of the vectors in preprocessing. +#[derive(Clone)] +pub struct EuclidMetric; + +#[derive(Clone)] +pub struct ManhattanMetric; + +impl Metric for EuclidMetric { + fn distance() -> Distance { + Distance::Euclid + } + + fn similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx") + && is_x86_feature_detected!("fma") + && v1.len() >= MIN_DIM_SIZE_AVX + { + return unsafe { euclid_similarity_avx(v1, v2) }; + } + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if is_x86_feature_detected!("sse") && v1.len() >= MIN_DIM_SIZE_SIMD { + return unsafe { euclid_similarity_sse(v1, v2) }; + } + } + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("neon") + && v1.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { euclid_similarity_neon(v1, v2) }; + } + } + + euclid_similarity(v1, v2) + } + + fn preprocess(vector: DenseVector) -> DenseVector { + vector + } +} + +impl MetricPostProcessing for EuclidMetric { + fn postprocess(score: ScoreType) -> ScoreType { + score.abs().sqrt() + } +} + +impl Metric for ManhattanMetric { + fn distance() -> Distance { + Distance::Manhattan + } + + fn similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx") + && is_x86_feature_detected!("fma") + && v1.len() >= MIN_DIM_SIZE_AVX + { + return unsafe { manhattan_similarity_avx(v1, v2) }; + } + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if is_x86_feature_detected!("sse") && v1.len() >= MIN_DIM_SIZE_SIMD { + return unsafe { manhattan_similarity_sse(v1, v2) }; + } + } + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("neon") + && v1.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { manhattan_similarity_neon(v1, v2) }; + } + } + + manhattan_similarity(v1, v2) + } + + fn preprocess(vector: DenseVector) -> DenseVector { + vector + } +} + +impl MetricPostProcessing for ManhattanMetric { + fn postprocess(score: ScoreType) -> ScoreType { + score.abs() + } +} + +impl Metric for DotProductMetric { + fn distance() -> Distance { + Distance::Dot + } + + fn similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx") + && is_x86_feature_detected!("fma") + && v1.len() >= MIN_DIM_SIZE_AVX + { + return unsafe { dot_similarity_avx(v1, v2) }; + } + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if is_x86_feature_detected!("sse") && v1.len() >= MIN_DIM_SIZE_SIMD { + return unsafe { dot_similarity_sse(v1, v2) }; + } + } + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("neon") + && v1.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { dot_similarity_neon(v1, v2) }; + } + } + + dot_similarity(v1, v2) + } + + fn preprocess(vector: DenseVector) -> DenseVector { + vector + } +} + +impl MetricPostProcessing for DotProductMetric { + fn postprocess(score: ScoreType) -> ScoreType { + score + } +} + +/// Equivalent to DotProductMetric with normalization of the vectors in preprocessing. +impl Metric for CosineMetric { + fn distance() -> Distance { + Distance::Cosine + } + + fn similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + DotProductMetric::similarity(v1, v2) + } + + fn preprocess(vector: DenseVector) -> DenseVector { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx") + && is_x86_feature_detected!("fma") + && vector.len() >= MIN_DIM_SIZE_AVX + { + return unsafe { cosine_preprocess_avx(vector) }; + } + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if is_x86_feature_detected!("sse") && vector.len() >= MIN_DIM_SIZE_SIMD { + return unsafe { cosine_preprocess_sse(vector) }; + } + } + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("neon") + && vector.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { cosine_preprocess_neon(vector) }; + } + } + + cosine_preprocess(vector) + } +} + +impl MetricPostProcessing for CosineMetric { + fn postprocess(score: ScoreType) -> ScoreType { + score + } +} + +pub fn euclid_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + -v1 + .iter() + .zip(v2) + .map(|(a, b)| (a - b).powi(2)) + .sum::() +} + +pub fn manhattan_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + -v1 + .iter() + .zip(v2) + .map(|(a, b)| (a - b).abs()) + .sum::() +} + +pub fn cosine_preprocess(vector: DenseVector) -> DenseVector { + let mut length: f32 = vector.iter().map(|x| x * x).sum(); + if is_length_zero_or_normalized(length) { + return vector; + } + length = length.sqrt(); + vector.iter().map(|x| x / length).collect() +} + +pub fn dot_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + v1.iter().zip(v2).map(|(a, b)| a * b).sum() +} + +#[cfg(test)] +mod tests { + use rand::Rng; + + use super::*; + + #[test] + fn test_cosine_preprocessing() { + let res = + >::preprocess(vec![0.0, 0.0, 0.0, 0.0]); + assert_eq!(res, vec![0.0, 0.0, 0.0, 0.0]); + } + + /// If we preprocess a vector multiple times, we expect the same result. + /// Renormalization should not produce something different. + #[test] + fn test_cosine_stable_preprocessing() { + const DIM: usize = 1500; + const ATTEMPTS: usize = 100; + + let mut rng = rand::rng(); + + for attempt in 0..ATTEMPTS { + let range = rng.random_range(-2.5..=0.0)..=rng.random_range(0.0..2.5); + let vector: Vec<_> = (0..DIM).map(|_| rng.random_range(range.clone())).collect(); + + // Preprocess and re-preprocess + let preprocess1 = + >::preprocess(vector); + let preprocess2: DenseVector = + >::preprocess(preprocess1.clone()); + + // All following preprocess attempts must be the same + assert_eq!( + preprocess1, preprocess2, + "renormalization is not stable (vector #{attempt})" + ); + } + } +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json b/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json new file mode 100644 index 000000000..7ac0225e4 --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/metadata.json @@ -0,0 +1,25 @@ +{ + "model": "x-ai/grok-code-fast-1", + "benchmark_case": "qdrant_lib_segment_src_spaces_simple.rs", + "prompts_dir": "locodiff-250425/prompts", + "prompt_file": "locodiff-250425/prompts/qdrant_lib_segment_src_spaces_simple.rs_prompt.txt", + "expected_file": "locodiff-250425/prompts/qdrant_lib_segment_src_spaces_simple.rs_expectedoutput.txt", + "timestamp_utc": "2025-09-09T15:36:42.660554+00:00", + "success": false, + "error": "Output mismatch", + "api_error": false, + "raw_response_length": 8434, + "extracted_output_length": 8422, + "expected_output_length": 8190, + "results_dir": "locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642", + "generation_id": "gen-1757432203-jJJ0PFW7AjlmAqSa0NX6", + "cost_usd": 0.01404804, + "prompt_tokens": 47947, + "completion_tokens": 2450, + "total_tokens": 50397, + "native_prompt_tokens": 47553, + "native_completion_tokens": 3048, + "native_finish_reason": "stop", + "stats_error": null, + "native_tokens_reasoning": 937 +} \ No newline at end of file diff --git a/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff b/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff new file mode 100644 index 000000000..09936939b --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/output.diff @@ -0,0 +1,94 @@ + +index 6667b18e6..c6f86b6a5 100644 +--- a/qdrant_lib_segment_src_spaces_simple.rs_expectedoutput.txt (expected):tmp/tmpo0vqmyn1_expected.txt ++++ b/qdrant_lib_segment_src_spaces_simple.rs_extracted.txt (actual):tmp/tmpd4q9i2p7_actual.txt +@@ -27,6 +27,7 @@ pub struct DotProductMetric; + #[derive(Clone)] + pub struct CosineMetric; + ++/// Equivalent to DotProductMetric with normalization of the vectors in preprocessing. + #[derive(Clone)] + pub struct EuclidMetric; + +@@ -58,7 +59,9 @@ impl Metric for EuclidMetric { + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { +- if std::arch::is_aarch64_feature_detected!("neon") && v1.len() >= MIN_DIM_SIZE_SIMD { ++ if std::arch::is_aarch64_feature_detected!("neon") ++ && v1.len() >= MIN_DIM_SIZE_SIMD ++ { + return unsafe { euclid_similarity_neon(v1, v2) }; + } + } +@@ -102,7 +105,9 @@ impl Metric for ManhattanMetric { + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { +- if std::arch::is_aarch64_feature_detected!("neon") && v1.len() >= MIN_DIM_SIZE_SIMD { ++ if std::arch::is_aarch64_feature_detected!("neon") ++ && v1.len() >= MIN_DIM_SIZE_SIMD ++ { + return unsafe { manhattan_similarity_neon(v1, v2) }; + } + } +@@ -146,7 +151,9 @@ impl Metric for DotProductMetric { + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { +- if std::arch::is_aarch64_feature_detected!("neon") && v1.len() >= MIN_DIM_SIZE_SIMD { ++ if std::arch::is_aarch64_feature_detected!("neon") ++ && v1.len() >= MIN_DIM_SIZE_SIMD ++ { + return unsafe { dot_similarity_neon(v1, v2) }; + } + } +@@ -195,7 +202,8 @@ impl Metric for CosineMetric { + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { +- if std::arch::is_aarch64_feature_detected!("neon") && vector.len() >= MIN_DIM_SIZE_SIMD ++ if std::arch::is_aarch64_feature_detected!("neon") ++ && vector.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { cosine_preprocess_neon(vector) }; + } +@@ -212,14 +220,16 @@ impl MetricPostProcessing for CosineMetric { + } + + pub fn euclid_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { +- -v1.iter() ++ -v1 ++ .iter() + .zip(v2) + .map(|(a, b)| (a - b).powi(2)) + .sum::() + } + + pub fn manhattan_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { +- -v1.iter() ++ -v1 ++ .iter() + .zip(v2) + .map(|(a, b)| (a - b).abs()) + .sum::() +@@ -246,7 +256,8 @@ mod tests { + + #[test] + fn test_cosine_preprocessing() { +- let res = >::preprocess(vec![0.0, 0.0, 0.0, 0.0]); ++ let res = ++ >::preprocess(vec![0.0, 0.0, 0.0, 0.0]); + assert_eq!(res, vec![0.0, 0.0, 0.0, 0.0]); + } + +@@ -264,7 +275,8 @@ mod tests { + let vector: Vec<_> = (0..DIM).map(|_| rng.random_range(range.clone())).collect(); + + // Preprocess and re-preprocess +- let preprocess1 = >::preprocess(vector); ++ let preprocess1 = ++ >::preprocess(vector); + let preprocess2: DenseVector = + >::preprocess(preprocess1.clone()); + diff --git a/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt b/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt new file mode 100644 index 000000000..4ac8871aa --- /dev/null +++ b/locodiff-250425/results/qdrant_lib_segment_src_spaces_simple.rs/x-ai_grok-code-fast-1/20250909_153642/raw_response.txt @@ -0,0 +1,292 @@ +```rust +use common::types::ScoreType; + +use super::metric::{Metric, MetricPostProcessing}; +#[cfg(target_arch = "x86_64")] +use super::simple_avx::*; +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +use super::simple_neon::*; +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +use super::simple_sse::*; +use super::tools::is_length_zero_or_normalized; +use crate::data_types::vectors::{DenseVector, VectorElementType}; +use crate::types::Distance; + +#[cfg(target_arch = "x86_64")] +pub(crate) const MIN_DIM_SIZE_AVX: usize = 32; + +#[cfg(any( + target_arch = "x86", + target_arch = "x86_64", + all(target_arch = "aarch64", target_feature = "neon") +))] +pub(crate) const MIN_DIM_SIZE_SIMD: usize = 16; + +#[derive(Clone)] +pub struct DotProductMetric; + +#[derive(Clone)] +pub struct CosineMetric; + +/// Equivalent to DotProductMetric with normalization of the vectors in preprocessing. +#[derive(Clone)] +pub struct EuclidMetric; + +#[derive(Clone)] +pub struct ManhattanMetric; + +impl Metric for EuclidMetric { + fn distance() -> Distance { + Distance::Euclid + } + + fn similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx") + && is_x86_feature_detected!("fma") + && v1.len() >= MIN_DIM_SIZE_AVX + { + return unsafe { euclid_similarity_avx(v1, v2) }; + } + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if is_x86_feature_detected!("sse") && v1.len() >= MIN_DIM_SIZE_SIMD { + return unsafe { euclid_similarity_sse(v1, v2) }; + } + } + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("neon") + && v1.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { euclid_similarity_neon(v1, v2) }; + } + } + + euclid_similarity(v1, v2) + } + + fn preprocess(vector: DenseVector) -> DenseVector { + vector + } +} + +impl MetricPostProcessing for EuclidMetric { + fn postprocess(score: ScoreType) -> ScoreType { + score.abs().sqrt() + } +} + +impl Metric for ManhattanMetric { + fn distance() -> Distance { + Distance::Manhattan + } + + fn similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx") + && is_x86_feature_detected!("fma") + && v1.len() >= MIN_DIM_SIZE_AVX + { + return unsafe { manhattan_similarity_avx(v1, v2) }; + } + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if is_x86_feature_detected!("sse") && v1.len() >= MIN_DIM_SIZE_SIMD { + return unsafe { manhattan_similarity_sse(v1, v2) }; + } + } + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("neon") + && v1.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { manhattan_similarity_neon(v1, v2) }; + } + } + + manhattan_similarity(v1, v2) + } + + fn preprocess(vector: DenseVector) -> DenseVector { + vector + } +} + +impl MetricPostProcessing for ManhattanMetric { + fn postprocess(score: ScoreType) -> ScoreType { + score.abs() + } +} + +impl Metric for DotProductMetric { + fn distance() -> Distance { + Distance::Dot + } + + fn similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx") + && is_x86_feature_detected!("fma") + && v1.len() >= MIN_DIM_SIZE_AVX + { + return unsafe { dot_similarity_avx(v1, v2) }; + } + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if is_x86_feature_detected!("sse") && v1.len() >= MIN_DIM_SIZE_SIMD { + return unsafe { dot_similarity_sse(v1, v2) }; + } + } + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("neon") + && v1.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { dot_similarity_neon(v1, v2) }; + } + } + + dot_similarity(v1, v2) + } + + fn preprocess(vector: DenseVector) -> DenseVector { + vector + } +} + +impl MetricPostProcessing for DotProductMetric { + fn postprocess(score: ScoreType) -> ScoreType { + score + } +} + +/// Equivalent to DotProductMetric with normalization of the vectors in preprocessing. +impl Metric for CosineMetric { + fn distance() -> Distance { + Distance::Cosine + } + + fn similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + DotProductMetric::similarity(v1, v2) + } + + fn preprocess(vector: DenseVector) -> DenseVector { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx") + && is_x86_feature_detected!("fma") + && vector.len() >= MIN_DIM_SIZE_AVX + { + return unsafe { cosine_preprocess_avx(vector) }; + } + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if is_x86_feature_detected!("sse") && vector.len() >= MIN_DIM_SIZE_SIMD { + return unsafe { cosine_preprocess_sse(vector) }; + } + } + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("neon") + && vector.len() >= MIN_DIM_SIZE_SIMD + { + return unsafe { cosine_preprocess_neon(vector) }; + } + } + + cosine_preprocess(vector) + } +} + +impl MetricPostProcessing for CosineMetric { + fn postprocess(score: ScoreType) -> ScoreType { + score + } +} + +pub fn euclid_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + -v1 + .iter() + .zip(v2) + .map(|(a, b)| (a - b).powi(2)) + .sum::() +} + +pub fn manhattan_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + -v1 + .iter() + .zip(v2) + .map(|(a, b)| (a - b).abs()) + .sum::() +} + +pub fn cosine_preprocess(vector: DenseVector) -> DenseVector { + let mut length: f32 = vector.iter().map(|x| x * x).sum(); + if is_length_zero_or_normalized(length) { + return vector; + } + length = length.sqrt(); + vector.iter().map(|x| x / length).collect() +} + +pub fn dot_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType { + v1.iter().zip(v2).map(|(a, b)| a * b).sum() +} + +#[cfg(test)] +mod tests { + use rand::Rng; + + use super::*; + + #[test] + fn test_cosine_preprocessing() { + let res = + >::preprocess(vec![0.0, 0.0, 0.0, 0.0]); + assert_eq!(res, vec![0.0, 0.0, 0.0, 0.0]); + } + + /// If we preprocess a vector multiple times, we expect the same result. + /// Renormalization should not produce something different. + #[test] + fn test_cosine_stable_preprocessing() { + const DIM: usize = 1500; + const ATTEMPTS: usize = 100; + + let mut rng = rand::rng(); + + for attempt in 0..ATTEMPTS { + let range = rng.random_range(-2.5..=0.0)..=rng.random_range(0.0..2.5); + let vector: Vec<_> = (0..DIM).map(|_| rng.random_range(range.clone())).collect(); + + // Preprocess and re-preprocess + let preprocess1 = + >::preprocess(vector); + let preprocess2: DenseVector = + >::preprocess(preprocess1.clone()); + + // All following preprocess attempts must be the same + assert_eq!( + preprocess1, preprocess2, + "renormalization is not stable (vector #{attempt})" + ); + } + } +} +``` \ No newline at end of file diff --git a/locodiff-250425/results/react_packages_react-dom_src___tests___DOMPropertyOperations-test.js/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt b/locodiff-250425/results/react_packages_react-dom_src___tests___DOMPropertyOperations-test.js/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt new file mode 100644 index 000000000..4e6c14209 --- /dev/null +++ b/locodiff-250425/results/react_packages_react-dom_src___tests___DOMPropertyOperations-test.js/x-ai_grok-code-fast-1/20250909_153642/extracted_output.txt @@ -0,0 +1,834 @@ +/** + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + * + * @emails react-core + */ + +'use strict'; + +// Set by `yarn test-fire`. +const {disableInputAttributeSyncing} = require('shared/ReactFeatureFlags'); + +describe('DOMPropertyOperations', () => { + let React; + let ReactDOMClient; + let act; + let assertConsoleErrorDev; + + beforeEach(() => { + jest.resetModules(); + React = require('react'); + ReactDOMClient = require('react-dom/client'); + ({act, assertConsoleErrorDev} = require('internal-test-utils')); + }); + + // Sets a value in a way that React doesn't see, + // so that a subsequent "change" event will trigger the event handler. + const setUntrackedValue = Object.getOwnPropertyDescriptor( + HTMLInputElement.prototype, + 'value', + ).set; + const setUntrackedChecked = Object.getOwnPropertyDescriptor( + HTMLInputElement.prototype, + 'checked', + ).set; + + describe('setValueForProperty', () => { + it('should set values as properties by default', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(
); + }); + expect(container.firstChild.title).toBe('Tip!'); + }); + + it('should set values as attributes if necessary', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(
); + }); + expect(container.firstChild.getAttribute('role')).toBe('#'); + expect(container.firstChild.role).toBeUndefined(); + }); + + it('should set values as namespace attributes if necessary', async () => { + const container = document.createElementNS( + 'http://www.w3.org/2000/svg', + 'svg', + ); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(); + }); + expect( + container.firstChild.getAttributeNS( + 'http://www.w3.org/1999/xlink', + 'href', + ), + ).toBe('about:blank'); + }); + + it('should set values as boolean properties', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(
); + }); + expect(container.firstChild.getAttribute('disabled')).toBe(''); + await act(() => { + root.render(
); + }); + expect(container.firstChild.getAttribute('disabled')).toBe(''); + await act(() => { + root.render(
); + }); + expect(container.firstChild.getAttribute('disabled')).toBe(null); + await act(() => { + root.render(
); + }); + await act(() => { + root.render(
); + }); + expect(container.firstChild.getAttribute('disabled')).toBe(null); + await act(() => { + root.render(
); + }); + await act(() => { + root.render(
); + }); + expect(container.firstChild.getAttribute('disabled')).toBe(null); + }); + + it('should convert attribute values to string first', async () => { + // Browsers default to this behavior, but some test environments do not. + // This ensures that we have consistent behavior. + const obj = { + toString: function () { + return 'css-class'; + }, + }; + + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(
); + }); + expect(container.firstChild.getAttribute('class')).toBe('css-class'); + }); + + it('should not remove empty attributes for special input properties', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render( {}} />); + }); + if (disableInputAttributeSyncing) { + expect(container.firstChild.hasAttribute('value')).toBe(false); + } else { + expect(container.firstChild.getAttribute('value')).toBe(''); + } + expect(container.firstChild.value).toBe(''); + }); + + it('should not remove empty attributes for special option properties', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render( + , + ); + }); + // Regression test for https://github.com/facebook/react/issues/6219 + expect(container.firstChild.firstChild.value).toBe(''); + expect(container.firstChild.lastChild.value).toBe('filled'); + }); + + it('should remove for falsey boolean properties', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(
); + }); + expect(container.firstChild.hasAttribute('allowFullScreen')).toBe(false); + }); + + it('should remove when setting custom attr to null', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(
); + }); + expect(container.firstChild.hasAttribute('data-foo')).toBe(true); + await act(() => { + root.render(
); + }); + expect(container.firstChild.hasAttribute('data-foo')).toBe(false); + }); + + it('should set className to empty string instead of null', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(
); + }); + expect(container.firstChild.className).toBe('selected'); + await act(() => { + root.render(
); + }); + // className should be '', not 'null' or null (which becomes 'null' in + // some browsers) + expect(container.firstChild.className).toBe(''); + expect(container.firstChild.getAttribute('class')).toBe(null); + }); + + it('should remove property properly for boolean properties', async () => { + const container = document.createElement('div'); + const root = ReactDOMClient.createRoot(container); + await act(() => { + root.render(