mirror of
https://github.com/qdrant/qdrant.git
synced 2026-09-21 13:37:46 -05:00
feat: UpdateOnlySegment / UpdateOnlyEdgeShard batch writer skeleton (#10021)
* feat: `UpdateOnlySegment` / `UpdateOnlyEdgeShard` batch writer skeleton Mirror image of the read-only pair, for the serverless updater: a shard/segment whose public surface is writes only, built for batches of many tiny operations against remote, append-only storage. Implemented: * `UpdateOnlySegment<S>` with a deliberately narrow open — id tracker, payload storage and one storage per named vector, all cold. No vector index, no quantized vectors, no payload index on the segments the writer only reads from. * `SegmentUpdateView`, the shared home of resolution logic, generic over the component traits (`VectorDataStorageRead` is a `VectorDataRead` without the index, so a segment that opens no index can produce the view). Batched `locate_points` / `point_versions` / `read_stored_points`. * `UpdateOnlyEdgeShard<S>::apply_batch`: fold the batch to one entry per point, locate the points, read only the ones that cannot be resolved from the batch alone, materialize `FullyQualifiedPoint`s, append them and tombstone the slots they replace. `todo!()`, pending the append-only components on the roadmap (appendable `DynamicStoredFlags` and `ChunkedVectors`, an appendable payload blobstore and field indexes): `store_points`, `tombstone_points`, `flush`, and creating the first appendable segment. Filter-selected operations, point sync, conditional upserts and the schema-level operations are rejected up front rather than silently skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix: codespell implementor → implementer in SegmentUpdateView docs Co-authored-by: Cursor <cursoragent@cursor.com> * docs: trim update-only writer docstrings to guarantees Less verbose throughout: state each function's contract — ordering, absent-value behavior, preconditions, durability — and drop narration about where types are used or why alternatives were rejected. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * refactor: fold SegmentUpdateView into UpdateOnlySegment as inherent methods The view was premature: it had exactly one producer, and its trait bounds bought an unexercised option. Resolution (locate / versions / read raw) now lives as inherent methods on UpdateOnlySegment, still generic over the backend. A shared view can be extracted when a second producer appears, e.g. batched CoW moves out of regular segments. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * refactor: split edge update_only batch.rs into a module Pure move: mutation.rs (PointMutation fold + materialize), plan.rs (UpdateBatchPlan operation intake), tests.rs. PointUpdates::new/push narrowed to pub(super). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat: parallel per-segment batch reads + tombstone every copy of a point locate_points and read_stored_points visit segments in parallel on a dedicated edge-update rayon pool (build_search_pool generalized to build_segment_pool with a thread-name prefix). locate_points now keeps every slot a point occupies, not just the newest copy: a rewrite or delete retires all of them. Tombstoning only the newest slot would let an older duplicate left by an interrupted move outlive the point — and resurrect it after a delete. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * refactor: point_versions returns a map keyed by internal id The id tracker's batch read is keyed by internal id already; returning AHashMap drops the positions_of reverse-lookup adapter. Absent key = unwritten slot, defaulted to version 0 at the caller. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat: accept a deferred threshold when opening UpdateOnlySegment Groundwork for an external rebuilder working the same directory: the cutoff loads slots at or above it into the appendable id tracker's deferred track (same appendable-only filter as ReadOnlySegment). It hides nothing from the writer — resolution runs WithDeferred, so every point still locates at its latest slot. The edge shard passes None until the rebuilder coordination exists. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat: preview_batch — resolve a batch without writing anything apply_batch and the new preview_batch share one resolution stage (resolve_batch: locate, read, materialize into per-point PointActions), so a dry-run reports exactly what an apply would do. Plus segment_configs(): per-segment configs with the write target marked. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat: prefetched + parallel segment opens for the update-only writer UpdateOnlySegment::open now mirrors ReadOnlySegment::open: a per-segment CachedFs primed by preopen, config parsed once and handed to open_via. The edge shard opens segments in parallel on its pool, keeping fail-hard semantics. With Populate::No throughout, prefetches transfer no data-file content — only configs, the id tracker and the deleted flags, whose opens consume them whole anyway. Also: ReadOnlyAppendableIdTracker::preopen now tolerates the not-yet-created mappings/versions files of an empty appendable segment, matching its open's contract — previously unreachable because followers skip appendable segments on error, while the writer must open them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat: edge-shard-update — dry-run batch upserts against a shard Counterpart of edge-shard-query for the write path: opens an UpdateOnlyEdgeShard over a local directory or S3/GCS object storage, generates random points shaped by the shard's own schema (segment config + payload-index schema), and logs what applying them would do — locations, versions, actions, tombstones — via preview_batch. Nothing is written: the write half is still todo!(). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix: box PointAction::Store to appease clippy::large_enum_variant A resolved point is ~384 bytes while every other variant is empty. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix: adapt to dev's dead-code sweep (#10030) Restore NamedVectors::remove_ref — removed as dead on dev, but the batch fold's DeleteVectors arm is now its first caller. Drop the allow(dead_code) on segment::update_only (no longer needed) and switch the writer's unread fs field to expect(dead_code), per the new ast-grep rule. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: root <111755117+qdrant-cloud-bot@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Andrey Vasnetsov
co-authored by
Claude Opus 5
root
Cursor
Claude Opus 5
root
Cursor
parent
50ddd1e3d9
commit
6aebf544d6
Generated
+17
@@ -2287,6 +2287,23 @@ dependencies = [
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "edge-shard-update"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"clap",
|
||||
"common",
|
||||
"edge",
|
||||
"env_logger",
|
||||
"fs-err",
|
||||
"io_bridge_object_store",
|
||||
"log",
|
||||
"object_store",
|
||||
"rand 0.10.2",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "either"
|
||||
version = "1.15.0"
|
||||
|
||||
@@ -363,6 +363,7 @@ members = [
|
||||
"lib/edge/ffi",
|
||||
"lib/edge/ffi/bindgen",
|
||||
"lib/edge/tools/shard_query",
|
||||
"lib/edge/tools/shard_update",
|
||||
"lib/blobstore",
|
||||
"lib/macros",
|
||||
"lib/posting_list",
|
||||
|
||||
@@ -25,7 +25,7 @@ use uuid::Uuid;
|
||||
|
||||
use crate::config::optimizers::EdgeOptimizersConfig;
|
||||
use crate::config::shard::{EDGE_CONFIG_FILE, EdgeConfig};
|
||||
use crate::read_view::build_search_pool;
|
||||
use crate::read_view::build_segment_pool;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct EdgeShard {
|
||||
@@ -66,7 +66,11 @@ impl EdgeShard {
|
||||
let mut segments = SegmentHolder::default();
|
||||
ensure_appendable_segment(&mut segments, path, &segments_path, &config)?;
|
||||
|
||||
let search_pool = build_search_pool(config.search_thread_count(), config.search_pool_core)?;
|
||||
let search_pool = build_segment_pool(
|
||||
"edge-search",
|
||||
config.search_thread_count(),
|
||||
config.search_pool_core,
|
||||
)?;
|
||||
|
||||
let config_path = path.join(EDGE_CONFIG_FILE);
|
||||
let config = Arc::new(
|
||||
@@ -139,7 +143,11 @@ impl EdgeShard {
|
||||
|
||||
ensure_appendable_segment(&mut segments, path, &segments_path, &config)?;
|
||||
|
||||
let search_pool = build_search_pool(config.search_thread_count(), config.search_pool_core)?;
|
||||
let search_pool = build_segment_pool(
|
||||
"edge-search",
|
||||
config.search_thread_count(),
|
||||
config.search_pool_core,
|
||||
)?;
|
||||
|
||||
let config_path = path.join(EDGE_CONFIG_FILE);
|
||||
let config = Arc::new(
|
||||
|
||||
@@ -7,6 +7,7 @@ mod read_view;
|
||||
mod reexports;
|
||||
mod requests;
|
||||
mod types;
|
||||
mod update_only;
|
||||
pub use types::*;
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -31,3 +32,7 @@ pub use requests::{
|
||||
ScrollRequest, SearchMatrixRequest, SearchRequest,
|
||||
};
|
||||
pub use shard::segment_manifest::{SegmentManifestState, SegmentsManifest};
|
||||
pub use update_only::{
|
||||
PointAction, PointCopy, PointPreview, PointUpdates, SegmentConfigInfo, UpdateBatchOutcome,
|
||||
UpdateBatchPlan, UpdateBatchPreview, UpdateOnlyEdgeShard,
|
||||
};
|
||||
|
||||
@@ -92,7 +92,8 @@ impl<S: UniversalReadExt + 'static> ReadOnlyEdgeShard<S> {
|
||||
|
||||
// Segments never carry `max_search_threads` / `search_pool_core`, so the pool is sized and
|
||||
// pinned from the caller-provided config alone: the CPU-derived default unless set.
|
||||
let search_pool = crate::read_view::build_search_pool(
|
||||
let search_pool = crate::read_view::build_segment_pool(
|
||||
"edge-search",
|
||||
provided_config.search_thread_count(),
|
||||
provided_config.search_pool_core,
|
||||
)?;
|
||||
|
||||
@@ -64,15 +64,12 @@ impl<H: ReadSegmentHandle> EdgeReadView<H> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a shard's search thread pool with `num_threads` worker threads — the pool behind
|
||||
/// [`EdgeReadView::par_map_segments`]. Both the read-write [`EdgeShard`](crate::EdgeShard) and the
|
||||
/// read-only [`ReadOnlyEdgeShard`](crate::ReadOnlyEdgeShard) build one at open and keep it for
|
||||
/// their lifetime, so per-segment reads and parallel segment opens don't spawn fresh threads per
|
||||
/// operation.
|
||||
/// Build a shard's per-segment thread pool with `num_threads` worker threads, its threads named
|
||||
/// `{thread_name_prefix}-{idx}`. Shards build one at open and keep it for their lifetime, so
|
||||
/// per-segment work doesn't spawn fresh threads per operation.
|
||||
///
|
||||
/// `num_threads` is the already-resolved thread count (see [`EdgeConfig::search_thread_count`]);
|
||||
/// callers pass `config.search_thread_count()` so a configured `0` is expanded to the CPU-derived
|
||||
/// default that matches the core search runtime.
|
||||
/// a configured `0` must be expanded by the caller.
|
||||
///
|
||||
/// `pin_core` pins every pool thread to the given CPU core ([`EdgeConfig::search_pool_core`]):
|
||||
/// the pool keeps its IO overlap but its compute is bounded to one core. Best-effort — an
|
||||
@@ -81,7 +78,8 @@ impl<H: ReadSegmentHandle> EdgeReadView<H> {
|
||||
/// Returns an error (rather than panicking) when the underlying thread spawn fails — this runs
|
||||
/// during shard open/load and follower open, so a transient resource failure must not abort the
|
||||
/// process.
|
||||
pub(crate) fn build_search_pool(
|
||||
pub(crate) fn build_segment_pool(
|
||||
thread_name_prefix: &'static str,
|
||||
num_threads: usize,
|
||||
pin_core: Option<usize>,
|
||||
) -> OperationResult<Arc<ThreadPool>> {
|
||||
@@ -91,23 +89,27 @@ pub(crate) fn build_search_pool(
|
||||
let available =
|
||||
core_affinity::get_core_ids().is_some_and(|ids| ids.iter().any(|c| c.id == *core));
|
||||
if !available {
|
||||
log::warn!("search pool core {core} is not available; leaving threads unpinned");
|
||||
log::warn!(
|
||||
"{thread_name_prefix} pool core {core} is not available; leaving threads unpinned"
|
||||
);
|
||||
}
|
||||
available
|
||||
});
|
||||
|
||||
let mut builder = ThreadPoolBuilder::new()
|
||||
.num_threads(num_threads)
|
||||
.thread_name(|idx| format!("edge-search-{idx}"));
|
||||
.thread_name(move |idx| format!("{thread_name_prefix}-{idx}"));
|
||||
if let Some(core) = pin_core {
|
||||
builder = builder.start_handler(move |idx| {
|
||||
if !core_affinity::set_for_current(core_affinity::CoreId { id: core }) {
|
||||
log::warn!("failed to pin edge search thread {idx} to core {core}");
|
||||
log::warn!("failed to pin edge {thread_name_prefix} thread {idx} to core {core}");
|
||||
}
|
||||
});
|
||||
}
|
||||
let pool = builder.build().map_err(|err| {
|
||||
OperationError::service_error(format!("failed to build edge search thread pool: {err}"))
|
||||
OperationError::service_error(format!(
|
||||
"failed to build edge {thread_name_prefix} thread pool: {err}"
|
||||
))
|
||||
})?;
|
||||
Ok(Arc::new(pool))
|
||||
}
|
||||
@@ -119,11 +121,11 @@ mod tests {
|
||||
/// Best-effort pinning: valid and out-of-range core ids must both yield a working pool.
|
||||
#[test]
|
||||
fn pinned_pool_builds_and_runs() {
|
||||
let pool = build_search_pool(2, Some(0)).unwrap();
|
||||
let pool = build_segment_pool("edge-search", 2, Some(0)).unwrap();
|
||||
let sum: i32 = pool.install(|| (0..4).sum());
|
||||
assert_eq!(sum, 6);
|
||||
|
||||
let pool = build_search_pool(1, Some(usize::MAX)).unwrap();
|
||||
let pool = build_segment_pool("edge-search", 1, Some(usize::MAX)).unwrap();
|
||||
assert_eq!(pool.install(|| 1 + 1), 2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
mod reexports_from_qdrant_crates {
|
||||
pub use segment::common::operation_error::{OperationError, OperationResult};
|
||||
pub use segment::data_types::facets::{FacetHit, FacetResponse, FacetValue, FacetValueHit};
|
||||
pub use segment::data_types::fully_qualified_point::FullyQualifiedPoint;
|
||||
pub use segment::data_types::index::{
|
||||
BoolIndexParams, DatetimeIndexParams, FloatIndexParams, GeoIndexParams, IntegerIndexParams,
|
||||
KeywordIndexParams, Language, SnowballLanguage, SnowballParams, StopwordsSet,
|
||||
@@ -14,8 +15,11 @@ mod reexports_from_qdrant_crates {
|
||||
pub use segment::data_types::vectors::{
|
||||
DEFAULT_VECTOR_NAME, NamedQuery, TypedMultiDenseVector,
|
||||
};
|
||||
pub use segment::index::payload_config::{PAYLOAD_INDEX_CONFIG_FILE, PayloadConfig};
|
||||
pub use segment::index::query_optimization::rescore_formula::parsed_formula::DecayKind;
|
||||
pub use segment::json_path::JsonPath;
|
||||
pub use segment::segment::SEGMENT_STATE_FILE;
|
||||
pub use segment::segment_constructor::get_payload_index_path;
|
||||
pub use segment::types::{
|
||||
AcornSearchParams, AnyVariants, BinaryQuantizationConfig, BinaryQuantizationEncoding,
|
||||
BinaryQuantizationQueryEncoding, CompressionRatio, Condition, DateTimeWrapper, Distance,
|
||||
@@ -27,9 +31,9 @@ mod reexports_from_qdrant_crates {
|
||||
NestedCondition, Payload, PayloadFieldSchema, PayloadIndexInfo, PayloadSchemaParams,
|
||||
PayloadSchemaType, PayloadSelector, PayloadSelectorExclude, PayloadSelectorInclude,
|
||||
ProductQuantizationConfig, QuantizationConfig, QuantizationSearchParams, Range,
|
||||
RangeInterface, ScalarQuantizationConfig, ScalarType, ScoredPoint, SearchParams, Slice,
|
||||
SliceCondition, ValueVariants, ValuesCount, VectorStorageDatatype, WithPayloadInterface,
|
||||
WithVector,
|
||||
RangeInterface, ScalarQuantizationConfig, ScalarType, ScoredPoint, SearchParams,
|
||||
SegmentConfig, SegmentState, Slice, SliceCondition, SparseVectorDataConfig, ValueVariants,
|
||||
ValuesCount, VectorDataConfig, VectorStorageDatatype, WithPayloadInterface, WithVector,
|
||||
};
|
||||
pub use segment::vector_storage::query::{
|
||||
ContextPair, ContextQuery, DiscoverQuery, FeedbackItem,
|
||||
|
||||
@@ -0,0 +1,273 @@
|
||||
//! Applying a folded batch: locate, resolve, materialize, append — each one
|
||||
//! batched pass over the whole point set, so a batch's cost scales with the
|
||||
//! points it touches, not the operations in it.
|
||||
|
||||
use ahash::AHashMap;
|
||||
use common::counter::hardware_counter::HardwareCounterCell;
|
||||
use common::types::PointOffsetType;
|
||||
use common::universal_io::UniversalRead;
|
||||
use rayon::ThreadPool;
|
||||
use rayon::prelude::*;
|
||||
use segment::common::operation_error::{OperationError, OperationResult};
|
||||
use segment::data_types::fully_qualified_point::{FullyQualifiedPoint, StoredPoint};
|
||||
use segment::types::{PointIdType, SeqNumberType};
|
||||
use shard::operations::CollectionUpdateOperations;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::update_only::UpdateOnlyEdgeShard;
|
||||
use crate::update_only::batch::UpdateBatchPlan;
|
||||
use crate::update_only::holder::UpdateOnlySegmentHolder;
|
||||
use crate::update_only::preview::{PointAction, PointPreview, resolve_batch};
|
||||
|
||||
/// What a batch did, counted per point rather than per operation: a point
|
||||
/// named by ten operations counts once.
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateBatchOutcome {
|
||||
/// Points written: created, or rewritten into a fresh slot.
|
||||
pub stored: usize,
|
||||
/// Points removed.
|
||||
pub deleted: usize,
|
||||
/// Points already at or beyond the batch's version, left untouched — what
|
||||
/// makes a replayed batch a no-op.
|
||||
pub skipped: usize,
|
||||
/// Points an operation named that no segment holds and the batch did not
|
||||
/// create — a payload update to a point that is not there.
|
||||
pub missing: usize,
|
||||
}
|
||||
|
||||
/// One copy of a point: where it lives, and at what version.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(super) struct PointLocation {
|
||||
pub(super) segment: Uuid,
|
||||
pub(super) internal_id: PointOffsetType,
|
||||
pub(super) version: SeqNumberType,
|
||||
/// Whether the holding segment accepts appends; breaks a version tie.
|
||||
appendable: bool,
|
||||
}
|
||||
|
||||
impl PointLocation {
|
||||
/// Whether this copy of the point supersedes `other`: the higher version
|
||||
/// wins, and on a tie the appendable copy is the live one (a point being
|
||||
/// moved between segments exists in both at the same version).
|
||||
fn supersedes(&self, other: &Self) -> bool {
|
||||
(self.version, self.appendable) > (other.version, other.appendable)
|
||||
}
|
||||
}
|
||||
|
||||
/// Every copy of one point across the shard's segments.
|
||||
pub(super) struct PointLocations {
|
||||
/// The live copy: its version decides whether the batch is already
|
||||
/// applied, and its slot is the one a resolve reads from.
|
||||
pub(super) newest: PointLocation,
|
||||
/// Every slot the point occupies, `newest`'s included. A rewrite or a
|
||||
/// delete retires them all — tombstoning only the newest slot would let
|
||||
/// an older duplicate (left by an interrupted move) outlive the point
|
||||
/// and, on a delete, resurrect it.
|
||||
pub(super) slots: Vec<(Uuid, PointOffsetType)>,
|
||||
}
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlyEdgeShard<S> {
|
||||
/// Apply a batch of update operations, each paired with the operation
|
||||
/// number to record as its version. Operations are expected in ascending
|
||||
/// operation-number order; see [`UpdateBatchPlan::build`] for what is
|
||||
/// rejected.
|
||||
///
|
||||
/// Atomic in the sense that matters without a WAL: applied in full or the
|
||||
/// error is returned, and re-applying a batch that partially landed skips
|
||||
/// the points that already carry its version.
|
||||
pub fn apply_batch(
|
||||
&self,
|
||||
operations: impl IntoIterator<Item = (SeqNumberType, CollectionUpdateOperations)>,
|
||||
) -> OperationResult<UpdateBatchOutcome> {
|
||||
let plan = UpdateBatchPlan::build(operations)?;
|
||||
if plan.is_empty() {
|
||||
return Ok(UpdateBatchOutcome::default());
|
||||
}
|
||||
|
||||
let hw_counter = HardwareCounterCell::disposable();
|
||||
let segments = self.segments.read();
|
||||
|
||||
// 1-3. Locate, read, materialize — the decision stage shared with
|
||||
// `preview_batch`, so a preview cannot drift from the real apply.
|
||||
let resolved = resolve_batch(&segments, plan, &self.pool)?;
|
||||
|
||||
let mut outcome = UpdateBatchOutcome::default();
|
||||
let mut to_store: Vec<FullyQualifiedPoint> = Vec::new();
|
||||
let mut to_tombstone: AHashMap<Uuid, Vec<PointOffsetType>> = AHashMap::new();
|
||||
|
||||
for point in resolved {
|
||||
let PointPreview {
|
||||
id: _,
|
||||
current: _,
|
||||
slots,
|
||||
action,
|
||||
} = point;
|
||||
|
||||
match action {
|
||||
PointAction::Skip => {
|
||||
outcome.skipped += 1;
|
||||
continue;
|
||||
}
|
||||
PointAction::Missing => {
|
||||
outcome.missing += 1;
|
||||
continue;
|
||||
}
|
||||
PointAction::Store(point) => {
|
||||
to_store.push(*point);
|
||||
outcome.stored += 1;
|
||||
}
|
||||
PointAction::Delete => outcome.deleted += 1,
|
||||
}
|
||||
|
||||
// Whatever happened to the point, every slot it occupied — in any
|
||||
// segment — is retired: a rewrite left its replacement elsewhere,
|
||||
// a delete left nothing, and an older duplicate must not outlive
|
||||
// either.
|
||||
for (segment, internal_id) in slots {
|
||||
to_tombstone.entry(segment).or_default().push(internal_id);
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Append the resolved points, then retire the slots they replaced.
|
||||
if !to_store.is_empty() {
|
||||
let write_target = segments.write_target()?;
|
||||
write_target.write().store_points(&to_store, &hw_counter)?;
|
||||
|
||||
// The new slots must be durable before the tombstones that retire
|
||||
// the old ones: the reverse order can lose a point outright if the
|
||||
// process dies in between.
|
||||
write_target.read().flush()?;
|
||||
}
|
||||
|
||||
for (uuid, internal_ids) in to_tombstone {
|
||||
let segment = segments.get(&uuid).ok_or_else(|| {
|
||||
OperationError::service_error(format!("Segment {uuid} disappeared mid-batch"))
|
||||
})?;
|
||||
segment.write().tombstone_points(&internal_ids)?;
|
||||
segment.read().flush()?;
|
||||
}
|
||||
|
||||
Ok(outcome)
|
||||
}
|
||||
}
|
||||
|
||||
/// Locate every point the batch touches: every slot it occupies, with the
|
||||
/// newest copy marked, when more than one segment holds the point. Segments
|
||||
/// are visited in parallel on `pool`.
|
||||
pub(super) fn locate_points<S: UniversalRead + 'static>(
|
||||
segments: &UpdateOnlySegmentHolder<S>,
|
||||
plan: &UpdateBatchPlan,
|
||||
pool: &ThreadPool,
|
||||
) -> OperationResult<AHashMap<PointIdType, PointLocations>> {
|
||||
let ids: Vec<PointIdType> = plan.point_ids().collect();
|
||||
|
||||
let per_segment: Vec<Vec<(PointIdType, PointLocation)>> = pool.install(|| {
|
||||
segments
|
||||
.iter()
|
||||
.collect::<Vec<_>>()
|
||||
.into_par_iter()
|
||||
.map(|(uuid, segment)| {
|
||||
let segment = segment.read();
|
||||
let appendable = segment.is_appendable();
|
||||
|
||||
let mut found_ids = Vec::new();
|
||||
let mut internal_ids = Vec::new();
|
||||
segment.locate_points(ids.iter().copied(), |id, internal_id| {
|
||||
found_ids.push(id);
|
||||
internal_ids.push(internal_id);
|
||||
})?;
|
||||
let versions = segment.point_versions(&internal_ids)?;
|
||||
|
||||
let located = found_ids
|
||||
.into_iter()
|
||||
.zip(internal_ids)
|
||||
.map(|(id, internal_id)| {
|
||||
let location = PointLocation {
|
||||
segment: uuid,
|
||||
internal_id,
|
||||
// A slot without a stored version is unwritten,
|
||||
// which compares as version 0.
|
||||
version: versions.get(&internal_id).copied().unwrap_or(0),
|
||||
appendable,
|
||||
};
|
||||
(id, location)
|
||||
})
|
||||
.collect();
|
||||
Ok(located)
|
||||
})
|
||||
.collect::<OperationResult<Vec<_>>>()
|
||||
})?;
|
||||
|
||||
let mut locations: AHashMap<PointIdType, PointLocations> = AHashMap::new();
|
||||
for (id, location) in per_segment.into_iter().flatten() {
|
||||
let slot = (location.segment, location.internal_id);
|
||||
locations
|
||||
.entry(id)
|
||||
.and_modify(|current| {
|
||||
current.slots.push(slot);
|
||||
if location.supersedes(¤t.newest) {
|
||||
current.newest = location;
|
||||
}
|
||||
})
|
||||
.or_insert_with(|| PointLocations {
|
||||
newest: location,
|
||||
slots: vec![slot],
|
||||
});
|
||||
}
|
||||
|
||||
Ok(locations)
|
||||
}
|
||||
|
||||
/// Read the stored form of the points whose mutations need it, one batched
|
||||
/// pass per segment; segments are read in parallel on `pool`.
|
||||
pub(super) fn read_stored_points<S: UniversalRead + 'static>(
|
||||
segments: &UpdateOnlySegmentHolder<S>,
|
||||
plan: &UpdateBatchPlan,
|
||||
locations: &AHashMap<PointIdType, PointLocations>,
|
||||
pool: &ThreadPool,
|
||||
) -> OperationResult<AHashMap<PointIdType, StoredPoint>> {
|
||||
let mut by_segment: AHashMap<Uuid, Vec<(PointIdType, PointOffsetType)>> = AHashMap::new();
|
||||
for id in plan.point_ids_needing_stored_point() {
|
||||
// A point no segment holds has nothing to read; its mutations either
|
||||
// create it outright or resolve to nothing. Only the newest copy is
|
||||
// read — older duplicates are stale.
|
||||
if let Some(location) = locations.get(&id) {
|
||||
by_segment
|
||||
.entry(location.newest.segment)
|
||||
.or_default()
|
||||
.push((id, location.newest.internal_id));
|
||||
}
|
||||
}
|
||||
|
||||
let per_segment: Vec<Vec<(PointIdType, StoredPoint)>> = pool.install(|| {
|
||||
by_segment
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>()
|
||||
.into_par_iter()
|
||||
.map(|(uuid, entries)| {
|
||||
let segment = segments.get(&uuid).ok_or_else(|| {
|
||||
OperationError::service_error(format!("Segment {uuid} disappeared mid-batch"))
|
||||
})?;
|
||||
let segment = segment.read();
|
||||
|
||||
let internal_ids: Vec<PointOffsetType> = entries
|
||||
.iter()
|
||||
.map(|(_, internal_id)| *internal_id)
|
||||
.collect();
|
||||
// Not shared with the caller's counter: `HardwareCounterCell`
|
||||
// is not `Sync`, and the writer's accounting is disposable.
|
||||
let hw_counter = HardwareCounterCell::disposable();
|
||||
let points = segment.read_stored_points(&internal_ids, &hw_counter)?;
|
||||
|
||||
Ok(entries.into_iter().map(|(id, _)| id).zip(points).collect())
|
||||
})
|
||||
.collect::<OperationResult<Vec<_>>>()
|
||||
})?;
|
||||
|
||||
let mut stored = AHashMap::new();
|
||||
for (id, point) in per_segment.into_iter().flatten() {
|
||||
stored.insert(id, point);
|
||||
}
|
||||
|
||||
Ok(stored)
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
//! Folding a batch of update operations into per-point work.
|
||||
//!
|
||||
//! All operations touching the same point collapse into one [`PointUpdates`]:
|
||||
//! a point is written at most once, however many operations named it, and read
|
||||
//! only if some surviving mutation needs the stored point — a batch that
|
||||
//! upserts a point never reads it.
|
||||
//!
|
||||
//! This stage is pure: nothing here touches storage.
|
||||
|
||||
mod mutation;
|
||||
mod plan;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
pub use self::mutation::PointUpdates;
|
||||
pub use self::plan::UpdateBatchPlan;
|
||||
@@ -0,0 +1,198 @@
|
||||
//! What operations do to a single point, and how a point's mutations fold
|
||||
//! onto its stored form.
|
||||
|
||||
use segment::common::operation_error::{OperationError, OperationResult};
|
||||
use segment::data_types::fully_qualified_point::{FullyQualifiedPoint, StoredPoint};
|
||||
use segment::data_types::named_vectors::NamedVectors;
|
||||
use segment::data_types::segment_record::NamedVectorBytesOwned;
|
||||
use segment::json_path::JsonPath;
|
||||
use segment::types::{Payload, PayloadKeyType, PointIdType, SeqNumberType, VectorNameBuf};
|
||||
|
||||
/// The vectors an operation carries, in the form the operation carried them:
|
||||
/// storage-native bytes travel to the new slot untouched, decoded vectors are
|
||||
/// encoded by the storage. Keeping the two apart avoids a decode/re-encode
|
||||
/// round-trip a quantized storage would not survive losslessly.
|
||||
pub enum OperationVectors {
|
||||
Decoded(NamedVectors<'static>),
|
||||
Raw(NamedVectorBytesOwned),
|
||||
}
|
||||
|
||||
/// What a single operation does to a single point; one variant per accepted
|
||||
/// operation.
|
||||
pub enum PointMutation {
|
||||
/// Whole-point replacement (an upsert): both vectors and payload come from
|
||||
/// the operation, and nothing of a previously stored point survives.
|
||||
Replace {
|
||||
vectors: OperationVectors,
|
||||
payload: Payload,
|
||||
},
|
||||
/// The point is removed.
|
||||
Delete,
|
||||
/// Replace the named vectors, leaving the rest of the point alone.
|
||||
UpdateVectors(NamedVectors<'static>),
|
||||
/// Drop the named vectors, leaving the rest of the point alone.
|
||||
DeleteVectors(Vec<VectorNameBuf>),
|
||||
/// Merge into the stored payload, at `key` when given.
|
||||
SetPayload {
|
||||
payload: Payload,
|
||||
key: Option<JsonPath>,
|
||||
},
|
||||
/// Replace the whole payload.
|
||||
OverwritePayload(Payload),
|
||||
/// Drop the listed payload keys.
|
||||
DeletePayload(Vec<PayloadKeyType>),
|
||||
/// Drop the whole payload.
|
||||
ClearPayload,
|
||||
}
|
||||
|
||||
impl PointMutation {
|
||||
/// Whether this mutation makes every mutation before it irrelevant:
|
||||
/// nothing of the point as it stood survives, so neither the earlier
|
||||
/// mutations nor the stored point itself need to be looked at.
|
||||
fn discards_stored_point(&self) -> bool {
|
||||
match self {
|
||||
Self::Replace { .. } | Self::Delete => true,
|
||||
Self::UpdateVectors(_)
|
||||
| Self::DeleteVectors(_)
|
||||
| Self::SetPayload { .. }
|
||||
| Self::OverwritePayload(_)
|
||||
| Self::DeletePayload(_)
|
||||
| Self::ClearPayload => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Everything a batch does to one point, in operation order.
|
||||
pub struct PointUpdates {
|
||||
/// Operation number of the last operation folded in — the version the
|
||||
/// rewritten point is stored at.
|
||||
version: SeqNumberType,
|
||||
/// Mutations to fold onto the stored point, oldest first. Never empty.
|
||||
mutations: Vec<PointMutation>,
|
||||
}
|
||||
|
||||
impl PointUpdates {
|
||||
pub(super) fn new(version: SeqNumberType, mutation: PointMutation) -> Self {
|
||||
Self {
|
||||
version,
|
||||
mutations: vec![mutation],
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn push(&mut self, version: SeqNumberType, mutation: PointMutation) {
|
||||
if mutation.discards_stored_point() {
|
||||
self.mutations.clear();
|
||||
}
|
||||
self.version = self.version.max(version);
|
||||
self.mutations.push(mutation);
|
||||
}
|
||||
|
||||
/// Version the rewritten point is stored at.
|
||||
pub fn version(&self) -> SeqNumberType {
|
||||
self.version
|
||||
}
|
||||
|
||||
/// Whether applying these mutations requires reading the point as it is
|
||||
/// stored today. False exactly when the first surviving mutation replaces
|
||||
/// or removes the point.
|
||||
pub fn needs_stored_point(&self) -> bool {
|
||||
self.mutations
|
||||
.first()
|
||||
.is_none_or(|mutation| !mutation.discards_stored_point())
|
||||
}
|
||||
|
||||
/// Fold the mutations onto `stored` — the point as it stands, absent when
|
||||
/// no segment holds it — into the point to store.
|
||||
///
|
||||
/// `Ok(None)` means the batch leaves nothing to store: the point ends up
|
||||
/// deleted, or an operation that can only modify an existing point named
|
||||
/// one that does not exist.
|
||||
pub fn materialize(
|
||||
self,
|
||||
id: PointIdType,
|
||||
stored: Option<StoredPoint>,
|
||||
) -> OperationResult<Option<FullyQualifiedPoint>> {
|
||||
let Self { version, mutations } = self;
|
||||
|
||||
let mut exists = stored.is_some();
|
||||
let (mut stored_vectors, mut payload) = match stored {
|
||||
Some(stored) => {
|
||||
let StoredPoint {
|
||||
internal_id: _,
|
||||
vectors,
|
||||
payload,
|
||||
} = stored;
|
||||
(vectors, payload)
|
||||
}
|
||||
None => (NamedVectorBytesOwned::new(), Payload::default()),
|
||||
};
|
||||
// Vectors the batch supplied, which override `stored_vectors` by name
|
||||
// (see `FullyQualifiedPoint`), so replacing one does not require
|
||||
// removing its carried-over counterpart.
|
||||
let mut updated_vectors = NamedVectors::default();
|
||||
|
||||
for mutation in mutations {
|
||||
match mutation {
|
||||
PointMutation::Replace {
|
||||
vectors,
|
||||
payload: replacement,
|
||||
} => {
|
||||
exists = true;
|
||||
stored_vectors.clear();
|
||||
updated_vectors = NamedVectors::default();
|
||||
match vectors {
|
||||
OperationVectors::Decoded(vectors) => updated_vectors = vectors,
|
||||
OperationVectors::Raw(vectors) => stored_vectors = vectors,
|
||||
}
|
||||
payload = replacement;
|
||||
}
|
||||
PointMutation::Delete => {
|
||||
exists = false;
|
||||
stored_vectors.clear();
|
||||
updated_vectors = NamedVectors::default();
|
||||
payload = Payload::default();
|
||||
}
|
||||
PointMutation::UpdateVectors(vectors) => {
|
||||
if !exists {
|
||||
return Err(OperationError::PointIdError {
|
||||
missed_point_id: id,
|
||||
});
|
||||
}
|
||||
updated_vectors.merge(vectors);
|
||||
}
|
||||
PointMutation::DeleteVectors(names) => {
|
||||
for name in &names {
|
||||
stored_vectors.retain(|(stored_name, _)| stored_name != name);
|
||||
updated_vectors.remove_ref(name.as_str());
|
||||
}
|
||||
}
|
||||
PointMutation::SetPayload {
|
||||
payload: values,
|
||||
key,
|
||||
} => match key {
|
||||
Some(key) => payload.merge_by_key(&values, &key),
|
||||
None => payload.merge(&values),
|
||||
},
|
||||
PointMutation::OverwritePayload(values) => payload = values,
|
||||
PointMutation::DeletePayload(keys) => {
|
||||
for key in &keys {
|
||||
payload.remove(key);
|
||||
}
|
||||
}
|
||||
PointMutation::ClearPayload => payload = Payload::default(),
|
||||
}
|
||||
}
|
||||
|
||||
if !exists {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(FullyQualifiedPoint {
|
||||
id,
|
||||
version,
|
||||
stored_vectors,
|
||||
updated_vectors,
|
||||
payload,
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
//! Collapsing a batch of operations into one [`PointUpdates`] entry per point.
|
||||
|
||||
use ahash::AHashMap;
|
||||
use segment::common::operation_error::{OperationError, OperationResult};
|
||||
use segment::data_types::named_vectors::NamedVectors;
|
||||
use segment::types::{PointIdType, SeqNumberType};
|
||||
use shard::operations::CollectionUpdateOperations;
|
||||
use shard::operations::payload_ops::PayloadOps;
|
||||
use shard::operations::point_ops::{
|
||||
PointOperations, PointStructPersisted, PointStructRawPersisted,
|
||||
};
|
||||
use shard::operations::vector_ops::{PointVectorsPersisted, VectorOperations};
|
||||
|
||||
use super::mutation::{OperationVectors, PointMutation, PointUpdates};
|
||||
|
||||
/// A batch of update operations, collapsed to one entry per touched point.
|
||||
pub struct UpdateBatchPlan {
|
||||
/// Points in the order the batch first touched them, so the writer's
|
||||
/// appends are deterministic for a given batch.
|
||||
order: Vec<PointIdType>,
|
||||
updates: AHashMap<PointIdType, PointUpdates>,
|
||||
}
|
||||
|
||||
impl UpdateBatchPlan {
|
||||
/// Fold `operations` — each paired with the operation number to record —
|
||||
/// into one entry per point. Operations are expected in ascending
|
||||
/// operation-number order; the fold is order-sensitive.
|
||||
///
|
||||
/// Rejects everything outside the writer's contract: operations that
|
||||
/// select points by filter, point sync, conditional upserts, and the
|
||||
/// schema-level operations.
|
||||
pub fn build(
|
||||
operations: impl IntoIterator<Item = (SeqNumberType, CollectionUpdateOperations)>,
|
||||
) -> OperationResult<Self> {
|
||||
let mut plan = Self {
|
||||
order: Vec::new(),
|
||||
updates: AHashMap::new(),
|
||||
};
|
||||
|
||||
for (op_num, operation) in operations {
|
||||
match operation {
|
||||
CollectionUpdateOperations::PointOperation(operation) => {
|
||||
plan.push_point_operation(op_num, operation)?;
|
||||
}
|
||||
CollectionUpdateOperations::VectorOperation(operation) => {
|
||||
plan.push_vector_operation(op_num, operation)?;
|
||||
}
|
||||
CollectionUpdateOperations::PayloadOperation(operation) => {
|
||||
plan.push_payload_operation(op_num, operation)?;
|
||||
}
|
||||
CollectionUpdateOperations::FieldIndexOperation(_) => {
|
||||
return Err(unsupported("payload index operations"));
|
||||
}
|
||||
CollectionUpdateOperations::VectorNameOperation(_) => {
|
||||
return Err(unsupported("vector name operations"));
|
||||
}
|
||||
#[cfg(feature = "staging")]
|
||||
CollectionUpdateOperations::StagingOperation(_) => {
|
||||
return Err(unsupported("staging operations"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(plan)
|
||||
}
|
||||
|
||||
fn push(&mut self, id: PointIdType, version: SeqNumberType, mutation: PointMutation) {
|
||||
match self.updates.entry(id) {
|
||||
std::collections::hash_map::Entry::Occupied(mut entry) => {
|
||||
entry.get_mut().push(version, mutation);
|
||||
}
|
||||
std::collections::hash_map::Entry::Vacant(entry) => {
|
||||
entry.insert(PointUpdates::new(version, mutation));
|
||||
self.order.push(id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn push_point_operation(
|
||||
&mut self,
|
||||
op_num: SeqNumberType,
|
||||
operation: PointOperations,
|
||||
) -> OperationResult<()> {
|
||||
match operation {
|
||||
PointOperations::UpsertPoints(operation) => {
|
||||
for point in operation.into_point_vec() {
|
||||
// Decode before destructuring: `get_vectors` reads the
|
||||
// still-owned `vector` field, and taking the payload by
|
||||
// value afterwards saves a clone of it.
|
||||
let vectors = OperationVectors::Decoded(point.get_vectors().into_owned());
|
||||
let PointStructPersisted {
|
||||
id,
|
||||
vector: _,
|
||||
payload,
|
||||
} = point;
|
||||
self.push(
|
||||
id,
|
||||
op_num,
|
||||
PointMutation::Replace {
|
||||
vectors,
|
||||
payload: payload.unwrap_or_default(),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
PointOperations::UpsertPointsRaw(points) => {
|
||||
for point in points {
|
||||
let PointStructRawPersisted {
|
||||
id,
|
||||
vectors,
|
||||
payload,
|
||||
} = point;
|
||||
self.push(
|
||||
id,
|
||||
op_num,
|
||||
PointMutation::Replace {
|
||||
vectors: OperationVectors::Raw(vectors),
|
||||
payload: payload.unwrap_or_default(),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
PointOperations::DeletePoints { ids } => {
|
||||
for id in ids {
|
||||
self.push(id, op_num, PointMutation::Delete);
|
||||
}
|
||||
}
|
||||
PointOperations::UpsertPointsConditional(_) => {
|
||||
return Err(unsupported("conditional upserts"));
|
||||
}
|
||||
PointOperations::DeletePointsByFilter(_) => {
|
||||
return Err(unsupported("deleting points by filter"));
|
||||
}
|
||||
PointOperations::SyncPoints(_) | PointOperations::SyncPointsRaw(_) => {
|
||||
return Err(unsupported("point sync"));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push_vector_operation(
|
||||
&mut self,
|
||||
op_num: SeqNumberType,
|
||||
operation: VectorOperations,
|
||||
) -> OperationResult<()> {
|
||||
match operation {
|
||||
VectorOperations::UpdateVectors(operation) => {
|
||||
if operation.update_filter.is_some() {
|
||||
return Err(unsupported("conditional vector updates"));
|
||||
}
|
||||
for point in operation.points {
|
||||
let PointVectorsPersisted { id, vector } = point;
|
||||
let vectors = NamedVectors::from(vector).into_owned();
|
||||
self.push(id, op_num, PointMutation::UpdateVectors(vectors));
|
||||
}
|
||||
}
|
||||
VectorOperations::DeleteVectors(points, vector_names) => {
|
||||
for id in points.points {
|
||||
self.push(
|
||||
id,
|
||||
op_num,
|
||||
PointMutation::DeleteVectors(vector_names.clone()),
|
||||
);
|
||||
}
|
||||
}
|
||||
VectorOperations::DeleteVectorsByFilter(_, _) => {
|
||||
return Err(unsupported("deleting vectors by filter"));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push_payload_operation(
|
||||
&mut self,
|
||||
op_num: SeqNumberType,
|
||||
operation: PayloadOps,
|
||||
) -> OperationResult<()> {
|
||||
match operation {
|
||||
PayloadOps::SetPayload(operation) => {
|
||||
let points = require_points(operation.points, operation.filter.is_some())?;
|
||||
for id in points {
|
||||
self.push(
|
||||
id,
|
||||
op_num,
|
||||
PointMutation::SetPayload {
|
||||
payload: operation.payload.clone(),
|
||||
key: operation.key.clone(),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
PayloadOps::OverwritePayload(operation) => {
|
||||
let points = require_points(operation.points, operation.filter.is_some())?;
|
||||
for id in points {
|
||||
self.push(
|
||||
id,
|
||||
op_num,
|
||||
PointMutation::OverwritePayload(operation.payload.clone()),
|
||||
);
|
||||
}
|
||||
}
|
||||
PayloadOps::DeletePayload(operation) => {
|
||||
let points = require_points(operation.points, operation.filter.is_some())?;
|
||||
for id in points {
|
||||
self.push(
|
||||
id,
|
||||
op_num,
|
||||
PointMutation::DeletePayload(operation.keys.clone()),
|
||||
);
|
||||
}
|
||||
}
|
||||
PayloadOps::ClearPayload { points } => {
|
||||
for id in points {
|
||||
self.push(id, op_num, PointMutation::ClearPayload);
|
||||
}
|
||||
}
|
||||
PayloadOps::ClearPayloadByFilter(_) => {
|
||||
return Err(unsupported("clearing payload by filter"));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.order.is_empty()
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.order.len()
|
||||
}
|
||||
|
||||
/// Every point the batch touches, in first-touched order.
|
||||
pub fn point_ids(&self) -> impl Iterator<Item = PointIdType> + '_ {
|
||||
self.order.iter().copied()
|
||||
}
|
||||
|
||||
/// The points whose stored form has to be read before they can be
|
||||
/// rewritten.
|
||||
pub fn point_ids_needing_stored_point(&self) -> impl Iterator<Item = PointIdType> + '_ {
|
||||
self.order
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|id| self.updates[id].needs_stored_point())
|
||||
}
|
||||
|
||||
/// Consume the plan, yielding one entry per point in first-touched order.
|
||||
pub fn into_point_updates(mut self) -> impl Iterator<Item = (PointIdType, PointUpdates)> {
|
||||
let order = std::mem::take(&mut self.order);
|
||||
order.into_iter().filter_map(move |id| {
|
||||
let updates = self.updates.remove(&id)?;
|
||||
Some((id, updates))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Point-selecting operations must name their points: resolving a filter means
|
||||
/// querying payload indexes, which the writer never fetches.
|
||||
fn require_points(
|
||||
points: Option<Vec<PointIdType>>,
|
||||
has_filter: bool,
|
||||
) -> OperationResult<Vec<PointIdType>> {
|
||||
match points {
|
||||
Some(points) => Ok(points),
|
||||
None if has_filter => Err(unsupported("selecting points by filter")),
|
||||
None => Err(OperationError::validation_error(
|
||||
"No points or filter specified",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn unsupported(what: &str) -> OperationError {
|
||||
OperationError::validation_error(format!("The update-only writer does not support {what}"))
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
use segment::payload_json;
|
||||
use segment::types::{Payload, PointIdType};
|
||||
use shard::operations::CollectionUpdateOperations;
|
||||
use shard::operations::payload_ops::{PayloadOps, SetPayloadOp};
|
||||
use shard::operations::point_ops::{PointOperations, PointStructPersisted, VectorStructPersisted};
|
||||
|
||||
use super::UpdateBatchPlan;
|
||||
|
||||
fn point_id(id: u64) -> PointIdType {
|
||||
PointIdType::NumId(id)
|
||||
}
|
||||
|
||||
fn upsert(id: u64, payload: Payload) -> CollectionUpdateOperations {
|
||||
CollectionUpdateOperations::PointOperation(PointOperations::UpsertPoints(
|
||||
vec![PointStructPersisted {
|
||||
id: point_id(id),
|
||||
vector: VectorStructPersisted::Single(vec![1.0, 0.0]),
|
||||
payload: Some(payload),
|
||||
}]
|
||||
.into(),
|
||||
))
|
||||
}
|
||||
|
||||
fn set_payload(id: u64, payload: Payload) -> CollectionUpdateOperations {
|
||||
CollectionUpdateOperations::PayloadOperation(PayloadOps::SetPayload(SetPayloadOp {
|
||||
payload,
|
||||
points: Some(vec![point_id(id)]),
|
||||
filter: None,
|
||||
key: None,
|
||||
}))
|
||||
}
|
||||
|
||||
fn delete(id: u64) -> CollectionUpdateOperations {
|
||||
CollectionUpdateOperations::PointOperation(PointOperations::DeletePoints {
|
||||
ids: vec![point_id(id)],
|
||||
})
|
||||
}
|
||||
|
||||
/// Operations on the same point collapse into one entry, and the merged
|
||||
/// payload is the fold of all of them.
|
||||
#[test]
|
||||
fn folds_operations_on_the_same_point() {
|
||||
let plan = UpdateBatchPlan::build([
|
||||
(1, upsert(7, payload_json! { "a": 1 })),
|
||||
(2, set_payload(7, payload_json! { "b": 2 })),
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(plan.len(), 1);
|
||||
// The upsert supplies the whole point, so nothing has to be read.
|
||||
assert_eq!(plan.point_ids_needing_stored_point().count(), 0);
|
||||
|
||||
let (id, updates) = plan.into_point_updates().next().unwrap();
|
||||
assert_eq!(id, point_id(7));
|
||||
assert_eq!(updates.version(), 2);
|
||||
|
||||
let point = updates.materialize(id, None).unwrap().unwrap();
|
||||
assert_eq!(point.version, 2);
|
||||
assert_eq!(point.payload, payload_json! { "a": 1, "b": 2 });
|
||||
}
|
||||
|
||||
/// A batch that only modifies a point has to read it first.
|
||||
#[test]
|
||||
fn modification_only_batch_needs_the_stored_point() {
|
||||
let plan = UpdateBatchPlan::build([(1, set_payload(7, payload_json! { "b": 2 }))]).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
plan.point_ids_needing_stored_point().collect::<Vec<_>>(),
|
||||
vec![point_id(7)],
|
||||
);
|
||||
}
|
||||
|
||||
/// A delete discards everything before it: the point is neither read nor
|
||||
/// written.
|
||||
#[test]
|
||||
fn delete_discards_preceding_operations() {
|
||||
let plan = UpdateBatchPlan::build([
|
||||
(1, set_payload(7, payload_json! { "b": 2 })),
|
||||
(2, delete(7)),
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(plan.point_ids_needing_stored_point().count(), 0);
|
||||
|
||||
let (id, updates) = plan.into_point_updates().next().unwrap();
|
||||
assert!(updates.materialize(id, None).unwrap().is_none());
|
||||
}
|
||||
|
||||
/// ... and an upsert after a delete brings the point back.
|
||||
#[test]
|
||||
fn upsert_after_delete_recreates_the_point() {
|
||||
let plan =
|
||||
UpdateBatchPlan::build([(1, delete(7)), (2, upsert(7, payload_json! { "a": 1 }))]).unwrap();
|
||||
|
||||
let (id, updates) = plan.into_point_updates().next().unwrap();
|
||||
let point = updates.materialize(id, None).unwrap().unwrap();
|
||||
assert_eq!(point.payload, payload_json! { "a": 1 });
|
||||
}
|
||||
|
||||
/// Filter-selected operations are rejected up front, not silently applied to
|
||||
/// nothing.
|
||||
#[test]
|
||||
fn rejects_filter_selected_operations() {
|
||||
let operation =
|
||||
CollectionUpdateOperations::PayloadOperation(PayloadOps::SetPayload(SetPayloadOp {
|
||||
payload: payload_json! { "a": 1 },
|
||||
points: None,
|
||||
filter: Some(Default::default()),
|
||||
key: None,
|
||||
}));
|
||||
|
||||
assert!(UpdateBatchPlan::build([(1, operation)]).is_err());
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::universal_io::UniversalRead;
|
||||
use parking_lot::RwLock;
|
||||
use segment::common::operation_error::{OperationError, OperationResult};
|
||||
use segment::segment::update_only::UpdateOnlySegment;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// In-memory inventory of the segments a writer updates, keyed by segment
|
||||
/// UUID, with at most one — the appendable one — as the write target.
|
||||
pub(crate) struct UpdateOnlySegmentHolder<S: UniversalRead + 'static> {
|
||||
by_uuid: HashMap<Uuid, Arc<RwLock<UpdateOnlySegment<S>>>>,
|
||||
/// UUID of the single segment that accepts appends.
|
||||
write_target: Option<Uuid>,
|
||||
}
|
||||
|
||||
impl<S: UniversalRead + 'static> Default for UpdateOnlySegmentHolder<S> {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
by_uuid: HashMap::new(),
|
||||
write_target: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlySegmentHolder<S> {
|
||||
pub(crate) fn insert(&mut self, uuid: Uuid, segment: UpdateOnlySegment<S>) {
|
||||
if segment.is_appendable() {
|
||||
self.write_target = Some(uuid);
|
||||
}
|
||||
self.by_uuid.insert(uuid, Arc::new(RwLock::new(segment)));
|
||||
}
|
||||
|
||||
pub(crate) fn len(&self) -> usize {
|
||||
self.by_uuid.len()
|
||||
}
|
||||
|
||||
pub(crate) fn is_empty(&self) -> bool {
|
||||
self.by_uuid.is_empty()
|
||||
}
|
||||
|
||||
/// Every segment, paired with its UUID. Order is unspecified.
|
||||
pub(crate) fn iter(
|
||||
&self,
|
||||
) -> impl Iterator<Item = (Uuid, &Arc<RwLock<UpdateOnlySegment<S>>>)> + '_ {
|
||||
self.by_uuid.iter().map(|(uuid, segment)| (*uuid, segment))
|
||||
}
|
||||
|
||||
pub(crate) fn get(&self, uuid: &Uuid) -> Option<&Arc<RwLock<UpdateOnlySegment<S>>>> {
|
||||
self.by_uuid.get(uuid)
|
||||
}
|
||||
|
||||
/// UUID of the single segment that accepts appends, if one exists.
|
||||
pub(crate) fn write_target_uuid(&self) -> Option<Uuid> {
|
||||
self.write_target
|
||||
}
|
||||
|
||||
/// The single segment that accepts appends; an error when none exists.
|
||||
pub(crate) fn write_target(&self) -> OperationResult<&Arc<RwLock<UpdateOnlySegment<S>>>> {
|
||||
self.write_target
|
||||
.as_ref()
|
||||
.and_then(|uuid| self.by_uuid.get(uuid))
|
||||
.ok_or_else(|| {
|
||||
OperationError::service_error("No appendable segment exists, expected exactly one")
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use common::universal_io::{MmapFile, MmapFs, UniversalRead, UniversalReadFs};
|
||||
use parking_lot::RwLock;
|
||||
use rayon::prelude::*;
|
||||
use segment::common::operation_error::OperationResult;
|
||||
use segment::segment::update_only::UpdateOnlySegment;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::read_only::{LocalSegmentEnumerator, SegmentEnumerator};
|
||||
use crate::read_view::build_segment_pool;
|
||||
use crate::update_only::UpdateOnlyEdgeShard;
|
||||
use crate::update_only::holder::UpdateOnlySegmentHolder;
|
||||
|
||||
impl UpdateOnlyEdgeShard<MmapFile> {
|
||||
/// Open a writer over local memory-mapped files, discovering segments by
|
||||
/// scanning the `segments/` directory — the writer owns the directory it
|
||||
/// writes to, so there is no manifest to agree with.
|
||||
pub fn open_mmap(path: &Path) -> OperationResult<Self> {
|
||||
Self::open(MmapFs, path, LocalSegmentEnumerator::new(path))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlyEdgeShard<S> {
|
||||
/// Open a writer over the shard directory at `path`, using `fs` as the
|
||||
/// read backend and `enumerator` to discover the segments.
|
||||
///
|
||||
/// Segments are opened in parallel on the shard's thread pool, each over
|
||||
/// its own prefetching [`CachedFs`](common::universal_io::CachedFs) (see
|
||||
/// [`UpdateOnlySegment::open`]) — the same shape as the read-only
|
||||
/// follower's load — and entirely cold: no point data is fetched until a
|
||||
/// batch reads a point. A segment that fails to load is an error, not a
|
||||
/// skip — a writer that misses a segment would resolve a point against a
|
||||
/// stale copy of itself, or duplicate it.
|
||||
pub fn open(
|
||||
fs: S::Fs,
|
||||
path: &Path,
|
||||
enumerator: impl SegmentEnumerator + 'static,
|
||||
) -> OperationResult<Self>
|
||||
where
|
||||
S::Fs: UniversalReadFs<File = S>,
|
||||
{
|
||||
// Sized like the search pools: over-provisioned relative to the CPU
|
||||
// count, since on a remote backend the threads mostly wait on IO.
|
||||
let pool = build_segment_pool(
|
||||
"edge-update",
|
||||
common::defaults::search_thread_count(0),
|
||||
None,
|
||||
)?;
|
||||
|
||||
let segments: Vec<(Uuid, PathBuf)> = enumerator.list_segments()?.into_iter().collect();
|
||||
let opened: Vec<(Uuid, UpdateOnlySegment<S>)> = pool.install(|| {
|
||||
segments
|
||||
.into_par_iter()
|
||||
.map(|(uuid, segment_path)| {
|
||||
// No deferred threshold yet: it belongs to the coordination
|
||||
// with an external rebuilder, which does not exist in this
|
||||
// iteration.
|
||||
let segment = UpdateOnlySegment::<S>::open(&fs, &segment_path, uuid, None)?;
|
||||
Ok((uuid, segment))
|
||||
})
|
||||
.collect::<OperationResult<Vec<_>>>()
|
||||
})?;
|
||||
|
||||
let mut holder = UpdateOnlySegmentHolder::default();
|
||||
for (uuid, segment) in opened {
|
||||
holder.insert(uuid, segment);
|
||||
}
|
||||
|
||||
if holder.is_empty() {
|
||||
// Creating the first appendable segment needs the append-only
|
||||
// components the writer cannot build yet, so an empty directory is
|
||||
// not something this iteration can bootstrap.
|
||||
todo!("creating the initial appendable segment needs the append-only components");
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
path: path.to_path_buf(),
|
||||
fs,
|
||||
segments: RwLock::new(holder),
|
||||
pool,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
//! Update-only shard: a batch writer over an edge-shard directory, the mirror
|
||||
//! image of [`ReadOnlyEdgeShard`](crate::ReadOnlyEdgeShard). Its whole public
|
||||
//! surface is [`apply_batch`].
|
||||
//!
|
||||
//! Built for the serverless updater's cost model — batches of many tiny
|
||||
//! operations, remote per-file reads, no long-lived process:
|
||||
//!
|
||||
//! * a batch is folded before it is applied: a point is read at most once and
|
||||
//! written at most once, however many operations named it;
|
||||
//! * only the components a write needs are opened, and every lookup is one
|
||||
//! batched pass per component over the whole point set;
|
||||
//! * there is no WAL: a batch is durable when the storages are flushed.
|
||||
//!
|
||||
//! Storage is append-only throughout. Updating a point appends it in full and
|
||||
//! tombstones its old slot; a deletion writes nothing but the deleted-points
|
||||
//! bitmask.
|
||||
//!
|
||||
//! [`apply_batch`]: UpdateOnlyEdgeShard::apply_batch
|
||||
|
||||
mod apply;
|
||||
mod batch;
|
||||
mod holder;
|
||||
mod lifecycle;
|
||||
mod preview;
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::universal_io::UniversalRead;
|
||||
use parking_lot::RwLock;
|
||||
use rayon::ThreadPool;
|
||||
use segment::types::SegmentConfig;
|
||||
use uuid::Uuid;
|
||||
|
||||
pub use self::apply::UpdateBatchOutcome;
|
||||
pub use self::batch::{PointUpdates, UpdateBatchPlan};
|
||||
use self::holder::UpdateOnlySegmentHolder;
|
||||
pub use self::preview::{PointAction, PointCopy, PointPreview, UpdateBatchPreview};
|
||||
|
||||
/// A batch writer over the segments of one shard directory, generic over the
|
||||
/// backend `S`.
|
||||
///
|
||||
/// Compared to [`EdgeShard`](crate::EdgeShard), there is no WAL, no
|
||||
/// optimizers, and no `EdgeConfig` — the write target's own segment config is
|
||||
/// the only configuration a write needs.
|
||||
pub struct UpdateOnlyEdgeShard<S: UniversalRead + 'static> {
|
||||
path: PathBuf,
|
||||
/// Backend the segments were opened on, and the one their appends go
|
||||
/// through. Unread until the writer can create the appendable segment a
|
||||
/// fresh directory needs.
|
||||
#[expect(dead_code)]
|
||||
fs: S::Fs,
|
||||
segments: RwLock<UpdateOnlySegmentHolder<S>>,
|
||||
/// Thread pool the per-segment work of a batch runs on: on a remote
|
||||
/// backend each segment's reads block on the network, so segments are
|
||||
/// visited in parallel.
|
||||
pool: Arc<ThreadPool>,
|
||||
}
|
||||
|
||||
/// One segment's schema, as reported by
|
||||
/// [`UpdateOnlyEdgeShard::segment_configs`].
|
||||
pub struct SegmentConfigInfo {
|
||||
pub uuid: Uuid,
|
||||
/// Whether this segment is the write target — the one every write in a
|
||||
/// batch is appended to.
|
||||
pub is_write_target: bool,
|
||||
pub config: SegmentConfig,
|
||||
}
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlyEdgeShard<S> {
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
|
||||
/// Number of segments the writer has open.
|
||||
pub fn segments_count(&self) -> usize {
|
||||
self.segments.read().len()
|
||||
}
|
||||
|
||||
/// Every segment's config, cloned out, with the write target marked.
|
||||
/// Order is unspecified. The write target's config is the schema a write
|
||||
/// must conform to: the named vectors a point carries, and their shapes.
|
||||
pub fn segment_configs(&self) -> Vec<SegmentConfigInfo> {
|
||||
let segments = self.segments.read();
|
||||
let write_target = segments.write_target_uuid();
|
||||
segments
|
||||
.iter()
|
||||
.map(|(uuid, segment)| SegmentConfigInfo {
|
||||
uuid,
|
||||
is_write_target: Some(uuid) == write_target,
|
||||
config: segment.read().segment_config.clone(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
//! Dry-run of a batch: the apply pipeline run up to — but not including — the
|
||||
//! writes.
|
||||
//!
|
||||
//! [`preview_batch`] and [`apply_batch`] share one resolution stage
|
||||
//! ([`resolve_batch`]), so a preview reports exactly what an apply would do.
|
||||
//!
|
||||
//! [`preview_batch`]: UpdateOnlyEdgeShard::preview_batch
|
||||
//! [`apply_batch`]: UpdateOnlyEdgeShard::apply_batch
|
||||
|
||||
use common::types::PointOffsetType;
|
||||
use common::universal_io::UniversalRead;
|
||||
use rayon::ThreadPool;
|
||||
use segment::common::operation_error::OperationResult;
|
||||
use segment::data_types::fully_qualified_point::FullyQualifiedPoint;
|
||||
use segment::types::{PointIdType, SeqNumberType};
|
||||
use shard::operations::CollectionUpdateOperations;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::update_only::UpdateOnlyEdgeShard;
|
||||
use crate::update_only::apply::{locate_points, read_stored_points};
|
||||
use crate::update_only::batch::UpdateBatchPlan;
|
||||
use crate::update_only::holder::UpdateOnlySegmentHolder;
|
||||
|
||||
/// A resolved batch: one entry per touched point, in first-touched order.
|
||||
pub struct UpdateBatchPreview {
|
||||
pub points: Vec<PointPreview>,
|
||||
}
|
||||
|
||||
/// What the batch does to one point.
|
||||
pub struct PointPreview {
|
||||
pub id: PointIdType,
|
||||
/// The newest stored copy of the point; `None` when no segment holds it.
|
||||
pub current: Option<PointCopy>,
|
||||
/// Every slot the point occupies across segments, the newest's included —
|
||||
/// all of them are tombstoned when the action stores or deletes the point.
|
||||
pub slots: Vec<(Uuid, PointOffsetType)>,
|
||||
pub action: PointAction,
|
||||
}
|
||||
|
||||
/// One stored copy of a point: which segment holds it, in which slot, at what
|
||||
/// version.
|
||||
pub struct PointCopy {
|
||||
pub segment: Uuid,
|
||||
pub internal_id: PointOffsetType,
|
||||
pub version: SeqNumberType,
|
||||
}
|
||||
|
||||
/// The write one point's folded mutations resolved to.
|
||||
pub enum PointAction {
|
||||
/// The point is appended to the write target in this fully qualified
|
||||
/// form, and every slot in [`PointPreview::slots`] is tombstoned.
|
||||
/// Boxed: a resolved point is hundreds of bytes, the other variants none.
|
||||
Store(Box<FullyQualifiedPoint>),
|
||||
/// The point is removed: every slot is tombstoned, nothing is stored.
|
||||
Delete,
|
||||
/// Left untouched: the stored copy is already at or beyond the batch's
|
||||
/// version, so re-applying would move the point backwards.
|
||||
Skip,
|
||||
/// An operation that can only modify an existing point named one that no
|
||||
/// segment holds; there is nothing to write.
|
||||
Missing,
|
||||
}
|
||||
|
||||
/// Resolve a folded batch against the segments: locate every touched point,
|
||||
/// read the ones whose mutations need the stored form, and materialize each
|
||||
/// into its [`PointAction`]. Reads only — the single decision stage behind
|
||||
/// both [`UpdateOnlyEdgeShard::preview_batch`] and
|
||||
/// [`UpdateOnlyEdgeShard::apply_batch`].
|
||||
pub(super) fn resolve_batch<S: UniversalRead + 'static>(
|
||||
segments: &UpdateOnlySegmentHolder<S>,
|
||||
plan: UpdateBatchPlan,
|
||||
pool: &ThreadPool,
|
||||
) -> OperationResult<Vec<PointPreview>> {
|
||||
let locations = locate_points(segments, &plan, pool)?;
|
||||
let mut stored = read_stored_points(segments, &plan, &locations, pool)?;
|
||||
|
||||
let mut points = Vec::with_capacity(plan.len());
|
||||
for (id, updates) in plan.into_point_updates() {
|
||||
let location = locations.get(&id);
|
||||
let current = location.map(|location| PointCopy {
|
||||
segment: location.newest.segment,
|
||||
internal_id: location.newest.internal_id,
|
||||
version: location.newest.version,
|
||||
});
|
||||
let slots = location
|
||||
.map(|location| location.slots.clone())
|
||||
.unwrap_or_default();
|
||||
|
||||
// Already applied: the stored point is at or beyond this batch's
|
||||
// version, so re-applying would move it backwards.
|
||||
let already_applied = current
|
||||
.as_ref()
|
||||
.is_some_and(|current| current.version >= updates.version());
|
||||
|
||||
let action = if already_applied {
|
||||
PointAction::Skip
|
||||
} else {
|
||||
match updates.materialize(id, stored.remove(&id))? {
|
||||
Some(point) => PointAction::Store(Box::new(point)),
|
||||
None if current.is_some() => PointAction::Delete,
|
||||
None => PointAction::Missing,
|
||||
}
|
||||
};
|
||||
|
||||
points.push(PointPreview {
|
||||
id,
|
||||
current,
|
||||
slots,
|
||||
action,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(points)
|
||||
}
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlyEdgeShard<S> {
|
||||
/// Resolve a batch without writing anything: what
|
||||
/// [`apply_batch`](Self::apply_batch) would do, reported per point.
|
||||
///
|
||||
/// Runs the same resolution code path as the real apply — the same
|
||||
/// operations are rejected, the same points read — so the report cannot
|
||||
/// drift from the apply's behavior.
|
||||
pub fn preview_batch(
|
||||
&self,
|
||||
operations: impl IntoIterator<Item = (SeqNumberType, CollectionUpdateOperations)>,
|
||||
) -> OperationResult<UpdateBatchPreview> {
|
||||
let plan = UpdateBatchPlan::build(operations)?;
|
||||
if plan.is_empty() {
|
||||
return Ok(UpdateBatchPreview { points: Vec::new() });
|
||||
}
|
||||
|
||||
let segments = self.segments.read();
|
||||
let points = resolve_batch(&segments, plan, &self.pool)?;
|
||||
Ok(UpdateBatchPreview { points })
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
[package]
|
||||
name = "edge-shard-update"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
publish = false
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "edge-shard-update"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
edge = { path = "../.." }
|
||||
common = { path = "../../../common/common" }
|
||||
io_bridge_object_store = { path = "../../../common/io_bridge_object_store" }
|
||||
|
||||
fs-err = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
clap = { workspace = true }
|
||||
env_logger = { workspace = true }
|
||||
log = { workspace = true }
|
||||
rand = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
@@ -0,0 +1,667 @@
|
||||
//! Experimental binary that opens an [`UpdateOnlyEdgeShard`] over a local
|
||||
//! shard directory — or directly over object storage (AWS S3 / S3-compatible,
|
||||
//! or Google Cloud Storage) — and dry-runs a batch of random upserts against
|
||||
//! it.
|
||||
//!
|
||||
//! The counterpart of `edge-shard-query` for the write path. Point ids to
|
||||
//! overwrite are taken from the command line; the *shape* of each generated
|
||||
//! point is derived from the shard's own schema — every dense/sparse vector
|
||||
//! named in the appendable segment's config, and one payload value per field
|
||||
//! in its payload-index schema — so the batch is exactly what a real client
|
||||
//! could have written.
|
||||
//!
|
||||
//! Since the writer's write half is not implemented yet
|
||||
//! (`store_points`/`tombstone_points` are `todo!()`), the tool stops at the
|
||||
//! resolution stage: it runs [`preview_batch`], which shares the decision
|
||||
//! pipeline with the real `apply_batch`, and logs what *would* happen — which
|
||||
//! segments hold each point today, in which slots and at what versions, which
|
||||
//! points would be stored/skipped, and which slots would be tombstoned.
|
||||
//! Nothing is written, on any backend.
|
||||
//!
|
||||
//! Example — local shard directory:
|
||||
//!
|
||||
//! ```sh
|
||||
//! cargo run -p edge-shard-update -- \
|
||||
//! --path ./qdrant_storage/collections/benchmark/0 \
|
||||
//! --ids 1,2,42 \
|
||||
//! --op-num 20000
|
||||
//! ```
|
||||
//!
|
||||
//! Example — shard on S3-compatible object storage (here: GCS via its S3
|
||||
//! interoperability endpoint). Segments are discovered from the leader's
|
||||
//! segment manifest, and reads go through a local disk cache, exactly like
|
||||
//! `edge-shard-query`:
|
||||
//!
|
||||
//! ```sh
|
||||
//! cargo run -p edge-shard-update -- \
|
||||
//! --backend aws \
|
||||
//! --bucket qdrant-benchmark-snapshots \
|
||||
//! --endpoint https://storage.googleapis.com \
|
||||
//! --region auto \
|
||||
//! --access-key xxxx \
|
||||
//! --secret-key xxxx \
|
||||
//! --prefix serverless/shard-100k \
|
||||
//! --ids 1,2,42 \
|
||||
//! --op-num 20000
|
||||
//! ```
|
||||
//!
|
||||
//! [`preview_batch`]: UpdateOnlyEdgeShard::preview_batch
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use clap::{Args as ClapArgs, Parser, ValueEnum};
|
||||
use common::universal_io::{
|
||||
DiskCache, DiskCacheConfig, DiskCacheFs, DiskCacheFsContext, MmapFile, OkNotFound as _,
|
||||
UniversalRead, UniversalReadFileOps, UniversalReadFs, read_json_via,
|
||||
};
|
||||
use edge::external::uuid::Uuid;
|
||||
use edge::{
|
||||
FullyQualifiedPoint, ManifestSegmentEnumerator, PAYLOAD_INDEX_CONFIG_FILE, Payload,
|
||||
PayloadConfig, PayloadFieldSchema, PayloadSchemaParams, PayloadSchemaType, PointAction,
|
||||
PointId, PointOperations, PointStructPersisted, SegmentConfig, SegmentConfigInfo, SparseVector,
|
||||
UpdateOnlyEdgeShard, UpdateOperation, VectorPersisted, VectorStructPersisted,
|
||||
get_payload_index_path,
|
||||
};
|
||||
use io_bridge_object_store::backends::aws::{AwsConfig, AwsCredentials};
|
||||
use io_bridge_object_store::backends::gcp::{GcsConfig, GcsCredentials};
|
||||
use io_bridge_object_store::{AsyncRead, BlobFile, ObjectStoreSource};
|
||||
use object_store::aws::AmazonS3;
|
||||
use object_store::gcp::GoogleCloudStorage;
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{RngExt as _, SeedableRng as _};
|
||||
|
||||
/// Storage backend to read the shard from.
|
||||
#[derive(Clone, Copy, Debug, ValueEnum)]
|
||||
enum Backend {
|
||||
/// A local shard directory (`--path`), read via memory-mapped files and
|
||||
/// discovered by scanning `segments/`.
|
||||
Local,
|
||||
/// AWS S3 or an S3-compatible store (MinIO, RustFS, GCS interop, ...).
|
||||
Aws,
|
||||
/// Google Cloud Storage.
|
||||
Gcs,
|
||||
}
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(
|
||||
about = "Open an UpdateOnlyEdgeShard over a local directory or S3/GCS object storage and \
|
||||
dry-run a batch of random upserts (nothing is written: the write half is not \
|
||||
implemented yet)"
|
||||
)]
|
||||
struct Cli {
|
||||
#[command(flatten)]
|
||||
connection: ConnectionArgs,
|
||||
|
||||
/// Point ids to overwrite with random points: comma-separated integers or
|
||||
/// UUIDs. Ids no segment holds are created rather than overwritten.
|
||||
#[arg(long, value_delimiter = ',', required = true)]
|
||||
ids: Vec<String>,
|
||||
|
||||
/// Operation number recorded as the new points' version. A point whose
|
||||
/// stored version is at or beyond this is reported as skipped — the
|
||||
/// writer's replay semantics — so pick a value above the versions the log
|
||||
/// reports to see the points overwritten.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
op_num: u64,
|
||||
|
||||
/// RNG seed for the generated points, so a run is reproducible.
|
||||
#[arg(long, default_value_t = 42)]
|
||||
seed: u64,
|
||||
}
|
||||
|
||||
/// How to reach the shard: a local path, or an object-storage location.
|
||||
#[derive(ClapArgs, Debug)]
|
||||
struct ConnectionArgs {
|
||||
/// Which backend to read the shard from.
|
||||
#[arg(long, value_enum, default_value = "local")]
|
||||
backend: Backend,
|
||||
|
||||
/// [local] Path to the shard root directory (the one containing
|
||||
/// `segments/`).
|
||||
#[arg(long)]
|
||||
path: Option<PathBuf>,
|
||||
|
||||
/// [AWS/GCS] Bucket name (without any scheme prefix).
|
||||
#[arg(long, env = "BLOB_BUCKET")]
|
||||
bucket: Option<String>,
|
||||
|
||||
/// [AWS] Custom S3 endpoint URL (MinIO / RustFS / GCS interop; omit for real AWS).
|
||||
#[arg(long, env = "S3_ENDPOINT")]
|
||||
endpoint: Option<String>,
|
||||
|
||||
/// [AWS] Region (e.g. `us-east-1`). Required for real AWS; optional otherwise.
|
||||
#[arg(long, env = "S3_REGION")]
|
||||
region: Option<String>,
|
||||
|
||||
/// [AWS] Access key id. If omitted, the AWS default credential chain is used.
|
||||
#[arg(long, env = "S3_ACCESS_KEY")]
|
||||
access_key: Option<String>,
|
||||
|
||||
/// [AWS] Secret access key. Required when `--access-key` is given.
|
||||
#[arg(long, env = "S3_SECRET_KEY")]
|
||||
secret_key: Option<String>,
|
||||
|
||||
/// [AWS] Optional session token for short-lived credentials.
|
||||
#[arg(long, env = "S3_SESSION_TOKEN")]
|
||||
session_token: Option<String>,
|
||||
|
||||
/// [AWS] Use S3 Express One Zone (directory buckets, named `*--x-s3`).
|
||||
#[arg(long, env = "S3_EXPRESS")]
|
||||
s3_express: bool,
|
||||
|
||||
/// [GCS] Path to a service-account JSON key file. Takes precedence over
|
||||
/// `--gcs-service-account-key`; if neither is set, application default
|
||||
/// credentials (ADC) are used.
|
||||
#[arg(long, env = "GCS_SERVICE_ACCOUNT_PATH")]
|
||||
gcs_service_account_path: Option<String>,
|
||||
|
||||
/// [GCS] Inline service-account JSON key contents (instead of a path).
|
||||
#[arg(long, env = "GCS_SERVICE_ACCOUNT_KEY")]
|
||||
gcs_service_account_key: Option<String>,
|
||||
|
||||
/// [AWS/GCS] Key prefix inside the bucket pointing at the edge-shard root
|
||||
/// (the directory containing `segments/` and the segment manifest). Empty
|
||||
/// means the bucket root.
|
||||
#[arg(long, default_value = "")]
|
||||
prefix: String,
|
||||
|
||||
/// [AWS/GCS] Local directory for the segment disk cache. Remote blocks are
|
||||
/// fetched once and mirrored here; later reads hit this directory instead.
|
||||
/// Defaults to a stable subdirectory of the system temp dir.
|
||||
#[arg(long)]
|
||||
cache_dir: Option<PathBuf>,
|
||||
}
|
||||
|
||||
/// The shard's write-facing schema, read off the appendable segment: the
|
||||
/// segment config names the vectors a point must carry, the payload-index
|
||||
/// config names the payload fields worth generating.
|
||||
struct ShardSchema {
|
||||
config: SegmentConfig,
|
||||
payload_fields: Vec<(String, PayloadSchemaType)>,
|
||||
}
|
||||
|
||||
/// The write-target schema off the opened shard: its segment config (already
|
||||
/// parsed during the open — no extra reads) plus the payload-index schema,
|
||||
/// which is the one file the writer deliberately never opens, so the tool
|
||||
/// reads it through `fs` itself.
|
||||
fn read_schema<S: UniversalRead + 'static, F: UniversalReadFs>(
|
||||
shard: &UpdateOnlyEdgeShard<S>,
|
||||
fs: &F,
|
||||
shard_path: &Path,
|
||||
) -> Result<ShardSchema> {
|
||||
let mut target = None;
|
||||
|
||||
for info in shard.segment_configs() {
|
||||
let SegmentConfigInfo {
|
||||
uuid,
|
||||
is_write_target,
|
||||
config,
|
||||
} = info;
|
||||
log::info!(
|
||||
"segment {uuid}: appendable={is_write_target}, {} dense vector(s) {:?}, \
|
||||
{} sparse vector(s) {:?}",
|
||||
config.vector_data.len(),
|
||||
config.vector_data.keys().collect::<Vec<_>>(),
|
||||
config.sparse_vector_data.len(),
|
||||
config.sparse_vector_data.keys().collect::<Vec<_>>(),
|
||||
);
|
||||
if is_write_target {
|
||||
target = Some((uuid, config));
|
||||
}
|
||||
}
|
||||
|
||||
let (uuid, config) = target
|
||||
.ok_or_else(|| anyhow!("no appendable segment found — the shard has no write target"))?;
|
||||
log::info!("write target: segment {uuid}");
|
||||
|
||||
let segment_path = shard_path.join("segments").join(uuid.to_string());
|
||||
let payload_config_path = get_payload_index_path(&segment_path).join(PAYLOAD_INDEX_CONFIG_FILE);
|
||||
let payload_config: Option<PayloadConfig> =
|
||||
match read_json_via(fs, &payload_config_path).ok_not_found() {
|
||||
Ok(Some(payload_config)) => Some(payload_config),
|
||||
Ok(None) => {
|
||||
log::info!(
|
||||
"no payload index config at {} — generating without payload",
|
||||
payload_config_path.display(),
|
||||
);
|
||||
None
|
||||
}
|
||||
// A dry-run generator should not die on an unreadable auxiliary
|
||||
// file, but the reason must stay visible — it may be a
|
||||
// credentials problem.
|
||||
Err(err) => {
|
||||
log::warn!(
|
||||
"could not read payload index config at {}: {err} — generating without payload",
|
||||
payload_config_path.display(),
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
let payload_fields = match payload_config {
|
||||
Some(payload_config) => {
|
||||
let mut fields: Vec<(String, PayloadSchemaType)> = payload_config
|
||||
.indices
|
||||
.to_schemas()
|
||||
.into_iter()
|
||||
.map(|(key, schema)| (key.to_string(), base_schema_type(&schema)))
|
||||
.collect();
|
||||
fields.sort_by(|(a, _), (b, _)| a.cmp(b));
|
||||
fields
|
||||
}
|
||||
None => Vec::new(),
|
||||
};
|
||||
|
||||
for (field, schema_type) in &payload_fields {
|
||||
log::info!("payload schema: {field:?} -> {schema_type:?}");
|
||||
}
|
||||
|
||||
Ok(ShardSchema {
|
||||
config,
|
||||
payload_fields,
|
||||
})
|
||||
}
|
||||
|
||||
/// The base value type behind a payload field schema, with or without index
|
||||
/// params.
|
||||
fn base_schema_type(schema: &PayloadFieldSchema) -> PayloadSchemaType {
|
||||
match schema {
|
||||
PayloadFieldSchema::FieldType(schema_type) => *schema_type,
|
||||
PayloadFieldSchema::FieldParams(params) => match params {
|
||||
PayloadSchemaParams::Keyword(_) => PayloadSchemaType::Keyword,
|
||||
PayloadSchemaParams::Integer(_) => PayloadSchemaType::Integer,
|
||||
PayloadSchemaParams::Float(_) => PayloadSchemaType::Float,
|
||||
PayloadSchemaParams::Geo(_) => PayloadSchemaType::Geo,
|
||||
PayloadSchemaParams::Text(_) => PayloadSchemaType::Text,
|
||||
PayloadSchemaParams::Bool(_) => PayloadSchemaType::Bool,
|
||||
PayloadSchemaParams::Datetime(_) => PayloadSchemaType::Datetime,
|
||||
PayloadSchemaParams::Uuid(_) => PayloadSchemaType::Uuid,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
const WORDS: &[&str] = &[
|
||||
"amber", "basalt", "cobalt", "dune", "ember", "fjord", "garnet", "harbor",
|
||||
];
|
||||
|
||||
/// One random point in the shape the schema prescribes: every named vector at
|
||||
/// its configured dimensionality, one payload value per indexed field.
|
||||
fn random_point(id: PointId, schema: &ShardSchema, rng: &mut StdRng) -> PointStructPersisted {
|
||||
let mut vectors = HashMap::new();
|
||||
|
||||
for (name, vector_config) in &schema.config.vector_data {
|
||||
let dense = |rng: &mut StdRng| {
|
||||
(0..vector_config.size)
|
||||
.map(|_| rng.random_range(-1.0..1.0))
|
||||
.collect::<Vec<f32>>()
|
||||
};
|
||||
let vector = if vector_config.multivector_config.is_some() {
|
||||
VectorPersisted::MultiDense(vec![dense(rng), dense(rng)])
|
||||
} else {
|
||||
VectorPersisted::Dense(dense(rng))
|
||||
};
|
||||
vectors.insert(name.clone(), vector);
|
||||
}
|
||||
|
||||
for name in schema.config.sparse_vector_data.keys() {
|
||||
// Cumulative random gaps: sorted, unique indices without a sampler.
|
||||
let mut index = 0u32;
|
||||
let mut indices = Vec::new();
|
||||
let mut values = Vec::new();
|
||||
for _ in 0..8 {
|
||||
index += rng.random_range(1..1000);
|
||||
indices.push(index);
|
||||
values.push(rng.random_range(0.0..1.0));
|
||||
}
|
||||
vectors.insert(
|
||||
name.clone(),
|
||||
VectorPersisted::Sparse(SparseVector { indices, values }),
|
||||
);
|
||||
}
|
||||
|
||||
let mut payload = serde_json::Map::new();
|
||||
for (field, schema_type) in &schema.payload_fields {
|
||||
payload.insert(field.clone(), random_payload_value(*schema_type, rng));
|
||||
}
|
||||
|
||||
PointStructPersisted {
|
||||
id,
|
||||
vector: VectorStructPersisted::Named(vectors),
|
||||
payload: Some(
|
||||
serde_json::from_value::<Payload>(serde_json::Value::Object(payload))
|
||||
.expect("a JSON object is always a valid payload"),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
PayloadSchemaType::Keyword => word(rng).into(),
|
||||
PayloadSchemaType::Integer => rng.random_range(0..1000).into(),
|
||||
PayloadSchemaType::Float => rng.random_range(0.0..100.0).into(),
|
||||
PayloadSchemaType::Bool => rng.random::<bool>().into(),
|
||||
PayloadSchemaType::Geo => serde_json::json!({
|
||||
"lon": rng.random_range(-180.0..180.0),
|
||||
"lat": rng.random_range(-85.0..85.0),
|
||||
}),
|
||||
PayloadSchemaType::Text => format!("{} {} {}", word(rng), word(rng), word(rng)).into(),
|
||||
PayloadSchemaType::Datetime => format!(
|
||||
"2026-{:02}-{:02}T{:02}:{:02}:{:02}Z",
|
||||
rng.random_range(1..=12),
|
||||
rng.random_range(1..=28),
|
||||
rng.random_range(0..24),
|
||||
rng.random_range(0..60),
|
||||
rng.random_range(0..60),
|
||||
)
|
||||
.into(),
|
||||
PayloadSchemaType::Uuid => Uuid::from_u128(rng.random()).to_string().into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a point id from the command line: a bare integer id or a UUID string.
|
||||
fn parse_point_id(raw: &str) -> Result<PointId> {
|
||||
// `PointId` deserializes a JSON number into a numeric id; a bare UUID is not
|
||||
// valid JSON, so quote it and retry as a JSON string.
|
||||
if let Ok(id) = serde_json::from_str::<PointId>(raw) {
|
||||
return Ok(id);
|
||||
}
|
||||
let quoted = serde_json::to_string(raw).expect("string is always serializable");
|
||||
serde_json::from_str::<PointId>("ed)
|
||||
.with_context(|| format!("invalid point id (not an integer or UUID): {raw:?}"))
|
||||
}
|
||||
|
||||
fn log_preview_point(point: &edge::PointPreview) {
|
||||
let edge::PointPreview {
|
||||
id,
|
||||
current,
|
||||
slots,
|
||||
action,
|
||||
} = point;
|
||||
|
||||
match current {
|
||||
Some(current) => log::info!(
|
||||
"point {id}: newest copy in segment {} slot {} at version {}",
|
||||
current.segment,
|
||||
current.internal_id,
|
||||
current.version,
|
||||
),
|
||||
None => log::info!("point {id}: not stored in any segment (would be created)"),
|
||||
}
|
||||
for (segment, internal_id) in slots {
|
||||
log::info!("point {id}: occupies segment {segment} slot {internal_id}");
|
||||
}
|
||||
|
||||
match action {
|
||||
PointAction::Store(resolved) => {
|
||||
let FullyQualifiedPoint {
|
||||
id: _,
|
||||
version,
|
||||
stored_vectors,
|
||||
updated_vectors,
|
||||
payload,
|
||||
} = resolved.as_ref();
|
||||
log::info!(
|
||||
"point {id}: would be stored at version {version}: \
|
||||
{} vector(s) carried over as raw bytes {:?}, \
|
||||
{} vector(s) supplied by the batch {:?}, payload {}",
|
||||
stored_vectors.len(),
|
||||
stored_vectors
|
||||
.iter()
|
||||
.map(|(name, bytes)| format!("{name}({} B)", bytes.len()))
|
||||
.collect::<Vec<_>>(),
|
||||
updated_vectors.len(),
|
||||
updated_vectors.keys().collect::<Vec<_>>(),
|
||||
serde_json::to_string(payload).unwrap_or_else(|err| err.to_string()),
|
||||
);
|
||||
if !slots.is_empty() {
|
||||
log::info!(
|
||||
"point {id}: its {} current slot(s) would be tombstoned",
|
||||
slots.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
PointAction::Delete => {
|
||||
log::info!(
|
||||
"point {id}: would be deleted, tombstoning {} slot(s)",
|
||||
slots.len()
|
||||
);
|
||||
}
|
||||
PointAction::Skip => log::info!(
|
||||
"point {id}: already at or beyond version — skipped (replay no-op); \
|
||||
re-run with a higher --op-num to overwrite",
|
||||
),
|
||||
PointAction::Missing => {
|
||||
log::info!("point {id}: names a point no segment holds — nothing to do");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate the random batch, resolve it against the open shard, and log what
|
||||
/// it would do. The backend is behind `S`, so this is the whole dry run for
|
||||
/// local and object-storage shards alike.
|
||||
fn dry_run<S: UniversalRead + 'static>(
|
||||
shard: &UpdateOnlyEdgeShard<S>,
|
||||
schema: &ShardSchema,
|
||||
ids: &[PointId],
|
||||
op_num: u64,
|
||||
seed: u64,
|
||||
) -> Result<()> {
|
||||
let mut rng = StdRng::seed_from_u64(seed);
|
||||
let points: Vec<PointStructPersisted> = ids
|
||||
.iter()
|
||||
.map(|&id| {
|
||||
let point = random_point(id, schema, &mut rng);
|
||||
log::info!(
|
||||
"generated point {id}: vectors {:?}, payload {}",
|
||||
match &point.vector {
|
||||
VectorStructPersisted::Single(v) => vec![format!("(default, {} dim)", v.len())],
|
||||
VectorStructPersisted::MultiDense(m) =>
|
||||
vec![format!("(default, {} multivector(s))", m.len())],
|
||||
VectorStructPersisted::Named(named) => named
|
||||
.iter()
|
||||
.map(|(name, vector)| match vector {
|
||||
VectorPersisted::Dense(v) => format!("{name}({} dim)", v.len()),
|
||||
VectorPersisted::Sparse(s) =>
|
||||
format!("{name}({} nnz)", s.indices.len()),
|
||||
VectorPersisted::MultiDense(m) =>
|
||||
format!("{name}({} multivector(s))", m.len()),
|
||||
})
|
||||
.collect(),
|
||||
},
|
||||
point
|
||||
.payload
|
||||
.as_ref()
|
||||
.map(|payload| serde_json::to_string(payload)
|
||||
.unwrap_or_else(|err| err.to_string()))
|
||||
.unwrap_or_else(|| "<none>".to_string()),
|
||||
);
|
||||
point
|
||||
})
|
||||
.collect();
|
||||
|
||||
let operation = UpdateOperation::PointOperation(PointOperations::UpsertPoints(points.into()));
|
||||
|
||||
let preview = shard
|
||||
.preview_batch([(op_num, operation)])
|
||||
.context("failed to resolve the batch")?;
|
||||
|
||||
let mut stored = 0usize;
|
||||
let mut skipped = 0usize;
|
||||
let mut tombstones: HashMap<Uuid, usize> = HashMap::new();
|
||||
for point in &preview.points {
|
||||
log_preview_point(point);
|
||||
match &point.action {
|
||||
PointAction::Store(_) => {
|
||||
stored += 1;
|
||||
for (segment, _) in &point.slots {
|
||||
*tombstones.entry(*segment).or_default() += 1;
|
||||
}
|
||||
}
|
||||
PointAction::Delete => {
|
||||
for (segment, _) in &point.slots {
|
||||
*tombstones.entry(*segment).or_default() += 1;
|
||||
}
|
||||
}
|
||||
PointAction::Skip => skipped += 1,
|
||||
PointAction::Missing => {}
|
||||
}
|
||||
}
|
||||
|
||||
log::info!(
|
||||
"summary: {stored} point(s) would be appended to the write target, {skipped} skipped",
|
||||
);
|
||||
for (segment, count) in &tombstones {
|
||||
log::info!("summary: segment {segment} would receive {count} tombstone(s)");
|
||||
}
|
||||
log::info!("dry run only: store_points/tombstone_points are not implemented — nothing written");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The bucket name, required by the object-storage backends.
|
||||
fn require_bucket(conn: &ConnectionArgs) -> Result<String> {
|
||||
conn.bucket
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow!("--bucket is required for the {:?} backend", conn.backend))
|
||||
}
|
||||
|
||||
fn build_aws_config(conn: &ConnectionArgs) -> Result<AwsConfig> {
|
||||
let credentials = match (&conn.access_key, &conn.secret_key) {
|
||||
(Some(access_key_id), Some(secret_access_key)) => AwsCredentials::Static {
|
||||
access_key_id: access_key_id.clone(),
|
||||
secret_access_key: secret_access_key.clone(),
|
||||
session_token: conn.session_token.clone(),
|
||||
},
|
||||
(None, None) => AwsCredentials::Default,
|
||||
_ => {
|
||||
return Err(anyhow!(
|
||||
"--access-key and --secret-key must be provided together"
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
Ok(AwsConfig {
|
||||
bucket: require_bucket(conn)?,
|
||||
region: conn.region.clone(),
|
||||
endpoint: conn.endpoint.clone(),
|
||||
s3_express: conn.s3_express,
|
||||
credentials,
|
||||
})
|
||||
}
|
||||
|
||||
fn build_gcs_config(conn: &ConnectionArgs) -> Result<GcsConfig> {
|
||||
let credentials = if let Some(key) = &conn.gcs_service_account_key {
|
||||
GcsCredentials::ServiceAccountKey(key.clone())
|
||||
} else if let Some(path) = &conn.gcs_service_account_path {
|
||||
GcsCredentials::ServiceAccountPath(path.clone())
|
||||
} else {
|
||||
GcsCredentials::Default
|
||||
};
|
||||
|
||||
Ok(GcsConfig {
|
||||
bucket: require_bucket(conn)?,
|
||||
credentials,
|
||||
})
|
||||
}
|
||||
|
||||
/// Build the disk-cache-backed filesystem used to read segment data: each
|
||||
/// remote block is fetched from object storage once and served from the local
|
||||
/// mirror afterwards.
|
||||
fn build_cached_fs<A>(
|
||||
remote_config: A::Config,
|
||||
remote_prefix: &Path,
|
||||
cache_dir: &Path,
|
||||
) -> Result<DiskCacheFs<BlobFile<A>>>
|
||||
where
|
||||
A: AsyncRead + Clone,
|
||||
{
|
||||
fs_err::create_dir_all(cache_dir)
|
||||
.with_context(|| format!("failed to create cache dir {}", cache_dir.display()))?;
|
||||
|
||||
let config = DiskCacheConfig::new(remote_prefix.to_path_buf(), cache_dir.to_path_buf())
|
||||
.context("failed to build disk cache config")?;
|
||||
|
||||
DiskCacheFs::<BlobFile<A>>::from_context(DiskCacheFsContext {
|
||||
config: Arc::new(config),
|
||||
remote: remote_config,
|
||||
})
|
||||
.context("failed to build disk-cache filesystem")
|
||||
}
|
||||
|
||||
/// Dry-run against a local shard directory: segments discovered by scanning
|
||||
/// `segments/`, read over memory-mapped files.
|
||||
fn run_local(cli: &Cli, ids: &[PointId]) -> Result<()> {
|
||||
let path = cli
|
||||
.connection
|
||||
.path
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow!("--path is required for the local backend"))?;
|
||||
|
||||
let shard = UpdateOnlyEdgeShard::<MmapFile>::open_mmap(&path)
|
||||
.context("failed to open update-only edge shard")?;
|
||||
log::info!(
|
||||
"opened update-only shard with {} segment(s)",
|
||||
shard.segments_count()
|
||||
);
|
||||
|
||||
let schema = read_schema(&shard, &common::universal_io::MmapFs, &path)?;
|
||||
dry_run(&shard, &schema, ids, cli.op_num, cli.seed)
|
||||
}
|
||||
|
||||
/// Dry-run against object storage: segments discovered from the leader's
|
||||
/// segment manifest, read through the local disk cache.
|
||||
fn run_remote<A>(cli: &Cli, ids: &[PointId], remote_config: A::Config) -> Result<()>
|
||||
where
|
||||
A: AsyncRead + Clone,
|
||||
{
|
||||
let prefix = PathBuf::from(&cli.connection.prefix);
|
||||
let cache_dir = cli
|
||||
.connection
|
||||
.cache_dir
|
||||
.clone()
|
||||
.unwrap_or_else(|| std::env::temp_dir().join("edge-shard-update-cache"));
|
||||
|
||||
let cached_fs = build_cached_fs::<A>(remote_config, &prefix, &cache_dir)?;
|
||||
log::info!("caching segment reads under {}", cache_dir.display());
|
||||
|
||||
let enumerator = ManifestSegmentEnumerator::new(cached_fs.clone(), &prefix);
|
||||
let shard =
|
||||
UpdateOnlyEdgeShard::<DiskCache<BlobFile<A>>>::open(cached_fs.clone(), &prefix, enumerator)
|
||||
.context("failed to open update-only edge shard over object storage")?;
|
||||
log::info!(
|
||||
"opened update-only shard with {} segment(s)",
|
||||
shard.segments_count()
|
||||
);
|
||||
|
||||
let schema = read_schema(&shard, &cached_fs, &prefix)?;
|
||||
dry_run(&shard, &schema, ids, cli.op_num, cli.seed)
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
|
||||
.format_timestamp_millis()
|
||||
.init();
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
let ids = cli
|
||||
.ids
|
||||
.iter()
|
||||
.map(|raw| parse_point_id(raw))
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
log::info!("dry-run upsert of {} random point(s): {ids:?}", ids.len());
|
||||
|
||||
let conn = &cli.connection;
|
||||
match conn.backend {
|
||||
Backend::Local => run_local(&cli, &ids),
|
||||
Backend::Aws => {
|
||||
run_remote::<ObjectStoreSource<AmazonS3>>(&cli, &ids, build_aws_config(conn)?)
|
||||
}
|
||||
Backend::Gcs => {
|
||||
run_remote::<ObjectStoreSource<GoogleCloudStorage>>(&cli, &ids, build_gcs_config(conn)?)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Point representations of the batch update path: [`StoredPoint`] is a point
|
||||
//! as read out of the segment that owns it, [`FullyQualifiedPoint`] is a point
|
||||
//! resolved in full, ready to be appended.
|
||||
|
||||
use common::types::PointOffsetType;
|
||||
|
||||
use crate::data_types::named_vectors::NamedVectors;
|
||||
use crate::data_types::segment_record::NamedVectorBytesOwned;
|
||||
use crate::types::{Payload, PointIdType, SeqNumberType};
|
||||
|
||||
/// The stored form of a point: the base a batch of mutations is folded onto.
|
||||
///
|
||||
/// Vectors are storage-native bytes, never decoded, so they can move to a new
|
||||
/// slot without the lossy decode/re-encode round-trip — the same contract as
|
||||
/// [`SegmentEntry::upsert_moved_point`].
|
||||
///
|
||||
/// Carries no version: versions are resolved separately, before the point is
|
||||
/// read (see `UpdateOnlySegment::point_versions`).
|
||||
///
|
||||
/// [`SegmentEntry::upsert_moved_point`]: crate::entry::entry_point::SegmentEntry::upsert_moved_point
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StoredPoint {
|
||||
/// Slot the point occupies in the segment it was read from.
|
||||
pub internal_id: PointOffsetType,
|
||||
/// Every named vector the point has, in storage-native bytes.
|
||||
pub vectors: NamedVectorBytesOwned,
|
||||
/// The point's complete payload; empty when it has none.
|
||||
pub payload: Payload,
|
||||
}
|
||||
|
||||
/// A point resolved to everything a segment needs in order to store it:
|
||||
/// storing a fully qualified point reads nothing back.
|
||||
///
|
||||
/// Vectors come in two halves: `stored_vectors` carried over verbatim as
|
||||
/// storage-native bytes, `updated_vectors` supplied by the batch and therefore
|
||||
/// decoded. A name present in both is taken from `updated_vectors`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FullyQualifiedPoint {
|
||||
pub id: PointIdType,
|
||||
/// Operation number to record as the point's version — the highest one
|
||||
/// among the batch operations folded into it.
|
||||
pub version: SeqNumberType,
|
||||
/// Vectors carried over from the point's previous slot, in storage-native
|
||||
/// bytes. Empty for a point the batch creates from scratch.
|
||||
pub stored_vectors: NamedVectorBytesOwned,
|
||||
/// Vectors supplied by the batch, overriding `stored_vectors` by name.
|
||||
pub updated_vectors: NamedVectors<'static>,
|
||||
/// The point's complete payload, already merged.
|
||||
pub payload: Payload,
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
pub mod build_index_result;
|
||||
pub mod collection_defaults;
|
||||
pub mod facets;
|
||||
pub mod fully_qualified_point;
|
||||
pub mod groups;
|
||||
pub mod index;
|
||||
pub mod load_profile;
|
||||
|
||||
@@ -263,6 +263,10 @@ impl<'a> NamedVectors<'a> {
|
||||
.insert(Cow::Borrowed(name), CowVector::from(vector));
|
||||
}
|
||||
|
||||
pub fn remove_ref(&mut self, key: &VectorName) {
|
||||
self.map.remove(key);
|
||||
}
|
||||
|
||||
pub fn contains_key(&self, key: &VectorName) -> bool {
|
||||
self.map.contains_key(key)
|
||||
}
|
||||
|
||||
@@ -25,11 +25,17 @@ impl<S: UniversalRead> ReadOnlyAppendableIdTracker<S> {
|
||||
|
||||
/// Schedule background prefetch of the mappings log and versions file that
|
||||
/// [`open`](Self::open) reads via [`live_reload`](Self::live_reload).
|
||||
///
|
||||
/// Either file may not exist yet — the writer only creates them once it
|
||||
/// flushes the first point, and [`open`](Self::open) treats a missing file
|
||||
/// as an empty storage — so absence is tolerated here too.
|
||||
pub fn preopen(fs: &impl CachedReadFs<File = S>, segment_path: &Path) -> OperationResult<()> {
|
||||
let options = Self::open_options();
|
||||
|
||||
fs.schedule_prefetch(&mappings_path(segment_path), Some(options), None)?;
|
||||
fs.schedule_prefetch(&versions_path(segment_path), Some(options), None)?;
|
||||
fs.schedule_prefetch(&mappings_path(segment_path), Some(options), None)
|
||||
.ok_not_found()?;
|
||||
fs.schedule_prefetch(&versions_path(segment_path), Some(options), None)
|
||||
.ok_not_found()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -15,6 +15,8 @@ mod tests;
|
||||
|
||||
pub mod read_only;
|
||||
|
||||
pub mod update_only;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::path::PathBuf;
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
//! The write half of an [`UpdateOnlySegment`]: storing resolved points and
|
||||
//! tombstoning the slots they replace.
|
||||
//!
|
||||
//! Not implemented in this iteration: every method needs append-only
|
||||
//! components that do not exist yet — an appendable `DynamicStoredFlags`
|
||||
//! (deleted-points bitmask), appendable `ChunkedVectors`, an appendable
|
||||
//! payload blobstore and field indexes. Today's equivalents all mutate at an
|
||||
//! offset, which an object store cannot do. The signatures fix the shape of
|
||||
//! the write; the bodies are `todo!()`.
|
||||
|
||||
use common::counter::hardware_counter::HardwareCounterCell;
|
||||
use common::types::PointOffsetType;
|
||||
use common::universal_io::UniversalRead;
|
||||
|
||||
use super::UpdateOnlySegment;
|
||||
use crate::common::operation_error::OperationResult;
|
||||
use crate::data_types::fully_qualified_point::FullyQualifiedPoint;
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlySegment<S> {
|
||||
/// Append `points` to this segment, each into a fresh slot, and repoint
|
||||
/// the id tracker at those slots. A point that already exists here is
|
||||
/// never rewritten in place: it is written anew and its previous slot is
|
||||
/// tombstoned.
|
||||
///
|
||||
/// Requires [`is_appendable`](UpdateOnlySegment::is_appendable).
|
||||
pub fn store_points(
|
||||
&mut self,
|
||||
points: &[FullyQualifiedPoint],
|
||||
hw_counter: &HardwareCounterCell,
|
||||
) -> OperationResult<()> {
|
||||
debug_assert!(self.is_appendable());
|
||||
let _ = (points, hw_counter);
|
||||
todo!("needs the append-only storages and field indexes of the write target")
|
||||
}
|
||||
|
||||
/// Mark `internal_ids` deleted in this segment. Nothing but the
|
||||
/// deleted-points bitmask is written: the payload row, the vectors and the
|
||||
/// field indexes at those slots are left untouched.
|
||||
pub fn tombstone_points(&mut self, internal_ids: &[PointOffsetType]) -> OperationResult<()> {
|
||||
let _ = internal_ids;
|
||||
todo!("needs an appendable deleted-points bitmask")
|
||||
}
|
||||
|
||||
/// Persist everything written since the last flush. There is no WAL:
|
||||
/// writes are durable only once this returns.
|
||||
pub fn flush(&self) -> OperationResult<()> {
|
||||
todo!("needs the append-only storages")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use atomic_refcell::AtomicRefCell;
|
||||
use common::storage_version::{StorageVersion, VERSION_FILE};
|
||||
use common::types::PointOffsetType;
|
||||
use common::universal_io::{
|
||||
CachedFs, CachedReadFs, OkNotFound as _, Populate, UniversalRead, UniversalReadFs,
|
||||
read_json_via,
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{UpdateOnlySegment, UpdateOnlyVectorData};
|
||||
use crate::common::operation_error::{OperationError, OperationResult};
|
||||
use crate::id_tracker::read_only_tracker_enum::ReadOnlyIdTrackerEnum;
|
||||
use crate::payload_storage::read_only::ReadOnlyPayloadStorage;
|
||||
use crate::segment::{SEGMENT_STATE_FILE, SegmentVersion};
|
||||
use crate::segment_constructor::get_vector_storage_path;
|
||||
use crate::types::{SegmentConfig, SegmentState};
|
||||
use crate::vector_storage::read_only::VectorStorageReadEnum;
|
||||
use crate::vector_storage::sparse::read_only::ReadOnlySparseVectorStorage;
|
||||
|
||||
/// Every component of an update-only segment is opened cold: a writer touches
|
||||
/// a handful of scattered points, so warming a storage would fetch far more
|
||||
/// than it reads.
|
||||
///
|
||||
/// On a remote backend this is also what keeps [`preopen`] cheap: a prefetch
|
||||
/// is an open, and an open with `Populate::No` transfers no content — so the
|
||||
/// data files (vectors, payload pages) are never downloaded, only the files
|
||||
/// whose opens consume them whole (the configs, the id tracker, the deleted
|
||||
/// flags).
|
||||
///
|
||||
/// [`preopen`]: UpdateOnlySegment::preopen
|
||||
const WRITER_POPULATE: Populate = Populate::No;
|
||||
|
||||
/// Build the per-segment [`CachedFs`] an open runs over: the version and
|
||||
/// state files are prefetched, and the directory listing snapshot is taken so
|
||||
/// probes for optional files resolve without inner-filesystem round-trips.
|
||||
/// Mirror of the read-only segment's `build_cached_fs`, minus the payload
|
||||
/// index config the writer never opens.
|
||||
fn build_cached_fs<Fs: UniversalReadFs>(
|
||||
fs: &Fs,
|
||||
segment_path: &Path,
|
||||
) -> OperationResult<CachedFs<Fs>> {
|
||||
let mut cached_fs = CachedFs::new(fs.clone(), segment_path)?;
|
||||
|
||||
// Absence is tolerated here: the subsequent read reports it gracefully.
|
||||
for file_name in [VERSION_FILE, SEGMENT_STATE_FILE] {
|
||||
cached_fs
|
||||
.schedule_prefetch(&segment_path.join(file_name), None, None)
|
||||
.ok_not_found()?;
|
||||
}
|
||||
|
||||
cached_fs.cache_file_info()?;
|
||||
|
||||
Ok(cached_fs)
|
||||
}
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlySegment<S> {
|
||||
/// Open the segment over a per-segment [`CachedFs`]: every file the
|
||||
/// components will read is prefetched concurrently
|
||||
/// ([`preopen`](Self::preopen)) before the component opens consume it, so
|
||||
/// a remote backend pays for the depth of the longest dependent chain
|
||||
/// rather than one blocking round-trip per file.
|
||||
///
|
||||
/// `fs` is the canonical backend: the caching wrapper lives only for this
|
||||
/// open, and it is `fs` that components keep for later re-opens and the
|
||||
/// segment keeps for its appends.
|
||||
///
|
||||
/// `deferred_internal_id` is the cutoff agreed with an external rebuilder
|
||||
/// working the same directory — see [`open_via`](Self::open_via).
|
||||
pub fn open(
|
||||
fs: &S::Fs,
|
||||
segment_path: &Path,
|
||||
uuid: Uuid,
|
||||
deferred_internal_id: Option<PointOffsetType>,
|
||||
) -> OperationResult<Self>
|
||||
where
|
||||
S::Fs: UniversalReadFs<File = S>,
|
||||
{
|
||||
let cached_fs = build_cached_fs(fs, segment_path)?;
|
||||
let config = Self::preopen(&cached_fs, segment_path)?;
|
||||
Self::open_via(
|
||||
&cached_fs,
|
||||
fs,
|
||||
segment_path,
|
||||
config,
|
||||
uuid,
|
||||
deferred_internal_id,
|
||||
)
|
||||
}
|
||||
|
||||
/// Open the segment's components: the id tracker, the payload storage and
|
||||
/// one storage per named vector — nothing else.
|
||||
///
|
||||
/// `fs` opens the component files (in production the [`CachedFs`] that
|
||||
/// [`open`](Self::open) primed); `raw_fs` is the canonical backend, kept
|
||||
/// by components that re-open files after this call and by the segment
|
||||
/// itself for its appends. `config` is the one [`preopen`](Self::preopen)
|
||||
/// already parsed, so the state file is not read twice.
|
||||
///
|
||||
/// `deferred_internal_id` is the cutoff agreed with an external rebuilder
|
||||
/// working the same directory: slots at or above it load into the id
|
||||
/// tracker's deferred track, marking them as outside the rebuild snapshot.
|
||||
/// It does not hide anything from the writer — resolution runs
|
||||
/// `WithDeferred`, so every point still locates at its latest slot. `None`
|
||||
/// (a writer running alone) keeps every mapping active. Only the
|
||||
/// appendable segment has a deferred track; on any other segment the
|
||||
/// cutoff is ignored.
|
||||
pub fn open_via(
|
||||
fs: &impl UniversalReadFs<File = S>,
|
||||
raw_fs: &S::Fs,
|
||||
segment_path: &Path,
|
||||
config: SegmentConfig,
|
||||
uuid: Uuid,
|
||||
deferred_internal_id: Option<PointOffsetType>,
|
||||
) -> OperationResult<Self> {
|
||||
if SegmentVersion::load_universal(fs, segment_path)?.is_none() {
|
||||
// `FileNotFound`, not a service error: the version file is written
|
||||
// last, so its absence means the segment vanished mid-open (or was
|
||||
// never completed).
|
||||
return Err(OperationError::FileNotFound {
|
||||
path: segment_path.join(VERSION_FILE),
|
||||
});
|
||||
}
|
||||
|
||||
let payload_storage = Arc::new(AtomicRefCell::new(ReadOnlyPayloadStorage::open(
|
||||
fs,
|
||||
segment_path.to_path_buf(),
|
||||
WRITER_POPULATE,
|
||||
)?));
|
||||
|
||||
let appendable = config.is_appendable();
|
||||
|
||||
// Detect the persisted format by attempting each format's open (no
|
||||
// per-file `exists` round-trips — important for object-storage
|
||||
// backends). The deferred threshold applies to the appendable tracker
|
||||
// only, mirroring `ReadOnlySegment::open_via`.
|
||||
let id_tracker = Arc::new(AtomicRefCell::new(ReadOnlyIdTrackerEnum::detect_and_load(
|
||||
fs,
|
||||
raw_fs,
|
||||
segment_path,
|
||||
deferred_internal_id.filter(|_| appendable),
|
||||
)?));
|
||||
|
||||
let mut vector_data = HashMap::new();
|
||||
for (vector_name, vector_config) in &config.vector_data {
|
||||
let path = get_vector_storage_path(segment_path, vector_name);
|
||||
let storage =
|
||||
VectorStorageReadEnum::open(fs, vector_config, &path, Some(WRITER_POPULATE))?
|
||||
.ok_or_else(|| {
|
||||
OperationError::service_error(format!(
|
||||
"Dense vector storage '{vector_name}' was not found, or is corrupted.",
|
||||
))
|
||||
})?;
|
||||
vector_data.insert(
|
||||
vector_name.clone(),
|
||||
UpdateOnlyVectorData {
|
||||
vector_storage: Arc::new(AtomicRefCell::new(storage)),
|
||||
},
|
||||
);
|
||||
}
|
||||
for vector_name in config.sparse_vector_data.keys() {
|
||||
let path = get_vector_storage_path(segment_path, vector_name);
|
||||
let storage = VectorStorageReadEnum::Sparse(Box::new(
|
||||
ReadOnlySparseVectorStorage::open(fs, &path, WRITER_POPULATE)?,
|
||||
));
|
||||
vector_data.insert(
|
||||
vector_name.clone(),
|
||||
UpdateOnlyVectorData {
|
||||
vector_storage: Arc::new(AtomicRefCell::new(storage)),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
uuid,
|
||||
segment_path: segment_path.to_path_buf(),
|
||||
fs: raw_fs.clone(),
|
||||
id_tracker,
|
||||
payload_storage,
|
||||
vector_data,
|
||||
segment_config: config,
|
||||
appendable,
|
||||
})
|
||||
}
|
||||
|
||||
/// Schedule the prefetch of every file [`open`](Self::open) will read, so a
|
||||
/// caching filesystem can fetch them concurrently instead of one blocking
|
||||
/// round-trip at a time. Returns the segment config parsed along the way,
|
||||
/// so the open does not read it twice.
|
||||
pub fn preopen(
|
||||
fs: &impl CachedReadFs<File = S>,
|
||||
segment_path: &Path,
|
||||
) -> OperationResult<SegmentConfig> {
|
||||
let SegmentState {
|
||||
initial_version: _,
|
||||
version: _,
|
||||
config,
|
||||
} = read_json_via(fs, segment_path.join(SEGMENT_STATE_FILE))?;
|
||||
|
||||
ReadOnlyPayloadStorage::preopen(fs, segment_path.to_path_buf(), WRITER_POPULATE)?;
|
||||
ReadOnlyIdTrackerEnum::preopen(fs, segment_path)?;
|
||||
|
||||
for (vector_name, vector_config) in &config.vector_data {
|
||||
let path = get_vector_storage_path(segment_path, vector_name);
|
||||
VectorStorageReadEnum::<S>::preopen(fs, vector_config, &path, Some(WRITER_POPULATE))?;
|
||||
}
|
||||
for vector_name in config.sparse_vector_data.keys() {
|
||||
let path = get_vector_storage_path(segment_path, vector_name);
|
||||
ReadOnlySparseVectorStorage::<S>::preopen(fs, &path, WRITER_POPULATE)?;
|
||||
}
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
//! Update-only segment: the write-side counterpart of
|
||||
//! [`ReadOnlySegment`](crate::segment::read_only::ReadOnlySegment).
|
||||
//!
|
||||
//! Its public surface is storing points and tombstoning slots; internally it
|
||||
//! still *reads*, because an operation like `set_payload` names a point but
|
||||
//! not its vectors. It opens exactly what that requires — the id tracker, the
|
||||
//! payload storage and the vector storages. No vector index, no quantized
|
||||
//! vectors, no payload index: on a remote backend those files are never
|
||||
//! fetched.
|
||||
|
||||
mod append;
|
||||
mod lifecycle;
|
||||
mod resolve;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use atomic_refcell::AtomicRefCell;
|
||||
use common::universal_io::UniversalRead;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::id_tracker::read_only_tracker_enum::ReadOnlyIdTrackerEnum;
|
||||
use crate::payload_storage::read_only::ReadOnlyPayloadStorage;
|
||||
use crate::types::{SegmentConfig, VectorNameBuf};
|
||||
use crate::vector_storage::read_only::VectorStorageReadEnum;
|
||||
|
||||
/// A segment open for updates: read components to resolve points with,
|
||||
/// append-only components to store them through. Generic over the backend `S`,
|
||||
/// like `ReadOnlySegment`.
|
||||
pub struct UpdateOnlySegment<S: UniversalRead + 'static> {
|
||||
pub uuid: Uuid,
|
||||
/// Path to the segment directory.
|
||||
pub segment_path: PathBuf,
|
||||
/// Backend the segment was opened on. Retained because appends need a
|
||||
/// filesystem handle of their own: the caching wrapper an open may go
|
||||
/// through only lives for that open.
|
||||
pub fs: S::Fs,
|
||||
|
||||
pub id_tracker: Arc<AtomicRefCell<ReadOnlyIdTrackerEnum<S>>>,
|
||||
pub payload_storage: Arc<AtomicRefCell<ReadOnlyPayloadStorage<S>>>,
|
||||
pub vector_data: HashMap<VectorNameBuf, UpdateOnlyVectorData<S>>,
|
||||
|
||||
pub segment_config: SegmentConfig,
|
||||
/// Whether this segment can accept appends and therefore be the target of
|
||||
/// a write.
|
||||
appendable: bool,
|
||||
}
|
||||
|
||||
/// A single named vector of an [`UpdateOnlySegment`]: storage only — no vector
|
||||
/// index, no quantized vectors.
|
||||
pub struct UpdateOnlyVectorData<S: UniversalRead + 'static> {
|
||||
pub vector_storage: Arc<AtomicRefCell<VectorStorageReadEnum<S>>>,
|
||||
}
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlySegment<S> {
|
||||
/// Whether this segment accepts appends, and can therefore be the target of
|
||||
/// a write.
|
||||
pub fn is_appendable(&self) -> bool {
|
||||
self.appendable
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
//! Resolving stored points: the read half of applying a batch of updates.
|
||||
//!
|
||||
//! Everything here is batched: each component is handed the whole set of ids
|
||||
//! at once, one pass per component rather than one round-trip per point.
|
||||
|
||||
use ahash::AHashMap;
|
||||
use common::counter::hardware_counter::HardwareCounterCell;
|
||||
use common::generic_consts::Random;
|
||||
use common::types::{DeferredBehavior, PointOffsetType};
|
||||
use common::universal_io::UniversalRead;
|
||||
|
||||
use super::UpdateOnlySegment;
|
||||
use crate::common::operation_error::OperationResult;
|
||||
use crate::data_types::fully_qualified_point::StoredPoint;
|
||||
use crate::data_types::segment_record::NamedVectorBytesOwned;
|
||||
use crate::id_tracker::IdTrackerRead;
|
||||
use crate::payload_storage::PayloadStorageRead;
|
||||
use crate::types::{Payload, PointIdType, SeqNumberType};
|
||||
use crate::vector_storage::VectorStorageRead;
|
||||
|
||||
impl<S: UniversalRead + 'static> UpdateOnlySegment<S> {
|
||||
/// Locate `point_ids` in this segment, streaming each `(point_id,
|
||||
/// internal_id)` pair that resolves; ids the segment does not hold are
|
||||
/// skipped. Deferred heads are included, so a point shadowed by an
|
||||
/// optimization in progress resolves to its latest slot.
|
||||
pub fn locate_points(
|
||||
&self,
|
||||
point_ids: impl IntoIterator<Item = PointIdType>,
|
||||
callback: impl FnMut(PointIdType, PointOffsetType),
|
||||
) -> OperationResult<()> {
|
||||
self.id_tracker.borrow().resolve_external_ids(
|
||||
point_ids,
|
||||
DeferredBehavior::WithDeferred,
|
||||
callback,
|
||||
)
|
||||
}
|
||||
|
||||
/// Versions of the points occupying `internal_ids`, keyed by internal id —
|
||||
/// the same key the id tracker's batch read yields.
|
||||
///
|
||||
/// A slot the tracker has no version for is absent from the map; it counts
|
||||
/// as `0`, the version an unwritten point compares as.
|
||||
pub fn point_versions(
|
||||
&self,
|
||||
internal_ids: &[PointOffsetType],
|
||||
) -> OperationResult<AHashMap<PointOffsetType, SeqNumberType>> {
|
||||
let mut versions = AHashMap::with_capacity(internal_ids.len());
|
||||
|
||||
self.id_tracker.borrow().internal_versions_batch(
|
||||
internal_ids.iter().copied(),
|
||||
|internal_id, version| {
|
||||
versions.insert(internal_id, version);
|
||||
},
|
||||
)?;
|
||||
|
||||
Ok(versions)
|
||||
}
|
||||
|
||||
/// Read the stored form of the points occupying `internal_ids`, returned
|
||||
/// in the same order — one batched pass per component, vectors as
|
||||
/// storage-native bytes.
|
||||
///
|
||||
/// A slot with no value in a given component contributes nothing: a point
|
||||
/// without payload gets an empty [`Payload`], and a vector name the point
|
||||
/// does not have (or has deleted) is absent from [`StoredPoint::vectors`].
|
||||
///
|
||||
/// `internal_ids` must be free of duplicates.
|
||||
pub fn read_stored_points(
|
||||
&self,
|
||||
internal_ids: &[PointOffsetType],
|
||||
hw_counter: &HardwareCounterCell,
|
||||
) -> OperationResult<Vec<StoredPoint>> {
|
||||
let mut stored: Vec<StoredPoint> = internal_ids
|
||||
.iter()
|
||||
.map(|&internal_id| StoredPoint {
|
||||
internal_id,
|
||||
vectors: NamedVectorBytesOwned::new(),
|
||||
payload: Payload::default(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
self.payload_storage
|
||||
.borrow()
|
||||
.read_payloads::<Random, usize>(
|
||||
internal_ids.iter().copied().enumerate(),
|
||||
|position, payload| {
|
||||
stored[position].payload = payload;
|
||||
Ok(())
|
||||
},
|
||||
hw_counter,
|
||||
)?;
|
||||
|
||||
for (vector_name, vector_data) in &self.vector_data {
|
||||
let vector_storage = vector_data.vector_storage.borrow();
|
||||
vector_storage.read_vector_bytes::<Random, usize>(
|
||||
internal_ids.iter().copied().enumerate(),
|
||||
|position, internal_id, bytes| {
|
||||
// A vector deleted on its own (`delete_vectors`) still has
|
||||
// bytes in the storage; carrying them over would resurrect
|
||||
// it in the rewritten point.
|
||||
if vector_storage.is_deleted_vector(internal_id) {
|
||||
return;
|
||||
}
|
||||
stored[position].vectors.push((vector_name.clone(), bytes));
|
||||
},
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(stored)
|
||||
}
|
||||
}
|
||||
@@ -28,10 +28,9 @@ mod vector_index;
|
||||
mod vector_storage;
|
||||
|
||||
pub(crate) use id_tracker::create_mutable_id_tracker;
|
||||
pub(crate) use paths::get_payload_index_path;
|
||||
pub use paths::{
|
||||
PAYLOAD_INDEX_PATH, VECTOR_INDEX_PATH, VECTOR_STORAGE_PATH, get_vector_index_path,
|
||||
get_vector_name_with_prefix, get_vector_storage_path,
|
||||
PAYLOAD_INDEX_PATH, VECTOR_INDEX_PATH, VECTOR_STORAGE_PATH, get_payload_index_path,
|
||||
get_vector_index_path, get_vector_name_with_prefix, get_vector_storage_path,
|
||||
};
|
||||
pub(crate) use payload_storage::create_payload_storage;
|
||||
pub use segment::{NewSegmentToken, build_segment, load_segment, normalize_segment_dir};
|
||||
|
||||
@@ -25,6 +25,6 @@ pub fn get_vector_index_path(segment_path: &Path, vector_name: &VectorName) -> P
|
||||
segment_path.join(get_vector_name_with_prefix(VECTOR_INDEX_PATH, vector_name))
|
||||
}
|
||||
|
||||
pub(crate) fn get_payload_index_path(segment_path: &Path) -> PathBuf {
|
||||
pub fn get_payload_index_path(segment_path: &Path) -> PathBuf {
|
||||
segment_path.join(PAYLOAD_INDEX_PATH)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user