diff --git a/docs/ja/src/concepts/indexing/vector_indexing.md b/docs/ja/src/concepts/indexing/vector_indexing.md index 2fc2613e..941a319a 100644 --- a/docs/ja/src/concepts/indexing/vector_indexing.md +++ b/docs/ja/src/concepts/indexing/vector_indexing.md @@ -486,7 +486,10 @@ laurus train pq-codebook --field embedding --input vectors.jsonl --update-schema ``` （`--input` の代わりに `--from-index` でインデックスにコミット済みの -ベクトルを直接サンプリングすることもできます、Issue #920）、または +ベクトルを直接サンプリングすることも、`laurus create index +--train-pq-codebook ` で学習をインデックス作成に畳み込んで +train-before-first-commit の順序ハザードを完全に解消することもできます +— いずれも Issue #920）、または プログラムから [`Engine::train_pq_codebook`](https://docs.rs/laurus/latest/laurus/struct.Engine.html) （`engine.train_pq_codebook("embedding", &vectors, None)`。from-index diff --git a/docs/ja/src/laurus-cli/commands.md b/docs/ja/src/laurus-cli/commands.md index afe76dd7..bdcad8ca 100644 --- a/docs/ja/src/laurus-cli/commands.md +++ b/docs/ja/src/laurus-cli/commands.md @@ -23,7 +23,7 @@ laurus --index-dir /var/data/my_index --format json search "title:rust" 新しいインデックスを作成します。`--schema` が指定された場合はその TOML ファイルを使用し、省略された場合は対話型スキーマウィザードが起動します。 ```bash -laurus create index [--schema ] +laurus create index [--schema ] [--train-pq-codebook ] ``` **引数:** @@ -31,6 +31,7 @@ laurus create index [--schema ] | フラグ | 必須 | 説明 | | :--- | :--- | :--- | | `--schema ` | いいえ | インデックススキーマを定義する TOML ファイルのパス。省略時はインデックスディレクトリに既存の `schema.toml` があればそれを使用し、なければ対話型ウィザードが起動します。 | +| `--train-pq-codebook ` | いいえ | インデックス作成の一部として共有 PQ codebook を学習します（Issue #920）。`ProductQuantization` + `pq_codebook_path` を設定したすべての HNSW フィールドを、作成直後にこの JSONL ファイル（`put docs` / `add docs` と同じ形式、事前計算済み `Vector` 値）から学習します。最初の commit がすぐに codebook でエンコードできるため、create → `train pq-codebook` → ingest の順序を手動で守る必要がなくなります。ファイル不在または対象フィールドなしの場合は、何も作成する前にエラーになります。 | **スキーマファイルの形式:** @@ -65,6 +66,13 @@ laurus --index-dir ./my_index create index # Field name: title # ... # Index created at ./my_index. + +# 作成と共有 PQ codebook の学習を1ステップで（Issue #920） +laurus --index-dir ./my_index create index --schema schema.toml \ + --train-pq-codebook train.jsonl +# Index created at ./my_index. +# Training PQ codebook for field 'embedding' on 300 vectors... +# Trained codebook 'embedding.pqcb' (m = 4, k = 256, sub_dim = 8, dimension = 32) from 300 vectors. ``` > **注意:** `schema.toml` と `store/` の両方が存在する場合はエラーが返されます。再作成するにはインデックスディレクトリを削除してください。`schema.toml` のみ存在する場合（作成が中断された場合など）は、`--schema` なしで `create index` を実行すると既存スキーマからストレージが復旧されます。 diff --git a/docs/src/concepts/indexing/vector_indexing.md b/docs/src/concepts/indexing/vector_indexing.md index d0fb2d11..cae9babb 100644 --- a/docs/src/concepts/indexing/vector_indexing.md +++ b/docs/src/concepts/indexing/vector_indexing.md @@ -509,7 +509,10 @@ laurus train pq-codebook --field embedding --input vectors.jsonl --update-schema ``` (or `--from-index` in place of `--input` to sample vectors already -committed to the index, Issue #920), or programmatically via +committed to the index; or fold training into index creation with +`laurus create index --train-pq-codebook `, which removes the +train-before-first-commit ordering hazard entirely — both Issue #920), +or programmatically via [`Engine::train_pq_codebook`](https://docs.rs/laurus/latest/laurus/struct.Engine.html) (`engine.train_pq_codebook("embedding", &vectors, None)`; pair with `engine.sample_committed_vectors("embedding", Some(n))` for the diff --git a/docs/src/laurus-cli/commands.md b/docs/src/laurus-cli/commands.md index 3ff86c40..84102107 100644 --- a/docs/src/laurus-cli/commands.md +++ b/docs/src/laurus-cli/commands.md @@ -23,7 +23,7 @@ laurus --index-dir /var/data/my_index --format json search "title:rust" Create a new index. If `--schema` is given, uses that TOML file; otherwise launches the interactive schema wizard. ```bash -laurus create index [--schema ] +laurus create index [--schema ] [--train-pq-codebook ] ``` **Arguments:** @@ -31,6 +31,7 @@ laurus create index [--schema ] | Flag | Required | Description | | :--- | :--- | :--- | | `--schema ` | No | Path to a TOML file defining the index schema. When omitted, the command checks if a `schema.toml` already exists in the index directory and uses it; otherwise the interactive wizard is launched. | +| `--train-pq-codebook ` | No | Train shared PQ codebooks as part of creation (Issue #920). Every HNSW field configuring `ProductQuantization` + `pq_codebook_path` is trained from this JSONL file (the `put docs` / `add docs` shape with pre-computed `Vector` values) immediately after the index is created, so the very first commit can already encode against the codebook — removing the create → `train pq-codebook` → ingest ordering the failure policy otherwise requires you to manage manually. Errors before creating anything if the file is missing or no field is eligible. | **Schema file format:** @@ -65,6 +66,13 @@ laurus --index-dir ./my_index create index # Field name: title # ... # Index created at ./my_index. + +# Create and train the shared PQ codebook in one step (Issue #920) +laurus --index-dir ./my_index create index --schema schema.toml \ + --train-pq-codebook train.jsonl +# Index created at ./my_index. +# Training PQ codebook for field 'embedding' on 300 vectors... +# Trained codebook 'embedding.pqcb' (m = 4, k = 256, sub_dim = 8, dimension = 32) from 300 vectors. ``` > **Note:** If both `schema.toml` and `store/` already exist, an error is returned. Delete the index directory to recreate. If only `schema.toml` exists (e.g. after an interrupted creation), running `create index` without `--schema` recovers the index by creating the missing storage from the existing schema. diff --git a/laurus-cli/src/cli.rs b/laurus-cli/src/cli.rs index 6cc20e12..d7b31680 100644 --- a/laurus-cli/src/cli.rs +++ b/laurus-cli/src/cli.rs @@ -67,12 +67,23 @@ pub struct CreateCommand { #[derive(Subcommand)] pub enum CreateResource { /// Create a new index. If --schema is given, uses that TOML file; - /// otherwise launches the interactive schema wizard. + /// otherwise launches the interactive schema wizard. With + /// --train-pq-codebook, shared PQ codebooks are trained as part of + /// creation (Issue #920), removing the train-before-first-commit + /// ordering hazard for fields that configure pq_codebook_path. Index { /// Path to an existing schema TOML file. When omitted, the /// interactive schema wizard is launched instead. #[arg(long)] schema: Option, + /// Path to a JSONL training file (the `put docs` / `add docs` + /// shape; pre-computed Vector values). When given, every HNSW + /// field that configures ProductQuantization + pq_codebook_path + /// gets its shared codebook trained from this file immediately + /// after creation, so the very first commit can already encode + /// against it. + #[arg(long)] + train_pq_codebook: Option, }, /// Interactively generate a schema TOML file. Schema { diff --git a/laurus-cli/src/commands/create.rs b/laurus-cli/src/commands/create.rs index f35770b9..f82810e4 100644 --- a/laurus-cli/src/commands/create.rs +++ b/laurus-cli/src/commands/create.rs @@ -27,10 +27,21 @@ use crate::context; /// schema file (the schema is persisted inside the index directory as /// `schema.toml`). /// +/// With `train_pq_codebook`, every HNSW field configuring +/// `ProductQuantization` + `pq_codebook_path` gets its shared codebook +/// trained from the given JSONL file immediately after creation (Issue +/// #920) — the very first commit can already encode against it, removing +/// the train-before-first-commit ordering hazard the #918 failure policy +/// otherwise leaves to the user. Validation (JSONL exists, at least one +/// eligible field) runs **before** anything is created so a failure never +/// leaves a half-initialized index behind. +/// /// # Arguments /// /// * `schema_path` - Optional path to a schema TOML file. When `None`, the /// interactive wizard is used instead. +/// * `train_pq_codebook` - Optional JSONL training file (bulk-ingest +/// shape); trains the shared codebook(s) as part of creation. /// * `index_dir` - Path to the index directory for the new index. /// /// # Errors @@ -38,27 +49,113 @@ use crate::context; /// Returns an error if: /// - The schema file cannot be read or parsed (when `schema_path` is given). /// - The interactive wizard fails (when `schema_path` is `None`). -/// - The index cannot be created. -pub async fn run_index(schema_path: Option<&Path>, index_dir: &Path) -> Result<()> { - match schema_path { - Some(path) => { - context::create_index(index_dir, path).await?; +/// - `train_pq_codebook` is given but the file does not exist, or the +/// effective schema has no `ProductQuantization` + `pq_codebook_path` +/// field. +/// - The index cannot be created, or codebook training fails. +pub async fn run_index( + schema_path: Option<&Path>, + train_pq_codebook: Option<&Path>, + index_dir: &Path, +) -> Result<()> { + // Determine the schema that will actually be persisted, replicating + // init_index's recovery rule: an existing schema.toml without store/ + // wins over the argument/wizard schema. + let schema = if index_dir.join("schema.toml").exists() && !index_dir.join("store").exists() { + context::read_schema(index_dir)? + } else { + match schema_path { + Some(path) => { + let content = + std::fs::read_to_string(path).context("Failed to read schema file")?; + toml::from_str(&content).context("Failed to parse schema TOML")? + } + None => build_schema_interactive()?, } - None => { - // If schema.toml already exists, use it directly instead of - // launching the wizard (recovery path for missing store/). - if index_dir.join("schema.toml").exists() { - let schema = context::read_schema(index_dir)?; - context::create_index_from_schema(index_dir, schema).await?; - } else { - let schema = build_schema_interactive()?; - context::create_index_from_schema(index_dir, schema).await?; + }; + + // Pre-creation validation for --train-pq-codebook: fail before + // creating anything so a bad invocation never leaves a created-but- + // untrained index behind. + let pq_fields = match train_pq_codebook { + Some(jsonl) => { + if !jsonl.exists() { + anyhow::bail!( + "--train-pq-codebook file '{}' does not exist", + jsonl.display() + ); + } + let fields = eligible_pq_fields(&schema); + if fields.is_empty() { + anyhow::bail!( + "--train-pq-codebook was given but no HNSW field configures \ + ProductQuantization + pq_codebook_path; nothing to train \ + (set pq_codebook_path on the field, or drop the flag)" + ); } + fields } - } + None => Vec::new(), + }; + + context::create_index_from_schema(index_dir, schema).await?; println!("Index created at {}.", index_dir.display()); + + if let Some(jsonl) = train_pq_codebook { + let engine = context::open_index(index_dir).await?; + for field in &pq_fields { + let vectors = crate::commands::train::collect_vectors_from_jsonl(field, jsonl, None)?; + if vectors.is_empty() { + anyhow::bail!( + "no training vectors found in '{}' for field '{field}'", + jsonl.display() + ); + } + println!( + "Training PQ codebook for field '{field}' on {} vectors...", + vectors.len() + ); + // output = None writes to the field's configured + // pq_codebook_path, so the schema persisted above and the + // trained file agree by construction. + let info = engine.train_pq_codebook(field, &vectors, None)?; + println!( + "Trained codebook '{}' (m = {}, k = {}, sub_dim = {}, dimension = {}) \ + from {} vectors.", + info.path, + info.subvector_count, + info.centroids, + info.sub_dimension, + info.dimension, + info.training_vectors + ); + } + } Ok(()) } + +/// Collect the fields eligible for create-time codebook training: HNSW +/// fields configuring `ProductQuantization` with a `pq_codebook_path`, +/// sorted by field name (`Schema::fields` is a HashMap, so iteration +/// order alone would be nondeterministic). +fn eligible_pq_fields(schema: &Schema) -> Vec { + use laurus::vector::core::quantization::QuantizationMethod; + let mut fields: Vec = schema + .fields + .iter() + .filter_map(|(name, option)| match option { + FieldOption::Hnsw(o) + if matches!(o.quantizer, QuantizationMethod::ProductQuantization { .. }) + && o.pq_codebook_path.is_some() => + { + Some(name.clone()) + } + _ => None, + }) + .collect(); + fields.sort(); + fields +} use laurus::lexical::core::field::{ BooleanOption, BytesOption, DateTimeOption, FloatOption, Geo3dOption, GeoOption, IntegerOption, TextOption, @@ -628,4 +725,157 @@ ef_construction = 16 other => panic!("expected Hnsw field, got {:?}", field_type_label(other)), } } + + // --- create index --train-pq-codebook (Issue #920) --- + + use laurus::{DataValue, Document}; + + const DIM: usize = 32; + + /// Write a schema TOML declaring `fields` as HNSW + PQ + + /// pq_codebook_path, and a JSONL training file carrying `count` + /// deterministic vectors for every one of those fields per line. + fn setup_pq_schema_and_jsonl( + dir: &Path, + fields: &[&str], + count: usize, + ) -> (std::path::PathBuf, std::path::PathBuf) { + let mut schema = String::new(); + for field in fields { + schema.push_str(&format!( + "[fields.{field}.Hnsw]\ndimension = 32\ndistance = \"Euclidean\"\n\ + m = 8\nef_construction = 32\npq_codebook_path = \"{field}.pqcb\"\n\n\ + [fields.{field}.Hnsw.quantizer.ProductQuantization]\nsubvector_count = 4\n\n" + )); + } + std::fs::create_dir_all(dir).unwrap(); + let schema_path = dir.join("input-schema.toml"); + std::fs::write(&schema_path, schema).unwrap(); + + let mut state = 0x2468_ACE0_u64; + let mut jsonl = String::new(); + for i in 0..count { + let mut cells = Vec::new(); + for field in fields { + let data: Vec = (0..DIM) + .map(|_| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + format!( + "{:.4}", + ((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0 + ) + }) + .collect(); + cells.push(format!( + "\"{field}\": {{\"Vector\": [{}]}}", + data.join(", ") + )); + } + jsonl.push_str(&format!( + "{{\"id\": \"doc{i}\", \"document\": {{\"fields\": {{{}}}}}}}\n", + cells.join(", ") + )); + } + let jsonl_path = dir.join("train.jsonl"); + std::fs::write(&jsonl_path, jsonl).unwrap(); + (schema_path, jsonl_path) + } + + /// Put one small document per field and commit on a freshly opened + /// engine — succeeds only when every PQ field's codebook is trained. + async fn ingest_and_commit(index_dir: &Path, fields: &[&str]) -> anyhow::Result<()> { + let engine = context::open_index(index_dir).await?; + let mut builder = Document::builder(); + for field in fields { + builder = builder.add_field( + *field, + DataValue::Vector((0..DIM).map(|j| j as f32 * 0.01).collect()), + ); + } + engine.put_document("probe", builder.build()).await?; + engine.commit().await?; + Ok(()) + } + + /// Issue #920: the flag trains the codebook as part of creation, so + /// the very first commit encodes against it — the direct regression + /// for the train-before-first-commit hazard. + #[tokio::test] + async fn create_with_train_flag_makes_first_commit_succeed() { + let dir = tempfile::tempdir().unwrap(); + let (schema, jsonl) = setup_pq_schema_and_jsonl(dir.path(), &["embedding"], 300); + + run_index(Some(&schema), Some(&jsonl), dir.path()) + .await + .unwrap(); + + assert!(dir.path().join("store/vector/embedding.pqcb").exists()); + ingest_and_commit(dir.path(), &["embedding"]) + .await + .expect("first commit must encode against the create-time codebook"); + } + + /// Control pinning the hazard itself: without the flag, the same + /// schema's first commit hard-errors per the #918 failure policy. + #[tokio::test] + async fn create_without_train_flag_leaves_first_commit_failing() { + let dir = tempfile::tempdir().unwrap(); + let (schema, _jsonl) = setup_pq_schema_and_jsonl(dir.path(), &["embedding"], 1); + + run_index(Some(&schema), None, dir.path()).await.unwrap(); + + assert!(!dir.path().join("store/vector/embedding.pqcb").exists()); + let err = ingest_and_commit(dir.path(), &["embedding"]) + .await + .expect_err("commit without a trained codebook must fail"); + assert!( + err.to_string().contains("pq_codebook_path"), + "the failure must be the untrained-codebook hard-error: {err}" + ); + } + + /// The flag with no eligible PQ field bails before creating anything. + #[tokio::test] + async fn create_with_train_flag_rejects_schema_without_pq_field() { + let dir = tempfile::tempdir().unwrap(); + let schema_path = dir.path().join("input-schema.toml"); + // HNSW but Scalar8Bit (default quantizer), no pq_codebook_path. + std::fs::write( + &schema_path, + "[fields.embedding.Hnsw]\ndimension = 32\ndistance = \"Euclidean\"\n\ + m = 8\nef_construction = 32\n", + ) + .unwrap(); + let jsonl = dir.path().join("train.jsonl"); + std::fs::write(&jsonl, "").unwrap(); + + let index_dir = dir.path().join("idx"); + let err = run_index(Some(&schema_path), Some(&jsonl), &index_dir) + .await + .unwrap_err(); + assert!( + err.to_string().contains("no HNSW field configures"), + "error must explain why nothing can be trained: {err}" + ); + assert!( + !index_dir.join("schema.toml").exists(), + "a rejected invocation must not leave a half-created index" + ); + } + + /// Multiple eligible fields all get their codebooks trained. + #[tokio::test] + async fn create_with_train_flag_trains_every_pq_field() { + let dir = tempfile::tempdir().unwrap(); + let (schema, jsonl) = setup_pq_schema_and_jsonl(dir.path(), &["emb_a", "emb_b"], 300); + + run_index(Some(&schema), Some(&jsonl), dir.path()) + .await + .unwrap(); + + assert!(dir.path().join("store/vector/emb_a.pqcb").exists()); + assert!(dir.path().join("store/vector/emb_b.pqcb").exists()); + } } diff --git a/laurus-cli/src/commands/train.rs b/laurus-cli/src/commands/train.rs index adda1a56..7e3341e6 100644 --- a/laurus-cli/src/commands/train.rs +++ b/laurus-cli/src/commands/train.rs @@ -21,6 +21,66 @@ use laurus::vector::Vector; use crate::commands::bulk::parse_entry; use crate::context; +/// Collect pre-computed training vectors for `field` from a JSONL file +/// (bulk-ingest shape: `{"id": "...", "document": {"fields": {...}}}`). +/// +/// Blank lines are skipped; every non-blank entry must carry a +/// pre-computed `Vector` value for `field`. Collection stops after the +/// first `sample_size` vectors when set (deterministic file order — no +/// random sampling). +/// +/// Shared between `train pq-codebook --input` and +/// `create index --train-pq-codebook` (Issue #920). +/// +/// # Arguments +/// +/// * `field` - The vector field whose values to extract. +/// * `input` - Path to the JSONL training file. +/// * `sample_size` - Optional cap: only the first N vectors are collected. +/// +/// # Errors +/// +/// Returns an error if the file cannot be opened or read, an entry fails +/// to parse, or an entry lacks a pre-computed vector for `field` (the +/// message names the line). +pub(crate) fn collect_vectors_from_jsonl( + field: &str, + input: &Path, + sample_size: Option, +) -> Result> { + let reader = std::io::BufReader::new( + std::fs::File::open(input) + .with_context(|| format!("failed to open JSONL file '{}'", input.display()))?, + ); + + let mut vectors: Vec = Vec::new(); + for (index, line) in reader.lines().enumerate() { + let line_no = index + 1; + let line = line.with_context(|| format!("line {line_no}: failed to read"))?; + if line.trim().is_empty() { + continue; + } + let (_, doc) = parse_entry(&line, line_no)?; + let Some(value) = doc.fields.get(field) else { + bail!("line {line_no}: entry has no '{field}' field"); + }; + let Some(data) = value.as_vector() else { + bail!( + "line {line_no}: field '{field}' is not a pre-computed vector \ + (embedder-generated training input is not supported; provide \ + `{{\"{field}\": {{\"Vector\": [..]}}}}` values)" + ); + }; + vectors.push(Vector::new(data.clone())); + if let Some(cap) = sample_size + && vectors.len() >= cap + { + break; + } + } + Ok(vectors) +} + /// Execute the `train pq-codebook` command. /// /// Collects training vectors from exactly one of two sources — a JSONL @@ -82,37 +142,7 @@ pub async fn run_pq_codebook( } else { // Checked above: when `from_index` is false, `input` is Some. let input = input.expect("validated: --input given when --from-index is not"); - let reader = std::io::BufReader::new( - std::fs::File::open(input) - .with_context(|| format!("failed to open JSONL file '{}'", input.display()))?, - ); - - let mut vectors: Vec = Vec::new(); - for (index, line) in reader.lines().enumerate() { - let line_no = index + 1; - let line = line.with_context(|| format!("line {line_no}: failed to read"))?; - if line.trim().is_empty() { - continue; - } - let (_, doc) = parse_entry(&line, line_no)?; - let Some(value) = doc.fields.get(field) else { - bail!("line {line_no}: entry has no '{field}' field"); - }; - let Some(data) = value.as_vector() else { - bail!( - "line {line_no}: field '{field}' is not a pre-computed vector \ - (embedder-generated training input is not supported; provide \ - `{{\"{field}\": {{\"Vector\": [..]}}}}` values)" - ); - }; - vectors.push(Vector::new(data.clone())); - if let Some(cap) = sample_size - && vectors.len() >= cap - { - break; - } - } - vectors + collect_vectors_from_jsonl(field, input, sample_size)? }; if vectors.is_empty() { if from_index { diff --git a/laurus-cli/src/main.rs b/laurus-cli/src/main.rs index ba540edb..d4ad2aac 100644 --- a/laurus-cli/src/main.rs +++ b/laurus-cli/src/main.rs @@ -39,8 +39,11 @@ async fn main() -> Result<()> { match cli.command { Command::Create(cmd) => match cmd.resource { - CreateResource::Index { schema } => { - create::run_index(schema.as_deref(), &index_dir).await + CreateResource::Index { + schema, + train_pq_codebook, + } => { + create::run_index(schema.as_deref(), train_pq_codebook.as_deref(), &index_dir).await } CreateResource::Schema { output } => create::run_schema(&output), },