diff --git a/benchmarks/full_report/report.ipynb b/benchmarks/full_report/report.ipynb index 11a4924c4c1..70eb810ab26 100644 --- a/benchmarks/full_report/report.ipynb +++ b/benchmarks/full_report/report.ipynb @@ -8,11 +8,12 @@ "outputs": [], "source": [ "import sys\n", + "\n", "sys.dont_write_bytecode = True\n", "\n", "import os\n", "\n", - "module_path = os.path.abspath(os.path.join('.'))\n", + "module_path = os.path.abspath(os.path.join(\".\"))\n", "if module_path not in sys.path:\n", " sys.path.append(module_path)" ] @@ -61,58 +62,91 @@ " gt_in_sample = knn(in_sample, data, metric, 10)\n", "\n", " print(\"generated gt\")\n", - " \n", + "\n", " with tempfile.TemporaryDirectory() as d:\n", " write_lance(d, data)\n", " ds = lance.dataset(d)\n", "\n", - " for q, target in zip(tqdm(in_sample, desc=\"checking brute force\"), gt_in_sample):\n", - " res = ds.to_table(nearest={\n", - " \"column\": \"vec\",\n", - " \"q\": q,\n", - " \"k\": 10,\n", - " \"metric\": metric,\n", - " }, columns=[\"id\"])\n", + " for q, target in zip(\n", + " tqdm(in_sample, desc=\"checking brute force\"), gt_in_sample\n", + " ):\n", + " res = ds.to_table(\n", + " nearest={\n", + " \"column\": \"vec\",\n", + " \"q\": q,\n", + " \"k\": 10,\n", + " \"metric\": metric,\n", + " },\n", + " columns=[\"id\"],\n", + " )\n", " assert len(np.intersect1d(res[\"id\"].to_numpy(), target)) == 10\n", - " \n", - " ds = ds.create_index(\"vec\", \"IVF_PQ\", metric=metric, num_partitions=num_partitions, num_sub_vectors=num_sub_vectors)\n", - " \n", + "\n", + " ds = ds.create_index(\n", + " \"vec\",\n", + " \"IVF_PQ\",\n", + " metric=metric,\n", + " num_partitions=num_partitions,\n", + " num_sub_vectors=num_sub_vectors,\n", + " )\n", + "\n", " recall_data = []\n", " for nprobes in nprobes_list:\n", " for refine_factor in refine_factor_list:\n", " hits = 0\n", " # check that brute force impl is correct\n", - " for q, target in zip(tqdm(query, desc=f\"out of sample, nprobes={nprobes}, refine={refine_factor}\"), gt):\n", - " res = ds.to_table(nearest={\n", - " \"column\": \"vec\",\n", - " \"q\": q,\n", - " \"k\": 10,\n", - " \"nprobes\": nprobes,\n", - " \"refine_factor\": refine_factor,\n", - " }, columns=[\"id\"])[\"id\"].to_numpy()\n", + " for q, target in zip(\n", + " tqdm(\n", + " query,\n", + " desc=f\"out of sample, nprobes={nprobes}, refine={refine_factor}\",\n", + " ),\n", + " gt,\n", + " ):\n", + " res = ds.to_table(\n", + " nearest={\n", + " \"column\": \"vec\",\n", + " \"q\": q,\n", + " \"k\": 10,\n", + " \"nprobes\": nprobes,\n", + " \"refine_factor\": refine_factor,\n", + " },\n", + " columns=[\"id\"],\n", + " )[\"id\"].to_numpy()\n", " hits += len(np.intersect1d(res, target))\n", - " recall_data.append([\n", - " \"out_of_sample\",\n", - " nprobes,\n", - " refine_factor,\n", - " hits / 10 / len(gt),\n", - " ])\n", + " recall_data.append(\n", + " [\n", + " \"out_of_sample\",\n", + " nprobes,\n", + " refine_factor,\n", + " hits / 10 / len(gt),\n", + " ]\n", + " )\n", " # check that brute force impl is correct\n", - " for q, target in zip(tqdm(in_sample, desc=f\"in sample nprobes={nprobes}, refine={refine_factor}\"), gt_in_sample):\n", - " res = ds.to_table(nearest={\n", - " \"column\": \"vec\",\n", - " \"q\": q,\n", - " \"k\": 10,\n", - " \"nprobes\": nprobes,\n", - " \"refine_factor\": refine_factor,\n", - " }, columns=[\"id\"])[\"id\"].to_numpy()\n", + " for q, target in zip(\n", + " tqdm(\n", + " in_sample,\n", + " desc=f\"in sample nprobes={nprobes}, refine={refine_factor}\",\n", + " ),\n", + " gt_in_sample,\n", + " ):\n", + " res = ds.to_table(\n", + " nearest={\n", + " \"column\": \"vec\",\n", + " \"q\": q,\n", + " \"k\": 10,\n", + " \"nprobes\": nprobes,\n", + " \"refine_factor\": refine_factor,\n", + " },\n", + " columns=[\"id\"],\n", + " )[\"id\"].to_numpy()\n", " hits += len(np.intersect1d(res, target))\n", - " recall_data.append([\n", - " \"in_sample\",\n", - " nprobes,\n", - " refine_factor,\n", - " hits / 10 / len(gt_in_sample),\n", - " ])\n", + " recall_data.append(\n", + " [\n", + " \"in_sample\",\n", + " nprobes,\n", + " refine_factor,\n", + " hits / 10 / len(gt_in_sample),\n", + " ]\n", + " )\n", " return recall_data" ] }, @@ -124,15 +158,19 @@ "outputs": [], "source": [ "def make_plot(recall_data):\n", - " df = pd.DataFrame(recall_data, columns=[\"case\", \"nprobes\", \"refine_factor\", \"recall\"])\n", - " \n", + " df = pd.DataFrame(\n", + " recall_data, columns=[\"case\", \"nprobes\", \"refine_factor\", \"recall\"]\n", + " )\n", + "\n", " num_cases = len(df[\"case\"].unique())\n", " (fig, axs) = plt.subplots(1, 2, figsize=(16, 8))\n", - " \n", + "\n", " for case, ax in zip(df[\"case\"].unique(), axs):\n", " current_case = df[df[\"case\"] == case]\n", " sns.heatmap(\n", - " current_case.drop(columns=[\"case\"]).set_index([\"nprobes\", \"refine_factor\"])[\"recall\"].unstack(),\n", + " current_case.drop(columns=[\"case\"])\n", + " .set_index([\"nprobes\", \"refine_factor\"])[\"recall\"]\n", + " .unstack(),\n", " annot=True,\n", " ax=ax,\n", " ).set(title=f\"Recall -- {case}\")" diff --git a/benchmarks/sift/Results.ipynb b/benchmarks/sift/Results.ipynb index 7758587255e..4062e9d45a9 100644 --- a/benchmarks/sift/Results.ipynb +++ b/benchmarks/sift/Results.ipynb @@ -35,6 +35,7 @@ "outputs": [], "source": [ "import pandas as pd\n", + "\n", "df = pd.read_csv(\"query.csv\")" ] }, @@ -46,6 +47,7 @@ "outputs": [], "source": [ "import seaborn as sns\n", + "\n", "sns.set_style(\"darkgrid\")" ] }, diff --git a/ci/check_breaking_changes.py b/ci/check_breaking_changes.py index aa83d1ae7ee..ce569ca637a 100644 --- a/ci/check_breaking_changes.py +++ b/ci/check_breaking_changes.py @@ -4,6 +4,7 @@ Can also be used as a library to detect breaking changes without version validation. """ + import argparse import os import sys @@ -40,10 +41,17 @@ def detect_breaking_changes(repo, base, head): parser = argparse.ArgumentParser() parser.add_argument("base", help="Base commit/tag for comparison") parser.add_argument("head", help="Head commit/tag for comparison") - parser.add_argument("last_stable_version", nargs="?", help="Last stable version (for validation)") - parser.add_argument("current_version", nargs="?", help="Current version (for validation)") - parser.add_argument("--detect-only", action="store_true", - help="Only detect breaking changes, don't validate version") + parser.add_argument( + "last_stable_version", nargs="?", help="Last stable version (for validation)" + ) + parser.add_argument( + "current_version", nargs="?", help="Current version (for validation)" + ) + parser.add_argument( + "--detect-only", + action="store_true", + help="Only detect breaking changes, don't validate version", + ) args = parser.parse_args() repo = Github(os.environ["GITHUB_TOKEN"]).get_repo(os.environ["GITHUB_REPOSITORY"]) diff --git a/ci/setup_version.py b/ci/setup_version.py index 4b248d8e527..b0f6f514d34 100644 --- a/ci/setup_version.py +++ b/ci/setup_version.py @@ -3,7 +3,7 @@ This script is used to set the pre-release version for beta releases (e.g. 0.10.17-beta.1) when the tag indicates a beta release. -With the new automated release process, stable versions are already +With the new automated release process, stable versions are already updated by bump-my-version during the release workflow. """ @@ -38,9 +38,9 @@ def main(): print(f"Setting beta version: {current_version} -> {args.version[1:]}") else: # For stable releases, version should already match - assert ( - parsed_version.release == current_version_parsed.release - ), f"Version mismatch for stable release: {parsed_version.release} != {current_version_parsed.release}" + assert parsed_version.release == current_version_parsed.release, ( + f"Version mismatch for stable release: {parsed_version.release} != {current_version_parsed.release}" + ) with open("python/Cargo.toml", "w") as f: f.writelines(lines) diff --git a/memtest/tests/integration_test.rs b/memtest/tests/integration_test.rs index b83b50cd3d9..a453efd42c3 100644 --- a/memtest/tests/integration_test.rs +++ b/memtest/tests/integration_test.rs @@ -2,7 +2,7 @@ use libc::{c_void, size_t}; use std::ptr; // Import from the library we're testing -use memtest::{memtest_get_stats, memtest_reset_stats, MemtestStats}; +use memtest::{MemtestStats, memtest_get_stats, memtest_reset_stats}; extern "C" { fn malloc(size: size_t) -> *mut c_void; diff --git a/notebooks/quickstart.ipynb b/notebooks/quickstart.ipynb index 5abd79db2d7..2704edfb34d 100644 --- a/notebooks/quickstart.ipynb +++ b/notebooks/quickstart.ipynb @@ -221,7 +221,7 @@ "shutil.rmtree(\"/tmp/test.lance\", ignore_errors=True)\n", "\n", "tbl = pa.Table.from_pandas(df)\n", - "pa.dataset.write_dataset(tbl, \"/tmp/test.parquet\", format='parquet')\n", + "pa.dataset.write_dataset(tbl, \"/tmp/test.parquet\", format=\"parquet\")\n", "\n", "parquet = pa.dataset.dataset(\"/tmp/test.parquet\")\n", "parquet.to_table().to_pandas()" @@ -542,7 +542,7 @@ } ], "source": [ - "lance.dataset('/tmp/test.lance', version=1).to_table().to_pandas()" + "lance.dataset(\"/tmp/test.lance\", version=1).to_table().to_pandas()" ] }, { @@ -600,7 +600,7 @@ } ], "source": [ - "lance.dataset('/tmp/test.lance', version=2).to_table().to_pandas()" + "lance.dataset(\"/tmp/test.lance\", version=2).to_table().to_pandas()" ] }, { @@ -698,7 +698,7 @@ } ], "source": [ - "lance.dataset('/tmp/test.lance', version=\"stable\").to_table().to_pandas()" + "lance.dataset(\"/tmp/test.lance\", version=\"stable\").to_table().to_pandas()" ] }, { @@ -796,11 +796,13 @@ "\n", "with open(\"sift/sift_base.fvecs\", mode=\"rb\") as fobj:\n", " buf = fobj.read()\n", - " data = np.array(struct.unpack(\"<128000000f\", buf[4 : 4 + 4 * 1000000 * 128])).reshape((1000000, 128))\n", + " data = np.array(\n", + " struct.unpack(\"<128000000f\", buf[4 : 4 + 4 * 1000000 * 128])\n", + " ).reshape((1000000, 128))\n", " dd = dict(zip(range(1000000), data))\n", "\n", "table = vec_to_table(dd)\n", - "lance.write_dataset(table, uri, max_rows_per_group=8192, max_rows_per_file=1024*1024)" + "lance.write_dataset(table, uri, max_rows_per_group=8192, max_rows_per_file=1024 * 1024)" ] }, { @@ -864,6 +866,7 @@ ], "source": [ "import duckdb\n", + "\n", "# if this segfaults make sure duckdb v0.7+ is installed\n", "samples = duckdb.query(\"SELECT vector FROM sift1m USING SAMPLE 100\").to_df().vector\n", "samples" @@ -910,10 +913,12 @@ "import time\n", "\n", "start = time.time()\n", - "tbl = sift1m.to_table(columns=[\"id\"], nearest={\"column\": \"vector\", \"q\": samples[0], \"k\": 10})\n", + "tbl = sift1m.to_table(\n", + " columns=[\"id\"], nearest={\"column\": \"vector\", \"q\": samples[0], \"k\": 10}\n", + ")\n", "end = time.time()\n", "\n", - "print(f\"Time(sec): {end-start}\")\n", + "print(f\"Time(sec): {end - start}\")\n", "print(tbl.to_pandas())" ] }, @@ -993,7 +998,7 @@ "\n", "sift1m.create_index(\n", " \"vector\",\n", - " index_type=\"IVF_PQ\", # IVF_PQ, IVF_HNSW_PQ and IVF_HNSW_SQ are supported\n", + " index_type=\"IVF_PQ\", # IVF_PQ, IVF_HNSW_PQ and IVF_HNSW_SQ are supported\n", " num_partitions=256, # IVF\n", " num_sub_vectors=16, # PQ\n", ")" @@ -1068,7 +1073,7 @@ " start = time.time()\n", " tbl = sift1m.to_table(nearest={\"column\": \"vector\", \"q\": q, \"k\": 10})\n", " end = time.time()\n", - " tot += (end - start)\n", + " tot += end - start\n", "\n", "print(f\"Avg(sec): {tot / len(samples)}\")\n", "print(tbl.to_pandas())" @@ -1425,7 +1430,7 @@ "source": [ "tbl = sift1m.to_table()\n", "tbl = tbl.append_column(\"item_id\", pa.array(range(len(tbl))))\n", - "tbl = tbl.append_column(\"revenue\", pa.array((np.random.randn(len(tbl))+5)*1000))\n", + "tbl = tbl.append_column(\"revenue\", pa.array((np.random.randn(len(tbl)) + 5) * 1000))\n", "tbl.to_pandas()" ] }, @@ -1564,7 +1569,9 @@ } ], "source": [ - "sift1m.to_table(columns=[\"revenue\"], nearest={\"column\": \"vector\", \"q\": samples[0], \"k\": 10}).to_pandas()" + "sift1m.to_table(\n", + " columns=[\"revenue\"], nearest={\"column\": \"vector\", \"q\": samples[0], \"k\": 10}\n", + ").to_pandas()" ] } ], diff --git a/notebooks/youtube_transcript_search.ipynb b/notebooks/youtube_transcript_search.ipynb index 42b18d6e557..32f269f02e8 100644 --- a/notebooks/youtube_transcript_search.ipynb +++ b/notebooks/youtube_transcript_search.ipynb @@ -70,7 +70,7 @@ "source": [ "from datasets import load_dataset\n", "\n", - "data = load_dataset('jamescalam/youtube-transcriptions', split='train')\n", + "data = load_dataset(\"jamescalam/youtube-transcriptions\", split=\"train\")\n", "data" ] }, @@ -132,14 +132,20 @@ " text = vid.text.values\n", " time_end = vid[\"end\"].values\n", " contexts = vid.iloc[:-window:stride, :].copy()\n", - " contexts[\"text\"] = [' '.join(text[start_i:start_i+window])\n", - " for start_i in range(0, len(vid)-window, stride)]\n", - " contexts[\"end\"] = [time_end[start_i+window-1]\n", - " for start_i in range(0, len(vid)-window, stride)] \n", + " contexts[\"text\"] = [\n", + " \" \".join(text[start_i : start_i + window])\n", + " for start_i in range(0, len(vid) - window, stride)\n", + " ]\n", + " contexts[\"end\"] = [\n", + " time_end[start_i + window - 1]\n", + " for start_i in range(0, len(vid) - window, stride)\n", + " ]\n", " return contexts\n", + "\n", " # concat result from all videos\n", " return pd.concat([process_video(vid) for _, vid in raw_df.groupby(\"title\")])\n", "\n", + "\n", "df = contextualize(data.to_pandas(), 20, 4)" ] }, @@ -180,7 +186,6 @@ "metadata": {}, "outputs": [], "source": [ - "import functools\n", "import openai\n", "import ratelimiter\n", "from retry import retry\n", @@ -190,12 +195,14 @@ "# API limit at 60/min == 1/sec\n", "limiter = ratelimiter.RateLimiter(max_calls=0.9, period=1.0)\n", "\n", + "\n", "# Get the embedding with retry\n", "@retry(tries=10, delay=1, max_delay=30, backoff=3, jitter=1)\n", - "def embed_func(c): \n", + "def embed_func(c):\n", " rs = openai.Embedding.create(input=c, engine=embed_model)\n", " return [record[\"embedding\"] for record in rs[\"data\"]]\n", "\n", + "\n", "rate_limited = limiter(embed_func)" ] }, @@ -226,15 +233,19 @@ "\n", "openai.api_key = \"sk-...\"\n", "\n", + "\n", "# We request in batches rather than 1 embedding at a time\n", "def to_batches(arr, batch_size):\n", " length = len(arr)\n", + "\n", " def _chunker(arr):\n", " for start_i in range(0, len(df), batch_size):\n", - " yield arr[start_i:start_i+batch_size]\n", + " yield arr[start_i : start_i + batch_size]\n", + "\n", " # add progress meter\n", " yield from tqdm(_chunker(arr), total=math.ceil(length / batch_size))\n", - " \n", + "\n", + "\n", "batch_size = 1000\n", "batches = to_batches(df.text.values.tolist(), batch_size)\n", "embeds = [emb for c in batches for emb in rate_limited(c)]" @@ -280,10 +291,12 @@ } ], "source": [ - "ds = ds.create_index(\"vector\",\n", - " index_type=\"IVF_PQ\", \n", - " num_partitions=64, # IVF\n", - " num_sub_vectors=96) # PQ" + "ds = ds.create_index(\n", + " \"vector\",\n", + " index_type=\"IVF_PQ\",\n", + " num_partitions=64, # IVF\n", + " num_sub_vectors=96,\n", + ") # PQ" ] }, { @@ -304,28 +317,17 @@ "def create_prompt(query, context):\n", " limit = 3750\n", "\n", - " prompt_start = (\n", - " \"Answer the question based on the context below.\\n\\n\"+\n", - " \"Context:\\n\"\n", - " )\n", - " prompt_end = (\n", - " f\"\\n\\nQuestion: {query}\\nAnswer:\"\n", - " )\n", + " prompt_start = \"Answer the question based on the context below.\\n\\n\" + \"Context:\\n\"\n", + " prompt_end = f\"\\n\\nQuestion: {query}\\nAnswer:\"\n", " # append contexts until hitting limit\n", " for i in range(1, len(context)):\n", " if len(\"\\n\\n---\\n\\n\".join(context.text[:i])) >= limit:\n", " prompt = (\n", - " prompt_start +\n", - " \"\\n\\n---\\n\\n\".join(context.text[:i-1]) +\n", - " prompt_end\n", + " prompt_start + \"\\n\\n---\\n\\n\".join(context.text[: i - 1]) + prompt_end\n", " )\n", " break\n", - " elif i == len(context)-1:\n", - " prompt = (\n", - " prompt_start +\n", - " \"\\n\\n---\\n\\n\".join(context.text) +\n", - " prompt_end\n", - " ) \n", + " elif i == len(context) - 1:\n", + " prompt = prompt_start + \"\\n\\n---\\n\\n\".join(context.text) + prompt_end\n", " return prompt" ] }, @@ -350,16 +352,17 @@ "def complete(prompt):\n", " # query text-davinci-003\n", " res = openai.Completion.create(\n", - " engine='text-davinci-003',\n", + " engine=\"text-davinci-003\",\n", " prompt=prompt,\n", " temperature=0,\n", " max_tokens=400,\n", " top_p=1,\n", " frequency_penalty=0,\n", " presence_penalty=0,\n", - " stop=None\n", + " stop=None,\n", " )\n", - " return res['choices'][0]['text'].strip()\n", + " return res[\"choices\"][0][\"text\"].strip()\n", + "\n", "\n", "# check that it works\n", "query = \"who was the 12th person on the moon and when did they land?\"\n", @@ -381,8 +384,9 @@ " \"k\": 3,\n", " \"q\": emb,\n", " \"nprobes\": 20,\n", - " \"refine_factor\": 100\n", - " }).to_pandas()\n", + " \"refine_factor\": 100,\n", + " }\n", + " ).to_pandas()\n", " prompt = create_prompt(question, context)\n", " return complete(prompt), context.reset_index()" ] @@ -443,8 +447,10 @@ } ], "source": [ - "query = (\"Which training method should I use for sentence transformers \"\n", - " \"when I only have pairs of related sentences?\")\n", + "query = (\n", + " \"Which training method should I use for sentence transformers \"\n", + " \"when I only have pairs of related sentences?\"\n", + ")\n", "completion, context = answer(query)\n", "\n", "print(completion)\n", diff --git a/pre-commit-config.yaml b/pre-commit-config.yaml new file mode 100644 index 00000000000..20cdae104a8 --- /dev/null +++ b/pre-commit-config.yaml @@ -0,0 +1,11 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v2.3.0 + hooks: + - id: check-yaml + - id: end-of-file-fixer + - id: trailing-whitespace + - repo: https://github.com/psf/black + rev: 22.10.0 + hooks: + - id: black \ No newline at end of file diff --git a/rust/lance-linalg/src/svd.rs b/rust/lance-linalg/src/svd.rs new file mode 100644 index 00000000000..56014c53d83 --- /dev/null +++ b/rust/lance-linalg/src/svd.rs @@ -0,0 +1,361 @@ +//Helpful for printing error messages to determine whether +//the input (matrix, m, n) for the svd() function is valid +use lance_core::{Error, Result}; +//Machine epsilon for f64:smallest value where 1.0 + EPS != 1.0 +const EPS: f64 = 2.220_446_049_250_313e-16; +//Maximum number of Jacobi sweeps before giving up, when not +//scaled by dimension +const MAX_JACOBI_SWEEPS: usize = 100; + +//Creates an nxn identity matrix (row-major flat vector) +fn create_identity_matrix(n: usize) -> Vec { + let mut m = vec![0f64; n * n]; + for i in 0..n { + m[i * n + i] = 1.0; + } + m +} + +//Applies a Givens rotation to a symmetric matrix S in-place +//Computes S ← G^T * S * G where G is the rotation in the (p,q) plane +fn jacobi_rotate(s: &mut [f64], n: usize, p: usize, q: usize, c: f64, sn: f64) { + //Reads the current values of the 2x2 submatrix at (p,p), (q,q), and (p,q) + let s_pp = s[p * n + p]; + let s_qq = s[q * n + q]; + let s_pq = s[p * n + q]; + + //Updates the 2x2 submatrix using the closed-form Jacobi rotation formulas + s[p * n + p] = c * c * s_pp - 2.0 * sn * c * s_pq + sn * sn * s_qq; + s[q * n + q] = sn * sn * s_pp + 2.0 * sn * c * s_pq + c * c * s_qq; + s[p * n + q] = 0.0; + s[q * n + p] = 0.0; + + //Updates all other rows and columns that interact with p or q + for r in 0..n { + //Skips the 2x2 block already handles above + if r == p || r == q { + continue; + } + let s_rp = s[r * n + p]; + let s_rq = s[r * n + q]; + let new_rp = c * s_rp - sn * s_rq; + let new_rq = sn * s_rp + c * s_rq; + + //Updates both (r,p) and (p,r) to maintain symmetry + s[r * n + p] = new_rp; + s[p * n + r] = new_rp; + + //Updates both (r,q) and (q,r) to maintain symmetry + s[r * n + q] = new_rq; + s[q * n + r] = new_rq; + } +} + +//Accumulates a Givens rotation into V from the right: V ← V * G +//Rotates columns p and q of V by angle (c, sn) +fn apply_givens_rotation_from_right(v: &mut [f64], n: usize, p: usize, q: usize, c: f64, sn: f64) { + for r in 0..n { + let vp = v[r * n + p]; + let vq = v[r * n + q]; + + v[r * n + p] = c * vp - sn * vq; + v[r * n + q] = sn * vp + c * vq; + } +} + +//Jacobi eigenvalue algorithm: Decomposes a symmetric matrix A into V * diag(eigenvalues) * V^T +//by repeatedly applying Givens rotation to zero out off-diagonal entries. +//Input: a - summetric nxn matrix (row-major flat vector) +//Output: (eigenvalues, V) where the columns of V are the eigenvectors +fn jacobi_eigen(a: &[f64], n: usize) -> Result<(Vec, Vec)> { + //Copy of the matrix; this copy will be diagonalized in-place. + let mut s = a.to_vec(); + //Accumulates the product of all Givens rotations + let mut v = create_identity_matrix(n); + + //There is nothing to do for 1x1 matrices due to there being no off-diagonal entries. + if n <= 1 { + let eigenvalues: Vec = (0..n).map(|i| s[i * n + i]).collect(); + return Ok((eigenvalues, v)); + } + + //Uses the Frobenius norm of the input matrix as a scale reference for convergence + let norm: f64 = s.iter().map(|&x| x * x).sum::().sqrt(); + if norm < EPS { + let eigenvalues: Vec = (0..n).map(|i| s[i * n + i]).collect(); + return Ok((eigenvalues, v)); + } + + //Relative tolerance: off-diagonal entries (and the overall off-diagonal + //Frobenius norm) smaller than this, relative to the matrix's own scale, + //are converged + let relative_tolerance = EPS * norm * (n as f64); + + //Scales the sweep budget with the matrix dimension + //Larger matrices need more sweeps and individual rotations to + //fully diagonalize. + let max_sweeps = MAX_JACOBI_SWEEPS.max(n * n); + + let mut converged = false; + for _sweep in 0..max_sweeps { + //This is one full cyclic sweep. + //It visits each (p, q) pair with p < q exactly once, + //applying a rotation when the entry is not already negligible. + for p in 0..n { + for q in p + 1..n { + let s_pq = s[p * n + q]; + if s_pq.abs() <= relative_tolerance { + continue; + } + + //Computes the Jacobi rotation angle theta that zeros out s[p,q] + let diff = s[q * n + q] - s[p * n + p]; + let theta = if diff.abs() < EPS { + std::f64::consts::FRAC_PI_4 + } else { + 0.5 * (2.0 * s_pq / diff).atan() + }; + let (sn, c) = theta.sin_cos(); + + //Applies the Givens rotation: S ← G^T * S * G + //This zeroes out s[p,q] and s[q,p] while updating the rest of the matrix. + jacobi_rotate(&mut s, n, p, q, c, sn); + + //Accumulates the rotation into V: V ← V * G + //At convergence, V's columns are the eigenvectors of the original matrix. + apply_givens_rotation_from_right(&mut v, n, p, q, c, sn); + } + } + + //Recomputes the off-diagonal Frobenius norm after the full sweep to + //determine whether another sweep is needed + let mut off_diag_sq = 0.0f64; + for p in 0..n { + for q in p + 1..n { + off_diag_sq += s[p * n + q] * s[p * n + q]; + } + } + if off_diag_sq.sqrt() < relative_tolerance { + converged = true; + break; + } + } + + //Checks whether the Jacobi eigensolver converged within the + //maximum number of sweeps + if !converged { + return Err(Error::invalid_input(format!( + "svd: Jacobi eigensolver failed to converge within {max_sweeps} sweeps for a {n}x{n} matrix" + ))); + } + + //Extracts eigenvalues from the diagonal of the now-diagonalized matrix + let eigenvalues: Vec = (0..n).map(|i| s[i * n + i]).collect(); + Ok((eigenvalues, v)) +} + +//Computes C = A^T A where A is a mxn row-major flat vector +//and C is an nxn row-major flat vector +fn compute_ata(a: &[f64], m: usize, n: usize) -> Vec { + let mut c = vec![0f64; n * n]; + for i in 0..n { + for l in 0..m { + let a_li = a[l * n + i]; + for j in 0..n { + c[i * n + j] += a_li * a[l * n + j]; + } + } + } + c +} + +//Modified Gram-Schmidt orthonormalization +//Takes a list of column vectors and makes them orthonormal +//If a column is zero (e.g., from a zero singular value), it is +//replaced with a standard basis vector orthogonal to all prior columns. +fn gram_schmidt(cols: &mut Vec>) { + let ncols = cols.len(); + for i in 0..ncols { + //Subtracts the projection of cols[i] onto each already-orthonormal column + for j in 0..i { + let dot: f64 = cols[i].iter().zip(cols[j].iter()).map(|(&a, &b)| a * b).sum(); + let cj = cols[j].clone(); + for (a, b) in cols[i].iter_mut().zip(cj.iter()) { + *a -= dot * b; + } + } + + //Normalizes the resulting vector + let norm = cols[i].iter().map(|&x| x * x).sum::().sqrt(); + if norm > EPS { + for x in cols[i].iter_mut() { + *x /= norm; + } + } else { + //If the column is zero, finds a standard basis vector not in the span yet + let dim = cols[i].len(); + 'search: for k in 0..dim { + let mut e = vec![0.0f64; dim]; + //Tries the k-th standard basis vector and + //orthogonalize it against all previous columns + e[k] = 1.0; + for j in 0..i { + let dot: f64 = e.iter().zip(cols[j].iter()).map(|(&a, &b)| a * b).sum(); + let cj = cols[j].clone(); + for (a, b) in e.iter_mut().zip(cj.iter()) { + *a -= dot * b; + } + } + let n2 = e.iter().map(|&x| x * x).sum::().sqrt(); + if n2 > EPS { + //Found a valid replacement vector. Normalizes and uses it. + for x in e.iter_mut() { + *x /= n2; + } + cols[i] = e; + break 'search; + } + } + } + } +} + +//Computes y = A * x where A is an mxn row-major flat vector, +//x is a vector of length n, and y is a vector of length m +fn multiply_a_by_vector(a: &[f64], x: &[f64], m: usize, n: usize) -> Vec { + (0..m).map(|i| (0..n).map(|j| a[i * n + j] * x[j]).sum()).collect() +} + +/// Computes the SVD of an `m x n` row-major matrix `a`, decomposing it into +/// `U * diag(sigma) * V^T`. +/// +/// Returns `(u, sigma, vt)` where `u` is `m x m` (row-major), `sigma` has length +/// `min(m, n)` sorted descending, and `vt` is `n x n` (row-major, rows are right +/// singular vectors). +/// +/// # Example +/// ``` +/// let (u, sigma, vt) = svd(&[1.0, 0.0, 0.0, 1.0], 2, 2)?; +/// assert_eq!(u, vec![1.0, 0.0, 0.0, 1.0]); +/// assert_eq!(sigma, vec![1.0, 1.0]); +/// assert_eq!(vt, vec![1.0, 0.0, 0.0, 1.0]); +/// ``` +pub fn svd(a: &[f64], m: usize, n: usize) -> Result<(Vec, Vec, Vec)> { + //Checks whether the input matrix A has at least 1 row and at least 1 column + if a.is_empty() || m == 0 || n == 0 { + return Err(Error::invalid_input(format!( + "svd: matrix must have >=1 row and column (m={m}, n={n}, len={})", + a.len() + ))); + } + + //Checks that m*n, m*m, and n*n don't overflow usize before they are used + let mn = m.checked_mul(n).ok_or_else(|| { + Error::invalid_input(format!( + "svd: m*n overflows usize (m={m}, n={n}, len={})", + a.len() + )) + })?; + let _mm = m.checked_mul(m).ok_or_else(|| { + Error::invalid_input(format!( + "svd: m*m overflows usize (m={m}, n={n}, len={})", + a.len() + )) + })?; + let _nn = n.checked_mul(n).ok_or_else(|| { + Error::invalid_input(format!( + "svd: n*n overflows usize (m={m}, n={n}, len={})", + a.len() + )) + })?; + + //Checks whether the data length of the input matrix A matches the + //product of the specified number of rows and number of columns + if a.len() != mn { + return Err(Error::invalid_input(format!( + "svd: data length {} != m*n ({m}*{n})", a.len() + ))); + } + //Checks whether the input matrix A contains null or infinite entries + if !a.iter().all(|x| x.is_finite()) { + return Err(Error::invalid_input( + "svd: matrix contains NaN or infinite entries" + )); + } + + //Step 1: Forms A^T A (nxn symmetric positive semi-definite matrix) + let ata = compute_ata(a, m, n); + //Step 2: Eigendecomposes A^T A into eigenvalues λ and + //eigenvectors V using Jacobi iteration + let (eigenvalues, v) = jacobi_eigen(&ata, n); + + //k is the number of singular values. + let k = m.min(n); + + //Sorts all n eigenvalues in descending order + let mut order: Vec = (0..n).collect(); + for i in 1..n { + let mut j = i; + while j > 0 && eigenvalues[order[j - 1]] < eigenvalues[order[j]] { + order.swap(j - 1, j); + j -= 1; + } + } + + //Step 3: Computes singular values σᵢ = √λᵢ for the top k eigenvalues + let mut sigma: Vec = order[..k] + .iter() + .map(|&i| if eigenvalues[i] > 0.0 { eigenvalues[i].sqrt() } else { 0.0 }) + .collect(); + + //Step 4: Computes the left singular vectors uᵢ = (1/σᵢ) * A * vᵢ + //Builds U column by column using the sorted eigenvectors + let mut u_cols: Vec> = Vec::with_capacity(m); + for index in 0..k { + //eᵢ is the index of the idx-th largest eigenvector in V + let ei = order[index]; + //Extracts the eigenvector eᵢ (column eᵢ of V, stored as a row-major flat nxn vector) + let vi: Vec = (0..n).map(|r| v[r * n + ei]).collect(); + //Computes A * vᵢ to get the unnormalized left singular vector + let av = multiply_a_by_vector(a, &vi, m, n); + if sigma[index] > EPS * 10.0 { + //Normalizes the left singular vector to get uᵢ + u_cols.push(av.iter().map(|&x| x / sigma[index]).collect()); + } else { + //The placeholder column of sigma is the zero singular value + //This column is filled by the Gram-Schmidt algorithm. + sigma[index] = 0.0; + u_cols.push(vec![0.0; m]); + } + } + + //Step 5: If m > k, pad U with zero columns. + //The Gram-Schmidt algorithm will fill these columns with + //orthonormal vectors that complete the basis of R^m. + for _ in k..m { + u_cols.push(vec![0.0; m]); + } + + //Orthonormalize all columns of U + gram_schmidt(&mut u_cols); + + //Pack the columns of U into an mxm flat row-major matrix + let mut u = vec![0f64; m * m]; + for (ci, col) in u_cols.iter().enumerate() { + for (ri, &value) in col.iter().enumerate() { + u[ri * m + ci] = value; + } + } + + //Step 6: Build V^T (nxn row-major flat vector) + //Row i of V^T = column order[i] of V = the i-th right singular vector + let mut vt = vec![0f64; n * n]; + for new_row in 0..n { + let old_col = order[new_row]; + for c in 0..n { + vt[new_row * n + c] = v[c * n + old_col]; + } + } + + Ok((u, sigma, vt)) +} \ No newline at end of file diff --git a/rust/lance/benches/mem_wal/vector/hnsw/disk_ann_compare.py b/rust/lance/benches/mem_wal/vector/hnsw/disk_ann_compare.py index e3b81123ece..1599c2d38ed 100644 --- a/rust/lance/benches/mem_wal/vector/hnsw/disk_ann_compare.py +++ b/rust/lance/benches/mem_wal/vector/hnsw/disk_ann_compare.py @@ -18,7 +18,11 @@ recall@10 vs p50/p99 latency and QPS. The Lance index is served fully cached (large index_cache_size_bytes). """ -import argparse, json, os, time + +import argparse +import json +import os +import time import numpy as np K = 10 @@ -27,7 +31,9 @@ DIM = 1536 EF_SWEEP = [16, 32, 64, 128, 256] HF_TREE = "https://huggingface.co/api/datasets/KShivendu/dbpedia-entities-openai-1M/tree/main/data" -HF_BASE = "https://huggingface.co/datasets/KShivendu/dbpedia-entities-openai-1M/resolve/main/" +HF_BASE = ( + "https://huggingface.co/datasets/KShivendu/dbpedia-entities-openai-1M/resolve/main/" +) def data_dir(base, rows): @@ -42,10 +48,13 @@ def normalize(x): # ---------------- prepare ---------------- def load_corpus(cache_dir, needed): - import requests, pyarrow.parquet as pq + import requests + import pyarrow.parquet as pq + os.makedirs(cache_dir, exist_ok=True) shards = sorted( - e["path"] for e in requests.get(HF_TREE, timeout=60).json() + e["path"] + for e in requests.get(HF_TREE, timeout=60).json() if e["type"] == "file" and e["path"].endswith(".parquet") ) out = np.empty((needed, DIM), dtype=np.float32) @@ -62,7 +71,7 @@ def load_corpus(cache_dir, needed): col = pq.read_table(local, columns=["openai"]).column("openai") arr = np.stack(col.to_pylist()).astype(np.float32) take = min(len(arr), needed - n) - out[n:n + take] = arr[:take] + out[n : n + take] = arr[:take] n += take print(f" shard {os.path.basename(rel)} -> {take} (cum {n})", flush=True) assert n == needed, f"only got {n}/{needed}" @@ -98,10 +107,13 @@ def cmd_prepare(args): rng = np.random.default_rng(SEED) qidx = rng.choice(args.rows, size=NUM_QUERIES, replace=False) queries = corpus[qidx].copy() - print(f"corpus={len(corpus)} queries={len(queries)} dim={DIM}; computing GT...", flush=True) + print( + f"corpus={len(corpus)} queries={len(queries)} dim={DIM}; computing GT...", + flush=True, + ) t = time.perf_counter() gt = numpy_ground_truth(corpus, queries) - print(f" GT in {time.perf_counter()-t:.1f}s", flush=True) + print(f" GT in {time.perf_counter() - t:.1f}s", flush=True) np.save(os.path.join(d, "corpus.npy"), corpus) np.save(os.path.join(d, "queries.npy"), queries) np.save(os.path.join(d, "gt.npy"), gt) @@ -110,7 +122,9 @@ def cmd_prepare(args): # ---------------- shared run helpers ---------------- def recall_at_k(gt, got): - return sum(len(set(g.tolist()) & set(r.tolist())) for g, r in zip(gt, got)) / (len(gt) * K) + return sum(len(set(g.tolist()) & set(r.tolist())) for g, r in zip(gt, got)) / ( + len(gt) * K + ) def latency_qps(query_fn, queries, repeats=3): @@ -133,56 +147,95 @@ def sweep(name, make_q, params, queries, gt): got = np.stack([qf(v) for v in queries]) rec = recall_at_k(gt, got) p50, p99, qps = latency_qps(qf, queries) - rows.append({"param": p, "recall": rec, "p50_us": p50, "p99_us": p99, "qps": qps}) - print(f" {name} param={p} recall={rec:.4f} p50={p50:.0f}us p99={p99:.0f}us qps={qps:.0f}", flush=True) + rows.append( + {"param": p, "recall": rec, "p50_us": p50, "p99_us": p99, "qps": qps} + ) + print( + f" {name} param={p} recall={rec:.4f} p50={p50:.0f}us p99={p99:.0f}us qps={qps:.0f}", + flush=True, + ) return rows # ---------------- systems ---------------- def run_lance(base, rows, corpus, queries, gt): - import lance, pyarrow as pa, shutil + import lance + import pyarrow as pa + import shutil + uri = os.path.join(base, f"lance_{rows}") shutil.rmtree(uri, ignore_errors=True) - vecs = pa.FixedSizeListArray.from_arrays(pa.array(corpus.reshape(-1), type=pa.float32()), DIM) + vecs = pa.FixedSizeListArray.from_arrays( + pa.array(corpus.reshape(-1), type=pa.float32()), DIM + ) tbl = pa.table({"id": pa.array(np.arange(rows, dtype=np.int64)), "vec": vecs}) ds = lance.write_dataset(tbl, uri, mode="overwrite") # The flushed memtable index is a SINGLE-partition HNSW+SQ, so model it with # num_partitions=1 (nprobes=1); ef is the search knob, like DiskANN/FAISS. t = time.perf_counter() - ds.create_index("vec", "IVF_HNSW_SQ", metric="cosine", num_partitions=1, - m=20, ef_construction=150) + ds.create_index( + "vec", + "IVF_HNSW_SQ", + metric="cosine", + num_partitions=1, + m=20, + ef_construction=150, + ) build_s = time.perf_counter() - t ds = lance.dataset(uri, index_cache_size_bytes=48 * 1024**3) def make_q(ef): def q(v): - return ds.to_table(nearest={"column": "vec", "q": v, "k": K, - "nprobes": 1, "ef": ef}, - columns=["id"]).column("id").to_numpy() + return ( + ds.to_table( + nearest={"column": "vec", "q": v, "k": K, "nprobes": 1, "ef": ef}, + columns=["id"], + ) + .column("id") + .to_numpy() + ) + return q - return {"build_s": build_s, "nlist": 1, "sweep": sweep("lance", make_q, None, queries, gt)} + + return { + "build_s": build_s, + "nlist": 1, + "sweep": sweep("lance", make_q, None, queries, gt), + } def run_lance_flushed(base, rows, corpus, queries, gt, lance_path, id_offset, column): # Open a flushed MemTable generation directly from its dataset path and # benchmark its on-disk IVF_HNSW_SQ index (single partition), fully cached. import lance + ds = lance.dataset(lance_path, index_cache_size_bytes=48 * 1024**3) def make_q(ef): def q(v): - ids = ds.to_table(nearest={"column": column, "q": v, "k": K, - "nprobes": 1, "ef": ef}, - columns=["id"]).column("id").to_numpy() + ids = ( + ds.to_table( + nearest={"column": column, "q": v, "k": K, "nprobes": 1, "ef": ef}, + columns=["id"], + ) + .column("id") + .to_numpy() + ) return ids - id_offset # map flushed-gen id -> corpus index + return q - return {"lance_path": lance_path, "id_offset": id_offset, - "sweep": sweep("lance_flushed", make_q, None, queries, gt)} + + return { + "lance_path": lance_path, + "id_offset": id_offset, + "sweep": sweep("lance_flushed", make_q, None, queries, gt), + } def run_faiss(base, rows, corpus, queries, gt): # Full-precision HNSW reference (shows what no quantization buys). import faiss + index = faiss.IndexHNSWFlat(DIM, 32, faiss.METRIC_INNER_PRODUCT) index.hnsw.efConstruction = 200 t = time.perf_counter() @@ -194,16 +247,20 @@ def make_q(ef): def q(v): index.hnsw.efSearch = ef return index.search(v.reshape(1, -1), K)[1][0] + return q + return {"build_s": build_s, "sweep": sweep("faiss", make_q, None, queries, gt)} def run_faiss_sq(base, rows, corpus, queries, gt): # HNSW + 8-bit scalar quantization — apples-to-apples with Lance IVF_HNSW_SQ. import faiss + try: - index = faiss.IndexHNSWSQ(DIM, faiss.ScalarQuantizer.QT_8bit, 32, - faiss.METRIC_INNER_PRODUCT) + index = faiss.IndexHNSWSQ( + DIM, faiss.ScalarQuantizer.QT_8bit, 32, faiss.METRIC_INNER_PRODUCT + ) except Exception: # Fall back to L2; on unit-normalized vectors L2 ranking == cosine. index = faiss.IndexHNSWSQ(DIM, faiss.ScalarQuantizer.QT_8bit, 32) @@ -218,28 +275,44 @@ def make_q(ef): def q(v): index.hnsw.efSearch = ef return index.search(v.reshape(1, -1), K)[1][0] + return q + return {"build_s": build_s, "sweep": sweep("faiss_sq", make_q, None, queries, gt)} def run_diskann(base, rows, corpus, queries, gt): import diskannpy as dap + idx_dir = os.path.join(base, f"diskann_{rows}") os.makedirs(idx_dir, exist_ok=True) t = time.perf_counter() dap.build_memory_index( - data=corpus, distance_metric="cosine", index_directory=idx_dir, - index_prefix="ann", complexity=150, graph_degree=64, - num_threads=0, alpha=1.2, use_pq_build=False, num_pq_bytes=0, + data=corpus, + distance_metric="cosine", + index_directory=idx_dir, + index_prefix="ann", + complexity=150, + graph_degree=64, + num_threads=0, + alpha=1.2, + use_pq_build=False, + num_pq_bytes=0, ) build_s = time.perf_counter() - t - idx = dap.StaticMemoryIndex(index_directory=idx_dir, index_prefix="ann", - num_threads=0, initial_search_complexity=256) + idx = dap.StaticMemoryIndex( + index_directory=idx_dir, + index_prefix="ann", + num_threads=0, + initial_search_complexity=256, + ) def make_q(L): def q(v): return idx.search(v, k_neighbors=K, complexity=max(L, K)).identifiers + return q + return {"build_s": build_s, "sweep": sweep("diskann", make_q, None, queries, gt)} @@ -250,11 +323,23 @@ def cmd_run(args): gt = np.load(os.path.join(d, "gt.npy")) print(f"=== {args.system} rows={args.rows} corpus={len(corpus)} ===", flush=True) if args.system == "lance_flushed": - res = run_lance_flushed(args.base, args.rows, corpus, queries, gt, - args.lance_path, args.id_offset, args.column) + res = run_lance_flushed( + args.base, + args.rows, + corpus, + queries, + gt, + args.lance_path, + args.id_offset, + args.column, + ) else: - fn = {"lance": run_lance, "faiss": run_faiss, "faiss_sq": run_faiss_sq, - "diskann": run_diskann}[args.system] + fn = { + "lance": run_lance, + "faiss": run_faiss, + "faiss_sq": run_faiss_sq, + "diskann": run_diskann, + }[args.system] res = fn(args.base, args.rows, corpus, queries, gt) res["rows"] = args.rows res["system"] = args.system @@ -267,9 +352,16 @@ def cmd_run(args): def main(): ap = argparse.ArgumentParser() sub = ap.add_subparsers(dest="cmd", required=True) - p = sub.add_parser("prepare"); p.add_argument("--rows", type=int, required=True); p.add_argument("--base", required=True) - r = sub.add_parser("run"); r.add_argument("--rows", type=int, required=True); r.add_argument("--base", required=True); r.add_argument("--system", required=True) - r.add_argument("--lance-path", default=None); r.add_argument("--id-offset", type=int, default=0); r.add_argument("--column", default="vector") + p = sub.add_parser("prepare") + p.add_argument("--rows", type=int, required=True) + p.add_argument("--base", required=True) + r = sub.add_parser("run") + r.add_argument("--rows", type=int, required=True) + r.add_argument("--base", required=True) + r.add_argument("--system", required=True) + r.add_argument("--lance-path", default=None) + r.add_argument("--id-offset", type=int, default=0) + r.add_argument("--column", default="vector") args = ap.parse_args() (cmd_prepare if args.cmd == "prepare" else cmd_run)(args) diff --git a/skills/lance-user-guide/scripts/python_end_to_end.py b/skills/lance-user-guide/scripts/python_end_to_end.py index ec2d02713c9..b366fe93151 100644 --- a/skills/lance-user-guide/scripts/python_end_to_end.py +++ b/skills/lance-user-guide/scripts/python_end_to_end.py @@ -11,16 +11,24 @@ import lance -def _build_fixed_size_vectors(num_rows: int, dim: int) -> tuple[pa.FixedSizeListArray, np.ndarray]: +def _build_fixed_size_vectors( + num_rows: int, dim: int +) -> tuple[pa.FixedSizeListArray, np.ndarray]: vectors = np.random.rand(num_rows, dim).astype("float32") flat = pa.array(vectors.reshape(-1), type=pa.float32()) return pa.FixedSizeListArray.from_arrays(flat, dim), vectors def main() -> None: - parser = argparse.ArgumentParser(description="Minimal Lance write/index/query example") - parser.add_argument("--uri", default="example.lance", help="Dataset URI (directory)") - parser.add_argument("--mode", default="overwrite", choices=["create", "append", "overwrite"]) + parser = argparse.ArgumentParser( + description="Minimal Lance write/index/query example" + ) + parser.add_argument( + "--uri", default="example.lance", help="Dataset URI (directory)" + ) + parser.add_argument( + "--mode", default="overwrite", choices=["create", "append", "overwrite"] + ) parser.add_argument("--rows", type=int, default=1000) parser.add_argument("--dim", type=int, default=32) @@ -40,7 +48,13 @@ def main() -> None: uri = str(Path(args.uri)) vec_arr, vec_np = _build_fixed_size_vectors(args.rows, args.dim) categories = pa.array(["a" if i % 2 == 0 else "b" for i in range(args.rows)]) - table = pa.table({"id": pa.array(range(args.rows), pa.int64()), "category": categories, "vector": vec_arr}) + table = pa.table( + { + "id": pa.array(range(args.rows), pa.int64()), + "category": categories, + "vector": vec_arr, + } + ) ds = lance.write_dataset(table, uri, mode=args.mode) ds = lance.dataset(uri) diff --git a/tall pylance b/tall pylance new file mode 100644 index 00000000000..933b1652748 --- /dev/null +++ b/tall pylance @@ -0,0 +1,173 @@ +diff --git a/benchmarks/full_report/report.ipynb b/benchmarks/full_report/report.ipynb +index 11a4924c..70eb810a 100644 +--- a/benchmarks/full_report/report.ipynb ++++ b/benchmarks/full_report/report.ipynb +@@ -8,11 +8,12 @@ + "outputs": [], + "source": [ + "import sys\n", ++ "\n", + "sys.dont_write_bytecode = True\n", + "\n", + "import os\n", + "\n", +- "module_path = os.path.abspath(os.path.join('.'))\n", ++ "module_path = os.path.abspath(os.path.join(\".\"))\n", + "if module_path not in sys.path:\n", + " sys.path.append(module_path)" + ] +@@ -61,58 +62,91 @@ + " gt_in_sample = knn(in_sample, data, metric, 10)\n", + "\n", + " print(\"generated gt\")\n", +- " \n", ++ "\n", + " with tempfile.TemporaryDirectory() as d:\n", + " write_lance(d, data)\n", + " ds = lance.dataset(d)\n", + "\n", +- " for q, target in zip(tqdm(in_sample, desc=\"checking brute force\"), gt_in_sample):\n", +- " res = ds.to_table(nearest={\n", +- " \"column\": \"vec\",\n", +- " \"q\": q,\n", +- " \"k\": 10,\n", +- " \"metric\": metric,\n", +- " }, columns=[\"id\"])\n", ++ " for q, target in zip(\n", ++ " tqdm(in_sample, desc=\"checking brute force\"), gt_in_sample\n", ++ " ):\n", ++ " res = ds.to_table(\n", ++ " nearest={\n", ++ " \"column\": \"vec\",\n", ++ " \"q\": q,\n", ++ " \"k\": 10,\n", ++ " \"metric\": metric,\n", ++ " },\n", ++ " columns=[\"id\"],\n", ++ " )\n", + " assert len(np.intersect1d(res[\"id\"].to_numpy(), target)) == 10\n", +- " \n", +- " ds = ds.create_index(\"vec\", \"IVF_PQ\", metric=metric, num_partitions=num_partitions, num_sub_vectors=num_sub_vectors)\n", +- " \n", ++ "\n", ++ " ds = ds.create_index(\n", ++ " \"vec\",\n", ++ " \"IVF_PQ\",\n", ++ " metric=metric,\n", ++ " num_partitions=num_partitions,\n", ++ " num_sub_vectors=num_sub_vectors,\n", ++ " )\n", ++ "\n", + " recall_data = []\n", + " for nprobes in nprobes_list:\n", + " for refine_factor in refine_factor_list:\n", + " hits = 0\n", + " # check that brute force impl is correct\n", +- " for q, target in zip(tqdm(query, desc=f\"out of sample, nprobes={nprobes}, refine={refine_factor}\"), gt):\n", +- " res = ds.to_table(nearest={\n", +- " \"column\": \"vec\",\n", +- " \"q\": q,\n", +- " \"k\": 10,\n", +- " \"nprobes\": nprobes,\n", +- " \"refine_factor\": refine_factor,\n", +- " }, columns=[\"id\"])[\"id\"].to_numpy()\n", ++ " for q, target in zip(\n", ++ " tqdm(\n", ++ " query,\n", ++ " desc=f\"out of sample, nprobes={nprobes}, refine={refine_factor}\",\n", ++ " ),\n", ++ " gt,\n", ++ " ):\n", ++ " res = ds.to_table(\n", ++ " nearest={\n", ++ " \"column\": \"vec\",\n", ++ " \"q\": q,\n", ++ " \"k\": 10,\n", ++ " \"nprobes\": nprobes,\n", ++ " \"refine_factor\": refine_factor,\n", ++ " },\n", ++ " columns=[\"id\"],\n", ++ " )[\"id\"].to_numpy()\n", + " hits += len(np.intersect1d(res, target))\n", +- " recall_data.append([\n", +- " \"out_of_sample\",\n", +- " nprobes,\n", +- " refine_factor,\n", +- " hits / 10 / len(gt),\n", +- " ])\n", ++ " recall_data.append(\n", ++ " [\n", ++ " \"out_of_sample\",\n", ++ " nprobes,\n", ++ " refine_factor,\n", ++ " hits / 10 / len(gt),\n", ++ " ]\n", ++ " )\n", + " # check that brute force impl is correct\n", +- " for q, target in zip(tqdm(in_sample, desc=f\"in sample nprobes={nprobes}, refine={refine_factor}\"), gt_in_sample):\n", +- " res = ds.to_table(nearest={\n", +- " \"column\": \"vec\",\n", +- " \"q\": q,\n", +- " \"k\": 10,\n", +- " \"nprobes\": nprobes,\n", +- " \"refine_factor\": refine_factor,\n", +- " }, columns=[\"id\"])[\"id\"].to_numpy()\n", ++ " for q, target in zip(\n", ++ " tqdm(\n", ++ " in_sample,\n", ++ " desc=f\"in sample nprobes={nprobes}, refine={refine_factor}\",\n", ++ " ),\n", ++ " gt_in_sample,\n", ++ " ):\n", ++ " res = ds.to_table(\n", ++ " nearest={\n", ++ " \"column\": \"vec\",\n", ++ " \"q\": q,\n", ++ " \"k\": 10,\n", ++ " \"nprobes\": nprobes,\n", ++ " \"refine_factor\": refine_factor,\n", ++ " },\n", ++ " columns=[\"id\"],\n", ++ " )[\"id\"].to_numpy()\n", + " hits += len(np.intersect1d(res, target))\n", +- " recall_data.append([\n", +- " \"in_sample\",\n", +- " nprobes,\n", +- " refine_factor,\n", +- " hits / 10 / len(gt_in_sample),\n", +- " ])\n", ++ " recall_data.append(\n", ++ " [\n", ++ " \"in_sample\",\n", ++ " nprobes,\n", ++ " refine_factor,\n", ++ " hits / 10 / len(gt_in_sample),\n", ++ " ]\n", ++ " )\n", + " return recall_data" + ] + }, +@@ -124,15 +158,19 @@ + "outputs": [], + "source": [ + "def make_plot(recall_data):\n", +- " df = pd.DataFrame(recall_data, columns=[\"case\", \"nprobes\", \"refine_factor\", \"recall\"])\n", +- " \n", ++ " df = pd.DataFrame(\n", ++ " recall_data, columns=[\"case\", \"nprobes\", \"refine_factor\", \"recall\"]\n", ++ " )\n", ++ "\n", + " num_cases = len(df[\"case\"].unique())\n", + " (fig, axs) = plt.subplots(1, 2, figsize=(16, 8))\n", +- " \n", ++ "\n", + " for case, ax in zip(df[\"case\"].unique(), axs):\n", + " current_case = df[df[\"case\"] == case]\n", + " sns.heatmap(\n", +- " current_case.drop(columns=[\"case\"]).set_index([\"nprobes\", \"refine_factor\"])[\"recall\"].unstack(),\n", ++ " current_case.drop(columns=[\"case\"])\n", ++ " .set_index([\"nprobes\", \"refine_factor\"])[\"recall\"]\n", ++ " .unstack(),\n", + " annot=True,\n", + " ax=ax,\n", + " ).set(title=f\"Recall -- {case}\")" diff --git a/test_data/v0.30.0_pre_created_at/datagen.py b/test_data/v0.30.0_pre_created_at/datagen.py index 3d610c1fedc..8a669c12cda 100644 --- a/test_data/v0.30.0_pre_created_at/datagen.py +++ b/test_data/v0.30.0_pre_created_at/datagen.py @@ -2,6 +2,7 @@ Generate test data with Lance 0.29.0 to create an index without created_at field. This tests backward compatibility for the created_at field added in 0.30.0. """ + import lance import pyarrow as pa import pyarrow.compute as pc @@ -13,12 +14,14 @@ ndims = 16 nvecs = 256 -data = pa.table({ - "id": pa.array(range(nvecs)), - "vec": pa.FixedSizeListArray.from_arrays( - pc.random(ndims * nvecs).cast(pa.float32()), ndims - ), -}) +data = pa.table( + { + "id": pa.array(range(nvecs)), + "vec": pa.FixedSizeListArray.from_arrays( + pc.random(ndims * nvecs).cast(pa.float32()), ndims + ), + } +) # Write dataset dataset = lance.write_dataset(data, "index_without_created_at") @@ -33,4 +36,4 @@ print("Created dataset with IVF_PQ index on 'vec' field") print(f"Dataset version: {dataset.version}") -print(f"Index created with Lance {lance.__version__}") \ No newline at end of file +print(f"Index created with Lance {lance.__version__}") diff --git a/test_debug.py b/test_debug.py index e42798d5ff3..3f4452f7bf4 100644 --- a/test_debug.py +++ b/test_debug.py @@ -5,46 +5,52 @@ from lance.namespace import DirectoryNamespace CONFIG = { - 'allow_http': 'true', - 'aws_access_key_id': 'ACCESS_KEY', - 'aws_secret_access_key': 'SECRET_KEY', - 'aws_endpoint': 'http://localhost:4566', - 'aws_region': 'us-east-1', + "allow_http": "true", + "aws_access_key_id": "ACCESS_KEY", + "aws_secret_access_key": "SECRET_KEY", + "aws_endpoint": "http://localhost:4566", + "aws_region": "us-east-1", } storage_options = copy.deepcopy(CONFIG) storage_options_with_refresh = dict(storage_options) -storage_options_with_refresh['refresh_offset_millis'] = '1000' +storage_options_with_refresh["refresh_offset_millis"] = "1000" -dir_props = {f'storage.{k}': v for k, v in storage_options_with_refresh.items()} -dir_props['root'] = 's3://lance-namespace-integtest/namespace_root' -dir_props['ops_metrics_enabled'] = 'true' -dir_props['vend_input_storage_options'] = 'true' -dir_props['vend_input_storage_options_refresh_interval_millis'] = '3600000' +dir_props = {f"storage.{k}": v for k, v in storage_options_with_refresh.items()} +dir_props["root"] = "s3://lance-namespace-integtest/namespace_root" +dir_props["ops_metrics_enabled"] = "true" +dir_props["vend_input_storage_options"] = "true" +dir_props["vend_input_storage_options_refresh_interval_millis"] = "3600000" namespace = DirectoryNamespace(**dir_props) -table1 = pa.Table.from_pylist([{'a': 1, 'b': 2}]) -table_name = 'debug_print3_' + uuid.uuid4().hex -table_id = ['test_ns', table_name] +table1 = pa.Table.from_pylist([{"a": 1, "b": 2}]) +table_name = "debug_print3_" + uuid.uuid4().hex +table_id = ["test_ns", table_name] -print('=== Creating table ===') +print("=== Creating table ===") ds = lance.write_dataset( table1, namespace_client=namespace, table_id=table_id, - mode='create', + mode="create", storage_options=storage_options, ) -print('=== Table created ===') -print('Describe count after write:', namespace.retrieve_ops_metrics().get('describe_table', 0)) +print("=== Table created ===") +print( + "Describe count after write:", + namespace.retrieve_ops_metrics().get("describe_table", 0), +) -print('') -print('=== Opening dataset via lance.dataset() ===') +print("") +print("=== Opening dataset via lance.dataset() ===") ds_from_namespace = lance.dataset( namespace_client=namespace, table_id=table_id, storage_options=storage_options, ) -print('=== Dataset opened ===') -print('Describe count after lance.dataset():', namespace.retrieve_ops_metrics().get('describe_table', 0)) +print("=== Dataset opened ===") +print( + "Describe count after lance.dataset():", + namespace.retrieve_ops_metrics().get("describe_table", 0), +)