diff --git a/CHANGELOG.md b/CHANGELOG.md index d0f8fd35..ad720660 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,30 +4,10 @@ All notable changes to this project are documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). See the [CONTRIBUTING guide](./CONTRIBUTING.md#Changelog) for instructions on how to add changelog entries. -## [Unreleased 3.2](https://github.com/opensearch-project/opensearch-jvector/compare/2.x...HEAD) -### Features -### Enhancements -* PQ refinement during merge [109](https://github.com/opensearch-project/opensearch-jvector/issues/109) -* Persistent Ordinal To docID Mapping [167](https://github.com/opensearch-project/opensearch-jvector/pull/167) -* Incremental Insertion With Leading Segment [167](https://github.com/opensearch-project/opensearch-jvector/pull/167) -* Remove Redundant FlatVectorFormat [167](https://github.com/opensearch-project/opensearch-jvector/pull/167) -* Remove Redundant DocValuesFormat [167](https://github.com/opensearch-project/opensearch-jvector/pull/167) -### Bug Fixes -* Fix for sorted indices [167](https://github.com/opensearch-project/opensearch-jvector/pull/167) -* Fix for missing fields [167](https://github.com/opensearch-project/opensearch-jvector/pull/167) -### Infrastructure -* Upgrade to JDK24 [165] (https://github.com/opensearch-project/opensearch-jvector/pull/165) -* Upgrade Gradle to 8.14 [165] (https://github.com/opensearch-project/opensearch-jvector/pull/165) -### Documentation -* Add docker instructions [163] (https://github.com/opensearch-project/opensearch-jvector/pull/163) -### Maintenance -* Fix documentation bugs [161] (https://github.com/opensearch-project/opensearch-jvector/pull/161) -### Refactoring -* Remove jVector Codec [167](https://github.com/opensearch-project/opensearch-jvector/pull/167) - -## [Unreleased 2.x](https://github.com/opensearch-project/opensearch-jvector/compare/2.18...2.x) +## [Unreleased 3.3](https://github.com/opensearch-project/opensearch-jvector/compare/3.2...HEAD) ### Features ### Enhancements +* Add script for loading vector data using Parquet [192](https://github.com/opensearch-project/opensearch-jvector/issues/192) ### Bug Fixes ### Infrastructure ### Documentation diff --git a/scripts/parquet-loader/README.md b/scripts/parquet-loader/README.md new file mode 100644 index 00000000..96157538 --- /dev/null +++ b/scripts/parquet-loader/README.md @@ -0,0 +1,49 @@ +# OpenSearch Parquet Loader + +This script efficiently loads data from Parquet files into an OpenSearch index, leveraging batch processing and multiple connections to maximize indexing speed. + +## Prerequisites + +- Python 3.8+ +- Access to an OpenSearch cluster + +## Setup + +1. **Create a python virtual environment:** + ```bash + sudo apt install python3.11-venv + # Using venv (Python 3.3+) + python3 -m venv .venv + + # Activate the virtual environment + # On Windows: + .venv\Scripts\activate + # On macOS/Linux: + source .venv/bin/activate + ``` + +2. **Install dependencies:** + ```bash + pip install -r requirements.txt + ``` + +3. **Configure environment variables:** + Create a `.env` file in the project root and add your OpenSearch connection details: + ``` + OPENSEARCH_HOSTS='["http://localhost:9200"]' + OPENSEARCH_INDEX="my-index" + ``` + Alternatively, set these variables on your local terminal + export OPENSEARCH_HOSTS='[{"host":"localhost","port":9200,"scheme":"http"}]' + export OPENSEARCH_INDEX="my-index" + +## Usage + +Run the script from the command line, providing the path to your Parquet file and the number of parallel workers: + +```bash +python osbench.py --path /path/to/your/data.parquet --procs 4 +``` + +- `/path/to/your/data.parquet`: The Parquet file to load. +- `--workers`: The number of parallel connections to use. diff --git a/scripts/parquet-loader/create_index.sh b/scripts/parquet-loader/create_index.sh new file mode 100644 index 00000000..4c8b9ee2 --- /dev/null +++ b/scripts/parquet-loader/create_index.sh @@ -0,0 +1,39 @@ +curl -X PUT "https://localhost:9200/jvector-index?pretty" --insecure -H 'Content-Type: application/json' -d' +{ + "settings": { + "index": { + "knn": true, + "refresh_interval": -1, + "number_of_replicas": 0, + "number_of_shards": 1, + "merge": { + "policy": { + "max_merged_segment": "50g" + } + } + } + }, + "mappings": { + "_source": { + "excludes": ["embeddings", "chunk_id"], + "recovery_source_excludes": ["embeddings", "chunk_id"] + }, + "properties": { + "chunk_id": {"type": "long"}, + "embeddings": { + "type": "knn_vector", + "method": { + "name": "disk_ann", + "space_type": "l2", + "engine": "jvector", + "parameters": { + "m": 32, + "ef_construction": 200, + "advanced.num_pq_subspaces": 48 + } + }, + "dimension": 384 + } + } + } +}' -u admin diff --git a/scripts/parquet-loader/osbench.py b/scripts/parquet-loader/osbench.py new file mode 100644 index 00000000..ee08cbc7 --- /dev/null +++ b/scripts/parquet-loader/osbench.py @@ -0,0 +1,541 @@ +#!/usr/bin/env python3 +import argparse +import json +import os +import sys +import time +import socket +import traceback +from pathlib import Path +from typing import List, Dict, Any, Tuple +from collections import Counter, defaultdict +from threading import Thread +from queue import Empty +import multiprocessing as mp + +import pyarrow as pa +import pyarrow.parquet as pq +import numpy as np +from tqdm import tqdm + +import urllib3 + +urllib3.disable_warnings() + +from opensearchpy import OpenSearch, Urllib3HttpConnection, TransportError +from opensearchpy.helpers import streaming_bulk + +from dotenv import load_dotenv + +# ------------------------------ +# Optional ultra-fast JSON (NumPy → JSON without .tolist()) +# ------------------------------ +try: + import orjson + from opensearchpy.serializer import JSONSerializer + + class ORJSONSerializer(JSONSerializer): + def dumps(self, data): + if isinstance(data, (bytes, str)): + return data + return orjson.dumps(data, option=orjson.OPT_SERIALIZE_NUMPY).decode("utf-8") +except Exception: + ORJSONSerializer = None + +PROG_ADD = "add" # increment docs by n +PROG_FILE = "file_done" # one file finished +PROG_DONE = "done" # one worker finished + +def start_progress_consumer(progress_q, procs: int = 1): + """ + Run a background thread that updates two bars: + - docs indexed (primary, by actions confirmed from workers) + - files done (secondary, optional) + """ + def _runner(progress_q, stop_event, *args, **kwargs): + import traceback, os, sys, queue + # ensure progress-bar variables exist and are safe-to-call even if not set up + class _NoopBar: + def update(self, n=1): pass + def set_description(self, *a, **k): pass + def close(self): pass + def refresh(self): pass + + file_bar = _NoopBar() + # other bars used by the runner (if any) can be initialized similarly: + files_bar = _NoopBar() + docs_bar = _NoopBar() + + try: + while not stop_event.is_set(): + try: + kind, value = progress_q.get(timeout=0.5) + except queue.Empty: + continue + if kind == PROG_ADD: + doc_bar.update(int(value)) + elif kind == PROG_FILE: + file_bar.update(1) + elif kind == PROG_DONE: + done += 1 + except Exception as e: + # Avoid using buffered stderr at interpreter shutdown — write raw bytes. + tb = traceback.format_exc() + try: + os.write(2, b"Exception in runner thread:\n") + os.write(2, tb.encode("utf-8", "backslashreplace")) + except Exception: + # best-effort: if os.write fails, silently ignore to avoid aborting + pass + + # try to notify the main process via the queue (best-effort) + try: + progress_q.put(("runner_error", tb)) + except Exception: + pass + + # ensure we set stop_event so main cleans up worker threads/processes + try: + stop_event.set() + t.join(timeout=10) + except Exception: + pass + + # do not re-raise here because re-raising can cause unclean interpreter shutdown + return + finally: + doc_bar.close() + file_bar.close() + + # create a stop event for the runner thread and pass it in + from threading import Event + stop_event = Event() + t = Thread(target=_runner, args=(progress_q, stop_event), daemon=False) + t.start() + return t, stop_event + + + +# ------------------------------ +# Args / CLI +# ------------------------------ +def parse_args(): + p = argparse.ArgumentParser("Parquet → OpenSearch (multiprocessing)") + p.add_argument("--path", required=True, help="Parquet file or directory") + p.add_argument("--recursive", action="store_true", help="Recurse subdirs for *.parquet") + p.add_argument("--index", required=False, default=os.getenv("OPENSEARCH_INDEX")) + p.add_argument("--procs", type=int, default=32, help="Worker processes") + p.add_argument("--chunk_size", type=int, default=8000, help="Rows per Arrow batch (read size)") + p.add_argument("--bulk_chunk", type=int, default=1000, help="Docs per bulk request (send size)") + p.add_argument("--max_chunk_bytes", type=int, default=10 * 1024 * 1024, help="Max bytes per bulk request") + p.add_argument("--pool_maxsize", type=int, default=8, help="per-process HTTP pool sockets") + p.add_argument("--max_retries", type=int, default=32, help="Transport retries (client-level)") + p.add_argument("--use_orjson", action="store_true", help="Use orjson serializer (handles NumPy)") + p.add_argument("--debug", action="store_true", help="Verbose stack traces on connection errors") + p.add_argument("--error_log", type=str, help="Write critical errors (JSONL) per process with suffix ..jsonl") + p.add_argument("--print_retryable", action="store_true", help="Also print retryable errors (noisy)") + return p.parse_args() + + +def _run_worker(args_tuple): + # args_tuple is exactly the tuple you were building in `work` + return worker_proc(*args_tuple) + +# ------------------------------ +# Utilities +# ------------------------------ +def list_parquet_files(path: str, recursive: bool) -> List[Path]: + p = Path(path) + if p.is_file(): + if p.suffix.lower() != ".parquet": + raise ValueError(f"Not a .parquet file: {p}") + return [p] + if not p.is_dir(): + raise ValueError(f"Path not found: {p}") + pattern = "**/*.parquet" if recursive else "*.parquet" + files = sorted(p.glob(pattern)) + if not files: + raise ValueError(f"No .parquet files under {p}") + return files + + +def split_round_robin(files: List[Path], k: int) -> List[List[Path]]: + if k <= 1: + return [files] + return [files[i::k] for i in range(k)] + + +def get_client(max_retries: int, pool_maxsize: int, use_orjson: bool) -> OpenSearch: + hosts = json.loads(os.getenv("OPENSEARCH_HOSTS", '[{"host":"localhost","port":9200,"scheme":"http"}]')) + user, pwd = os.getenv("OPENSEARCH_USER"), os.getenv("OPENSEARCH_PASS") + http_auth = (user, pwd) if user and pwd else None + + kwargs = dict( + hosts=hosts, + http_auth=http_auth, + http_compress=True, + retry_on_timeout=True, + verify_certs=False, + max_retries=max_retries, + timeout=60, + connection_class=Urllib3HttpConnection, + pool_maxsize=pool_maxsize, + ) + if use_orjson and ORJSONSerializer: + kwargs["serializer"] = ORJSONSerializer() + return OpenSearch(**kwargs) + + +def check_connectivity(client: OpenSearch, debug: bool): + # DNS preflight + try: + hosts = json.loads(os.getenv("OPENSEARCH_HOSTS", "[]")) + for h in hosts: + host = h.get("host") + if not host: + continue + socket.getaddrinfo(host, None) + except Exception as e: + print(f"❌ DNS resolution failed: {e}", file=sys.stderr) + if debug: + traceback.print_exc() + sys.exit(2) + # ping + info + try: + if not client.ping(params={"request_timeout": 5}): + raise RuntimeError("Ping returned False (URL/creds/cluster?)") + _ = client.info() + except Exception as e: + print(f"❌ Connection check failed: {e}", file=sys.stderr) + if debug: + traceback.print_exc() + sys.exit(2) + + +# ------------------------------ +# Error classification +# ------------------------------ +RETRYABLE_STATUS = {408, 409, 429, 500, 502, 503, 504} +RETRYABLE_TYPES = { + "es_rejected_execution_exception", + "cluster_block_exception", + "master_not_discovered_exception", + "too_many_requests", + "receive_timeout_transport_exception", + "process_cluster_event_timeout_exception", + "timeout_exception", +} + +def classify_bulk_error(item: Dict[str, Any]) -> str: + op, detail = next(iter(item.items())) + status = detail.get("status") + err = detail.get("error") or {} + et = err.get("type") + if status in RETRYABLE_STATUS or (et and et in RETRYABLE_TYPES): + return "retryable" + return "critical" + + +def print_retryable(item, pid): + op, d = next(iter(item.items())) + status = d.get("status") + err = d.get("error") or {} + rid = d.get("_id", "") + reason = (err.get("reason") or "").splitlines()[0] + print(f"[{pid}] [retryable] {op} status={status} _id={rid} :: {reason}", file=sys.stderr, flush=True) + + +def print_critical(item, pid): + op, d = next(iter(item.items())) + status = d.get("status") + err = d.get("error") or {} + rid = d.get("_id", "") + reason = (err.get("reason") or "").splitlines()[0] + et = err.get("type", "unknown_error") + print(f"[{pid}] [CRITICAL] {op} status={status} type={et} _id={rid} :: {reason}", file=sys.stderr, flush=True) + + +# ------------------------------ +# Batch → actions (fast path for FixedSizeList) +# ------------------------------ +def actions_from_batch_numpy(batch: pa.RecordBatch, index_name: str, numpy_ok: bool) -> List[Dict[str, Any]]: + names = batch.schema.names + # tolerate missing columns: fallback to synthetic row ids if chunk_id absent + i_chunk = names.index("chunk_id") if "chunk_id" in names else None + + # find embeddings column (accept common alternatives or detect by type) + i_emb = None + if "embeddings" in names: + i_emb = names.index("embeddings") + else: + for alt in ("embedding", "vector", "vec", "features", "feat"): + if alt in names: + i_emb = names.index(alt) + break + + if i_emb is None: + # type-based detection: FixedSizeList + for idx, fname in enumerate(names): + f = batch.schema.field(idx) + if pa.types.is_fixed_size_list(f.type) and pa.types.is_floating(f.type.value_type): + i_emb = idx + break + + if i_emb is None: + # no embeddings found -> skip this batch instead of crashing worker + print(f"[worker] Missing required column 'embeddings' in batch — available columns: {names}", file=sys.stderr) + return [] + + n = batch.num_rows + + if i_chunk is not None: + chunk_ids = batch.column(i_chunk).to_pylist() + else: + chunk_ids = [str(i) for i in range(n)] + emb_field = batch.schema.field(i_emb) + emb_col = batch.column(i_emb) + emb_type = emb_field.type + + out: List[Dict[str, Any]] = [] + + if pa.types.is_fixed_size_list(emb_type) and pa.types.is_floating(emb_type.value_type): + dim = emb_type.list_size + vals = emb_col.values.to_numpy(zero_copy_only=False) # len n*dim + mat = vals.reshape(n, dim) + if numpy_ok and ORJSONSerializer: + for cid, vec in zip(chunk_ids, mat): + if cid is None: continue + out.append({"_index": index_name, "_id": cid, "chunk_id": cid, "embeddings": vec}) + else: + vecs = mat.tolist() + for cid, vec in zip(chunk_ids, vecs): + if cid is None: continue + out.append({"_index": index_name, "_id": cid, "chunk_id": cid, "embeddings": vec}) + return out + + # Fallback (variable-length lists) + emb_list = emb_col.to_pylist() + for cid, vec in zip(chunk_ids, emb_list): + if cid is None: continue + out.append({"_index": index_name, "_id": cid, "chunk_id": cid, "embeddings": (np.asarray(vec) if numpy_ok and ORJSONSerializer else vec)}) + return out + + +def fake_streaming_bulk( + client: Any, + actions: Any, + chunk_size: int = 500, + max_chunk_bytes: int = 100 * 1024 * 1024, + raise_on_error: bool = True, + raise_on_exception: bool = True, + max_retries: int = 0, + initial_backoff: int = 2, + max_backoff: int = 600, + yield_ok: bool = True, + ignore_status: Any = (), + *args: Any, + **kwargs: Any, +) -> Any: + for a in actions: + yield True, a + +# ------------------------------ +# Worker process +# ------------------------------ +def worker_proc(shard_id: int, + files: List[str], + index_name: str, + chunk_rows: int, + bulk_chunk: int, + max_chunk_bytes: int, + pool_maxsize: int, + max_retries: int, + use_orjson: bool, + print_retryable_flag: bool, + error_log_path: str | None, + progress_q, + progress_tick: int = 1000 ) -> Dict[str, Any]: + pid = os.getpid() + stats = defaultdict(int) + critical_sig_counter: Counter = Counter() + + # per-process error log file (avoid contention) + log_fp = None + if error_log_path: + log_fp = open(f"{error_log_path}.{pid}.jsonl", "w", encoding="utf-8") + + client = get_client(max_retries=max_retries, pool_maxsize=pool_maxsize, use_orjson=use_orjson) + + wanted_cols = ["chunk_id", "embeddings"] + start = time.time() + try: + for f in files: + pf = pq.ParquetFile(f) + for batch in pf.iter_batches(batch_size=chunk_rows, use_threads=True, columns=wanted_cols): + actions = actions_from_batch_numpy(batch, index_name, numpy_ok=use_orjson) + sent_since_last = 0 + # stream this batch with streaming_bulk (single process = 1 inflight request) + for ok, item in streaming_bulk( + client, + (a for a in actions), + chunk_size=bulk_chunk, + max_chunk_bytes=max_chunk_bytes, + raise_on_error=False, + ): + if ok: + stats["success"] += 1 + else: + stats["failed"] += 1 + cls = classify_bulk_error(item) + if cls == "retryable": + stats["retryable"] += 1 + if print_retryable_flag: + print_retryable(item, pid) + else: + stats["critical"] += 1 + # count signature + op, d = next(iter(item.items())) + status = d.get("status", -1) + err = d.get("error") or {} + et = err.get("type", "unknown_error") + reason = (err.get("reason") or "").splitlines()[0][:200] + critical_sig_counter[(et, status, reason)] += 1 + print_critical(item, pid) + if log_fp: + log_fp.write(json.dumps(item, ensure_ascii=False) + "\n") + # progress heartbeat every N docs + sent_since_last += 1 + if sent_since_last >= progress_tick: + try: + progress_q.put_nowait((PROG_ADD, sent_since_last)) + except Exception: + pass + sent_since_last = 0 + # flush any remainder for this batch + if sent_since_last: + try: + progress_q.put_nowait((PROG_ADD, sent_since_last)) + except Exception: + pass + + except Exception as e: + # propagate a clear failure line to parent via stats + stats["exception"] = 1 + stats["exception_msg"] = str(e) + stats["exception_tb"] = traceback.format_exc() + finally: + if log_fp: + log_fp.close() + stats["elapsed"] = time.time() - start + stats["pid"] = pid + stats["shard"] = shard_id + stats["files"] = len(files) + stats["crit_sigs"] = dict(critical_sig_counter) + # finished a file + try: + progress_q.put_nowait((PROG_FILE, 1)) + except Exception: + pass + return stats + + +# ------------------------------ +# Main +# ------------------------------ +def main(): + load_dotenv() + + args = parse_args() + if not args.index: + print("Missing --index or OPENSEARCH_INDEX", file=sys.stderr); sys.exit(1) + + files = list_parquet_files(args.path, args.recursive) + total_rows = 0 + for f in files: + try: + pf = pq.ParquetFile(str(f)) + total_rows += pf.metadata.num_rows + except Exception: + pass + + manager = mp.Manager() + progress_q = manager.Queue(maxsize=10000) + + # start the background consumer (2 progress bars) + progress_thread, progress_stop_event = start_progress_consumer( + progress_q, + procs=args.procs + ) + + # quick one-time connectivity check in parent + parent_client = get_client(args.max_retries, args.pool_maxsize, args.use_orjson) + check_connectivity(parent_client, args.debug) + + shards = split_round_robin(files, args.procs) + work: List[Tuple] = [] + for sid, shard in enumerate(shards): + work.append(( + sid, + [str(p) for p in shard], + args.index, + args.chunk_size, + args.bulk_chunk, + args.max_chunk_bytes, + args.pool_maxsize, + args.max_retries, + args.use_orjson, + args.print_retryable, + args.error_log, + progress_q, + 1000 + )) + + print(f"Starting {len(work)} processes over {len(files)} files; " + f"batch={args.chunk_size} bulk_chunk={args.bulk_chunk} max_chunk_bytes={args.max_chunk_bytes}", flush=True) + + start = time.time() + agg = defaultdict(int) + crit_counter = Counter() + + # Use spawn for safety across platforms + ctx = mp.get_context("spawn") + with ctx.Pool(processes=args.procs) as pool: + #mp.set_start_method("spawn", force=True) + #with mp.Pool(processes=len(work)) as pool: + # tqdm over processes finishing + for stats in tqdm(pool.imap_unordered(_run_worker, work, chunksize=1), total=len(work), desc="Workers", position=2): + # accumulate + for k, v in stats.items(): + if k == "crit_sigs": + crit_counter.update({tuple(sig): count for sig, count in v.items()}) + elif k in ("exception_msg", "exception_tb"): # printed below if exists + continue + elif isinstance(v, (int, float)): + agg[k] += v + if stats.get("exception", 0): + print(f"\n[worker {stats.get('pid')}] FAILED:\n{stats.get('exception_msg')}\n{stats.get('exception_tb')}", file=sys.stderr) + + elapsed = time.time() - start + progress_thread.join(timeout=1) + + # Summary + print("\nIndexing finished.") + print(f"Files processed: {len(files)}") + if total_rows: + print(f"Approx rows (metadata): {total_rows}") + print(f"Processes: {args.procs}") + print(f"Successfully indexed: {agg.get('success',0)}") + print(f"Failed: {agg.get('failed',0)}") + print(f" • retryable (suppressed):{agg.get('retryable',0)}") + print(f" • critical (shown): {sum(crit_counter.values())}") + print(f"Total time: {elapsed:.2f}s") + + if crit_counter: + print("\nTop critical error signatures:") + for (et, status, reason), c in crit_counter.most_common(10): + print(f" [{status}] {et} x{c} :: {reason}") + + +if __name__ == "__main__": + main() + diff --git a/scripts/parquet-loader/osrecall.py b/scripts/parquet-loader/osrecall.py new file mode 100644 index 00000000..e604c296 --- /dev/null +++ b/scripts/parquet-loader/osrecall.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +import argparse +import json +import os +import socket +import statistics +import sys +import time +from typing import Dict, List, Set + +from dotenv import load_dotenv +import urllib3 + +urllib3.disable_warnings() + +import numpy as np +from opensearchpy import OpenSearch, Urllib3HttpConnection +from tqdm import tqdm + +# ------------- recall math ------------- +def calculate_recall(truth: Set[str], recalled: Set[str]) -> float: + if not truth: + return 0.0 + missed = truth - recalled + return (1.0 - (float(len(missed)) / float(len(truth)))) * 100.0 + + +# ------------- OpenSearch client ------------- +def get_opensearch_client(max_retries: int = 5, pool_maxsize: int = 4): + """ + Configure via env or args: + OPENSEARCH_HOSTS='[{"host":"localhost","port":9200,"scheme":"http"}]' + OPENSEARCH_USER=... OPENSEARCH_PASS=... + OPENSEARCH_SSL_VERIFY=true|false + """ + hosts = json.loads(os.getenv("OPENSEARCH_HOSTS", '[{"host":"localhost","port":9200,"scheme":"http"}]')) + user = os.getenv("OPENSEARCH_USER") + pwd = os.getenv("OPENSEARCH_PASS") + http_auth = (user, pwd) if user and pwd else None + verify_certs = os.getenv("OPENSEARCH_SSL_VERIFY", "false").lower() == "true" + + return OpenSearch( + hosts=hosts, + http_auth=http_auth, + verify_certs=verify_certs, + http_compress=True, + retry_on_timeout=True, + max_retries=max_retries, + timeout=30, + connection_class=Urllib3HttpConnection, + pool_maxsize=pool_maxsize, + ) + + +def connection_check(client: OpenSearch, debug: bool = False): + # DNS preflight to catch typos fast + try: + hosts = json.loads(os.getenv("OPENSEARCH_HOSTS", "[]")) + for h in hosts: + host = h.get("host") + if host: + socket.getaddrinfo(host, None) + except Exception as e: + print(f"❌ DNS resolution failed: {e}", file=sys.stderr) + if debug: + import traceback; traceback.print_exc() + sys.exit(2) + # Ping + info + try: + if not client.ping(params={"request_timeout": 5}): + raise RuntimeError("Ping returned False (URL/creds/cluster?)") + _ = client.info() + except Exception as e: + print(f"❌ Connection check failed: {e}", file=sys.stderr) + if debug: + import traceback; traceback.print_exc() + sys.exit(2) + + +# ------------- OpenSearch queries ------------- +def do_search(client: OpenSearch, index: str, field: str, query_vec: List[float], k: int, overquery_factor: int): + """ + POST /{index}/_search with a knn query + """ + body = { + "size": k, + "query": { + "knn": { + field: { + "vector": query_vec, + "k": k * overquery_factor, +# "method_parameters": { +# "overquery_factor": overquery_factor, +# "advanced.threshold": 0.0, +# "advanced.rerank_floor": 0.0 +# } + } + } + } + } + resp = client.search(index=index, body=body) + hits = resp.get("hits", {}).get("hits", []) + return hits + + + + +# ------------- Runner ------------- +def run(args): + # Load truth + with open(args.truth_file) as fd: + truth_json = json.load(fd) + + query_points: Dict[int, List[float]] = {} + nn_ids: Dict[int, List[str]] = {} + for i, qp in enumerate(truth_json.get("query_points", [])): + query_points[i] = qp["point"] + nn_ids[i] = qp["nn_id"] + + if not query_points: + print("No query_points found in truth file.", file=sys.stderr) + sys.exit(1) + + print(f"Loaded {len(query_points)} queries; first has {len(nn_ids[0])} truth neighbors") + + + # OS client + client = get_opensearch_client(max_retries=args.max_retries, pool_maxsize=args.pool_maxsize) + connection_check(client, debug=args.debug) + + + # Warm-up + print("Warming up...") + warm_keys = list(query_points.keys())[: min(5, len(query_points))] + for qpidx in warm_keys: + _ = do_search(client, args.index, args.field, query_points[qpidx], args.k, args.refine) + + # Measure + print("Running measured queries...") + recalls: List[float] = [] + latencies: List[float] = [] + + iterator = query_points.keys() + if args.limit and args.limit > 0: + iterator = list(iterator)[: args.limit] + + for qpidx in tqdm(iterator, total=(args.limit or len(query_points)), desc="Queries"): + qp = query_points[qpidx] + t0 = time.perf_counter() + hits = do_search(client, args.index, args.field, qp, args.k, args.refine) + t1 = time.perf_counter() + latencies.append(t1 - t0) + + # gather returned ids + got_ids = {h.get("_id") for h in hits if h.get("_id") is not None} + + # truth ids for this query + truth_ids = set(map(str, nn_ids[qpidx])) # ensure comparable types + + r = calculate_recall(truth_ids, got_ids) + recalls.append(r) + + if args.verbose: + print(f"q={qpidx} recall={r:.2f} hits={len(got_ids)} latency={(t1-t0)*1000:.1f}ms") + + # Summary + if recalls: + print(f"\nmean recall: {np.mean(recalls):.2f} %") + print(f"p50 recall: {np.percentile(recalls, 50):.2f} %") + print(f"p95 recall: {np.percentile(recalls, 95):.2f} %") + if latencies: + ms = [x * 1000.0 for x in latencies] + print(f"mean latency: {statistics.mean(ms):.1f} ms") + print(f"p50 latency: {np.percentile(ms, 50):.1f} ms") + print(f"p95 latency: {np.percentile(ms, 95):.1f} ms") + print(f"qps (approx): {len(latencies) / sum(latencies):.1f} q/s") + + +# ------------- CLI ------------- +if __name__ == "__main__": + load_dotenv() + + ap = argparse.ArgumentParser("Recall tester for OpenSearch k-NN") + ap.add_argument("--truth_file", type=str, required=True, help="Ground-truth JSON") + ap.add_argument("--index", type=str, default=os.getenv("OPENSEARCH_INDEX"), help="OpenSearch index name") + ap.add_argument("--field", type=str, default="embeddings", help="Vector field name") + ap.add_argument("-k", type=int, default=100, help="Top-K to return") + ap.add_argument("--num_candidates", type=int, default=None, help="Approx. candidates (>= k). If unset, uses k*refine") + ap.add_argument("--refine", type=int, default=20, help="Fallback multiplier for num_candidates (k*refine)") + ap.add_argument("--api", choices=["search", "knn_search"], default="search", help="Which API to use") + ap.add_argument("--limit", type=int, default=0, help="Limit number of queries (0 = all)") + ap.add_argument("--max_retries", type=int, default=5, help="Client transport retries") + ap.add_argument("--pool_maxsize", type=int, default=4, help="HTTP connection pool per process") + ap.add_argument("--verbose", action="store_true") + ap.add_argument("--debug", action="store_true") + args = ap.parse_args() + + if not args.index: + print("❌ Missing --index or OPENSEARCH_INDEX", file=sys.stderr) + sys.exit(1) + + run(args) diff --git a/scripts/parquet-loader/requirements.txt b/scripts/parquet-loader/requirements.txt new file mode 100644 index 00000000..dcfcbf8c --- /dev/null +++ b/scripts/parquet-loader/requirements.txt @@ -0,0 +1,5 @@ +opensearch-py +pyarrow +pandas +tqdm +python-dotenv