#![allow(clippy::expect_used, clippy::indexing_slicing)]
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use tempfile::{TempDir, tempdir};
use zeppelin_embed::epoch::{
ComputeUnits, EmbeddingEpoch, EmbeddingRuntime, EmbeddingTower, EpochTransitionError,
Normalization, StoreEpoch,
};
use zeppelin_embed::fts::tokenizer::TokenizerConfig;
use zeppelin_embed::ingest::{DocId, Revision, SearchRequest};
use zeppelin_embed::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
use zeppelin_embed::lifecycle::{CancelToken, OpenOptions, QueryControl, Store};
use zeppelin_embed::manifest::io::commit_manifest;
use zeppelin_embed::manifest::{EpochMeta, Manifest};
use zeppelin_embed::meta::{AliveSet, ColumnStoreBuilder, Schema};
use zeppelin_embed::quant::{Bit4Factors, quantize_bit4};
use zeppelin_embed::scan::ScanOptions;
use zeppelin_embed::segment::writer::{
SegmentBuild, SegmentDocumentVersions, SegmentFactors, write_segment_with_documents,
};
use zeppelin_embed::segment::{SegmentId, SegmentMeta};
use zeppelin_embed::vfs::{StdVfs, SyncKind, Vfs, VfsFile};
const DIMS: usize = 4;
struct EpochFixture {
directory: TempDir,
epoch_a: StoreEpoch,
epoch_b: StoreEpoch,
segment_a: SegmentId,
}
fn policy() -> DurabilityPolicy {
DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None).expect("fixture policy")
}
fn store_epoch(version: &str, weight: u8) -> StoreEpoch {
let document = EmbeddingTower {
model_id: "epoch-fixture".to_owned(),
model_version: version.to_owned(),
weights_digest: vec![weight],
dims: DIMS as u32,
normalization: Normalization::None,
prompt_prefix: "document: ".to_owned(),
max_tokens: 32,
runtime: EmbeddingRuntime::CpuReference,
compute_units: ComputeUnits::Cpu,
os_build: None,
};
let mut query = document.clone();
query.prompt_prefix = "query: ".to_owned();
StoreEpoch {
embedding: EmbeddingEpoch {
document,
query,
alignment_digest: vec![weight, 0xa1],
},
tokenizer: TokenizerConfig::text_default().epoch(),
}
}
fn write_epoch_segment(
directory: &Path,
id: SegmentId,
epoch: &StoreEpoch,
vectors: &[f32],
) -> SegmentMeta {
let rows = vectors.len() / DIMS;
let mut codes = vec![0_u8; rows * DIMS.div_ceil(2)];
let mut factors = Vec::<Bit4Factors>::with_capacity(rows);
for (vector, encoded) in vectors
.chunks_exact(DIMS)
.zip(codes.chunks_exact_mut(DIMS.div_ceil(2)))
{
factors.push(quantize_bit4(vector, encoded).expect("finite fixture vector"));
}
let mut columns = ColumnStoreBuilder::new(Schema::timestamp_only());
for row in 0..rows {
columns
.push_row(row as i64, &[])
.expect("fixture timestamp row");
}
let columns = columns.finish().expect("fixture columns");
let alive = AliveSet::new(rows as u32);
let all_doc_ids = [DocId::new(10), DocId::new(20), DocId::new(30)];
let doc_ids = &all_doc_ids[..rows];
let revisions = vec![Revision::new(1); rows];
let mut meta = write_segment_with_documents(
&StdVfs,
directory,
SegmentBuild {
id,
scheme: 4,
dims: DIMS as u32,
codes: &codes,
factors: SegmentFactors::Bit4(&factors),
rescore: vectors,
columns: &columns,
alive: &alive,
},
SegmentDocumentVersions {
doc_ids,
revisions: &revisions,
},
policy(),
)
.expect("epoch fixture segment");
meta.epoch_id = Some(epoch.identity().embedding);
meta
}
fn publish_incomplete_target_fixture() -> EpochFixture {
let directory = tempdir().expect("incomplete epoch fixture directory");
let epoch_a = store_epoch("A", 0xa0);
let epoch_b = store_epoch("B", 0xb0);
let segment_a = SegmentId::new(0x0021_0000_0011, [0xa1; 10]);
let meta_a = write_epoch_segment(
directory.path(),
segment_a,
&epoch_a,
&[0.0, 1.0, 2.0, 3.0, 2.0, 1.0, 0.0, -1.0, 4.0, 3.0, 2.0, 1.0],
);
let meta_b = write_epoch_segment(
directory.path(),
SegmentId::new(0x0021_0000_0012, [0xb2; 10]),
&epoch_b,
&[4.0, 0.0, 1.0, 2.0, 0.0, -1.0, -2.0, -3.0],
);
commit_manifest(
&StdVfs,
directory.path(),
&Manifest {
#[cfg(feature = "graph-cypher")]
graph: None,
generation: 1,
log_seq: 0,
segments: vec![meta_a, meta_b],
epochs: vec![EpochMeta::from(&epoch_a), EpochMeta::from(&epoch_b)],
epoch_alias: Some(epoch_a.identity()),
schema: Schema::timestamp_only(),
},
policy(),
)
.expect("incomplete two-epoch manifest");
EpochFixture {
directory,
epoch_a,
epoch_b,
segment_a,
}
}
fn publish_two_epoch_fixture() -> EpochFixture {
let directory = tempdir().expect("epoch fixture directory");
let epoch_a = store_epoch("A", 0xa0);
let epoch_b = store_epoch("B", 0xb0);
let segment_a = SegmentId::new(0x0021_0000_0001, [0xa1; 10]);
let segment_b = SegmentId::new(0x0021_0000_0002, [0xb2; 10]);
let meta_a = write_epoch_segment(
directory.path(),
segment_a,
&epoch_a,
&[0.0, 1.0, 2.0, 3.0, 2.0, 1.0, 0.0, -1.0, 4.0, 3.0, 2.0, 1.0],
);
let meta_b = write_epoch_segment(
directory.path(),
segment_b,
&epoch_b,
&[
4.0, 0.0, 1.0, 2.0, 0.0, -1.0, -2.0, -3.0, 1.0, 1.0, 1.0, 1.0,
],
);
commit_manifest(
&StdVfs,
directory.path(),
&Manifest {
#[cfg(feature = "graph-cypher")]
graph: None,
generation: 1,
log_seq: 0,
segments: vec![meta_a, meta_b],
epochs: vec![EpochMeta::from(&epoch_a), EpochMeta::from(&epoch_b)],
epoch_alias: Some(epoch_a.identity()),
schema: Schema::timestamp_only(),
},
policy(),
)
.expect("two-epoch manifest");
EpochFixture {
directory,
epoch_a,
epoch_b,
segment_a,
}
}
fn search_bits(store: &Store) -> Vec<(u128, u64, u32)> {
store
.search(
SearchRequest::new(&[0.5, 0.5, 0.5, 0.5]),
3,
ScanOptions { thread_budget: 1 },
QueryControl::Cancel(CancelToken::new()),
)
.expect("epoch search")
.candidates
.into_iter()
.map(|candidate| {
let version = candidate.document().expect("fixture document identity");
(
version.doc_id().get(),
version.revision().get(),
candidate.score().to_bits(),
)
})
.collect()
}
#[test]
fn rollback_preserves_old_epoch_results_exactly_and_rollback_after_drop_epoch_is_a_typed_error() {
let fixture = publish_two_epoch_fixture();
let identity_a = fixture.epoch_a.identity();
let identity_b = fixture.epoch_b.identity();
let store = Store::open(
fixture.directory.path(),
OpenOptions::default().with_epoch(fixture.epoch_a.clone()),
)
.expect("open epoch A");
let epoch_a_before = search_bits(&store);
store
.switch_epoch_alias(identity_b)
.expect("publish epoch B alias");
let epoch_b = search_bits(&store);
assert_ne!(epoch_b, epoch_a_before, "fixture epochs must differ");
store
.switch_epoch_alias(identity_a)
.expect("roll back to epoch A");
assert_eq!(search_bits(&store), epoch_a_before);
store
.switch_epoch_alias(identity_b)
.expect("restore epoch B before dropping A");
let report = store
.drop_epoch(identity_a.embedding)
.expect("drop old epoch A");
assert_eq!(report.segments_dropped(), &[fixture.segment_a]);
assert!(
!fixture
.directory
.path()
.join(fixture.segment_a.file_name())
.exists()
);
let error = store
.switch_epoch_alias(identity_a)
.expect_err("a dropped epoch cannot be restored");
assert!(matches!(
error,
EpochTransitionError::EpochUnavailable { target } if target == identity_a
));
assert_eq!(search_bits(&store), epoch_b);
}
#[test]
fn an_alias_cannot_publish_an_incomplete_target_epoch() {
let fixture = publish_incomplete_target_fixture();
let identity_b = fixture.epoch_b.identity();
let store = Store::open(
fixture.directory.path(),
OpenOptions::default().with_epoch(fixture.epoch_a.clone()),
)
.expect("open epoch A");
let error = store
.switch_epoch_alias(identity_b)
.expect_err("an N-1 target must never publish");
assert!(matches!(
error,
EpochTransitionError::IncompleteEpoch {
target,
missing_documents: 1,
unexpected_documents: 0,
} if target == identity_b
));
let unchanged = store
.switch_epoch_alias(fixture.epoch_a.identity())
.expect("epoch A remains published");
assert!(!unchanged.manifest_committed());
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum DropEvent {
ManifestCommit,
Unlink,
}
#[derive(Clone, Default)]
struct EventVfs {
events: Arc<Mutex<Vec<DropEvent>>>,
}
impl EventVfs {
fn events(&self) -> Vec<DropEvent> {
self.events.lock().expect("event mutex").clone()
}
}
impl Vfs for EventVfs {
fn ensure_directory(&self, path: &Path, create: bool) -> std::io::Result<bool> {
StdVfs.ensure_directory(path, create)
}
fn open(&self, path: &Path) -> std::io::Result<u64> {
StdVfs.open(path)
}
fn open_for_map(&self, path: &Path) -> std::io::Result<std::fs::File> {
StdVfs.open_for_map(path)
}
fn read(&self, path: &Path) -> std::io::Result<Vec<u8>> {
StdVfs.read(path)
}
fn read_range(&self, path: &Path, offset: u64, length: usize) -> std::io::Result<Vec<u8>> {
StdVfs.read_range(path, offset, length)
}
fn write(&self, path: &Path, bytes: &[u8]) -> std::io::Result<()> {
StdVfs.write(path, bytes)
}
fn open_append(&self, path: &Path) -> std::io::Result<Box<dyn VfsFile>> {
StdVfs.open_append(path)
}
fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
StdVfs.rename(from, to)?;
if to.file_name().is_some_and(|name| name == "manifest.ze") {
self.events
.lock()
.expect("event mutex")
.push(DropEvent::ManifestCommit);
}
Ok(())
}
fn sync(&self, path: &Path, kind: SyncKind) -> std::io::Result<()> {
StdVfs.sync(path, kind)
}
fn list(&self, directory: &Path) -> std::io::Result<Vec<PathBuf>> {
StdVfs.list(directory)
}
fn for_each_direct_child(
&self,
directory: &Path,
visitor: &mut dyn FnMut(&Path) -> std::io::Result<()>,
) -> std::io::Result<()> {
StdVfs.for_each_direct_child(directory, visitor)
}
fn delete(&self, path: &Path) -> std::io::Result<()> {
StdVfs.delete(path)?;
self.events
.lock()
.expect("event mutex")
.push(DropEvent::Unlink);
Ok(())
}
}
#[test]
fn drop_epoch_commits_the_manifest_before_unlinking() {
let fixture = publish_two_epoch_fixture();
let identity_a = fixture.epoch_a.identity();
let identity_b = fixture.epoch_b.identity();
let store = Store::open(
fixture.directory.path(),
OpenOptions::default().with_epoch(fixture.epoch_a.clone()),
)
.expect("open epoch A");
store
.switch_epoch_alias(identity_b)
.expect("publish epoch B alias");
let vfs = EventVfs::default();
store
.drop_epoch_on_vfs(identity_a.embedding, &vfs)
.expect("drop old epoch A");
assert_eq!(
vfs.events(),
vec![DropEvent::ManifestCommit, DropEvent::Unlink]
);
}
#[test]
fn reindex_preserves_retained_embedding_epochs() {
let fixture = publish_two_epoch_fixture();
let store = Store::open(
fixture.directory.path(),
OpenOptions::default().with_epoch(fixture.epoch_a.clone()),
)
.expect("open");
let before = search_bits(&store);
store.reindex_text().expect("reindex all epochs");
assert_eq!(search_bits(&store), before);
store
.switch_epoch_alias(fixture.epoch_b.identity())
.expect("retained epoch B remains complete");
store
.switch_epoch_alias(fixture.epoch_a.identity())
.expect("epoch A remains complete");
assert_eq!(search_bits(&store), before);
}
#[cfg(feature = "graph-cypher")]
#[test]
fn graph_epoch_alias_refuses_without_recording_a_generation() {
let fixture = publish_two_epoch_fixture();
let options = OpenOptions::default().with_epoch(fixture.epoch_a.clone());
let store = Store::open(fixture.directory.path(), options.clone()).expect("open");
let generation = store.enable_graph().expect("graph");
let before = std::fs::read(fixture.directory.path().join("manifest.ze")).expect("manifest");
assert!(matches!(
store.switch_epoch_alias(fixture.epoch_b.identity()),
Err(EpochTransitionError::GraphEpochTransition)
));
assert_eq!(
std::fs::read(fixture.directory.path().join("manifest.ze")).expect("manifest"),
before
);
drop(store);
let reopened = Store::open(fixture.directory.path(), options).expect("reopen A");
assert_eq!(
reopened.snapshot().expect("snapshot").generation(),
generation
);
}
#[cfg(feature = "graph-cypher")]
#[test]
fn graph_epoch_drop_refuses_without_recording_a_generation() {
let fixture = publish_two_epoch_fixture();
let options = OpenOptions::default().with_epoch(fixture.epoch_a.clone());
let store = Store::open(fixture.directory.path(), options.clone()).expect("open");
let generation = store.enable_graph().expect("graph");
let before = std::fs::read(fixture.directory.path().join("manifest.ze")).expect("manifest");
assert!(matches!(
store.drop_epoch(fixture.epoch_b.identity().embedding),
Err(EpochTransitionError::GraphEpochTransition)
));
assert_eq!(
std::fs::read(fixture.directory.path().join("manifest.ze")).expect("manifest"),
before
);
drop(store);
let reopened = Store::open(fixture.directory.path(), options).expect("reopen A");
assert_eq!(
reopened.snapshot().expect("snapshot").generation(),
generation
);
}
#[test]
fn graph_free_epoch_snapshot_and_refusal_bytes_match_main() {
use zeppelin_embed::ingest::{DocumentVersion, IngestBatch, IngestDocument, IngestError};
use zeppelin_embed::lifecycle::StoreError;
let fixture = publish_two_epoch_fixture();
let store = Store::open(
fixture.directory.path(),
OpenOptions::default().with_epoch(fixture.epoch_a.clone()),
)
.expect("open");
store
.switch_epoch_alias(fixture.epoch_b.identity())
.expect("switch B");
store
.drop_epoch(fixture.epoch_a.identity().embedding)
.expect("drop A");
let large = IngestDocument::new(
DocumentVersion::new(DocId::new(99), Revision::new(1)),
vec![1.0; DIMS],
)
.with_text("x".repeat(2 * 1024 * 1024));
let error = store
.ingest(IngestBatch::new(vec![large]).with_epoch(fixture.epoch_b.identity()))
.expect_err("group cap");
assert!(matches!(
error,
IngestError::Store(StoreError::WalWrite(
zeppelin_embed::wal::WalWriteError::GroupTooLarge { .. }
))
));
store
.ingest(
IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(99), Revision::new(1)),
vec![1.0; DIMS],
)
.with_text("small"),
])
.with_epoch(fixture.epoch_b.identity()),
)
.expect("continued service");
store
.delete(zeppelin_embed::ingest::DeleteBatch::new(vec![DocId::new(
99,
)]))
.expect("delete");
let copies = tempdir().expect("copy root");
let copy = copies.path().join("snapshot");
store.write_snapshot(©).expect("snapshot");
if let Ok(destination) = std::env::var("ZE_GRAPH_FREE_PROBE_DIR") {
let root = Path::new(&destination);
for (name, source) in [
("source", fixture.directory.path()),
("snapshot", copy.as_path()),
] {
let target = root.join(name);
std::fs::create_dir_all(&target).expect("probe directory");
for entry in std::fs::read_dir(source).expect("probe files") {
let entry = entry.expect("entry");
std::fs::copy(entry.path(), target.join(entry.file_name())).expect("probe bytes");
}
}
}
drop(store);
let reopened = Store::open(©, OpenOptions::default().with_epoch(fixture.epoch_b))
.expect("reopen copy");
assert_eq!(
reopened.count_documents(None, None).expect("count").count,
3
);
}