edge-tool: bound upload and upsert memory, skip the WAL, fix generated payload paths (#10189)

* fix: bound edge-tool resource use and fix generated payload paths

* fix: keep sibling array elements when merging generated payload paths

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<rayon::ThreadPool>,
 }
 
-const WAL_PATH: &str = "wal";
 impl EdgeShard {
     /// Create a new edge shard at `path` with the given configuration.
     ///

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,

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.

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,

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<Vec<(String, usize)>> {
-    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<Vec<(String, usize)>> {
                 .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<String> {

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<String>,
-    pub payload: Vec<(String, PayloadSchemaType)>,
+    pub payload: Vec<(JsonPath, PayloadSchemaType)>,
 }
 
 impl Schema {
@@ -31,9 +32,9 @@ impl Schema {
         let mut sparse: Vec<String> = 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<String, serde_json::Value>,
+    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<String, serde_json::Value> {
+        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:?}",
+            );
+        }
+    }
+}

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<O: ObjectStore>(
     Ok(removed)
 }
 
-fn collect_files(root: &Path) -> Result<Vec<PathBuf>> {
+/// 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<Vec<PathBuf>> {
     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<O: ObjectStore>(
     store: &O,
     source: &Path,
@@ -140,9 +170,8 @@ async fn upload_file<O: ObjectStore>(
         .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<O: ObjectStore>(
         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

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::<Vec<_>>();
+    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::<Vec<_>>();
 
-    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!(