use std::sync::Arc;
use tracing::{debug, error, info, warn};
use ailake_catalog::{
make_data_file_entry, make_data_file_entry_indexing, CatalogProvider, DataFileEntry,
NewSnapshot, SnapshotOperation, TableIdent, VectorIndexInfo,
};
use ailake_core::{AilakeResult, RowId, VectorStoragePolicy};
use ailake_file::{AilakeFileReader, AilakeFileWriter};
use ailake_store::Store;
use ailake_vec::compute_centroid_and_radius;
use arrow_array::RecordBatch;
use arrow_schema::SchemaRef;
use bytes::Bytes;
use futures::future::try_join_all;
use crate::writer::build_and_patch_index;
#[derive(Debug, Clone, Default)]
pub enum CompactionIndexStrategy {
#[default]
Auto,
ForceHnsw,
ForceIvfPq,
}
#[derive(Debug, Clone)]
pub struct CompactionConfig {
pub min_files_to_compact: usize,
pub target_file_size_bytes: u64,
pub index_strategy: CompactionIndexStrategy,
pub max_files_per_pass: usize,
}
impl Default for CompactionConfig {
fn default() -> Self {
Self {
min_files_to_compact: 4,
target_file_size_bytes: 128 * 1024 * 1024, index_strategy: CompactionIndexStrategy::Auto,
max_files_per_pass: 20,
}
}
}
#[derive(Debug, Clone, Copy)]
pub enum CompactionMode {
Full, Partial, }
pub struct CompactionPlanner {
config: CompactionConfig,
}
impl CompactionPlanner {
pub fn new(config: CompactionConfig) -> Self {
Self { config }
}
pub fn plan(&self, files: &[DataFileEntry]) -> Vec<DataFileEntry> {
let (mut foreign, natives): (Vec<DataFileEntry>, Vec<DataFileEntry>) =
files.iter().cloned().partition(DataFileEntry::is_foreign);
if !foreign.is_empty() {
warn!(
"ailake: compaction plan — {} file(s) with no AI-Lake index detected \
(likely a foreign/external rewrite) — prioritizing for reindex: {:?}",
foreign.len(),
foreign.iter().map(|f| f.path.as_str()).collect::<Vec<_>>()
);
}
let mut size_candidates: Vec<DataFileEntry> = natives
.into_iter()
.filter(|f| f.file_size_bytes < self.config.target_file_size_bytes)
.collect();
if size_candidates.len() < self.config.min_files_to_compact {
if foreign.is_empty() {
debug!(
"ailake: compaction skipped — {} eligible files < min_files_to_compact={}",
size_candidates.len(),
self.config.min_files_to_compact
);
return vec![];
}
size_candidates.clear();
}
foreign.sort_unstable_by_key(|f| f.file_size_bytes);
size_candidates.sort_unstable_by_key(|f| f.file_size_bytes);
foreign.extend(size_candidates);
foreign.truncate(self.config.max_files_per_pass);
let candidates = foreign;
let total_bytes: u64 = candidates.iter().map(|f| f.file_size_bytes).sum();
info!(
"ailake: compaction plan — {} files ({} bytes) → 1 merged file",
candidates.len(),
total_bytes
);
candidates
}
}
#[derive(Clone)]
pub struct CompactionExecutor {
store: Arc<dyn Store>,
policy: VectorStoragePolicy,
index_strategy: CompactionIndexStrategy,
fts_config: Option<ailake_fts::FtsConfig>,
}
impl CompactionExecutor {
pub fn new(store: Arc<dyn Store>, policy: VectorStoragePolicy) -> Self {
Self {
store,
policy,
index_strategy: CompactionIndexStrategy::Auto,
fts_config: None,
}
}
pub fn with_index_strategy(mut self, strategy: CompactionIndexStrategy) -> Self {
self.index_strategy = strategy;
self
}
pub fn with_fts_config(mut self, cfg: ailake_fts::FtsConfig) -> Self {
self.fts_config = Some(cfg);
self
}
fn with_effective_fts(
&self,
table_props: &std::collections::HashMap<String, String>,
) -> std::borrow::Cow<'_, Self> {
if self.fts_config.is_some() {
return std::borrow::Cow::Borrowed(self);
}
match ailake_fts::FtsConfig::from_table_props(table_props) {
Some(cfg) => {
let mut cloned = self.clone();
cloned.fts_config = Some(cfg);
std::borrow::Cow::Owned(cloned)
}
None => std::borrow::Cow::Borrowed(self),
}
}
async fn read_files_parallel(
&self,
files: &[DataFileEntry],
) -> AilakeResult<Vec<(RecordBatch, Vec<Vec<f32>>)>> {
let futs = files.iter().map(|entry| {
let store = self.store.clone();
let path = entry.path.clone();
let column = self.policy.column_name.clone();
let dim = self.policy.dim;
let dv = entry.deletion_vector.clone();
async move {
let bytes: Bytes = store.get(&path).await?;
let reader = AilakeFileReader::new(bytes, &column, dim);
if !reader.is_ailake_file() {
debug!(
"ailake: compaction reading {} without an AI-Lake index \
(external write or indexing in progress) — data still merged",
path
);
}
let (batch, embeddings) = reader.read_parquet()?;
let pair = if let Some(dv) = dv {
let bitmap = crate::dv::load_deletion_vector(&store, &dv).await?;
crate::dv::filter_deleted_rows(batch, embeddings, &bitmap)?
} else {
(batch, embeddings)
};
Ok::<(RecordBatch, Vec<Vec<f32>>), ailake_core::AilakeError>(pair)
}
});
try_join_all(futs).await
}
pub async fn compact(
&self,
files: &[DataFileEntry],
output_path: &str,
) -> AilakeResult<DataFileEntry> {
if files.is_empty() {
return Err(ailake_core::AilakeError::Catalog(
"compact: no files provided".into(),
));
}
let pairs = self.read_files_parallel(files).await?;
let schema: SchemaRef = pairs[0].0.schema();
let (all_batches, all_embeddings): (Vec<_>, Vec<_>) = pairs.into_iter().unzip();
let all_embeddings: Vec<Vec<f32>> = all_embeddings.into_iter().flatten().collect();
let merged_batch = concat_batches(schema, &all_batches)?;
let record_count = merged_batch.num_rows() as u64;
let writer = {
let base = AilakeFileWriter::new(self.policy.clone());
let base = match &self.index_strategy {
CompactionIndexStrategy::Auto => base.with_auto_index(),
CompactionIndexStrategy::ForceHnsw => base,
CompactionIndexStrategy::ForceIvfPq => {
let cfg = ailake_index::IvfPqConfig::for_dataset(
self.policy.dim as usize,
all_embeddings.len(),
);
base.with_ivf_pq(cfg)
}
};
if let Some(ref fts_cfg) = self.fts_config {
match ailake_fts::merge_fts_blobs(fts_cfg, &merged_batch) {
Ok(blob) => base.with_prebuilt_fts_blob(blob),
Err(e) => {
warn!("ailake: FTS re-index during compaction failed: {e}");
base
}
}
} else {
base
}
};
let file_bytes = writer.write(&merged_batch, &all_embeddings)?;
let file_size = file_bytes.len() as u64;
let column_stats =
ailake_catalog::extract_column_stats(&file_bytes, &[self.policy.column_name.as_str()])
.and_then(|m| serde_json::to_string(&m).ok());
self.store.put(output_path, file_bytes.clone()).await?;
let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
let header = reader.read_header()?;
let ailk_start = reader.ailk_offset()?;
reader.verify_integrity()?;
let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
let mut entry = make_data_file_entry(
output_path,
record_count,
file_size,
¢roid,
VectorIndexInfo {
column: &self.policy.column_name,
dim: self.policy.dim,
hnsw_offset: ailk_start + header.hnsw_offset,
hnsw_len: header.hnsw_len,
},
);
entry.first_row_id = source_first_row_id;
entry.batch_id = DataFileEntry::merge_batch_ids(files);
entry.column_stats = column_stats;
Ok(entry)
}
pub async fn compact_incremental(
&self,
files: &[DataFileEntry],
output_path: &str,
) -> AilakeResult<DataFileEntry> {
const DOMINANT_RATIO: f64 = 0.40;
if files.is_empty() {
return Err(ailake_core::AilakeError::Catalog(
"compact_incremental: no files provided".into(),
));
}
if matches!(self.index_strategy, CompactionIndexStrategy::ForceIvfPq) {
debug!(
"ailake: compact_incremental — index_strategy=ForceIvfPq, which this method \
can never produce (HNSW graph extension only); falling back to full rebuild"
);
return self.compact(files, output_path).await;
}
let total_rows: u64 = files.iter().map(|f| f.record_count).sum();
let dom_idx = files
.iter()
.enumerate()
.max_by_key(|(_, f)| f.record_count)
.map(|(i, _)| i)
.unwrap_or(0);
let dom_rows = files[dom_idx].record_count;
if files[dom_idx].deletion_vector.is_some() {
debug!(
"ailake: compact_incremental — dominant file {} has DV-masked rows, \
falling back to full rebuild (graph reuse would desync row positions)",
files[dom_idx].path
);
return self.compact(files, output_path).await;
}
if (dom_rows as f64 / total_rows as f64) < DOMINANT_RATIO {
debug!(
"ailake: compact_incremental — no dominant file ({}/{} rows < {:.0}% threshold), \
falling back to full rebuild",
dom_rows,
total_rows,
DOMINANT_RATIO * 100.0
);
return self.compact(files, output_path).await;
}
let column = self.policy.column_name.clone();
let dim = self.policy.dim;
let dom_path = files[dom_idx].path.clone();
let futs: Vec<_> =
files
.iter()
.map(|entry| {
let store = self.store.clone();
let path = entry.path.clone();
let col = column.clone();
let is_dom = path == dom_path;
let dv = entry.deletion_vector.clone();
async move {
let bytes: Bytes = store.get(&path).await?;
let reader = AilakeFileReader::new(bytes.clone(), &col, dim);
let has_index = reader.is_ailake_file();
if is_dom && !has_index {
debug!(
"ailake: compact_incremental — dominant candidate {} has no \
AI-Lake index; will fall back to full rebuild if no HNSW to reuse",
path
);
}
let (raw_batch, raw_vecs) = reader.read_parquet()?;
let (batch, vecs) = if let Some(dv) = dv {
let bitmap = crate::dv::load_deletion_vector(&store, &dv).await?;
crate::dv::filter_deleted_rows(raw_batch, raw_vecs, &bitmap)?
} else {
(raw_batch, raw_vecs)
};
let retained = if is_dom && has_index {
Some(bytes)
} else {
None
};
Ok::<
(RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>),
ailake_core::AilakeError,
>((batch, vecs, is_dom, retained))
}
})
.collect();
#[allow(clippy::type_complexity)]
let raw: Vec<(RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>)> =
try_join_all(futs).await?;
let mut dom_batch: Option<RecordBatch> = None;
let mut dom_vecs: Vec<Vec<f32>> = Vec::new();
let mut dom_bytes_found: Option<Bytes> = None;
let mut other_batches: Vec<RecordBatch> = Vec::new();
let mut other_vecs: Vec<Vec<f32>> = Vec::new();
for (batch, vecs, is_dom, retained) in raw {
if is_dom {
dom_batch = Some(batch);
dom_vecs = vecs;
dom_bytes_found = retained;
} else {
other_batches.push(batch);
other_vecs.extend(vecs);
}
}
let (dom_batch, dom_bytes) = match (dom_batch, dom_bytes_found) {
(Some(b), Some(byt)) => (b, byt),
_ => {
debug!(
"ailake: compact_incremental — dominant file missing from read results, \
falling back to full rebuild"
);
return self.compact(files, output_path).await;
}
};
let dom_reader = AilakeFileReader::new(dom_bytes, &column, dim);
let mut hnsw = match dom_reader.load_index() {
Ok(idx) => idx,
Err(e) => {
debug!(
"ailake: compact_incremental — cannot load dominant HNSW ({}), \
falling back to full rebuild",
e
);
return self.compact(files, output_path).await;
}
};
let dom_count = dom_batch.num_rows() as u64;
for (j, vec) in other_vecs.iter().enumerate() {
hnsw.insert_node(RowId::new(dom_count + j as u64), vec.clone());
}
hnsw.quantize_to_f16();
let schema: SchemaRef = dom_batch.schema();
let mut all_batches = vec![dom_batch];
all_batches.extend(other_batches);
let merged_batch = concat_batches(schema, &all_batches)?;
let record_count = merged_batch.num_rows() as u64;
let mut all_embeddings = dom_vecs;
all_embeddings.extend(other_vecs);
let writer = {
let base = AilakeFileWriter::new(self.policy.clone());
if let Some(ref fts_cfg) = self.fts_config {
match ailake_fts::merge_fts_blobs(fts_cfg, &merged_batch) {
Ok(blob) => base.with_prebuilt_fts_blob(blob),
Err(e) => {
warn!("ailake: FTS re-index during incremental compaction failed: {e}");
base
}
}
} else {
base
}
};
let file_bytes = writer.write_with_prebuilt_hnsw(&merged_batch, &all_embeddings, &hnsw)?;
let file_size = file_bytes.len() as u64;
let column_stats =
ailake_catalog::extract_column_stats(&file_bytes, &[self.policy.column_name.as_str()])
.and_then(|m| serde_json::to_string(&m).ok());
self.store.put(output_path, file_bytes.clone()).await?;
let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
let header = reader.read_header()?;
let ailk_start = reader.ailk_offset()?;
reader.verify_integrity()?;
let source_first_row_id = files[dom_idx].first_row_id;
let mut entry = make_data_file_entry(
output_path,
record_count,
file_size,
¢roid,
VectorIndexInfo {
column: &self.policy.column_name,
dim: self.policy.dim,
hnsw_offset: ailk_start + header.hnsw_offset,
hnsw_len: header.hnsw_len,
},
);
entry.first_row_id = source_first_row_id;
entry.batch_id = DataFileEntry::merge_batch_ids(files);
entry.column_stats = column_stats;
info!(
"ailake: compact_incremental — merged {} files into {} \
({} rows from dominant + {} inserted incrementally)",
files.len(),
output_path,
dom_count,
record_count - dom_count
);
Ok(entry)
}
pub async fn compact_deferred(
&self,
files: &[DataFileEntry],
output_path: &str,
catalog: Arc<dyn CatalogProvider>,
table: &TableIdent,
) -> AilakeResult<DataFileEntry> {
if files.is_empty() {
return Err(ailake_core::AilakeError::Catalog(
"compact_deferred: no files provided".into(),
));
}
let pairs = self.read_files_parallel(files).await?;
let schema: SchemaRef = pairs[0].0.schema();
let (all_batches, all_embeddings): (Vec<_>, Vec<_>) = pairs.into_iter().unzip();
let all_embeddings: Vec<Vec<f32>> = all_embeddings.into_iter().flatten().collect();
let merged_batch = concat_batches(schema, &all_batches)?;
let record_count = merged_batch.num_rows() as u64;
let file_writer = AilakeFileWriter::new(self.policy.clone());
let parquet_bytes = file_writer.write_parquet_only(&merged_batch, &all_embeddings)?;
let file_size = parquet_bytes.len() as u64;
let column_stats = ailake_catalog::extract_column_stats(
&parquet_bytes,
&[self.policy.column_name.as_str()],
)
.and_then(|m| serde_json::to_string(&m).ok());
self.store.put(output_path, parquet_bytes).await?;
let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
let mut entry = make_data_file_entry_indexing(
output_path,
record_count,
file_size,
¢roid,
&self.policy.column_name,
self.policy.dim,
);
entry.first_row_id = source_first_row_id;
entry.batch_id = DataFileEntry::merge_batch_ids(files);
entry.column_stats = column_stats;
let store = self.store.clone();
let policy = self.policy.clone();
let table_id = table.clone();
let fp = output_path.to_string();
tokio::spawn(async move {
if let Err(e) = build_and_patch_index(store, catalog, policy, table_id, fp).await {
error!(
"ailake: compaction deferred HNSW build failed — file indexed as \
Parquet-only until next compaction rebuilds the index: {}",
e
);
}
});
Ok(entry)
}
pub async fn run(
&self,
planner: &CompactionPlanner,
table: &TableIdent,
catalog: Arc<dyn CatalogProvider>,
output_prefix: &str,
) -> AilakeResult<Option<DataFileEntry>> {
let all_files = catalog.list_files(table, None).await?;
let to_compact = planner.plan(&all_files);
if to_compact.is_empty() {
return Ok(None);
}
let meta_props = catalog
.load_table(table)
.await
.map(|m| m.properties)
.unwrap_or_default();
let executor = self.with_effective_fts(&meta_props);
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_else(|e| e.duration())
.as_millis();
let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
let merged = executor
.compact_incremental(&to_compact, &output_path)
.await?;
let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
let parent_snapshot_id = catalog
.load_table(table)
.await
.ok()
.and_then(|m| m.current_snapshot_id);
let snapshot = NewSnapshot {
snapshot_id: ailake_catalog::new_snapshot_id(),
parent_snapshot_id,
files,
operation: SnapshotOperation::Replace,
iceberg_schema: None,
extra_properties: std::collections::HashMap::new(),
bloom_filters: vec![],
equality_delete_files: vec![],
};
catalog.commit_snapshot(table, snapshot).await?;
info!(
"ailake: compaction committed — merged {} files into {}",
to_compact.len(),
output_path
);
if catalog.retires_files_physically() {
delete_old_files(&self.store, &to_compact).await;
} else {
info!(
"ailake: compaction — leaving {} retired file(s) in place; catalog backend \
manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
to_compact.len()
);
}
Ok(Some(merged))
}
pub async fn run_deferred(
&self,
planner: &CompactionPlanner,
table: &TableIdent,
catalog: Arc<dyn CatalogProvider>,
output_prefix: &str,
) -> AilakeResult<Option<DataFileEntry>> {
if !catalog.supports_in_place_rewrite() {
return Err(ailake_core::AilakeError::Catalog(
"deferred compaction is not supported with this catalog backend: the \
background index build patches the merged file in place at its committed \
path, which this catalog cannot re-register — run a blocking compact"
.into(),
));
}
let all_files = catalog.list_files(table, None).await?;
let to_compact = planner.plan(&all_files);
if to_compact.is_empty() {
return Ok(None);
}
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_else(|e| e.duration())
.as_millis();
let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
let merged = self
.compact_deferred(&to_compact, &output_path, catalog.clone(), table)
.await?;
let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
let parent_snapshot_id = catalog
.load_table(table)
.await
.ok()
.and_then(|m| m.current_snapshot_id);
let snapshot = NewSnapshot {
snapshot_id: ailake_catalog::new_snapshot_id(),
parent_snapshot_id,
files,
operation: SnapshotOperation::Replace,
iceberg_schema: None,
extra_properties: std::collections::HashMap::new(),
bloom_filters: vec![],
equality_delete_files: vec![],
};
catalog.commit_snapshot(table, snapshot).await?;
info!(
"ailake: compaction committed (deferred) — merged {} files into {} \
(index building in background)",
to_compact.len(),
output_path
);
if catalog.retires_files_physically() {
delete_old_files(&self.store, &to_compact).await;
} else {
info!(
"ailake: compaction — leaving {} retired file(s) in place; catalog backend \
manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
to_compact.len()
);
}
Ok(Some(merged))
}
}
async fn build_replace_file_list(
catalog: &Arc<dyn CatalogProvider>,
table: &TableIdent,
to_compact: &[DataFileEntry],
merged: DataFileEntry,
) -> AilakeResult<Vec<DataFileEntry>> {
let current_files = catalog.list_files(table, None).await?;
let compacted_paths: std::collections::HashSet<&str> =
to_compact.iter().map(|f| f.path.as_str()).collect();
let mut files: Vec<DataFileEntry> = current_files
.into_iter()
.filter(|f| !compacted_paths.contains(f.path.as_str()))
.collect();
files.push(merged);
Ok(files)
}
async fn delete_old_files(store: &Arc<dyn Store>, files: &[DataFileEntry]) {
for entry in files {
if let Err(e) = store.delete(&entry.path).await {
error!(
"ailake: compaction cleanup failed — could not delete {}: {} \
(orphan file in object store after successful catalog commit; \
delete manually to reclaim storage)",
entry.path, e
);
}
}
}
fn concat_batches(schema: SchemaRef, batches: &[RecordBatch]) -> AilakeResult<RecordBatch> {
arrow_select::concat::concat_batches(&schema, batches)
.map_err(|e| ailake_core::AilakeError::Arrow(e.to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
use ailake_catalog::IndexStatus;
#[test]
fn plan_returns_empty_if_too_few_files() {
let planner = CompactionPlanner::new(CompactionConfig {
min_files_to_compact: 4,
target_file_size_bytes: 1024 * 1024,
..Default::default()
});
let files: Vec<DataFileEntry> = (0..3)
.map(|i| DataFileEntry {
path: format!("file-{i}.parquet"),
record_count: 10,
file_size_bytes: 100,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
})
.collect();
assert!(planner.plan(&files).is_empty());
}
#[test]
fn plan_selects_small_files() {
let planner = CompactionPlanner::new(CompactionConfig {
min_files_to_compact: 2,
target_file_size_bytes: 1000,
..Default::default()
});
let files = vec![
DataFileEntry {
path: "small.parquet".into(),
record_count: 5,
file_size_bytes: 500,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "large.parquet".into(),
record_count: 5000,
file_size_bytes: 200_000_000,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "also-small.parquet".into(),
record_count: 5,
file_size_bytes: 800,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
];
let selected = planner.plan(&files);
assert_eq!(selected.len(), 2);
assert!(selected.iter().any(|f| f.path == "small.parquet"));
assert!(selected.iter().any(|f| f.path == "also-small.parquet"));
}
#[test]
fn plan_respects_max_files_per_pass() {
let planner = CompactionPlanner::new(CompactionConfig {
min_files_to_compact: 2,
target_file_size_bytes: 1_000_000,
max_files_per_pass: 3,
..Default::default()
});
let files: Vec<DataFileEntry> = (0..5)
.map(|i| DataFileEntry {
path: format!("f{i}.parquet"),
record_count: 10,
file_size_bytes: 100 + i as u64 * 100,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
})
.collect();
let selected = planner.plan(&files);
assert_eq!(selected.len(), 3);
assert_eq!(selected[0].file_size_bytes, 100);
assert_eq!(selected[1].file_size_bytes, 200);
assert_eq!(selected[2].file_size_bytes, 300);
}
#[test]
fn plan_sorts_smallest_first() {
let planner = CompactionPlanner::new(CompactionConfig {
min_files_to_compact: 2,
target_file_size_bytes: 10_000,
max_files_per_pass: 4,
..Default::default()
});
let files = vec![
DataFileEntry {
path: "c.parquet".into(),
record_count: 1,
file_size_bytes: 300,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "a.parquet".into(),
record_count: 1,
file_size_bytes: 100,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "b.parquet".into(),
record_count: 1,
file_size_bytes: 200,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
];
let selected = planner.plan(&files);
assert_eq!(selected[0].file_size_bytes, 100);
assert_eq!(selected[1].file_size_bytes, 200);
assert_eq!(selected[2].file_size_bytes, 300);
}
fn make_plan_entry(path: &str, size: u64, centroid_b64: Option<String>) -> DataFileEntry {
DataFileEntry {
path: path.to_string(),
record_count: 10,
file_size_bytes: size,
centroid_b64,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
}
}
#[test]
fn plan_prioritizes_foreign_written_files_regardless_of_size() {
let planner = CompactionPlanner::new(CompactionConfig {
min_files_to_compact: 4, target_file_size_bytes: 1000,
max_files_per_pass: 20,
..Default::default()
});
let files = vec![
make_plan_entry("native_small.parquet", 500, Some("AAAA".into())),
make_plan_entry("foreign_big.parquet", 200_000_000, None),
];
let selected = planner.plan(&files);
assert_eq!(
selected.len(),
1,
"foreign file alone should trigger a pass even below min_files_to_compact"
);
assert_eq!(selected[0].path, "foreign_big.parquet");
}
#[test]
fn plan_merges_foreign_and_size_candidates_together() {
let planner = CompactionPlanner::new(CompactionConfig {
min_files_to_compact: 2,
target_file_size_bytes: 1000,
max_files_per_pass: 20,
..Default::default()
});
let files = vec![
make_plan_entry("native_a.parquet", 300, Some("AAAA".into())),
make_plan_entry("native_b.parquet", 400, Some("AAAA".into())),
make_plan_entry("foreign.parquet", 200_000_000, None),
];
let selected = planner.plan(&files);
let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
assert_eq!(selected.len(), 3, "paths={paths:?}");
assert_eq!(selected[0].path, "foreign.parquet", "paths={paths:?}");
}
#[test]
fn plan_sorts_foreign_files_smallest_first_too() {
let planner = CompactionPlanner::new(CompactionConfig {
min_files_to_compact: 1,
target_file_size_bytes: 1000,
max_files_per_pass: 2,
..Default::default()
});
let files = vec![
make_plan_entry("foreign_huge.parquet", 500_000_000, None),
make_plan_entry("foreign_small.parquet", 100, None),
make_plan_entry("foreign_medium.parquet", 10_000_000, None),
];
let selected = planner.plan(&files);
let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
assert_eq!(selected.len(), 2, "max_files_per_pass=2, paths={paths:?}");
assert_eq!(
paths,
vec!["foreign_small.parquet", "foreign_medium.parquet"],
"foreign files must be size-sorted before truncation, cheapest first"
);
}
#[tokio::test]
async fn compact_merges_two_files() {
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
let batch_a = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
)
.unwrap();
let batch_b = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
)
.unwrap();
let writer_a = AilakeFileWriter::new(policy.clone());
let bytes_a = writer_a.write(&batch_a, &embs_a).unwrap();
let writer_b = AilakeFileWriter::new(policy.clone());
let bytes_b = writer_b.write(&batch_b, &embs_b).unwrap();
store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
let entries = vec![
DataFileEntry {
path: "data/a.parquet".into(),
record_count: 2,
file_size_bytes: bytes_a.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "data/b.parquet".into(),
record_count: 2,
file_size_bytes: bytes_b.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
];
let executor = CompactionExecutor::new(store.clone(), policy.clone());
let merged = executor
.compact(&entries, "data/merged.parquet")
.await
.unwrap();
assert_eq!(merged.record_count, 4);
assert_eq!(merged.path, "data/merged.parquet");
let merged_bytes = store.get("data/merged.parquet").await.unwrap();
let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
reader.verify_integrity().unwrap();
let (batch, embs) = reader.read_parquet().unwrap();
assert_eq!(batch.num_rows(), 4);
assert_eq!(embs.len(), 4);
}
#[tokio::test]
async fn compact_drops_deletion_vector_masked_rows() {
use ailake_catalog::provider::DeletionVector;
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use roaring::RoaringBitmap;
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
let batch_a = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
)
.unwrap();
let batch_b =
RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![2i32]))])
.unwrap();
let bytes_a = AilakeFileWriter::new(policy.clone())
.write(&batch_a, &embs_a)
.unwrap();
let bytes_b = AilakeFileWriter::new(policy.clone())
.write(&batch_b, &embs_b)
.unwrap();
store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
let mut bitmap = RoaringBitmap::new();
bitmap.insert(0);
let (puffin_bytes, offset, length) =
crate::delete::PuffinWriter::write_single_dv(&bitmap, 1).unwrap();
store.put("metadata/dv-1.dvd", puffin_bytes).await.unwrap();
let entry_a = DataFileEntry {
path: "data/a.parquet".into(),
record_count: 2,
file_size_bytes: bytes_a.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: Some(DeletionVector {
path: "metadata/dv-1.dvd".into(),
offset,
length,
cardinality: 1,
}),
first_row_id: None,
column_stats: None,
sequence_number: 0,
};
let entry_b = DataFileEntry {
path: "data/b.parquet".into(),
record_count: 1,
file_size_bytes: bytes_b.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
};
let executor = CompactionExecutor::new(store.clone(), policy.clone());
let merged = executor
.compact(&[entry_a, entry_b], "data/merged.parquet")
.await
.unwrap();
assert_eq!(merged.record_count, 2);
let merged_bytes = store.get("data/merged.parquet").await.unwrap();
let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
let (batch, embs) = reader.read_parquet().unwrap();
assert_eq!(batch.num_rows(), 2);
assert_eq!(embs.len(), 2);
let ids: Vec<i32> = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap()
.values()
.to_vec();
assert_eq!(
ids,
vec![1, 2],
"id=0 (deleted) must not survive compaction"
);
}
#[tokio::test]
async fn compact_with_ivf_pq_strategy_does_not_crash_on_verify_integrity() {
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let dim = 8;
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let n_per_file = 30usize;
let make_file = |path: &str, offset: i32| {
let ids: Vec<i32> = (offset..offset + n_per_file as i32).collect();
let embs: Vec<Vec<f32>> = ids
.iter()
.map(|&i| {
(0..dim as i32)
.map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
.collect()
})
.collect();
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
.unwrap();
let bytes = AilakeFileWriter::new(policy.clone())
.write(&batch, &embs)
.unwrap();
(path.to_string(), bytes)
};
let (path_a, bytes_a) = make_file("data/ivfpq_a.parquet", 0);
let (path_b, bytes_b) = make_file("data/ivfpq_b.parquet", n_per_file as i32);
for (path, bytes) in [(&path_a, &bytes_a), (&path_b, &bytes_b)] {
store.put(path, bytes.clone()).await.unwrap();
}
let make_entry = |path: &str, size: u64| DataFileEntry {
path: path.to_string(),
record_count: n_per_file as u64,
file_size_bytes: size,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
};
let entries = vec![
make_entry(&path_a, bytes_a.len() as u64),
make_entry(&path_b, bytes_b.len() as u64),
];
let executor = CompactionExecutor::new(store.clone(), policy.clone())
.with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
let merged = executor
.compact(&entries, "data/ivfpq_merged.parquet")
.await
.expect("compact() with ForceIvfPq must not fail in verify_integrity()");
assert_eq!(merged.record_count, 2 * n_per_file as u64);
let merged_bytes = store.get("data/ivfpq_merged.parquet").await.unwrap();
let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
reader
.verify_integrity()
.expect("verify_integrity() must handle an IVF-PQ-indexed file");
}
#[tokio::test]
async fn compact_incremental_merges_dominant_plus_small() {
use ailake_core::{RowId, VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let embs_dom: Vec<Vec<f32>> = vec![
vec![1.0, 0.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0, 0.0],
vec![0.0, 0.0, 1.0, 0.0],
vec![0.7, 0.7, 0.0, 0.0],
vec![0.0, 0.7, 0.7, 0.0],
vec![0.0, 0.0, 0.7, 0.7],
];
let batch_dom = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![0i32, 1, 2, 3, 4, 5]))],
)
.unwrap();
let embs_small: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 0.0, 1.0], vec![0.5, 0.5, 0.5, 0.5]];
let batch_small = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![6i32, 7]))],
)
.unwrap();
let bytes_dom = AilakeFileWriter::new(policy.clone())
.write(&batch_dom, &embs_dom)
.unwrap();
let bytes_small = AilakeFileWriter::new(policy.clone())
.write(&batch_small, &embs_small)
.unwrap();
store
.put("data/dominant.parquet", bytes_dom.clone())
.await
.unwrap();
store
.put("data/small.parquet", bytes_small.clone())
.await
.unwrap();
let entries = vec![
DataFileEntry {
path: "data/dominant.parquet".into(),
record_count: 6,
file_size_bytes: bytes_dom.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "data/small.parquet".into(),
record_count: 2,
file_size_bytes: bytes_small.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
];
let executor = CompactionExecutor::new(store.clone(), policy.clone());
let merged = executor
.compact_incremental(&entries, "data/merged.parquet")
.await
.unwrap();
assert_eq!(merged.record_count, 8);
assert_eq!(merged.path, "data/merged.parquet");
let merged_bytes = store.get("data/merged.parquet").await.unwrap();
let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
reader.verify_integrity().unwrap();
let (batch, embs) = reader.read_parquet().unwrap();
assert_eq!(batch.num_rows(), 8);
assert_eq!(embs.len(), 8);
for f in &embs[..6] {
assert_eq!(f.len(), 4);
}
let hnsw = reader.load_index().unwrap();
assert_eq!(hnsw.node_count(), 8);
let results = hnsw.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
assert_eq!(results[0].0, RowId::new(0));
let results = hnsw.search(&[0.0, 0.0, 0.0, 1.0], 1, 50);
assert_eq!(results[0].0, RowId::new(6));
}
#[tokio::test]
async fn compact_incremental_respects_force_ivf_pq_even_with_dominant_file() {
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let dim = 8;
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let make_file = |path: &str, offset: i32, n: usize| {
let ids: Vec<i32> = (offset..offset + n as i32).collect();
let embs: Vec<Vec<f32>> = ids
.iter()
.map(|&i| {
(0..dim as i32)
.map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
.collect()
})
.collect();
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
.unwrap();
let bytes = AilakeFileWriter::new(policy.clone())
.write(&batch, &embs)
.unwrap();
(path.to_string(), bytes, n as u64)
};
let (path_dom, bytes_dom, n_dom) = make_file("data/dom.parquet", 0, 90);
let (path_small, bytes_small, n_small) = make_file("data/small.parquet", 999, 30);
for (path, bytes) in [(&path_dom, &bytes_dom), (&path_small, &bytes_small)] {
store.put(path, bytes.clone()).await.unwrap();
}
let make_entry = |path: &str, record_count: u64, size: u64| DataFileEntry {
path: path.to_string(),
record_count,
file_size_bytes: size,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
};
let entries = vec![
make_entry(&path_dom, n_dom, bytes_dom.len() as u64),
make_entry(&path_small, n_small, bytes_small.len() as u64),
];
let executor = CompactionExecutor::new(store.clone(), policy.clone())
.with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
let merged = executor
.compact_incremental(&entries, "data/merged_force_ivfpq.parquet")
.await
.expect("compact_incremental() with ForceIvfPq must fall back to compact() cleanly");
let merged_bytes = store.get("data/merged_force_ivfpq.parquet").await.unwrap();
let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
reader.verify_integrity().unwrap();
match reader.load_any_index().unwrap() {
ailake_index::AnyIndex::IvfPq(_) => {}
ailake_index::AnyIndex::Hnsw(_) => panic!(
"ForceIvfPq was silently satisfied with HNSW instead — merged.record_count={}",
merged.record_count
),
}
}
#[tokio::test]
async fn compact_incremental_falls_back_when_no_dominant() {
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let make_batch = |ids: Vec<i32>, embs: Vec<Vec<f32>>| {
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
.unwrap();
AilakeFileWriter::new(policy.clone())
.write(&batch, &embs)
.unwrap()
};
let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
let bytes_a = make_batch(vec![0, 1], embs_a);
let bytes_b = make_batch(vec![2, 3], embs_b);
store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
let entries = vec![
DataFileEntry {
path: "data/a.parquet".into(),
record_count: 2,
file_size_bytes: bytes_a.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "data/b.parquet".into(),
record_count: 2,
file_size_bytes: bytes_b.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
];
let executor = CompactionExecutor::new(store.clone(), policy.clone());
let merged = executor
.compact_incremental(&entries, "data/merged.parquet")
.await
.unwrap();
assert_eq!(merged.record_count, 4);
let merged_bytes = store.get("data/merged.parquet").await.unwrap();
let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
reader.verify_integrity().unwrap();
}
#[tokio::test]
async fn compact_deferred_produces_parquet_only_file() {
use ailake_catalog::HadoopCatalog;
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let catalog_dir = TempDir::new().unwrap();
let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
let table = TableIdent {
namespace: "ns".into(),
name: "tbl".into(),
};
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
use ailake_catalog::TableProperties;
catalog
.create_table(
&table,
&TableProperties {
policy: policy.clone(),
extra: std::collections::HashMap::new(),
format_version: 2,
partition_column_type: None,
},
)
.await
.unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
let batch_a = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
)
.unwrap();
let bytes_a = AilakeFileWriter::new(policy.clone())
.write(&batch_a, &embs_a)
.unwrap();
store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
let batch_b = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
)
.unwrap();
let bytes_b = AilakeFileWriter::new(policy.clone())
.write(&batch_b, &embs_b)
.unwrap();
store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
let entries = vec![
DataFileEntry {
path: "data/a.parquet".into(),
record_count: 2,
file_size_bytes: bytes_a.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "data/b.parquet".into(),
record_count: 2,
file_size_bytes: bytes_b.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
];
let executor = CompactionExecutor::new(store.clone(), policy.clone());
let entry = executor
.compact_deferred(&entries, "data/merged.parquet", catalog.clone(), &table)
.await
.unwrap();
assert_eq!(entry.index_status, IndexStatus::Indexing);
assert_eq!(entry.record_count, 4);
let merged_bytes = store.get("data/merged.parquet").await.unwrap();
let pq_reader = ailake_parquet::ParquetVectorReader::new(merged_bytes, "embedding");
let count = pq_reader.record_count().unwrap();
assert_eq!(count, 4);
}
#[tokio::test]
async fn compact_aggregates_batch_ids_from_sources() {
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0]];
let batch_a =
RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![0i32]))])
.unwrap();
let bytes_a = AilakeFileWriter::new(policy.clone())
.write(&batch_a, &embs_a)
.unwrap();
store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 1.0, 0.0, 0.0]];
let batch_b =
RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1i32]))])
.unwrap();
let bytes_b = AilakeFileWriter::new(policy.clone())
.write(&batch_b, &embs_b)
.unwrap();
store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
let embs_c: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
let batch_c =
RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![2i32]))]).unwrap();
let bytes_c = AilakeFileWriter::new(policy.clone())
.write(&batch_c, &embs_c)
.unwrap();
store.put("data/c.parquet", bytes_c.clone()).await.unwrap();
fn entry(path: &str, size: u64, batch_id: Option<&str>) -> DataFileEntry {
DataFileEntry {
path: path.into(),
record_count: 1,
file_size_bytes: size,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: batch_id.map(String::from),
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
}
}
let entries = vec![
entry("data/a.parquet", bytes_a.len() as u64, Some("k-a")),
entry("data/b.parquet", bytes_b.len() as u64, Some("k-b")),
entry("data/c.parquet", bytes_c.len() as u64, None),
];
let executor = CompactionExecutor::new(store.clone(), policy);
let merged = executor
.compact(&entries, "data/merged.parquet")
.await
.unwrap();
assert_eq!(merged.record_count, 3);
assert_eq!(
merged.batch_ids(),
vec!["k-a".to_string(), "k-b".to_string()],
"merged file must carry every source's idempotency key, none invented"
);
}
#[tokio::test]
async fn run_preserves_untouched_files_outside_compaction_pass() {
use ailake_catalog::HadoopCatalog;
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let catalog_dir = TempDir::new().unwrap();
let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
let table = TableIdent {
namespace: "ns".into(),
name: "tbl".into(),
};
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
use ailake_catalog::TableProperties;
catalog
.create_table(
&table,
&TableProperties {
policy: policy.clone(),
extra: std::collections::HashMap::new(),
format_version: 2,
partition_column_type: None,
},
)
.await
.unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let write_file = |path: &str, ids: Vec<i32>, embs: Vec<Vec<f32>>| {
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
.unwrap();
let bytes = AilakeFileWriter::new(policy.clone())
.write(&batch, &embs)
.unwrap();
(path.to_string(), bytes)
};
let (path_a, bytes_a) = write_file(
"data/small_a.parquet",
vec![0, 1],
vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
);
let (path_b, bytes_b) = write_file(
"data/small_b.parquet",
vec![2, 3],
vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]],
);
let (path_big, bytes_big) = write_file(
"data/big.parquet",
vec![4, 5],
vec![vec![1.0, 1.0, 0.0, 0.0], vec![0.0, 1.0, 1.0, 0.0]],
);
for (path, bytes) in [
(&path_a, &bytes_a),
(&path_b, &bytes_b),
(&path_big, &bytes_big),
] {
store.put(path, bytes.clone()).await.unwrap();
}
let make_entry = |path: &str, size: u64| DataFileEntry {
path: path.to_string(),
record_count: 2,
file_size_bytes: size,
centroid_b64: Some("AAAA".into()),
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
};
let initial_snap_id = ailake_catalog::new_snapshot_id();
let initial_snapshot = NewSnapshot {
snapshot_id: initial_snap_id,
parent_snapshot_id: None,
files: vec![
make_entry(&path_a, 500),
make_entry(&path_b, 500),
make_entry(&path_big, 200_000_000), ],
operation: SnapshotOperation::Append,
iceberg_schema: None,
extra_properties: std::collections::HashMap::new(),
bloom_filters: vec![],
equality_delete_files: vec![],
};
catalog
.commit_snapshot(&table, initial_snapshot)
.await
.unwrap();
let planner = CompactionPlanner::new(CompactionConfig {
min_files_to_compact: 2,
target_file_size_bytes: 1000,
index_strategy: CompactionIndexStrategy::ForceHnsw,
max_files_per_pass: 20,
});
let executor = CompactionExecutor::new(store.clone(), policy.clone());
let merged = executor
.run(&planner, &table, catalog.clone(), "data")
.await
.unwrap()
.expect("compaction should have run — 2 eligible small files");
let files_after = catalog.list_files(&table, None).await.unwrap();
let paths_after: Vec<&str> = files_after.iter().map(|f| f.path.as_str()).collect();
assert!(
paths_after.contains(&path_big.as_str()),
"BUG: untouched 'big.parquet' vanished after run() — files_after={paths_after:?}"
);
assert!(
paths_after.contains(&merged.path.as_str()),
"merged output file must be present — files_after={paths_after:?}"
);
let meta_dir = catalog_dir.path().join("ns/tbl/metadata");
let latest_metadata = std::fs::read_dir(&meta_dir)
.unwrap()
.filter_map(|e| e.ok())
.filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
.max_by_key(|e| e.metadata().unwrap().modified().unwrap())
.expect("metadata.json must exist after commit");
let json: serde_json::Value =
serde_json::from_slice(&std::fs::read(latest_metadata.path()).unwrap()).unwrap();
let last_snapshot = json["snapshots"].as_array().unwrap().last().unwrap();
assert_eq!(
last_snapshot["parent-snapshot-id"].as_i64(),
Some(initial_snap_id),
"compaction's Replace snapshot must chain to the pre-compaction snapshot, not be orphaned"
);
assert!(
!paths_after.contains(&path_a.as_str()) && !paths_after.contains(&path_b.as_str()),
"compacted input files must no longer be listed — files_after={paths_after:?}"
);
assert_eq!(
files_after.len(),
2,
"expected exactly [big.parquet, merged] — files_after={paths_after:?}"
);
}
#[tokio::test]
async fn compact_preserves_rows_from_footerless_file() {
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let embs_native: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
let batch_native = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
)
.unwrap();
let bytes_native = AilakeFileWriter::new(policy.clone())
.write(&batch_native, &embs_native)
.unwrap();
store
.put("data/native.parquet", bytes_native.clone())
.await
.unwrap();
let embs_foreign: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
let batch_foreign = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
)
.unwrap();
let bytes_foreign = AilakeFileWriter::new(policy.clone())
.write_parquet_only(&batch_foreign, &embs_foreign)
.unwrap();
store
.put("data/foreign.parquet", bytes_foreign.clone())
.await
.unwrap();
let reader_foreign = AilakeFileReader::new(bytes_foreign.clone(), "embedding", 4);
assert!(
!reader_foreign.is_ailake_file(),
"sanity: write_parquet_only must not embed an AILK footer"
);
let entries = vec![
DataFileEntry {
path: "data/native.parquet".into(),
record_count: 2,
file_size_bytes: bytes_native.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
DataFileEntry {
path: "data/foreign.parquet".into(),
record_count: 2,
file_size_bytes: bytes_foreign.len() as u64,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: None,
vector_dim: None,
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
},
];
let executor = CompactionExecutor::new(store.clone(), policy.clone());
let merged = executor
.compact(&entries, "data/merged_full.parquet")
.await
.unwrap();
assert_eq!(
merged.record_count, 4,
"compact() must include all 4 rows — 2 native + 2 from the footerless file"
);
let merged_bytes = store.get("data/merged_full.parquet").await.unwrap();
let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
reader.verify_integrity().unwrap();
let merged_inc = executor
.compact_incremental(&entries, "data/merged_inc.parquet")
.await
.unwrap();
assert_eq!(
merged_inc.record_count, 4,
"compact_incremental() must include all 4 rows — 2 native + 2 from the footerless file"
);
}
#[tokio::test]
async fn compact_preserves_search_results() {
use crate::scanner::SearchConfig;
use ailake_catalog::HadoopCatalog;
use ailake_core::{VectorMetric, VectorPrecision};
use ailake_store::LocalStore;
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
use tempfile::TempDir;
let dir = TempDir::new().unwrap();
let store = Arc::new(LocalStore::new(dir.path()));
let policy = VectorStoragePolicy {
column_name: "embedding".into(),
dim: 4,
metric: VectorMetric::Cosine,
precision: VectorPrecision::F16,
pq: None,
keep_raw_for_reranking: true,
pre_normalize: false,
hnsw_m: None,
hnsw_ef_construction: None,
ivf_residual: false,
embedding_model: None,
modality: None,
partition_by: None,
partition_value: None,
partition_column_type: None,
partition_fields: vec![],
};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let embs_a: Vec<Vec<f32>> = vec![
vec![1.0, 0.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0, 0.0],
vec![0.0, 0.0, 1.0, 0.0],
];
let batch_a = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![0i32, 1, 2]))],
)
.unwrap();
let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 0.0, 1.0], vec![0.7, 0.7, 0.0, 0.0]];
let batch_b = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![3i32, 4]))],
)
.unwrap();
let bytes_a = AilakeFileWriter::new(policy.clone())
.write(&batch_a, &embs_a)
.unwrap();
let bytes_b = AilakeFileWriter::new(policy.clone())
.write(&batch_b, &embs_b)
.unwrap();
store.put("data/a.parquet", bytes_a).await.unwrap();
store.put("data/b.parquet", bytes_b).await.unwrap();
let file_size_a = store.file_size("data/a.parquet").await.unwrap();
let file_size_b = store.file_size("data/b.parquet").await.unwrap();
let catalog: Arc<dyn CatalogProvider> =
Arc::new(HadoopCatalog::new(store.clone(), "warehouse"));
let table = TableIdent::new("default", "compact_search_test");
catalog
.create_table(
&table,
&ailake_catalog::TableProperties {
policy: policy.clone(),
extra: std::collections::HashMap::new(),
format_version: 2,
partition_column_type: None,
},
)
.await
.unwrap();
let meta = catalog.load_table(&table).await.unwrap();
let parent_snapshot_id = meta.current_snapshot_id;
let entry_a = DataFileEntry {
path: "data/a.parquet".into(),
record_count: 3,
file_size_bytes: file_size_a,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: Some("embedding".into()),
vector_dim: Some(4),
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
};
let entry_b = DataFileEntry {
path: "data/b.parquet".into(),
record_count: 2,
file_size_bytes: file_size_b,
centroid_b64: None,
radius: None,
hnsw_offset: None,
hnsw_len: None,
vector_column: Some("embedding".into()),
vector_dim: Some(4),
extra_vector_indexes: vec![],
index_status: IndexStatus::Ready,
index_error: None,
batch_id: None,
embedding_model: None,
partition_value: None,
deletion_vector: None,
first_row_id: None,
column_stats: None,
sequence_number: 0,
};
catalog
.commit_snapshot(
&table,
NewSnapshot {
snapshot_id: 200,
parent_snapshot_id,
files: vec![entry_a, entry_b],
operation: SnapshotOperation::Append,
iceberg_schema: None,
extra_properties: std::collections::HashMap::new(),
bloom_filters: vec![],
equality_delete_files: vec![],
},
)
.await
.unwrap();
let query_before = vec![1.0f32, 0.0, 0.0, 0.0];
let config = SearchConfig {
top_k: 5,
ef_search: 100,
pruning_threshold: f32::INFINITY,
..Default::default()
};
let before = crate::scanner::search(
&table,
&query_before,
config.clone(),
"embedding",
4,
catalog.clone(),
store.clone(),
)
.await
.unwrap();
assert!(
!before.is_empty(),
"search before compaction must return results"
);
let before_dists: Vec<f32> = before.iter().map(|r| r.distance).collect();
let executor = CompactionExecutor::new(store.clone(), policy.clone());
let files_to_compact = catalog.list_files(&table, None).await.unwrap();
let merged = executor
.compact(&files_to_compact, "data/merged.parquet")
.await
.unwrap();
let meta2 = catalog.load_table(&table).await.unwrap();
let merged_entry = DataFileEntry {
path: merged.path.clone(),
record_count: merged.record_count,
file_size_bytes: merged.file_size_bytes,
centroid_b64: merged.centroid_b64,
radius: merged.radius,
hnsw_offset: merged.hnsw_offset,
hnsw_len: merged.hnsw_len,
vector_column: merged.vector_column,
vector_dim: merged.vector_dim,
extra_vector_indexes: merged.extra_vector_indexes,
index_status: merged.index_status,
index_error: merged.index_error,
batch_id: merged.batch_id,
embedding_model: merged.embedding_model,
partition_value: merged.partition_value,
deletion_vector: merged.deletion_vector,
first_row_id: None,
column_stats: None,
sequence_number: 0,
};
catalog
.commit_snapshot(
&table,
NewSnapshot {
snapshot_id: 201,
parent_snapshot_id: meta2.current_snapshot_id,
files: vec![merged_entry],
operation: SnapshotOperation::Replace,
iceberg_schema: None,
extra_properties: std::collections::HashMap::new(),
bloom_filters: vec![],
equality_delete_files: vec![],
},
)
.await
.unwrap();
let after = crate::scanner::search(
&table,
&query_before,
config.clone(),
"embedding",
4,
catalog,
store,
)
.await
.unwrap();
assert!(
!after.is_empty(),
"search after compaction must return results"
);
let after_dists: Vec<f32> = after.iter().map(|r| r.distance).collect();
assert_eq!(
before.len(),
after.len(),
"number of results should match before ({}) and after ({}) compaction",
before.len(),
after.len(),
);
let mut total_diff = 0.0f32;
for (i, (bd, ad)) in before_dists.iter().zip(after_dists.iter()).enumerate() {
let diff = (bd - ad).abs();
total_diff += diff;
assert!(
diff < 0.1,
"distance mismatch at position {i}: before={bd}, after={ad}, diff={diff}"
);
}
let avg_diff = total_diff / before_dists.len() as f32;
assert!(
avg_diff < 0.05,
"average distance error across all results too high: {avg_diff}"
);
}
}