1use std::sync::Arc;
3use tracing::{debug, error, info, warn};
4
5use ailake_catalog::{
6 make_data_file_entry, make_data_file_entry_indexing, CatalogProvider, DataFileEntry,
7 NewSnapshot, SnapshotOperation, TableIdent, VectorIndexInfo,
8};
9use ailake_core::{AilakeResult, RowId, VectorStoragePolicy};
10use ailake_file::{AilakeFileReader, AilakeFileWriter};
11use ailake_store::Store;
12use ailake_vec::compute_centroid_and_radius;
13use arrow_array::RecordBatch;
14use arrow_schema::SchemaRef;
15use bytes::Bytes;
16use futures::future::try_join_all;
17
18use crate::writer::build_and_patch_index;
19
20#[derive(Debug, Clone, Default)]
22pub enum CompactionIndexStrategy {
23 #[default]
26 Auto,
27 ForceHnsw,
29 ForceIvfPq,
34}
35
36#[derive(Debug, Clone)]
37pub struct CompactionConfig {
38 pub min_files_to_compact: usize,
40 pub target_file_size_bytes: u64,
42 pub index_strategy: CompactionIndexStrategy,
44 pub max_files_per_pass: usize,
51}
52
53impl Default for CompactionConfig {
54 fn default() -> Self {
55 Self {
56 min_files_to_compact: 4,
57 target_file_size_bytes: 128 * 1024 * 1024, index_strategy: CompactionIndexStrategy::Auto,
59 max_files_per_pass: 20,
60 }
61 }
62}
63
64#[derive(Debug, Clone, Copy)]
65pub enum CompactionMode {
66 Full, Partial, }
69
70pub struct CompactionPlanner {
71 config: CompactionConfig,
72}
73
74impl CompactionPlanner {
75 pub fn new(config: CompactionConfig) -> Self {
76 Self { config }
77 }
78
79 pub fn plan(&self, files: &[DataFileEntry]) -> Vec<DataFileEntry> {
97 let (mut foreign, natives): (Vec<DataFileEntry>, Vec<DataFileEntry>) =
99 files.iter().cloned().partition(DataFileEntry::is_foreign);
100 if !foreign.is_empty() {
101 warn!(
102 "ailake: compaction plan — {} file(s) with no AI-Lake index detected \
103 (likely a foreign/external rewrite) — prioritizing for reindex: {:?}",
104 foreign.len(),
105 foreign.iter().map(|f| f.path.as_str()).collect::<Vec<_>>()
106 );
107 }
108
109 let mut size_candidates: Vec<DataFileEntry> = natives
110 .into_iter()
111 .filter(|f| f.file_size_bytes < self.config.target_file_size_bytes)
112 .collect();
113
114 if size_candidates.len() < self.config.min_files_to_compact {
115 if foreign.is_empty() {
116 debug!(
117 "ailake: compaction skipped — {} eligible files < min_files_to_compact={}",
118 size_candidates.len(),
119 self.config.min_files_to_compact
120 );
121 return vec![];
122 }
123 size_candidates.clear();
126 }
127
128 foreign.sort_unstable_by_key(|f| f.file_size_bytes);
135 size_candidates.sort_unstable_by_key(|f| f.file_size_bytes);
136 foreign.extend(size_candidates);
138 foreign.truncate(self.config.max_files_per_pass);
139 let candidates = foreign;
140
141 let total_bytes: u64 = candidates.iter().map(|f| f.file_size_bytes).sum();
142 info!(
143 "ailake: compaction plan — {} files ({} bytes) → 1 merged file",
144 candidates.len(),
145 total_bytes
146 );
147 candidates
148 }
149}
150
151#[derive(Clone)]
162pub struct CompactionExecutor {
163 store: Arc<dyn Store>,
164 policy: VectorStoragePolicy,
165 index_strategy: CompactionIndexStrategy,
166 fts_config: Option<ailake_fts::FtsConfig>,
168}
169
170impl CompactionExecutor {
171 pub fn new(store: Arc<dyn Store>, policy: VectorStoragePolicy) -> Self {
172 Self {
173 store,
174 policy,
175 index_strategy: CompactionIndexStrategy::Auto,
176 fts_config: None,
177 }
178 }
179
180 pub fn with_index_strategy(mut self, strategy: CompactionIndexStrategy) -> Self {
182 self.index_strategy = strategy;
183 self
184 }
185
186 pub fn with_fts_config(mut self, cfg: ailake_fts::FtsConfig) -> Self {
188 self.fts_config = Some(cfg);
189 self
190 }
191
192 fn with_effective_fts(
199 &self,
200 table_props: &std::collections::HashMap<String, String>,
201 ) -> std::borrow::Cow<'_, Self> {
202 if self.fts_config.is_some() {
203 return std::borrow::Cow::Borrowed(self);
204 }
205 match ailake_fts::FtsConfig::from_table_props(table_props) {
206 Some(cfg) => {
207 let mut cloned = self.clone();
208 cloned.fts_config = Some(cfg);
209 std::borrow::Cow::Owned(cloned)
210 }
211 None => std::borrow::Cow::Borrowed(self),
212 }
213 }
214
215 async fn read_files_parallel(
223 &self,
224 files: &[DataFileEntry],
225 ) -> AilakeResult<Vec<(RecordBatch, Vec<Vec<f32>>)>> {
226 let futs = files.iter().map(|entry| {
227 let store = self.store.clone();
228 let path = entry.path.clone();
229 let column = self.policy.column_name.clone();
230 let dim = self.policy.dim;
231 let dv = entry.deletion_vector.clone();
232 async move {
233 let bytes: Bytes = store.get(&path).await?;
234 let reader = AilakeFileReader::new(bytes, &column, dim);
235 if !reader.is_ailake_file() {
236 debug!(
237 "ailake: compaction reading {} without an AI-Lake index \
238 (external write or indexing in progress) — data still merged",
239 path
240 );
241 }
242 let (batch, embeddings) = reader.read_parquet()?;
243 let pair = if let Some(dv) = dv {
247 let bitmap = crate::dv::load_deletion_vector(&store, &dv).await?;
248 crate::dv::filter_deleted_rows(batch, embeddings, &bitmap)?
249 } else {
250 (batch, embeddings)
251 };
252 Ok::<(RecordBatch, Vec<Vec<f32>>), ailake_core::AilakeError>(pair)
253 }
254 });
255 try_join_all(futs).await
256 }
257
258 pub async fn compact(
267 &self,
268 files: &[DataFileEntry],
269 output_path: &str,
270 ) -> AilakeResult<DataFileEntry> {
271 if files.is_empty() {
272 return Err(ailake_core::AilakeError::Catalog(
273 "compact: no files provided".into(),
274 ));
275 }
276
277 let pairs = self.read_files_parallel(files).await?;
281
282 let schema: SchemaRef = pairs[0].0.schema();
283 let (all_batches, all_embeddings): (Vec<_>, Vec<_>) = pairs.into_iter().unzip();
284 let all_embeddings: Vec<Vec<f32>> = all_embeddings.into_iter().flatten().collect();
285
286 let merged_batch = concat_batches(schema, &all_batches)?;
288 let record_count = merged_batch.num_rows() as u64;
289
290 let writer = {
292 let base = AilakeFileWriter::new(self.policy.clone());
293 let base = match &self.index_strategy {
294 CompactionIndexStrategy::Auto => base.with_auto_index(),
295 CompactionIndexStrategy::ForceHnsw => base,
296 CompactionIndexStrategy::ForceIvfPq => {
297 let cfg = ailake_index::IvfPqConfig::for_dataset(
298 self.policy.dim as usize,
299 all_embeddings.len(),
300 );
301 base.with_ivf_pq(cfg)
302 }
303 };
304 if let Some(ref fts_cfg) = self.fts_config {
305 match ailake_fts::merge_fts_blobs(fts_cfg, &merged_batch) {
306 Ok(blob) => base.with_prebuilt_fts_blob(blob),
307 Err(e) => {
308 warn!("ailake: FTS re-index during compaction failed: {e}");
309 base
310 }
311 }
312 } else {
313 base
314 }
315 };
316 let file_bytes = writer.write(&merged_batch, &all_embeddings)?;
317 let file_size = file_bytes.len() as u64;
318 let column_stats =
319 ailake_catalog::extract_column_stats(&file_bytes, &[self.policy.column_name.as_str()])
320 .and_then(|m| serde_json::to_string(&m).ok());
321 self.store.put(output_path, file_bytes.clone()).await?;
322
323 let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
325 let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
326 let header = reader.read_header()?;
327 let ailk_start = reader.ailk_offset()?;
328
329 reader.verify_integrity()?;
334
335 let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
338
339 let mut entry = make_data_file_entry(
340 output_path,
341 record_count,
342 file_size,
343 ¢roid,
344 VectorIndexInfo {
345 column: &self.policy.column_name,
346 dim: self.policy.dim,
347 hnsw_offset: ailk_start + header.hnsw_offset,
348 hnsw_len: header.hnsw_len,
349 },
350 );
351 entry.first_row_id = source_first_row_id;
352 entry.batch_id = DataFileEntry::merge_batch_ids(files);
356 entry.column_stats = column_stats;
357 Ok(entry)
358 }
359
360 pub async fn compact_incremental(
379 &self,
380 files: &[DataFileEntry],
381 output_path: &str,
382 ) -> AilakeResult<DataFileEntry> {
383 const DOMINANT_RATIO: f64 = 0.40;
384
385 if files.is_empty() {
386 return Err(ailake_core::AilakeError::Catalog(
387 "compact_incremental: no files provided".into(),
388 ));
389 }
390
391 if matches!(self.index_strategy, CompactionIndexStrategy::ForceIvfPq) {
400 debug!(
401 "ailake: compact_incremental — index_strategy=ForceIvfPq, which this method \
402 can never produce (HNSW graph extension only); falling back to full rebuild"
403 );
404 return self.compact(files, output_path).await;
405 }
406
407 let total_rows: u64 = files.iter().map(|f| f.record_count).sum();
409 let dom_idx = files
410 .iter()
411 .enumerate()
412 .max_by_key(|(_, f)| f.record_count)
413 .map(|(i, _)| i)
414 .unwrap_or(0);
415 let dom_rows = files[dom_idx].record_count;
416
417 if files[dom_idx].deletion_vector.is_some() {
424 debug!(
425 "ailake: compact_incremental — dominant file {} has DV-masked rows, \
426 falling back to full rebuild (graph reuse would desync row positions)",
427 files[dom_idx].path
428 );
429 return self.compact(files, output_path).await;
430 }
431
432 if (dom_rows as f64 / total_rows as f64) < DOMINANT_RATIO {
433 debug!(
434 "ailake: compact_incremental — no dominant file ({}/{} rows < {:.0}% threshold), \
435 falling back to full rebuild",
436 dom_rows,
437 total_rows,
438 DOMINANT_RATIO * 100.0
439 );
440 return self.compact(files, output_path).await;
441 }
442
443 let column = self.policy.column_name.clone();
444 let dim = self.policy.dim;
445 let dom_path = files[dom_idx].path.clone();
446
447 let futs: Vec<_> =
456 files
457 .iter()
458 .map(|entry| {
459 let store = self.store.clone();
460 let path = entry.path.clone();
461 let col = column.clone();
462 let is_dom = path == dom_path;
463 let dv = entry.deletion_vector.clone();
464 async move {
465 let bytes: Bytes = store.get(&path).await?;
466 let reader = AilakeFileReader::new(bytes.clone(), &col, dim);
467 let has_index = reader.is_ailake_file();
468 if is_dom && !has_index {
469 debug!(
470 "ailake: compact_incremental — dominant candidate {} has no \
471 AI-Lake index; will fall back to full rebuild if no HNSW to reuse",
472 path
473 );
474 }
475 let (raw_batch, raw_vecs) = reader.read_parquet()?;
476 let (batch, vecs) = if let Some(dv) = dv {
480 let bitmap = crate::dv::load_deletion_vector(&store, &dv).await?;
481 crate::dv::filter_deleted_rows(raw_batch, raw_vecs, &bitmap)?
482 } else {
483 (raw_batch, raw_vecs)
484 };
485 let retained = if is_dom && has_index {
486 Some(bytes)
487 } else {
488 None
489 };
490 Ok::<
491 (RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>),
492 ailake_core::AilakeError,
493 >((batch, vecs, is_dom, retained))
494 }
495 })
496 .collect();
497
498 #[allow(clippy::type_complexity)]
502 let raw: Vec<(RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>)> =
503 try_join_all(futs).await?;
504
505 let mut dom_batch: Option<RecordBatch> = None;
507 let mut dom_vecs: Vec<Vec<f32>> = Vec::new();
508 let mut dom_bytes_found: Option<Bytes> = None;
509 let mut other_batches: Vec<RecordBatch> = Vec::new();
510 let mut other_vecs: Vec<Vec<f32>> = Vec::new();
511
512 for (batch, vecs, is_dom, retained) in raw {
513 if is_dom {
514 dom_batch = Some(batch);
515 dom_vecs = vecs;
516 dom_bytes_found = retained;
517 } else {
518 other_batches.push(batch);
519 other_vecs.extend(vecs);
520 }
521 }
522
523 let (dom_batch, dom_bytes) = match (dom_batch, dom_bytes_found) {
524 (Some(b), Some(byt)) => (b, byt),
525 _ => {
526 debug!(
527 "ailake: compact_incremental — dominant file missing from read results, \
528 falling back to full rebuild"
529 );
530 return self.compact(files, output_path).await;
531 }
532 };
533
534 let dom_reader = AilakeFileReader::new(dom_bytes, &column, dim);
536 let mut hnsw = match dom_reader.load_index() {
537 Ok(idx) => idx,
538 Err(e) => {
539 debug!(
540 "ailake: compact_incremental — cannot load dominant HNSW ({}), \
541 falling back to full rebuild",
542 e
543 );
544 return self.compact(files, output_path).await;
545 }
546 };
547
548 let dom_count = dom_batch.num_rows() as u64;
549
550 for (j, vec) in other_vecs.iter().enumerate() {
553 hnsw.insert_node(RowId::new(dom_count + j as u64), vec.clone());
554 }
555 hnsw.quantize_to_f16();
556
557 let schema: SchemaRef = dom_batch.schema();
559 let mut all_batches = vec![dom_batch];
560 all_batches.extend(other_batches);
561 let merged_batch = concat_batches(schema, &all_batches)?;
562 let record_count = merged_batch.num_rows() as u64;
563
564 let mut all_embeddings = dom_vecs;
565 all_embeddings.extend(other_vecs);
566
567 let writer = {
570 let base = AilakeFileWriter::new(self.policy.clone());
571 if let Some(ref fts_cfg) = self.fts_config {
572 match ailake_fts::merge_fts_blobs(fts_cfg, &merged_batch) {
573 Ok(blob) => base.with_prebuilt_fts_blob(blob),
574 Err(e) => {
575 warn!("ailake: FTS re-index during incremental compaction failed: {e}");
576 base
577 }
578 }
579 } else {
580 base
581 }
582 };
583 let file_bytes = writer.write_with_prebuilt_hnsw(&merged_batch, &all_embeddings, &hnsw)?;
584 let file_size = file_bytes.len() as u64;
585 let column_stats =
586 ailake_catalog::extract_column_stats(&file_bytes, &[self.policy.column_name.as_str()])
587 .and_then(|m| serde_json::to_string(&m).ok());
588 self.store.put(output_path, file_bytes.clone()).await?;
589
590 let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
591 let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
592 let header = reader.read_header()?;
593 let ailk_start = reader.ailk_offset()?;
594
595 reader.verify_integrity()?;
598
599 let source_first_row_id = files[dom_idx].first_row_id;
603
604 let mut entry = make_data_file_entry(
605 output_path,
606 record_count,
607 file_size,
608 ¢roid,
609 VectorIndexInfo {
610 column: &self.policy.column_name,
611 dim: self.policy.dim,
612 hnsw_offset: ailk_start + header.hnsw_offset,
613 hnsw_len: header.hnsw_len,
614 },
615 );
616 entry.first_row_id = source_first_row_id;
617 entry.batch_id = DataFileEntry::merge_batch_ids(files);
618 entry.column_stats = column_stats;
619
620 info!(
621 "ailake: compact_incremental — merged {} files into {} \
622 ({} rows from dominant + {} inserted incrementally)",
623 files.len(),
624 output_path,
625 dom_count,
626 record_count - dom_count
627 );
628
629 Ok(entry)
630 }
631
632 pub async fn compact_deferred(
642 &self,
643 files: &[DataFileEntry],
644 output_path: &str,
645 catalog: Arc<dyn CatalogProvider>,
646 table: &TableIdent,
647 ) -> AilakeResult<DataFileEntry> {
648 if files.is_empty() {
649 return Err(ailake_core::AilakeError::Catalog(
650 "compact_deferred: no files provided".into(),
651 ));
652 }
653
654 let pairs = self.read_files_parallel(files).await?;
657
658 let schema: SchemaRef = pairs[0].0.schema();
659 let (all_batches, all_embeddings): (Vec<_>, Vec<_>) = pairs.into_iter().unzip();
660 let all_embeddings: Vec<Vec<f32>> = all_embeddings.into_iter().flatten().collect();
661
662 let merged_batch = concat_batches(schema, &all_batches)?;
663 let record_count = merged_batch.num_rows() as u64;
664
665 let file_writer = AilakeFileWriter::new(self.policy.clone());
667 let parquet_bytes = file_writer.write_parquet_only(&merged_batch, &all_embeddings)?;
668 let file_size = parquet_bytes.len() as u64;
669 let column_stats = ailake_catalog::extract_column_stats(
670 &parquet_bytes,
671 &[self.policy.column_name.as_str()],
672 )
673 .and_then(|m| serde_json::to_string(&m).ok());
674 self.store.put(output_path, parquet_bytes).await?;
675
676 let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
678 let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
679 let mut entry = make_data_file_entry_indexing(
680 output_path,
681 record_count,
682 file_size,
683 ¢roid,
684 &self.policy.column_name,
685 self.policy.dim,
686 );
687 entry.first_row_id = source_first_row_id;
688 entry.batch_id = DataFileEntry::merge_batch_ids(files);
689 entry.column_stats = column_stats;
690
691 let store = self.store.clone();
693 let policy = self.policy.clone();
694 let table_id = table.clone();
695 let fp = output_path.to_string();
696 tokio::spawn(async move {
697 if let Err(e) = build_and_patch_index(store, catalog, policy, table_id, fp).await {
698 error!(
699 "ailake: compaction deferred HNSW build failed — file indexed as \
700 Parquet-only until next compaction rebuilds the index: {}",
701 e
702 );
703 }
704 });
705
706 Ok(entry)
707 }
708
709 pub async fn run(
712 &self,
713 planner: &CompactionPlanner,
714 table: &TableIdent,
715 catalog: Arc<dyn CatalogProvider>,
716 output_prefix: &str,
717 ) -> AilakeResult<Option<DataFileEntry>> {
718 let all_files = catalog.list_files(table, None).await?;
719 let to_compact = planner.plan(&all_files);
720 if to_compact.is_empty() {
721 return Ok(None);
722 }
723
724 let meta_props = catalog
727 .load_table(table)
728 .await
729 .map(|m| m.properties)
730 .unwrap_or_default();
731 let executor = self.with_effective_fts(&meta_props);
732
733 let ts = std::time::SystemTime::now()
734 .duration_since(std::time::UNIX_EPOCH)
735 .unwrap_or_else(|e| e.duration())
736 .as_millis();
737 let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
738
739 let merged = executor
741 .compact_incremental(&to_compact, &output_path)
742 .await?;
743
744 let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
746 let parent_snapshot_id = catalog
749 .load_table(table)
750 .await
751 .ok()
752 .and_then(|m| m.current_snapshot_id);
753 let snapshot = NewSnapshot {
754 snapshot_id: ailake_catalog::new_snapshot_id(),
755 parent_snapshot_id,
756 files,
757 operation: SnapshotOperation::Replace,
758 iceberg_schema: None,
759 extra_properties: std::collections::HashMap::new(),
760 bloom_filters: vec![],
761 equality_delete_files: vec![],
762 };
763 catalog.commit_snapshot(table, snapshot).await?;
764
765 info!(
766 "ailake: compaction committed — merged {} files into {}",
767 to_compact.len(),
768 output_path
769 );
770
771 if catalog.retires_files_physically() {
772 delete_old_files(&self.store, &to_compact).await;
773 } else {
774 info!(
775 "ailake: compaction — leaving {} retired file(s) in place; catalog backend \
776 manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
777 to_compact.len()
778 );
779 }
780
781 Ok(Some(merged))
782 }
783
784 pub async fn run_deferred(
794 &self,
795 planner: &CompactionPlanner,
796 table: &TableIdent,
797 catalog: Arc<dyn CatalogProvider>,
798 output_prefix: &str,
799 ) -> AilakeResult<Option<DataFileEntry>> {
800 if !catalog.supports_in_place_rewrite() {
805 return Err(ailake_core::AilakeError::Catalog(
806 "deferred compaction is not supported with this catalog backend: the \
807 background index build patches the merged file in place at its committed \
808 path, which this catalog cannot re-register — run a blocking compact"
809 .into(),
810 ));
811 }
812 let all_files = catalog.list_files(table, None).await?;
813 let to_compact = planner.plan(&all_files);
814 if to_compact.is_empty() {
815 return Ok(None);
816 }
817
818 let ts = std::time::SystemTime::now()
819 .duration_since(std::time::UNIX_EPOCH)
820 .unwrap_or_else(|e| e.duration())
821 .as_millis();
822 let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
823
824 let merged = self
825 .compact_deferred(&to_compact, &output_path, catalog.clone(), table)
826 .await?;
827
828 let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
830 let parent_snapshot_id = catalog
833 .load_table(table)
834 .await
835 .ok()
836 .and_then(|m| m.current_snapshot_id);
837 let snapshot = NewSnapshot {
838 snapshot_id: ailake_catalog::new_snapshot_id(),
839 parent_snapshot_id,
840 files,
841 operation: SnapshotOperation::Replace,
842 iceberg_schema: None,
843 extra_properties: std::collections::HashMap::new(),
844 bloom_filters: vec![],
845 equality_delete_files: vec![],
846 };
847 catalog.commit_snapshot(table, snapshot).await?;
848
849 info!(
850 "ailake: compaction committed (deferred) — merged {} files into {} \
851 (index building in background)",
852 to_compact.len(),
853 output_path
854 );
855
856 if catalog.retires_files_physically() {
857 delete_old_files(&self.store, &to_compact).await;
858 } else {
859 info!(
860 "ailake: compaction — leaving {} retired file(s) in place; catalog backend \
861 manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
862 to_compact.len()
863 );
864 }
865
866 Ok(Some(merged))
867 }
868}
869
870async fn build_replace_file_list(
881 catalog: &Arc<dyn CatalogProvider>,
882 table: &TableIdent,
883 to_compact: &[DataFileEntry],
884 merged: DataFileEntry,
885) -> AilakeResult<Vec<DataFileEntry>> {
886 let current_files = catalog.list_files(table, None).await?;
887 let compacted_paths: std::collections::HashSet<&str> =
888 to_compact.iter().map(|f| f.path.as_str()).collect();
889 let mut files: Vec<DataFileEntry> = current_files
890 .into_iter()
891 .filter(|f| !compacted_paths.contains(f.path.as_str()))
892 .collect();
893 files.push(merged);
894 Ok(files)
895}
896
897async fn delete_old_files(store: &Arc<dyn Store>, files: &[DataFileEntry]) {
898 for entry in files {
899 if let Err(e) = store.delete(&entry.path).await {
900 error!(
901 "ailake: compaction cleanup failed — could not delete {}: {} \
902 (orphan file in object store after successful catalog commit; \
903 delete manually to reclaim storage)",
904 entry.path, e
905 );
906 }
907 }
908}
909
910fn concat_batches(schema: SchemaRef, batches: &[RecordBatch]) -> AilakeResult<RecordBatch> {
911 arrow_select::concat::concat_batches(&schema, batches)
912 .map_err(|e| ailake_core::AilakeError::Arrow(e.to_string()))
913}
914
915#[cfg(test)]
916mod tests {
917 use super::*;
918 use ailake_catalog::IndexStatus;
919
920 #[test]
921 fn plan_returns_empty_if_too_few_files() {
922 let planner = CompactionPlanner::new(CompactionConfig {
923 min_files_to_compact: 4,
924 target_file_size_bytes: 1024 * 1024,
925 ..Default::default()
926 });
927 let files: Vec<DataFileEntry> = (0..3)
928 .map(|i| DataFileEntry {
929 path: format!("file-{i}.parquet"),
930 record_count: 10,
931 file_size_bytes: 100,
932 centroid_b64: Some("AAAA".into()),
933 radius: None,
934 hnsw_offset: None,
935 hnsw_len: None,
936 vector_column: None,
937 vector_dim: None,
938 extra_vector_indexes: vec![],
939 index_status: IndexStatus::Ready,
940 index_error: None,
941 batch_id: None,
942 embedding_model: None,
943 partition_value: None,
944 deletion_vector: None,
945 first_row_id: None,
946 column_stats: None,
947 sequence_number: 0,
948 })
949 .collect();
950 assert!(planner.plan(&files).is_empty());
951 }
952
953 #[test]
954 fn plan_selects_small_files() {
955 let planner = CompactionPlanner::new(CompactionConfig {
956 min_files_to_compact: 2,
957 target_file_size_bytes: 1000,
958 ..Default::default()
959 });
960 let files = vec![
961 DataFileEntry {
962 path: "small.parquet".into(),
963 record_count: 5,
964 file_size_bytes: 500,
965 centroid_b64: Some("AAAA".into()),
966 radius: None,
967 hnsw_offset: None,
968 hnsw_len: None,
969 vector_column: None,
970 vector_dim: None,
971 extra_vector_indexes: vec![],
972 index_status: IndexStatus::Ready,
973 index_error: None,
974 batch_id: None,
975 embedding_model: None,
976 partition_value: None,
977 deletion_vector: None,
978 first_row_id: None,
979 column_stats: None,
980 sequence_number: 0,
981 },
982 DataFileEntry {
983 path: "large.parquet".into(),
984 record_count: 5000,
985 file_size_bytes: 200_000_000,
986 centroid_b64: Some("AAAA".into()),
987 radius: None,
988 hnsw_offset: None,
989 hnsw_len: None,
990 vector_column: None,
991 vector_dim: None,
992 extra_vector_indexes: vec![],
993 index_status: IndexStatus::Ready,
994 index_error: None,
995 batch_id: None,
996 embedding_model: None,
997 partition_value: None,
998 deletion_vector: None,
999 first_row_id: None,
1000 column_stats: None,
1001 sequence_number: 0,
1002 },
1003 DataFileEntry {
1004 path: "also-small.parquet".into(),
1005 record_count: 5,
1006 file_size_bytes: 800,
1007 centroid_b64: Some("AAAA".into()),
1008 radius: None,
1009 hnsw_offset: None,
1010 hnsw_len: None,
1011 vector_column: None,
1012 vector_dim: None,
1013 extra_vector_indexes: vec![],
1014 index_status: IndexStatus::Ready,
1015 index_error: None,
1016 batch_id: None,
1017 embedding_model: None,
1018 partition_value: None,
1019 deletion_vector: None,
1020 first_row_id: None,
1021 column_stats: None,
1022 sequence_number: 0,
1023 },
1024 ];
1025 let selected = planner.plan(&files);
1026 assert_eq!(selected.len(), 2);
1027 assert!(selected.iter().any(|f| f.path == "small.parquet"));
1028 assert!(selected.iter().any(|f| f.path == "also-small.parquet"));
1029 }
1030
1031 #[test]
1032 fn plan_respects_max_files_per_pass() {
1033 let planner = CompactionPlanner::new(CompactionConfig {
1034 min_files_to_compact: 2,
1035 target_file_size_bytes: 1_000_000,
1036 max_files_per_pass: 3,
1037 ..Default::default()
1038 });
1039 let files: Vec<DataFileEntry> = (0..5)
1040 .map(|i| DataFileEntry {
1041 path: format!("f{i}.parquet"),
1042 record_count: 10,
1043 file_size_bytes: 100 + i as u64 * 100,
1044 centroid_b64: Some("AAAA".into()),
1045 radius: None,
1046 hnsw_offset: None,
1047 hnsw_len: None,
1048 vector_column: None,
1049 vector_dim: None,
1050 extra_vector_indexes: vec![],
1051 index_status: IndexStatus::Ready,
1052 index_error: None,
1053 batch_id: None,
1054 embedding_model: None,
1055 partition_value: None,
1056 deletion_vector: None,
1057 first_row_id: None,
1058 column_stats: None,
1059 sequence_number: 0,
1060 })
1061 .collect();
1062 let selected = planner.plan(&files);
1063 assert_eq!(selected.len(), 3);
1064 assert_eq!(selected[0].file_size_bytes, 100);
1065 assert_eq!(selected[1].file_size_bytes, 200);
1066 assert_eq!(selected[2].file_size_bytes, 300);
1067 }
1068
1069 #[test]
1070 fn plan_sorts_smallest_first() {
1071 let planner = CompactionPlanner::new(CompactionConfig {
1072 min_files_to_compact: 2,
1073 target_file_size_bytes: 10_000,
1074 max_files_per_pass: 4,
1075 ..Default::default()
1076 });
1077 let files = vec![
1078 DataFileEntry {
1079 path: "c.parquet".into(),
1080 record_count: 1,
1081 file_size_bytes: 300,
1082 centroid_b64: Some("AAAA".into()),
1083 radius: None,
1084 hnsw_offset: None,
1085 hnsw_len: None,
1086 vector_column: None,
1087 vector_dim: None,
1088 extra_vector_indexes: vec![],
1089 index_status: IndexStatus::Ready,
1090 index_error: None,
1091 batch_id: None,
1092 embedding_model: None,
1093 partition_value: None,
1094 deletion_vector: None,
1095 first_row_id: None,
1096 column_stats: None,
1097 sequence_number: 0,
1098 },
1099 DataFileEntry {
1100 path: "a.parquet".into(),
1101 record_count: 1,
1102 file_size_bytes: 100,
1103 centroid_b64: Some("AAAA".into()),
1104 radius: None,
1105 hnsw_offset: None,
1106 hnsw_len: None,
1107 vector_column: None,
1108 vector_dim: None,
1109 extra_vector_indexes: vec![],
1110 index_status: IndexStatus::Ready,
1111 index_error: None,
1112 batch_id: None,
1113 embedding_model: None,
1114 partition_value: None,
1115 deletion_vector: None,
1116 first_row_id: None,
1117 column_stats: None,
1118 sequence_number: 0,
1119 },
1120 DataFileEntry {
1121 path: "b.parquet".into(),
1122 record_count: 1,
1123 file_size_bytes: 200,
1124 centroid_b64: Some("AAAA".into()),
1125 radius: None,
1126 hnsw_offset: None,
1127 hnsw_len: None,
1128 vector_column: None,
1129 vector_dim: None,
1130 extra_vector_indexes: vec![],
1131 index_status: IndexStatus::Ready,
1132 index_error: None,
1133 batch_id: None,
1134 embedding_model: None,
1135 partition_value: None,
1136 deletion_vector: None,
1137 first_row_id: None,
1138 column_stats: None,
1139 sequence_number: 0,
1140 },
1141 ];
1142 let selected = planner.plan(&files);
1143 assert_eq!(selected[0].file_size_bytes, 100);
1144 assert_eq!(selected[1].file_size_bytes, 200);
1145 assert_eq!(selected[2].file_size_bytes, 300);
1146 }
1147
1148 fn make_plan_entry(path: &str, size: u64, centroid_b64: Option<String>) -> DataFileEntry {
1149 DataFileEntry {
1150 path: path.to_string(),
1151 record_count: 10,
1152 file_size_bytes: size,
1153 centroid_b64,
1154 radius: None,
1155 hnsw_offset: None,
1156 hnsw_len: None,
1157 vector_column: None,
1158 vector_dim: None,
1159 extra_vector_indexes: vec![],
1160 index_status: IndexStatus::Ready,
1161 index_error: None,
1162 batch_id: None,
1163 embedding_model: None,
1164 partition_value: None,
1165 deletion_vector: None,
1166 first_row_id: None,
1167 column_stats: None,
1168 sequence_number: 0,
1169 }
1170 }
1171
1172 #[test]
1173 fn plan_prioritizes_foreign_written_files_regardless_of_size() {
1174 let planner = CompactionPlanner::new(CompactionConfig {
1175 min_files_to_compact: 4, target_file_size_bytes: 1000,
1177 max_files_per_pass: 20,
1178 ..Default::default()
1179 });
1180 let files = vec![
1181 make_plan_entry("native_small.parquet", 500, Some("AAAA".into())),
1184 make_plan_entry("foreign_big.parquet", 200_000_000, None),
1187 ];
1188 let selected = planner.plan(&files);
1189 assert_eq!(
1190 selected.len(),
1191 1,
1192 "foreign file alone should trigger a pass even below min_files_to_compact"
1193 );
1194 assert_eq!(selected[0].path, "foreign_big.parquet");
1195 }
1196
1197 #[test]
1198 fn plan_merges_foreign_and_size_candidates_together() {
1199 let planner = CompactionPlanner::new(CompactionConfig {
1200 min_files_to_compact: 2,
1201 target_file_size_bytes: 1000,
1202 max_files_per_pass: 20,
1203 ..Default::default()
1204 });
1205 let files = vec![
1206 make_plan_entry("native_a.parquet", 300, Some("AAAA".into())),
1207 make_plan_entry("native_b.parquet", 400, Some("AAAA".into())),
1208 make_plan_entry("foreign.parquet", 200_000_000, None),
1209 ];
1210 let selected = planner.plan(&files);
1211 let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1212 assert_eq!(selected.len(), 3, "paths={paths:?}");
1213 assert_eq!(selected[0].path, "foreign.parquet", "paths={paths:?}");
1215 }
1216
1217 #[test]
1218 fn plan_sorts_foreign_files_smallest_first_too() {
1219 let planner = CompactionPlanner::new(CompactionConfig {
1223 min_files_to_compact: 1,
1224 target_file_size_bytes: 1000,
1225 max_files_per_pass: 2,
1226 ..Default::default()
1227 });
1228 let files = vec![
1229 make_plan_entry("foreign_huge.parquet", 500_000_000, None),
1230 make_plan_entry("foreign_small.parquet", 100, None),
1231 make_plan_entry("foreign_medium.parquet", 10_000_000, None),
1232 ];
1233 let selected = planner.plan(&files);
1234 let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1235 assert_eq!(selected.len(), 2, "max_files_per_pass=2, paths={paths:?}");
1236 assert_eq!(
1237 paths,
1238 vec!["foreign_small.parquet", "foreign_medium.parquet"],
1239 "foreign files must be size-sorted before truncation, cheapest first"
1240 );
1241 }
1242
1243 #[tokio::test]
1244 async fn compact_merges_two_files() {
1245 use ailake_core::{VectorMetric, VectorPrecision};
1246 use ailake_store::LocalStore;
1247 use arrow_array::{Int32Array, RecordBatch};
1248 use arrow_schema::{DataType, Field, Schema};
1249 use std::sync::Arc;
1250 use tempfile::TempDir;
1251
1252 let dir = TempDir::new().unwrap();
1253 let store = Arc::new(LocalStore::new(dir.path()));
1254 let policy = VectorStoragePolicy {
1255 column_name: "embedding".into(),
1256 dim: 4,
1257 metric: VectorMetric::Cosine,
1258 precision: VectorPrecision::F16,
1259 pq: None,
1260 keep_raw_for_reranking: true,
1261 pre_normalize: false,
1262 hnsw_m: None,
1263 hnsw_ef_construction: None,
1264 ivf_residual: false,
1265 embedding_model: None,
1266 modality: None,
1267 partition_by: None,
1268 partition_value: None,
1269 partition_column_type: None,
1270 partition_fields: vec![],
1271 };
1272
1273 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1274 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]];
1275 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]];
1276
1277 let batch_a = RecordBatch::try_new(
1278 schema.clone(),
1279 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1280 )
1281 .unwrap();
1282 let batch_b = RecordBatch::try_new(
1283 schema.clone(),
1284 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1285 )
1286 .unwrap();
1287
1288 let writer_a = AilakeFileWriter::new(policy.clone());
1289 let bytes_a = writer_a.write(&batch_a, &embs_a).unwrap();
1290 let writer_b = AilakeFileWriter::new(policy.clone());
1291 let bytes_b = writer_b.write(&batch_b, &embs_b).unwrap();
1292
1293 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1294 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1295
1296 let entries = vec![
1297 DataFileEntry {
1298 path: "data/a.parquet".into(),
1299 record_count: 2,
1300 file_size_bytes: bytes_a.len() as u64,
1301 centroid_b64: None,
1302 radius: None,
1303 hnsw_offset: None,
1304 hnsw_len: None,
1305 vector_column: None,
1306 vector_dim: None,
1307 extra_vector_indexes: vec![],
1308 index_status: IndexStatus::Ready,
1309 index_error: None,
1310 batch_id: None,
1311 embedding_model: None,
1312 partition_value: None,
1313 deletion_vector: None,
1314 first_row_id: None,
1315 column_stats: None,
1316 sequence_number: 0,
1317 },
1318 DataFileEntry {
1319 path: "data/b.parquet".into(),
1320 record_count: 2,
1321 file_size_bytes: bytes_b.len() as u64,
1322 centroid_b64: None,
1323 radius: None,
1324 hnsw_offset: None,
1325 hnsw_len: None,
1326 vector_column: None,
1327 vector_dim: None,
1328 extra_vector_indexes: vec![],
1329 index_status: IndexStatus::Ready,
1330 index_error: None,
1331 batch_id: None,
1332 embedding_model: None,
1333 partition_value: None,
1334 deletion_vector: None,
1335 first_row_id: None,
1336 column_stats: None,
1337 sequence_number: 0,
1338 },
1339 ];
1340
1341 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1342 let merged = executor
1343 .compact(&entries, "data/merged.parquet")
1344 .await
1345 .unwrap();
1346
1347 assert_eq!(merged.record_count, 4);
1348 assert_eq!(merged.path, "data/merged.parquet");
1349
1350 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1351 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1352 reader.verify_integrity().unwrap();
1353 let (batch, embs) = reader.read_parquet().unwrap();
1354 assert_eq!(batch.num_rows(), 4);
1355 assert_eq!(embs.len(), 4);
1356 }
1357
1358 #[tokio::test]
1363 async fn compact_drops_deletion_vector_masked_rows() {
1364 use ailake_catalog::provider::DeletionVector;
1365 use ailake_core::{VectorMetric, VectorPrecision};
1366 use ailake_store::LocalStore;
1367 use arrow_array::{Int32Array, RecordBatch};
1368 use arrow_schema::{DataType, Field, Schema};
1369 use roaring::RoaringBitmap;
1370 use std::sync::Arc;
1371 use tempfile::TempDir;
1372
1373 let dir = TempDir::new().unwrap();
1374 let store = Arc::new(LocalStore::new(dir.path()));
1375 let policy = VectorStoragePolicy {
1376 column_name: "embedding".into(),
1377 dim: 4,
1378 metric: VectorMetric::Cosine,
1379 precision: VectorPrecision::F16,
1380 pq: None,
1381 keep_raw_for_reranking: true,
1382 pre_normalize: false,
1383 hnsw_m: None,
1384 hnsw_ef_construction: None,
1385 ivf_residual: false,
1386 embedding_model: None,
1387 modality: None,
1388 partition_by: None,
1389 partition_value: None,
1390 partition_column_type: None,
1391 partition_fields: vec![],
1392 };
1393
1394 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1395 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]];
1397 let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
1398
1399 let batch_a = RecordBatch::try_new(
1400 schema.clone(),
1401 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1402 )
1403 .unwrap();
1404 let batch_b =
1405 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![2i32]))])
1406 .unwrap();
1407
1408 let bytes_a = AilakeFileWriter::new(policy.clone())
1409 .write(&batch_a, &embs_a)
1410 .unwrap();
1411 let bytes_b = AilakeFileWriter::new(policy.clone())
1412 .write(&batch_b, &embs_b)
1413 .unwrap();
1414 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1415 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1416
1417 let mut bitmap = RoaringBitmap::new();
1419 bitmap.insert(0);
1420 let (puffin_bytes, offset, length) =
1421 crate::delete::PuffinWriter::write_single_dv(&bitmap, 1).unwrap();
1422 store.put("metadata/dv-1.dvd", puffin_bytes).await.unwrap();
1423
1424 let entry_a = DataFileEntry {
1425 path: "data/a.parquet".into(),
1426 record_count: 2,
1427 file_size_bytes: bytes_a.len() as u64,
1428 centroid_b64: None,
1429 radius: None,
1430 hnsw_offset: None,
1431 hnsw_len: None,
1432 vector_column: None,
1433 vector_dim: None,
1434 extra_vector_indexes: vec![],
1435 index_status: IndexStatus::Ready,
1436 index_error: None,
1437 batch_id: None,
1438 embedding_model: None,
1439 partition_value: None,
1440 deletion_vector: Some(DeletionVector {
1441 path: "metadata/dv-1.dvd".into(),
1442 offset,
1443 length,
1444 cardinality: 1,
1445 }),
1446 first_row_id: None,
1447 column_stats: None,
1448 sequence_number: 0,
1449 };
1450 let entry_b = DataFileEntry {
1451 path: "data/b.parquet".into(),
1452 record_count: 1,
1453 file_size_bytes: bytes_b.len() as u64,
1454 centroid_b64: None,
1455 radius: None,
1456 hnsw_offset: None,
1457 hnsw_len: None,
1458 vector_column: None,
1459 vector_dim: None,
1460 extra_vector_indexes: vec![],
1461 index_status: IndexStatus::Ready,
1462 index_error: None,
1463 batch_id: None,
1464 embedding_model: None,
1465 partition_value: None,
1466 deletion_vector: None,
1467 first_row_id: None,
1468 column_stats: None,
1469 sequence_number: 0,
1470 };
1471
1472 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1473 let merged = executor
1474 .compact(&[entry_a, entry_b], "data/merged.parquet")
1475 .await
1476 .unwrap();
1477
1478 assert_eq!(merged.record_count, 2);
1480 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1481 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1482 let (batch, embs) = reader.read_parquet().unwrap();
1483 assert_eq!(batch.num_rows(), 2);
1484 assert_eq!(embs.len(), 2);
1485 let ids: Vec<i32> = batch
1486 .column_by_name("id")
1487 .unwrap()
1488 .as_any()
1489 .downcast_ref::<Int32Array>()
1490 .unwrap()
1491 .values()
1492 .to_vec();
1493 assert_eq!(
1494 ids,
1495 vec![1, 2],
1496 "id=0 (deleted) must not survive compaction"
1497 );
1498 }
1499
1500 #[tokio::test]
1505 async fn compact_with_ivf_pq_strategy_does_not_crash_on_verify_integrity() {
1506 use ailake_core::{VectorMetric, VectorPrecision};
1507 use ailake_store::LocalStore;
1508 use arrow_array::{Int32Array, RecordBatch};
1509 use arrow_schema::{DataType, Field, Schema};
1510 use std::sync::Arc;
1511 use tempfile::TempDir;
1512
1513 let dir = TempDir::new().unwrap();
1514 let store = Arc::new(LocalStore::new(dir.path()));
1515 let dim = 8;
1516 let policy = VectorStoragePolicy {
1517 column_name: "embedding".into(),
1518 dim,
1519 metric: VectorMetric::Cosine,
1520 precision: VectorPrecision::F16,
1521 pq: None,
1522 keep_raw_for_reranking: true,
1523 pre_normalize: false,
1524 hnsw_m: None,
1525 hnsw_ef_construction: None,
1526 ivf_residual: false,
1527 embedding_model: None,
1528 modality: None,
1529 partition_by: None,
1530 partition_value: None,
1531 partition_column_type: None,
1532 partition_fields: vec![],
1533 };
1534
1535 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1536 let n_per_file = 30usize;
1537 let make_file = |path: &str, offset: i32| {
1538 let ids: Vec<i32> = (offset..offset + n_per_file as i32).collect();
1539 let embs: Vec<Vec<f32>> = ids
1540 .iter()
1541 .map(|&i| {
1542 (0..dim as i32)
1543 .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1544 .collect()
1545 })
1546 .collect();
1547 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1548 .unwrap();
1549 let bytes = AilakeFileWriter::new(policy.clone())
1550 .write(&batch, &embs)
1551 .unwrap();
1552 (path.to_string(), bytes)
1553 };
1554
1555 let (path_a, bytes_a) = make_file("data/ivfpq_a.parquet", 0);
1556 let (path_b, bytes_b) = make_file("data/ivfpq_b.parquet", n_per_file as i32);
1557 for (path, bytes) in [(&path_a, &bytes_a), (&path_b, &bytes_b)] {
1558 store.put(path, bytes.clone()).await.unwrap();
1559 }
1560
1561 let make_entry = |path: &str, size: u64| DataFileEntry {
1562 path: path.to_string(),
1563 record_count: n_per_file as u64,
1564 file_size_bytes: size,
1565 centroid_b64: None,
1566 radius: None,
1567 hnsw_offset: None,
1568 hnsw_len: None,
1569 vector_column: None,
1570 vector_dim: None,
1571 extra_vector_indexes: vec![],
1572 index_status: IndexStatus::Ready,
1573 index_error: None,
1574 batch_id: None,
1575 embedding_model: None,
1576 partition_value: None,
1577 deletion_vector: None,
1578 first_row_id: None,
1579 column_stats: None,
1580 sequence_number: 0,
1581 };
1582 let entries = vec![
1583 make_entry(&path_a, bytes_a.len() as u64),
1584 make_entry(&path_b, bytes_b.len() as u64),
1585 ];
1586
1587 let executor = CompactionExecutor::new(store.clone(), policy.clone())
1588 .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1589 let merged = executor
1590 .compact(&entries, "data/ivfpq_merged.parquet")
1591 .await
1592 .expect("compact() with ForceIvfPq must not fail in verify_integrity()");
1593
1594 assert_eq!(merged.record_count, 2 * n_per_file as u64);
1595
1596 let merged_bytes = store.get("data/ivfpq_merged.parquet").await.unwrap();
1598 let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1599 reader
1600 .verify_integrity()
1601 .expect("verify_integrity() must handle an IVF-PQ-indexed file");
1602 }
1603
1604 #[tokio::test]
1605 async fn compact_incremental_merges_dominant_plus_small() {
1606 use ailake_core::{RowId, VectorMetric, VectorPrecision};
1607 use ailake_store::LocalStore;
1608 use arrow_array::{Int32Array, RecordBatch};
1609 use arrow_schema::{DataType, Field, Schema};
1610 use std::sync::Arc;
1611 use tempfile::TempDir;
1612
1613 let dir = TempDir::new().unwrap();
1614 let store = Arc::new(LocalStore::new(dir.path()));
1615 let policy = VectorStoragePolicy {
1616 column_name: "embedding".into(),
1617 dim: 4,
1618 metric: VectorMetric::Cosine,
1619 precision: VectorPrecision::F16,
1620 pq: None,
1621 keep_raw_for_reranking: true,
1622 pre_normalize: false,
1623 hnsw_m: None,
1624 hnsw_ef_construction: None,
1625 ivf_residual: false,
1626 embedding_model: None,
1627 modality: None,
1628 partition_by: None,
1629 partition_value: None,
1630 partition_column_type: None,
1631 partition_fields: vec![],
1632 };
1633
1634 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1635
1636 let embs_dom: Vec<Vec<f32>> = vec![
1638 vec![1.0, 0.0, 0.0, 0.0],
1639 vec![0.0, 1.0, 0.0, 0.0],
1640 vec![0.0, 0.0, 1.0, 0.0],
1641 vec![0.7, 0.7, 0.0, 0.0],
1642 vec![0.0, 0.7, 0.7, 0.0],
1643 vec![0.0, 0.0, 0.7, 0.7],
1644 ];
1645 let batch_dom = RecordBatch::try_new(
1646 schema.clone(),
1647 vec![Arc::new(Int32Array::from(vec![0i32, 1, 2, 3, 4, 5]))],
1648 )
1649 .unwrap();
1650
1651 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]];
1653 let batch_small = RecordBatch::try_new(
1654 schema.clone(),
1655 vec![Arc::new(Int32Array::from(vec![6i32, 7]))],
1656 )
1657 .unwrap();
1658
1659 let bytes_dom = AilakeFileWriter::new(policy.clone())
1660 .write(&batch_dom, &embs_dom)
1661 .unwrap();
1662 let bytes_small = AilakeFileWriter::new(policy.clone())
1663 .write(&batch_small, &embs_small)
1664 .unwrap();
1665
1666 store
1667 .put("data/dominant.parquet", bytes_dom.clone())
1668 .await
1669 .unwrap();
1670 store
1671 .put("data/small.parquet", bytes_small.clone())
1672 .await
1673 .unwrap();
1674
1675 let entries = vec![
1676 DataFileEntry {
1677 path: "data/dominant.parquet".into(),
1678 record_count: 6,
1679 file_size_bytes: bytes_dom.len() as u64,
1680 centroid_b64: None,
1681 radius: None,
1682 hnsw_offset: None,
1683 hnsw_len: None,
1684 vector_column: None,
1685 vector_dim: None,
1686 extra_vector_indexes: vec![],
1687 index_status: IndexStatus::Ready,
1688 index_error: None,
1689 batch_id: None,
1690 embedding_model: None,
1691 partition_value: None,
1692 deletion_vector: None,
1693 first_row_id: None,
1694 column_stats: None,
1695 sequence_number: 0,
1696 },
1697 DataFileEntry {
1698 path: "data/small.parquet".into(),
1699 record_count: 2,
1700 file_size_bytes: bytes_small.len() as u64,
1701 centroid_b64: None,
1702 radius: None,
1703 hnsw_offset: None,
1704 hnsw_len: None,
1705 vector_column: None,
1706 vector_dim: None,
1707 extra_vector_indexes: vec![],
1708 index_status: IndexStatus::Ready,
1709 index_error: None,
1710 batch_id: None,
1711 embedding_model: None,
1712 partition_value: None,
1713 deletion_vector: None,
1714 first_row_id: None,
1715 column_stats: None,
1716 sequence_number: 0,
1717 },
1718 ];
1719
1720 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1721 let merged = executor
1722 .compact_incremental(&entries, "data/merged.parquet")
1723 .await
1724 .unwrap();
1725
1726 assert_eq!(merged.record_count, 8);
1728 assert_eq!(merged.path, "data/merged.parquet");
1729
1730 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1732 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1733 reader.verify_integrity().unwrap();
1734
1735 let (batch, embs) = reader.read_parquet().unwrap();
1736 assert_eq!(batch.num_rows(), 8);
1737 assert_eq!(embs.len(), 8);
1738
1739 for f in &embs[..6] {
1741 assert_eq!(f.len(), 4);
1742 }
1743
1744 let hnsw = reader.load_index().unwrap();
1746 assert_eq!(hnsw.node_count(), 8);
1747
1748 let results = hnsw.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
1750 assert_eq!(results[0].0, RowId::new(0));
1751
1752 let results = hnsw.search(&[0.0, 0.0, 0.0, 1.0], 1, 50);
1755 assert_eq!(results[0].0, RowId::new(6));
1756 }
1757
1758 #[tokio::test]
1768 async fn compact_incremental_respects_force_ivf_pq_even_with_dominant_file() {
1769 use ailake_core::{VectorMetric, VectorPrecision};
1770 use ailake_store::LocalStore;
1771 use arrow_array::{Int32Array, RecordBatch};
1772 use arrow_schema::{DataType, Field, Schema};
1773 use std::sync::Arc;
1774 use tempfile::TempDir;
1775
1776 let dir = TempDir::new().unwrap();
1777 let store = Arc::new(LocalStore::new(dir.path()));
1778 let dim = 8;
1779 let policy = VectorStoragePolicy {
1780 column_name: "embedding".into(),
1781 dim,
1782 metric: VectorMetric::Cosine,
1783 precision: VectorPrecision::F16,
1784 pq: None,
1785 keep_raw_for_reranking: true,
1786 pre_normalize: false,
1787 hnsw_m: None,
1788 hnsw_ef_construction: None,
1789 ivf_residual: false,
1790 embedding_model: None,
1791 modality: None,
1792 partition_by: None,
1793 partition_value: None,
1794 partition_column_type: None,
1795 partition_fields: vec![],
1796 };
1797
1798 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1799 let make_file = |path: &str, offset: i32, n: usize| {
1800 let ids: Vec<i32> = (offset..offset + n as i32).collect();
1801 let embs: Vec<Vec<f32>> = ids
1802 .iter()
1803 .map(|&i| {
1804 (0..dim as i32)
1805 .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1806 .collect()
1807 })
1808 .collect();
1809 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1810 .unwrap();
1811 let bytes = AilakeFileWriter::new(policy.clone())
1814 .write(&batch, &embs)
1815 .unwrap();
1816 (path.to_string(), bytes, n as u64)
1817 };
1818
1819 let (path_dom, bytes_dom, n_dom) = make_file("data/dom.parquet", 0, 90);
1822 let (path_small, bytes_small, n_small) = make_file("data/small.parquet", 999, 30);
1823
1824 for (path, bytes) in [(&path_dom, &bytes_dom), (&path_small, &bytes_small)] {
1825 store.put(path, bytes.clone()).await.unwrap();
1826 }
1827
1828 let make_entry = |path: &str, record_count: u64, size: u64| DataFileEntry {
1829 path: path.to_string(),
1830 record_count,
1831 file_size_bytes: size,
1832 centroid_b64: None,
1833 radius: None,
1834 hnsw_offset: None,
1835 hnsw_len: None,
1836 vector_column: None,
1837 vector_dim: None,
1838 extra_vector_indexes: vec![],
1839 index_status: IndexStatus::Ready,
1840 index_error: None,
1841 batch_id: None,
1842 embedding_model: None,
1843 partition_value: None,
1844 deletion_vector: None,
1845 first_row_id: None,
1846 column_stats: None,
1847 sequence_number: 0,
1848 };
1849 let entries = vec![
1850 make_entry(&path_dom, n_dom, bytes_dom.len() as u64),
1851 make_entry(&path_small, n_small, bytes_small.len() as u64),
1852 ];
1853
1854 let executor = CompactionExecutor::new(store.clone(), policy.clone())
1855 .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1856 let merged = executor
1857 .compact_incremental(&entries, "data/merged_force_ivfpq.parquet")
1858 .await
1859 .expect("compact_incremental() with ForceIvfPq must fall back to compact() cleanly");
1860
1861 let merged_bytes = store.get("data/merged_force_ivfpq.parquet").await.unwrap();
1862 let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1863 reader.verify_integrity().unwrap();
1864
1865 match reader.load_any_index().unwrap() {
1866 ailake_index::AnyIndex::IvfPq(_) => {}
1867 ailake_index::AnyIndex::Hnsw(_) => panic!(
1868 "ForceIvfPq was silently satisfied with HNSW instead — merged.record_count={}",
1869 merged.record_count
1870 ),
1871 }
1872 }
1873
1874 #[tokio::test]
1875 async fn compact_incremental_falls_back_when_no_dominant() {
1876 use ailake_core::{VectorMetric, VectorPrecision};
1877 use ailake_store::LocalStore;
1878 use arrow_array::{Int32Array, RecordBatch};
1879 use arrow_schema::{DataType, Field, Schema};
1880 use std::sync::Arc;
1881 use tempfile::TempDir;
1882
1883 let dir = TempDir::new().unwrap();
1884 let store = Arc::new(LocalStore::new(dir.path()));
1885 let policy = VectorStoragePolicy {
1886 column_name: "embedding".into(),
1887 dim: 4,
1888 metric: VectorMetric::Cosine,
1889 precision: VectorPrecision::F16,
1890 pq: None,
1891 keep_raw_for_reranking: true,
1892 pre_normalize: false,
1893 hnsw_m: None,
1894 hnsw_ef_construction: None,
1895 ivf_residual: false,
1896 embedding_model: None,
1897 modality: None,
1898 partition_by: None,
1899 partition_value: None,
1900 partition_column_type: None,
1901 partition_fields: vec![],
1902 };
1903
1904 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1905
1906 let make_batch = |ids: Vec<i32>, embs: Vec<Vec<f32>>| {
1908 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1909 .unwrap();
1910 AilakeFileWriter::new(policy.clone())
1911 .write(&batch, &embs)
1912 .unwrap()
1913 };
1914
1915 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]];
1916 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]];
1917 let bytes_a = make_batch(vec![0, 1], embs_a);
1918 let bytes_b = make_batch(vec![2, 3], embs_b);
1919
1920 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1921 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1922
1923 let entries = vec![
1924 DataFileEntry {
1925 path: "data/a.parquet".into(),
1926 record_count: 2,
1927 file_size_bytes: bytes_a.len() as u64,
1928 centroid_b64: None,
1929 radius: None,
1930 hnsw_offset: None,
1931 hnsw_len: None,
1932 vector_column: None,
1933 vector_dim: None,
1934 extra_vector_indexes: vec![],
1935 index_status: IndexStatus::Ready,
1936 index_error: None,
1937 batch_id: None,
1938 embedding_model: None,
1939 partition_value: None,
1940 deletion_vector: None,
1941 first_row_id: None,
1942 column_stats: None,
1943 sequence_number: 0,
1944 },
1945 DataFileEntry {
1946 path: "data/b.parquet".into(),
1947 record_count: 2,
1948 file_size_bytes: bytes_b.len() as u64,
1949 centroid_b64: None,
1950 radius: None,
1951 hnsw_offset: None,
1952 hnsw_len: None,
1953 vector_column: None,
1954 vector_dim: None,
1955 extra_vector_indexes: vec![],
1956 index_status: IndexStatus::Ready,
1957 index_error: None,
1958 batch_id: None,
1959 embedding_model: None,
1960 partition_value: None,
1961 deletion_vector: None,
1962 first_row_id: None,
1963 column_stats: None,
1964 sequence_number: 0,
1965 },
1966 ];
1967
1968 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1969 let merged = executor
1971 .compact_incremental(&entries, "data/merged.parquet")
1972 .await
1973 .unwrap();
1974
1975 assert_eq!(merged.record_count, 4);
1976
1977 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1978 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1979 reader.verify_integrity().unwrap();
1980 }
1981
1982 #[tokio::test]
1983 async fn compact_deferred_produces_parquet_only_file() {
1984 use ailake_catalog::HadoopCatalog;
1985 use ailake_core::{VectorMetric, VectorPrecision};
1986 use ailake_store::LocalStore;
1987 use arrow_array::{Int32Array, RecordBatch};
1988 use arrow_schema::{DataType, Field, Schema};
1989 use std::sync::Arc;
1990 use tempfile::TempDir;
1991
1992 let dir = TempDir::new().unwrap();
1993 let store = Arc::new(LocalStore::new(dir.path()));
1994 let catalog_dir = TempDir::new().unwrap();
1995 let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
1996 let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
1997 let table = TableIdent {
1998 namespace: "ns".into(),
1999 name: "tbl".into(),
2000 };
2001
2002 let policy = VectorStoragePolicy {
2003 column_name: "embedding".into(),
2004 dim: 4,
2005 metric: VectorMetric::Cosine,
2006 precision: VectorPrecision::F16,
2007 pq: None,
2008 keep_raw_for_reranking: true,
2009 pre_normalize: false,
2010 hnsw_m: None,
2011 hnsw_ef_construction: None,
2012 ivf_residual: false,
2013 embedding_model: None,
2014 modality: None,
2015 partition_by: None,
2016 partition_value: None,
2017 partition_column_type: None,
2018 partition_fields: vec![],
2019 };
2020
2021 use ailake_catalog::TableProperties;
2022 catalog
2023 .create_table(
2024 &table,
2025 &TableProperties {
2026 policy: policy.clone(),
2027 extra: std::collections::HashMap::new(),
2028 format_version: 2,
2029 partition_column_type: None,
2030 },
2031 )
2032 .await
2033 .unwrap();
2034
2035 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2036 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]];
2037 let batch_a = RecordBatch::try_new(
2038 schema.clone(),
2039 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2040 )
2041 .unwrap();
2042 let bytes_a = AilakeFileWriter::new(policy.clone())
2043 .write(&batch_a, &embs_a)
2044 .unwrap();
2045 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
2046
2047 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]];
2048 let batch_b = RecordBatch::try_new(
2049 schema.clone(),
2050 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2051 )
2052 .unwrap();
2053 let bytes_b = AilakeFileWriter::new(policy.clone())
2054 .write(&batch_b, &embs_b)
2055 .unwrap();
2056 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
2057
2058 let entries = vec![
2059 DataFileEntry {
2060 path: "data/a.parquet".into(),
2061 record_count: 2,
2062 file_size_bytes: bytes_a.len() as u64,
2063 centroid_b64: None,
2064 radius: None,
2065 hnsw_offset: None,
2066 hnsw_len: None,
2067 vector_column: None,
2068 vector_dim: None,
2069 extra_vector_indexes: vec![],
2070 index_status: IndexStatus::Ready,
2071 index_error: None,
2072 batch_id: None,
2073 embedding_model: None,
2074 partition_value: None,
2075 deletion_vector: None,
2076 first_row_id: None,
2077 column_stats: None,
2078 sequence_number: 0,
2079 },
2080 DataFileEntry {
2081 path: "data/b.parquet".into(),
2082 record_count: 2,
2083 file_size_bytes: bytes_b.len() as u64,
2084 centroid_b64: None,
2085 radius: None,
2086 hnsw_offset: None,
2087 hnsw_len: None,
2088 vector_column: None,
2089 vector_dim: None,
2090 extra_vector_indexes: vec![],
2091 index_status: IndexStatus::Ready,
2092 index_error: None,
2093 batch_id: None,
2094 embedding_model: None,
2095 partition_value: None,
2096 deletion_vector: None,
2097 first_row_id: None,
2098 column_stats: None,
2099 sequence_number: 0,
2100 },
2101 ];
2102
2103 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2104 let entry = executor
2105 .compact_deferred(&entries, "data/merged.parquet", catalog.clone(), &table)
2106 .await
2107 .unwrap();
2108
2109 assert_eq!(entry.index_status, IndexStatus::Indexing);
2111 assert_eq!(entry.record_count, 4);
2112
2113 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
2115 let pq_reader = ailake_parquet::ParquetVectorReader::new(merged_bytes, "embedding");
2116 let count = pq_reader.record_count().unwrap();
2117 assert_eq!(count, 4);
2118 }
2119
2120 #[tokio::test]
2128 async fn compact_aggregates_batch_ids_from_sources() {
2129 use ailake_core::{VectorMetric, VectorPrecision};
2130 use ailake_store::LocalStore;
2131 use arrow_array::{Int32Array, RecordBatch};
2132 use arrow_schema::{DataType, Field, Schema};
2133 use std::sync::Arc;
2134 use tempfile::TempDir;
2135
2136 let dir = TempDir::new().unwrap();
2137 let store = Arc::new(LocalStore::new(dir.path()));
2138 let policy = VectorStoragePolicy {
2139 column_name: "embedding".into(),
2140 dim: 4,
2141 metric: VectorMetric::Cosine,
2142 precision: VectorPrecision::F16,
2143 pq: None,
2144 keep_raw_for_reranking: true,
2145 pre_normalize: false,
2146 hnsw_m: None,
2147 hnsw_ef_construction: None,
2148 ivf_residual: false,
2149 embedding_model: None,
2150 modality: None,
2151 partition_by: None,
2152 partition_value: None,
2153 partition_column_type: None,
2154 partition_fields: vec![],
2155 };
2156
2157 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2158 let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0]];
2159 let batch_a =
2160 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![0i32]))])
2161 .unwrap();
2162 let bytes_a = AilakeFileWriter::new(policy.clone())
2163 .write(&batch_a, &embs_a)
2164 .unwrap();
2165 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
2166
2167 let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 1.0, 0.0, 0.0]];
2168 let batch_b =
2169 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1i32]))])
2170 .unwrap();
2171 let bytes_b = AilakeFileWriter::new(policy.clone())
2172 .write(&batch_b, &embs_b)
2173 .unwrap();
2174 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
2175
2176 let embs_c: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
2179 let batch_c =
2180 RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![2i32]))]).unwrap();
2181 let bytes_c = AilakeFileWriter::new(policy.clone())
2182 .write(&batch_c, &embs_c)
2183 .unwrap();
2184 store.put("data/c.parquet", bytes_c.clone()).await.unwrap();
2185
2186 fn entry(path: &str, size: u64, batch_id: Option<&str>) -> DataFileEntry {
2187 DataFileEntry {
2188 path: path.into(),
2189 record_count: 1,
2190 file_size_bytes: size,
2191 centroid_b64: None,
2192 radius: None,
2193 hnsw_offset: None,
2194 hnsw_len: None,
2195 vector_column: None,
2196 vector_dim: None,
2197 extra_vector_indexes: vec![],
2198 index_status: IndexStatus::Ready,
2199 index_error: None,
2200 batch_id: batch_id.map(String::from),
2201 embedding_model: None,
2202 partition_value: None,
2203 deletion_vector: None,
2204 first_row_id: None,
2205 column_stats: None,
2206 sequence_number: 0,
2207 }
2208 }
2209 let entries = vec![
2210 entry("data/a.parquet", bytes_a.len() as u64, Some("k-a")),
2211 entry("data/b.parquet", bytes_b.len() as u64, Some("k-b")),
2212 entry("data/c.parquet", bytes_c.len() as u64, None),
2213 ];
2214
2215 let executor = CompactionExecutor::new(store.clone(), policy);
2216 let merged = executor
2217 .compact(&entries, "data/merged.parquet")
2218 .await
2219 .unwrap();
2220
2221 assert_eq!(merged.record_count, 3);
2222 assert_eq!(
2223 merged.batch_ids(),
2224 vec!["k-a".to_string(), "k-b".to_string()],
2225 "merged file must carry every source's idempotency key, none invented"
2226 );
2227 }
2228
2229 #[tokio::test]
2236 async fn run_preserves_untouched_files_outside_compaction_pass() {
2237 use ailake_catalog::HadoopCatalog;
2238 use ailake_core::{VectorMetric, VectorPrecision};
2239 use ailake_store::LocalStore;
2240 use arrow_array::{Int32Array, RecordBatch};
2241 use arrow_schema::{DataType, Field, Schema};
2242 use std::sync::Arc;
2243 use tempfile::TempDir;
2244
2245 let dir = TempDir::new().unwrap();
2246 let store = Arc::new(LocalStore::new(dir.path()));
2247 let catalog_dir = TempDir::new().unwrap();
2248 let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
2249 let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
2250 let table = TableIdent {
2251 namespace: "ns".into(),
2252 name: "tbl".into(),
2253 };
2254
2255 let policy = VectorStoragePolicy {
2256 column_name: "embedding".into(),
2257 dim: 4,
2258 metric: VectorMetric::Cosine,
2259 precision: VectorPrecision::F16,
2260 pq: None,
2261 keep_raw_for_reranking: true,
2262 pre_normalize: false,
2263 hnsw_m: None,
2264 hnsw_ef_construction: None,
2265 ivf_residual: false,
2266 embedding_model: None,
2267 modality: None,
2268 partition_by: None,
2269 partition_value: None,
2270 partition_column_type: None,
2271 partition_fields: vec![],
2272 };
2273
2274 use ailake_catalog::TableProperties;
2275 catalog
2276 .create_table(
2277 &table,
2278 &TableProperties {
2279 policy: policy.clone(),
2280 extra: std::collections::HashMap::new(),
2281 format_version: 2,
2282 partition_column_type: None,
2283 },
2284 )
2285 .await
2286 .unwrap();
2287
2288 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2289
2290 let write_file = |path: &str, ids: Vec<i32>, embs: Vec<Vec<f32>>| {
2292 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
2293 .unwrap();
2294 let bytes = AilakeFileWriter::new(policy.clone())
2295 .write(&batch, &embs)
2296 .unwrap();
2297 (path.to_string(), bytes)
2298 };
2299
2300 let (path_a, bytes_a) = write_file(
2301 "data/small_a.parquet",
2302 vec![0, 1],
2303 vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
2304 );
2305 let (path_b, bytes_b) = write_file(
2306 "data/small_b.parquet",
2307 vec![2, 3],
2308 vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]],
2309 );
2310 let (path_big, bytes_big) = write_file(
2313 "data/big.parquet",
2314 vec![4, 5],
2315 vec![vec![1.0, 1.0, 0.0, 0.0], vec![0.0, 1.0, 1.0, 0.0]],
2316 );
2317
2318 for (path, bytes) in [
2319 (&path_a, &bytes_a),
2320 (&path_b, &bytes_b),
2321 (&path_big, &bytes_big),
2322 ] {
2323 store.put(path, bytes.clone()).await.unwrap();
2324 }
2325
2326 let make_entry = |path: &str, size: u64| DataFileEntry {
2327 path: path.to_string(),
2328 record_count: 2,
2329 file_size_bytes: size,
2330 centroid_b64: Some("AAAA".into()),
2335 radius: None,
2336 hnsw_offset: None,
2337 hnsw_len: None,
2338 vector_column: None,
2339 vector_dim: None,
2340 extra_vector_indexes: vec![],
2341 index_status: IndexStatus::Ready,
2342 index_error: None,
2343 batch_id: None,
2344 embedding_model: None,
2345 partition_value: None,
2346 deletion_vector: None,
2347 first_row_id: None,
2348 column_stats: None,
2349 sequence_number: 0,
2350 };
2351
2352 let initial_snap_id = ailake_catalog::new_snapshot_id();
2353 let initial_snapshot = NewSnapshot {
2354 snapshot_id: initial_snap_id,
2355 parent_snapshot_id: None,
2356 files: vec![
2357 make_entry(&path_a, 500),
2358 make_entry(&path_b, 500),
2359 make_entry(&path_big, 200_000_000), ],
2361 operation: SnapshotOperation::Append,
2362 iceberg_schema: None,
2363 extra_properties: std::collections::HashMap::new(),
2364 bloom_filters: vec![],
2365 equality_delete_files: vec![],
2366 };
2367 catalog
2368 .commit_snapshot(&table, initial_snapshot)
2369 .await
2370 .unwrap();
2371
2372 let planner = CompactionPlanner::new(CompactionConfig {
2373 min_files_to_compact: 2,
2374 target_file_size_bytes: 1000,
2375 index_strategy: CompactionIndexStrategy::ForceHnsw,
2376 max_files_per_pass: 20,
2377 });
2378 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2379
2380 let merged = executor
2381 .run(&planner, &table, catalog.clone(), "data")
2382 .await
2383 .unwrap()
2384 .expect("compaction should have run — 2 eligible small files");
2385
2386 let files_after = catalog.list_files(&table, None).await.unwrap();
2387 let paths_after: Vec<&str> = files_after.iter().map(|f| f.path.as_str()).collect();
2388
2389 assert!(
2390 paths_after.contains(&path_big.as_str()),
2391 "BUG: untouched 'big.parquet' vanished after run() — files_after={paths_after:?}"
2392 );
2393 assert!(
2394 paths_after.contains(&merged.path.as_str()),
2395 "merged output file must be present — files_after={paths_after:?}"
2396 );
2397
2398 let meta_dir = catalog_dir.path().join("ns/tbl/metadata");
2405 let latest_metadata = std::fs::read_dir(&meta_dir)
2406 .unwrap()
2407 .filter_map(|e| e.ok())
2408 .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
2409 .max_by_key(|e| e.metadata().unwrap().modified().unwrap())
2410 .expect("metadata.json must exist after commit");
2411 let json: serde_json::Value =
2412 serde_json::from_slice(&std::fs::read(latest_metadata.path()).unwrap()).unwrap();
2413 let last_snapshot = json["snapshots"].as_array().unwrap().last().unwrap();
2414 assert_eq!(
2415 last_snapshot["parent-snapshot-id"].as_i64(),
2416 Some(initial_snap_id),
2417 "compaction's Replace snapshot must chain to the pre-compaction snapshot, not be orphaned"
2418 );
2419 assert!(
2420 !paths_after.contains(&path_a.as_str()) && !paths_after.contains(&path_b.as_str()),
2421 "compacted input files must no longer be listed — files_after={paths_after:?}"
2422 );
2423 assert_eq!(
2424 files_after.len(),
2425 2,
2426 "expected exactly [big.parquet, merged] — files_after={paths_after:?}"
2427 );
2428 }
2429
2430 #[tokio::test]
2436 async fn compact_preserves_rows_from_footerless_file() {
2437 use ailake_core::{VectorMetric, VectorPrecision};
2438 use ailake_store::LocalStore;
2439 use arrow_array::{Int32Array, RecordBatch};
2440 use arrow_schema::{DataType, Field, Schema};
2441 use std::sync::Arc;
2442 use tempfile::TempDir;
2443
2444 let dir = TempDir::new().unwrap();
2445 let store = Arc::new(LocalStore::new(dir.path()));
2446 let policy = VectorStoragePolicy {
2447 column_name: "embedding".into(),
2448 dim: 4,
2449 metric: VectorMetric::Cosine,
2450 precision: VectorPrecision::F16,
2451 pq: None,
2452 keep_raw_for_reranking: true,
2453 pre_normalize: false,
2454 hnsw_m: None,
2455 hnsw_ef_construction: None,
2456 ivf_residual: false,
2457 embedding_model: None,
2458 modality: None,
2459 partition_by: None,
2460 partition_value: None,
2461 partition_column_type: None,
2462 partition_fields: vec![],
2463 };
2464
2465 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2466
2467 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]];
2469 let batch_native = RecordBatch::try_new(
2470 schema.clone(),
2471 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2472 )
2473 .unwrap();
2474 let bytes_native = AilakeFileWriter::new(policy.clone())
2475 .write(&batch_native, &embs_native)
2476 .unwrap();
2477 store
2478 .put("data/native.parquet", bytes_native.clone())
2479 .await
2480 .unwrap();
2481
2482 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]];
2486 let batch_foreign = RecordBatch::try_new(
2487 schema.clone(),
2488 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2489 )
2490 .unwrap();
2491 let bytes_foreign = AilakeFileWriter::new(policy.clone())
2492 .write_parquet_only(&batch_foreign, &embs_foreign)
2493 .unwrap();
2494 store
2495 .put("data/foreign.parquet", bytes_foreign.clone())
2496 .await
2497 .unwrap();
2498
2499 let reader_foreign = AilakeFileReader::new(bytes_foreign.clone(), "embedding", 4);
2500 assert!(
2501 !reader_foreign.is_ailake_file(),
2502 "sanity: write_parquet_only must not embed an AILK footer"
2503 );
2504
2505 let entries = vec![
2506 DataFileEntry {
2507 path: "data/native.parquet".into(),
2508 record_count: 2,
2509 file_size_bytes: bytes_native.len() as u64,
2510 centroid_b64: None,
2511 radius: None,
2512 hnsw_offset: None,
2513 hnsw_len: None,
2514 vector_column: None,
2515 vector_dim: None,
2516 extra_vector_indexes: vec![],
2517 index_status: IndexStatus::Ready,
2518 index_error: None,
2519 batch_id: None,
2520 embedding_model: None,
2521 partition_value: None,
2522 deletion_vector: None,
2523 first_row_id: None,
2524 column_stats: None,
2525 sequence_number: 0,
2526 },
2527 DataFileEntry {
2528 path: "data/foreign.parquet".into(),
2529 record_count: 2,
2530 file_size_bytes: bytes_foreign.len() as u64,
2531 centroid_b64: None,
2532 radius: None,
2533 hnsw_offset: None,
2534 hnsw_len: None,
2535 vector_column: None,
2536 vector_dim: None,
2537 extra_vector_indexes: vec![],
2538 index_status: IndexStatus::Ready,
2539 index_error: None,
2540 batch_id: None,
2541 embedding_model: None,
2542 partition_value: None,
2543 deletion_vector: None,
2544 first_row_id: None,
2545 column_stats: None,
2546 sequence_number: 0,
2547 },
2548 ];
2549
2550 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2552 let merged = executor
2553 .compact(&entries, "data/merged_full.parquet")
2554 .await
2555 .unwrap();
2556 assert_eq!(
2557 merged.record_count, 4,
2558 "compact() must include all 4 rows — 2 native + 2 from the footerless file"
2559 );
2560 let merged_bytes = store.get("data/merged_full.parquet").await.unwrap();
2561 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
2562 reader.verify_integrity().unwrap();
2563
2564 let merged_inc = executor
2567 .compact_incremental(&entries, "data/merged_inc.parquet")
2568 .await
2569 .unwrap();
2570 assert_eq!(
2571 merged_inc.record_count, 4,
2572 "compact_incremental() must include all 4 rows — 2 native + 2 from the footerless file"
2573 );
2574 }
2575
2576 #[tokio::test]
2579 async fn compact_preserves_search_results() {
2580 use crate::scanner::SearchConfig;
2581 use ailake_catalog::HadoopCatalog;
2582 use ailake_core::{VectorMetric, VectorPrecision};
2583 use ailake_store::LocalStore;
2584 use arrow_array::{Int32Array, RecordBatch};
2585 use arrow_schema::{DataType, Field, Schema};
2586 use std::sync::Arc;
2587 use tempfile::TempDir;
2588 let dir = TempDir::new().unwrap();
2589 let store = Arc::new(LocalStore::new(dir.path()));
2590 let policy = VectorStoragePolicy {
2591 column_name: "embedding".into(),
2592 dim: 4,
2593 metric: VectorMetric::Cosine,
2594 precision: VectorPrecision::F16,
2595 pq: None,
2596 keep_raw_for_reranking: true,
2597 pre_normalize: false,
2598 hnsw_m: None,
2599 hnsw_ef_construction: None,
2600 ivf_residual: false,
2601 embedding_model: None,
2602 modality: None,
2603 partition_by: None,
2604 partition_value: None,
2605 partition_column_type: None,
2606 partition_fields: vec![],
2607 };
2608
2609 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2610
2611 let embs_a: Vec<Vec<f32>> = vec![
2613 vec![1.0, 0.0, 0.0, 0.0],
2614 vec![0.0, 1.0, 0.0, 0.0],
2615 vec![0.0, 0.0, 1.0, 0.0],
2616 ];
2617 let batch_a = RecordBatch::try_new(
2618 schema.clone(),
2619 vec![Arc::new(Int32Array::from(vec![0i32, 1, 2]))],
2620 )
2621 .unwrap();
2622
2623 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]];
2625 let batch_b = RecordBatch::try_new(
2626 schema.clone(),
2627 vec![Arc::new(Int32Array::from(vec![3i32, 4]))],
2628 )
2629 .unwrap();
2630
2631 let bytes_a = AilakeFileWriter::new(policy.clone())
2632 .write(&batch_a, &embs_a)
2633 .unwrap();
2634 let bytes_b = AilakeFileWriter::new(policy.clone())
2635 .write(&batch_b, &embs_b)
2636 .unwrap();
2637
2638 store.put("data/a.parquet", bytes_a).await.unwrap();
2639 store.put("data/b.parquet", bytes_b).await.unwrap();
2640 let file_size_a = store.file_size("data/a.parquet").await.unwrap();
2641 let file_size_b = store.file_size("data/b.parquet").await.unwrap();
2642
2643 let catalog: Arc<dyn CatalogProvider> =
2646 Arc::new(HadoopCatalog::new(store.clone(), "warehouse"));
2647 let table = TableIdent::new("default", "compact_search_test");
2648
2649 catalog
2651 .create_table(
2652 &table,
2653 &ailake_catalog::TableProperties {
2654 policy: policy.clone(),
2655 extra: std::collections::HashMap::new(),
2656 format_version: 2,
2657 partition_column_type: None,
2658 },
2659 )
2660 .await
2661 .unwrap();
2662
2663 let meta = catalog.load_table(&table).await.unwrap();
2665 let parent_snapshot_id = meta.current_snapshot_id;
2666
2667 let entry_a = DataFileEntry {
2668 path: "data/a.parquet".into(),
2669 record_count: 3,
2670 file_size_bytes: file_size_a,
2671 centroid_b64: None,
2672 radius: None,
2673 hnsw_offset: None,
2674 hnsw_len: None,
2675 vector_column: Some("embedding".into()),
2676 vector_dim: Some(4),
2677 extra_vector_indexes: vec![],
2678 index_status: IndexStatus::Ready,
2679 index_error: None,
2680 batch_id: None,
2681 embedding_model: None,
2682 partition_value: None,
2683 deletion_vector: None,
2684 first_row_id: None,
2685 column_stats: None,
2686 sequence_number: 0,
2687 };
2688 let entry_b = DataFileEntry {
2689 path: "data/b.parquet".into(),
2690 record_count: 2,
2691 file_size_bytes: file_size_b,
2692 centroid_b64: None,
2693 radius: None,
2694 hnsw_offset: None,
2695 hnsw_len: None,
2696 vector_column: Some("embedding".into()),
2697 vector_dim: Some(4),
2698 extra_vector_indexes: vec![],
2699 index_status: IndexStatus::Ready,
2700 index_error: None,
2701 batch_id: None,
2702 embedding_model: None,
2703 partition_value: None,
2704 deletion_vector: None,
2705 first_row_id: None,
2706 column_stats: None,
2707 sequence_number: 0,
2708 };
2709
2710 catalog
2711 .commit_snapshot(
2712 &table,
2713 NewSnapshot {
2714 snapshot_id: 200,
2715 parent_snapshot_id,
2716 files: vec![entry_a, entry_b],
2717 operation: SnapshotOperation::Append,
2718 iceberg_schema: None,
2719 extra_properties: std::collections::HashMap::new(),
2720 bloom_filters: vec![],
2721 equality_delete_files: vec![],
2722 },
2723 )
2724 .await
2725 .unwrap();
2726
2727 let query_before = vec![1.0f32, 0.0, 0.0, 0.0];
2729 let config = SearchConfig {
2730 top_k: 5,
2731 ef_search: 100,
2732 pruning_threshold: f32::INFINITY,
2733 ..Default::default()
2734 };
2735 let before = crate::scanner::search(
2736 &table,
2737 &query_before,
2738 config.clone(),
2739 "embedding",
2740 4,
2741 catalog.clone(),
2742 store.clone(),
2743 )
2744 .await
2745 .unwrap();
2746 assert!(
2747 !before.is_empty(),
2748 "search before compaction must return results"
2749 );
2750 let before_dists: Vec<f32> = before.iter().map(|r| r.distance).collect();
2751
2752 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2754 let files_to_compact = catalog.list_files(&table, None).await.unwrap();
2755 let merged = executor
2756 .compact(&files_to_compact, "data/merged.parquet")
2757 .await
2758 .unwrap();
2759
2760 let meta2 = catalog.load_table(&table).await.unwrap();
2762 let merged_entry = DataFileEntry {
2763 path: merged.path.clone(),
2764 record_count: merged.record_count,
2765 file_size_bytes: merged.file_size_bytes,
2766 centroid_b64: merged.centroid_b64,
2767 radius: merged.radius,
2768 hnsw_offset: merged.hnsw_offset,
2769 hnsw_len: merged.hnsw_len,
2770 vector_column: merged.vector_column,
2771 vector_dim: merged.vector_dim,
2772 extra_vector_indexes: merged.extra_vector_indexes,
2773 index_status: merged.index_status,
2774 index_error: merged.index_error,
2775 batch_id: merged.batch_id,
2776 embedding_model: merged.embedding_model,
2777 partition_value: merged.partition_value,
2778 deletion_vector: merged.deletion_vector,
2779 first_row_id: None,
2780 column_stats: None,
2781 sequence_number: 0,
2782 };
2783 catalog
2784 .commit_snapshot(
2785 &table,
2786 NewSnapshot {
2787 snapshot_id: 201,
2788 parent_snapshot_id: meta2.current_snapshot_id,
2789 files: vec![merged_entry],
2790 operation: SnapshotOperation::Replace,
2791 iceberg_schema: None,
2792 extra_properties: std::collections::HashMap::new(),
2793 bloom_filters: vec![],
2794 equality_delete_files: vec![],
2795 },
2796 )
2797 .await
2798 .unwrap();
2799
2800 let after = crate::scanner::search(
2802 &table,
2803 &query_before,
2804 config.clone(),
2805 "embedding",
2806 4,
2807 catalog,
2808 store,
2809 )
2810 .await
2811 .unwrap();
2812 assert!(
2813 !after.is_empty(),
2814 "search after compaction must return results"
2815 );
2816 let after_dists: Vec<f32> = after.iter().map(|r| r.distance).collect();
2817
2818 assert_eq!(
2820 before.len(),
2821 after.len(),
2822 "number of results should match before ({}) and after ({}) compaction",
2823 before.len(),
2824 after.len(),
2825 );
2826 let mut total_diff = 0.0f32;
2831 for (i, (bd, ad)) in before_dists.iter().zip(after_dists.iter()).enumerate() {
2832 let diff = (bd - ad).abs();
2833 total_diff += diff;
2834 assert!(
2835 diff < 0.1,
2836 "distance mismatch at position {i}: before={bd}, after={ad}, diff={diff}"
2837 );
2838 }
2839 let avg_diff = total_diff / before_dists.len() as f32;
2840 assert!(
2841 avg_diff < 0.05,
2842 "average distance error across all results too high: {avg_diff}"
2843 );
2844 }
2845}