diff --git a/lib/edge/src/edge_shard/mod.rs b/lib/edge/src/edge_shard/mod.rs index 27fa9276e3..fc7d3aa69b 100644 --- a/lib/edge/src/edge_shard/mod.rs +++ b/lib/edge/src/edge_shard/mod.rs @@ -15,7 +15,7 @@ use parking_lot::Mutex; use segment::common::operation_error::{OperationError, OperationResult}; use segment::entry::ReadSegmentEntry as _; use segment::segment_constructor::{load_segment, normalize_segment_dir}; -use shard::files::{PAYLOAD_INDEX_CONFIG_FILE, SEGMENTS_PATH, segment_manifest_path}; +use shard::files::{PAYLOAD_INDEX_CONFIG_FILE, SEGMENTS_PATH, WAL_PATH, segment_manifest_path}; use shard::operations::CollectionUpdateOperations; use shard::segment_holder::locked::LockedSegmentHolder; use shard::segment_holder::{FlushMode, SegmentHolder}; @@ -44,7 +44,6 @@ pub struct EdgeShard { search_pool: Arc, } -const WAL_PATH: &str = "wal"; impl EdgeShard { /// Create a new edge shard at `path` with the given configuration. /// diff --git a/lib/edge/src/lib.rs b/lib/edge/src/lib.rs index d42d96c071..e1fcd8e11a 100644 --- a/lib/edge/src/lib.rs +++ b/lib/edge/src/lib.rs @@ -31,6 +31,7 @@ pub use requests::{ CountRequest, FacetRequest, GroupRequest, Prefetch, QueryRequest, RetrieveRequest, ScrollRequest, SearchMatrixRequest, SearchRequest, }; +pub use shard::files::WAL_PATH; pub use shard::segment_manifest::{SegmentManifestState, SegmentsManifest}; pub use update_only::{ PointAction, PointCopy, PointPreview, PointUpdates, SegmentConfigInfo, UpdateBatchOutcome, diff --git a/lib/edge/tools/edge_tool/README.md b/lib/edge/tools/edge_tool/README.md index 88246d48c0..f9a3f36a17 100644 --- a/lib/edge/tools/edge_tool/README.md +++ b/lib/edge/tools/edge_tool/README.md @@ -58,6 +58,15 @@ this tool). Point ids are sequential and default to starting at the collection's (approximate) point count, so repeated `upsert` calls append rather than overwrite; override with `--start-id`. `--seed` controls the RNG (default `42`). +Points are written `--batch-size` at a time (default `1000`), one update operation per batch. Each +batch is generated, written and dropped before the next one starts, so peak memory follows the +batch rather than `--num`. Batching pays off twice, because an update operation is also +CBOR-serialized into one WAL record: writing everything in one operation holds the points once as +structs and again as that record. + +The batches are not one transaction. A failure partway through leaves the batches before it +applied — re-run to continue, since `--start-id` defaults to the collection's current point count. + ### `optimize` — run the shard optimizers ```sh @@ -77,11 +86,18 @@ Recursively uploads every file under the local `SOURCE` directory to `DESTINATIO inside the bucket, e.g. `my_collection/0`), preserving the relative directory structure — so the result is byte-for-byte the same layout `edge-shard-query`/`edge-shard-update` expect via `--prefix`. Every file is streamed through `object_store`'s multipart upload API (even small ones, -as a single final part), so upload memory use stays bounded regardless of segment size. +as a single final part), at most four parts of a file in flight at a time, so upload memory use +stays bounded regardless of segment size. + +The `wal/` directory is skipped: the read paths open segments and the config only, `EdgeShard` +recreates an empty WAL, and WAL segments are preallocated to their full capacity and never +truncated — a freshly created collection is 72 KB of segments next to 64 MB of WAL. - `--clean` — delete every existing object under `DESTINATION` before uploading, so a re-upload doesn't leave stale files behind from a previous run with a different shape (e.g. a collection - re-created with a different segment UUID). + re-created with a different segment UUID). Refused for an empty `DESTINATION`, which would match + every object in the bucket. +- `--include-wal` — upload `wal/` as well. - `--aws` (default) — AWS S3 or an S3-compatible store (MinIO, RustFS, ...). - `--gcs` — Google Cloud Storage. - `--bucket` [`BLOB_BUCKET`] — required. diff --git a/lib/edge/tools/edge_tool/src/args.rs b/lib/edge/tools/edge_tool/src/args.rs index 59112e6f46..37903677d1 100644 --- a/lib/edge/tools/edge_tool/src/args.rs +++ b/lib/edge/tools/edge_tool/src/args.rs @@ -110,6 +110,11 @@ pub struct UpsertArgs { #[arg(short = 'n', long, default_value_t = 100)] pub num: usize, + /// Points per update operation. Each batch is generated, written and dropped before the + /// next one starts, so peak memory follows the batch rather than `--num`. + #[arg(long, default_value_t = 1000)] + pub batch_size: usize, + /// First point id to use. Defaults to the collection's current (approximate) /// point count, so repeated `upsert` calls append rather than overwrite. #[arg(long)] @@ -136,10 +141,15 @@ pub struct UploadArgs { /// Delete every existing object under the destination prefix before uploading, so a /// re-upload doesn't leave stale files behind from a previous run with a different - /// shape (e.g. different segment UUIDs). + /// shape (e.g. different segment UUIDs). Refuses an empty destination. #[arg(long, default_value_t = false)] pub clean: bool, + /// Upload the `wal/` directory too. Skipped by default: the read paths never open it, + /// and its segments are preallocated to their full capacity and never truncated. + #[arg(long, default_value_t = false)] + pub include_wal: bool, + /// Upload to AWS S3 or an S3-compatible store (MinIO, RustFS, ...). Default backend. #[arg(long, conflicts_with = "gcs")] pub aws: bool, diff --git a/lib/edge/tools/edge_tool/src/create.rs b/lib/edge/tools/edge_tool/src/create.rs index fc19faf68f..8ed17160b8 100644 --- a/lib/edge/tools/edge_tool/src/create.rs +++ b/lib/edge/tools/edge_tool/src/create.rs @@ -98,13 +98,14 @@ pub fn run(args: CreateArgs) -> Result<()> { /// [`DEFAULT_DENSE_NAME`]; with more than one `--dense`, every one must be /// `NAME:SIZE`. fn parse_dense_specs(specs: &[String]) -> Result> { - if let [size] = specs - && !size.contains(':') + if let [spec] = specs + && !spec.contains(':') { - let size: usize = size + let size: usize = spec .trim() .parse() - .with_context(|| format!("invalid --dense size: {size:?}"))?; + .with_context(|| format!("invalid --dense size: {spec:?}"))?; + check_dense_size(size, spec)?; return Ok(vec![(DEFAULT_DENSE_NAME.to_string(), size)]); } @@ -121,11 +122,19 @@ fn parse_dense_specs(specs: &[String]) -> Result> { .trim() .parse() .with_context(|| format!("invalid dense vector size in {spec:?}"))?; + check_dense_size(size, spec)?; Ok((name.trim().to_string(), size)) }) .collect() } +fn check_dense_size(size: usize, spec: &str) -> Result<()> { + if size == 0 { + bail!("dense vector size must be greater than zero (got {spec:?})"); + } + Ok(()) +} + /// Parse `--sparse` specs: `default_missing_value = ""` turns a bare /// `--sparse` into [`DEFAULT_SPARSE_NAME`]; `--sparse=NAME` names it. fn parse_sparse_specs(specs: &[String]) -> Vec { diff --git a/lib/edge/tools/edge_tool/src/points.rs b/lib/edge/tools/edge_tool/src/points.rs index 309a03499d..9348887669 100644 --- a/lib/edge/tools/edge_tool/src/points.rs +++ b/lib/edge/tools/edge_tool/src/points.rs @@ -10,13 +10,14 @@ use edge::{ }; use rand::RngExt as _; use rand::rngs::StdRng; +use segment::json_path::JsonPathItem; /// The shard's write-facing schema: every named vector a point must carry, /// and every payload field worth generating a value for. pub struct Schema { pub dense: Vec<(String, usize)>, pub sparse: Vec, - pub payload: Vec<(String, PayloadSchemaType)>, + pub payload: Vec<(JsonPath, PayloadSchemaType)>, } impl Schema { @@ -31,9 +32,9 @@ impl Schema { let mut sparse: Vec = sparse.into_iter().collect(); sparse.sort(); - let mut payload: Vec<(String, PayloadSchemaType)> = payload_schema + let mut payload: Vec<(JsonPath, PayloadSchemaType)> = payload_schema .iter() - .map(|(key, info)| (key.to_string(), info.data_type)) + .map(|(key, info)| (key.clone(), info.data_type)) .collect(); payload.sort_by(|a, b| a.0.cmp(&b.0)); @@ -76,7 +77,7 @@ pub fn random_point(id: PointId, schema: &Schema, rng: &mut StdRng) -> PointStru let mut payload = serde_json::Map::new(); for (field, schema_type) in &schema.payload { - payload.insert(field.clone(), random_payload_value(*schema_type, rng)); + set_at_path(&mut payload, field, random_payload_value(*schema_type, rng)); } PointStructPersisted { @@ -86,6 +87,75 @@ pub fn random_point(id: PointId, schema: &Schema, rng: &mut StdRng) -> PointStru } } +/// Wrap `value` in the objects and arrays `path` walks through, so an index on +/// a nested or array field actually sees it — a flat `"a.b"` key leaves that +/// index empty. Paths sharing a prefix merge rather than overwrite. +fn set_at_path( + payload: &mut serde_json::Map, + path: &JsonPath, + value: serde_json::Value, +) { + let nested = path + .rest + .iter() + .rev() + .fold(value, |value, item| match item { + JsonPathItem::Key(key) => serde_json::json!({ key.clone(): value }), + JsonPathItem::WildcardIndex => serde_json::json!([value]), + JsonPathItem::Index(index) => { + let mut array = vec![serde_json::Value::Null; *index]; + array.push(value); + array.into() + } + }); + + match payload.get_mut(&path.first_key) { + Some(existing) => merge_value(existing, nested), + None => { + payload.insert(path.first_key.clone(), nested); + } + } +} + +/// Deep-merge objects and arrays so sibling paths under one prefix coexist +/// (`tags[].name` and `tags[].score` land on the same element); anything else +/// is replaced. +fn merge_value(dest: &mut serde_json::Value, source: serde_json::Value) { + match (dest, source) { + (serde_json::Value::Object(dest), serde_json::Value::Object(source)) => { + for (key, value) in source { + match dest.get_mut(&key) { + Some(existing) => merge_value(existing, value), + None => { + dest.insert(key, value); + } + } + } + } + (serde_json::Value::Array(dest), serde_json::Value::Array(source)) => { + for (index, value) in source.into_iter().enumerate() { + // Null is only ever the padding `set_at_path` puts ahead of an + // explicit index — no generated value is null — so it must not + // overwrite an element a sibling path already placed. + if value.is_null() { + if dest.len() <= index { + dest.resize(index + 1, serde_json::Value::Null); + } + continue; + } + match dest.get_mut(index) { + Some(existing) => merge_value(existing, value), + None => { + dest.resize(index, serde_json::Value::Null); + dest.push(value); + } + } + } + } + (dest, source) => *dest = source, + } +} + fn random_payload_value(schema_type: PayloadSchemaType, rng: &mut StdRng) -> serde_json::Value { let word = |rng: &mut StdRng| WORDS[rng.random_range(0..WORDS.len())].to_string(); match schema_type { @@ -112,3 +182,79 @@ fn random_payload_value(schema_type: PayloadSchemaType, rng: &mut StdRng) -> ser .into(), } } + +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + + use super::*; + + fn schema(paths: &[(&str, PayloadSchemaType)]) -> Schema { + Schema { + dense: Vec::new(), + sparse: Vec::new(), + payload: paths + .iter() + .map(|(path, kind)| (path.parse().unwrap(), *kind)) + .collect(), + } + } + + fn payload_of(schema: &Schema) -> serde_json::Map { + let mut rng = StdRng::seed_from_u64(1); + let point = random_point(PointId::NumId(0), schema, &mut rng); + point.payload.unwrap().0 + } + + #[test] + fn nested_paths_are_reachable_by_the_index_they_were_generated_for() { + let schema = schema(&[ + ("meta.city", PayloadSchemaType::Keyword), + ("meta.age", PayloadSchemaType::Integer), + ("flat", PayloadSchemaType::Float), + ]); + let payload = payload_of(&schema); + + for (path, _) in &schema.payload { + assert!( + !path.value_get(&payload).is_empty(), + "index on {path} sees no value in {payload:?}", + ); + } + assert!( + payload["meta"].is_object(), + "a nested path must not become a flat key: {payload:?}", + ); + } + + #[test] + fn sibling_paths_at_different_explicit_indexes_both_survive() { + let schema = schema(&[ + ("tags[0].name", PayloadSchemaType::Keyword), + ("tags[1].score", PayloadSchemaType::Float), + ]); + let payload = payload_of(&schema); + for (path, _) in &schema.payload { + assert!( + !path.value_get(&payload).is_empty(), + "index on {path} sees no value in {payload:?}", + ); + } + } + + #[test] + fn array_paths_are_reachable_too() { + let schema = schema(&[ + ("tags[].name", PayloadSchemaType::Keyword), + ("tags[].score", PayloadSchemaType::Float), + ]); + let payload = payload_of(&schema); + + for (path, _) in &schema.payload { + assert!( + !path.value_get(&payload).is_empty(), + "index on {path} sees no value in {payload:?}", + ); + } + } +} diff --git a/lib/edge/tools/edge_tool/src/upload.rs b/lib/edge/tools/edge_tool/src/upload.rs index 97e983cc87..66e8317a6f 100644 --- a/lib/edge/tools/edge_tool/src/upload.rs +++ b/lib/edge/tools/edge_tool/src/upload.rs @@ -1,9 +1,11 @@ //! Recursively upload a local collection directory to S3/GCS object storage. -use std::io::{BufReader, Read as _}; +use std::ffi::OsStr; +use std::io::Read as _; use std::path::{Path, PathBuf}; use anyhow::{Context, Result, anyhow, bail}; +use edge::WAL_PATH; use futures::StreamExt as _; use io_bridge_object_store::backends::aws::{AwsConfig, AwsCredentials}; use io_bridge_object_store::backends::gcp::{GcsConfig, GcsCredentials}; @@ -19,12 +21,20 @@ use crate::args::UploadArgs; /// 5 MiB on S3/GCS, so this stays comfortably above that. const MULTIPART_CHUNK_SIZE: usize = 8 * 1024 * 1024; +/// Parts of one file in flight at once. `WriteMultipart::write` starts an +/// upload the moment a chunk fills, however many are already running, so +/// without this a whole file is resident as spawned parts. +const MAX_PARTS_IN_FLIGHT: usize = 4; + pub fn run(args: UploadArgs) -> Result<()> { if !args.source.is_dir() { bail!("source {} is not a directory", args.source.display()); } + if args.clean && args.destination.trim_matches('/').is_empty() { + bail!("--clean with an empty destination would delete every object in the bucket"); + } - let files = collect_files(&args.source)?; + let files = collect_files(&args.source, args.include_wal)?; if files.is_empty() && !args.clean { log::warn!("{} has no files to upload", args.source.display()); return Ok(()); @@ -92,18 +102,38 @@ async fn clean_destination( Ok(removed) } -fn collect_files(root: &Path) -> Result> { +/// The WAL is write-path state: the read paths open segments and the config +/// only, and an [`edge::EdgeShard`] recreates an empty one. Its segments are +/// preallocated to their full capacity and never truncated, so it dwarfs a +/// freshly seeded collection. +fn collect_files(root: &Path, include_wal: bool) -> Result> { let mut files = Vec::new(); + let mut skipped_wal = 0usize; for entry in WalkDir::new(root) { let entry = entry.with_context(|| format!("failed to walk {}", root.display()))?; - if entry.file_type().is_file() { - files.push(entry.into_path()); + if !entry.file_type().is_file() { + continue; } + if !include_wal && entry.path().strip_prefix(root).is_ok_and(is_wal_path) { + skipped_wal += 1; + continue; + } + files.push(entry.into_path()); + } + if skipped_wal > 0 { + log::info!("skipped {skipped_wal} WAL file(s); pass --include-wal to upload them"); } files.sort(); Ok(files) } +fn is_wal_path(relative: &Path) -> bool { + relative + .components() + .next() + .is_some_and(|first| first.as_os_str() == OsStr::new(WAL_PATH)) +} + async fn upload_all( store: &O, source: &Path, @@ -140,9 +170,8 @@ async fn upload_file( .with_context(|| format!("failed to start upload of {}", file.display()))?; let mut write = WriteMultipart::new_with_chunk_size(upload, MULTIPART_CHUNK_SIZE); - let handle = + let mut reader = fs_err::File::open(file).with_context(|| format!("failed to open {}", file.display()))?; - let mut reader = BufReader::new(handle); let mut buffer = vec![0u8; MULTIPART_CHUNK_SIZE]; loop { let bytes_read = reader @@ -151,7 +180,10 @@ async fn upload_file( if bytes_read == 0 { break; } - // Sync but spawns an internal worker thread; `finish` waits for it. + write + .wait_for_capacity(MAX_PARTS_IN_FLIGHT) + .await + .with_context(|| format!("failed to upload a part of {}", file.display()))?; write.write(&buffer[..bytes_read]); } write diff --git a/lib/edge/tools/edge_tool/src/upsert.rs b/lib/edge/tools/edge_tool/src/upsert.rs index 4eab50fc98..e52f83f4c5 100644 --- a/lib/edge/tools/edge_tool/src/upsert.rs +++ b/lib/edge/tools/edge_tool/src/upsert.rs @@ -42,15 +42,31 @@ pub fn run(args: UpsertArgs) -> Result<()> { ); let mut rng = StdRng::seed_from_u64(args.seed); - let points = (0..args.num as u64) - .map(|offset| random_point(PointId::NumId(start_id + offset), &schema, &mut rng)) - .collect::>(); + let batch_size = args.batch_size.max(1); + let mut written = 0usize; + while written < args.num { + let batch = batch_size.min(args.num - written); + let points = (0..batch as u64) + .map(|offset| { + random_point( + PointId::NumId(start_id + written as u64 + offset), + &schema, + &mut rng, + ) + }) + .collect::>(); - shard - .update(UpdateOperation::PointOperation( - PointOperations::UpsertPoints(PointInsertOperations::PointsList(points)), - )) - .context("failed to upsert points")?; + shard + .update(UpdateOperation::PointOperation( + PointOperations::UpsertPoints(PointInsertOperations::PointsList(points)), + )) + .context("failed to upsert points")?; + + written += batch; + if written < args.num { + log::debug!("upserted {written}/{} point(s)", args.num); + } + } shard.flush().context("failed to flush the collection")?; log::info!(