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 })
948 .collect();
949 assert!(planner.plan(&files).is_empty());
950 }
951
952 #[test]
953 fn plan_selects_small_files() {
954 let planner = CompactionPlanner::new(CompactionConfig {
955 min_files_to_compact: 2,
956 target_file_size_bytes: 1000,
957 ..Default::default()
958 });
959 let files = vec![
960 DataFileEntry {
961 path: "small.parquet".into(),
962 record_count: 5,
963 file_size_bytes: 500,
964 centroid_b64: Some("AAAA".into()),
965 radius: None,
966 hnsw_offset: None,
967 hnsw_len: None,
968 vector_column: None,
969 vector_dim: None,
970 extra_vector_indexes: vec![],
971 index_status: IndexStatus::Ready,
972 index_error: None,
973 batch_id: None,
974 embedding_model: None,
975 partition_value: None,
976 deletion_vector: None,
977 first_row_id: None,
978 column_stats: None,
979 },
980 DataFileEntry {
981 path: "large.parquet".into(),
982 record_count: 5000,
983 file_size_bytes: 200_000_000,
984 centroid_b64: Some("AAAA".into()),
985 radius: None,
986 hnsw_offset: None,
987 hnsw_len: None,
988 vector_column: None,
989 vector_dim: None,
990 extra_vector_indexes: vec![],
991 index_status: IndexStatus::Ready,
992 index_error: None,
993 batch_id: None,
994 embedding_model: None,
995 partition_value: None,
996 deletion_vector: None,
997 first_row_id: None,
998 column_stats: None,
999 },
1000 DataFileEntry {
1001 path: "also-small.parquet".into(),
1002 record_count: 5,
1003 file_size_bytes: 800,
1004 centroid_b64: Some("AAAA".into()),
1005 radius: None,
1006 hnsw_offset: None,
1007 hnsw_len: None,
1008 vector_column: None,
1009 vector_dim: None,
1010 extra_vector_indexes: vec![],
1011 index_status: IndexStatus::Ready,
1012 index_error: None,
1013 batch_id: None,
1014 embedding_model: None,
1015 partition_value: None,
1016 deletion_vector: None,
1017 first_row_id: None,
1018 column_stats: None,
1019 },
1020 ];
1021 let selected = planner.plan(&files);
1022 assert_eq!(selected.len(), 2);
1023 assert!(selected.iter().any(|f| f.path == "small.parquet"));
1024 assert!(selected.iter().any(|f| f.path == "also-small.parquet"));
1025 }
1026
1027 #[test]
1028 fn plan_respects_max_files_per_pass() {
1029 let planner = CompactionPlanner::new(CompactionConfig {
1030 min_files_to_compact: 2,
1031 target_file_size_bytes: 1_000_000,
1032 max_files_per_pass: 3,
1033 ..Default::default()
1034 });
1035 let files: Vec<DataFileEntry> = (0..5)
1036 .map(|i| DataFileEntry {
1037 path: format!("f{i}.parquet"),
1038 record_count: 10,
1039 file_size_bytes: 100 + i as u64 * 100,
1040 centroid_b64: Some("AAAA".into()),
1041 radius: None,
1042 hnsw_offset: None,
1043 hnsw_len: None,
1044 vector_column: None,
1045 vector_dim: None,
1046 extra_vector_indexes: vec![],
1047 index_status: IndexStatus::Ready,
1048 index_error: None,
1049 batch_id: None,
1050 embedding_model: None,
1051 partition_value: None,
1052 deletion_vector: None,
1053 first_row_id: None,
1054 column_stats: None,
1055 })
1056 .collect();
1057 let selected = planner.plan(&files);
1058 assert_eq!(selected.len(), 3);
1059 assert_eq!(selected[0].file_size_bytes, 100);
1060 assert_eq!(selected[1].file_size_bytes, 200);
1061 assert_eq!(selected[2].file_size_bytes, 300);
1062 }
1063
1064 #[test]
1065 fn plan_sorts_smallest_first() {
1066 let planner = CompactionPlanner::new(CompactionConfig {
1067 min_files_to_compact: 2,
1068 target_file_size_bytes: 10_000,
1069 max_files_per_pass: 4,
1070 ..Default::default()
1071 });
1072 let files = vec![
1073 DataFileEntry {
1074 path: "c.parquet".into(),
1075 record_count: 1,
1076 file_size_bytes: 300,
1077 centroid_b64: Some("AAAA".into()),
1078 radius: None,
1079 hnsw_offset: None,
1080 hnsw_len: None,
1081 vector_column: None,
1082 vector_dim: None,
1083 extra_vector_indexes: vec![],
1084 index_status: IndexStatus::Ready,
1085 index_error: None,
1086 batch_id: None,
1087 embedding_model: None,
1088 partition_value: None,
1089 deletion_vector: None,
1090 first_row_id: None,
1091 column_stats: None,
1092 },
1093 DataFileEntry {
1094 path: "a.parquet".into(),
1095 record_count: 1,
1096 file_size_bytes: 100,
1097 centroid_b64: Some("AAAA".into()),
1098 radius: None,
1099 hnsw_offset: None,
1100 hnsw_len: None,
1101 vector_column: None,
1102 vector_dim: None,
1103 extra_vector_indexes: vec![],
1104 index_status: IndexStatus::Ready,
1105 index_error: None,
1106 batch_id: None,
1107 embedding_model: None,
1108 partition_value: None,
1109 deletion_vector: None,
1110 first_row_id: None,
1111 column_stats: None,
1112 },
1113 DataFileEntry {
1114 path: "b.parquet".into(),
1115 record_count: 1,
1116 file_size_bytes: 200,
1117 centroid_b64: Some("AAAA".into()),
1118 radius: None,
1119 hnsw_offset: None,
1120 hnsw_len: None,
1121 vector_column: None,
1122 vector_dim: None,
1123 extra_vector_indexes: vec![],
1124 index_status: IndexStatus::Ready,
1125 index_error: None,
1126 batch_id: None,
1127 embedding_model: None,
1128 partition_value: None,
1129 deletion_vector: None,
1130 first_row_id: None,
1131 column_stats: None,
1132 },
1133 ];
1134 let selected = planner.plan(&files);
1135 assert_eq!(selected[0].file_size_bytes, 100);
1136 assert_eq!(selected[1].file_size_bytes, 200);
1137 assert_eq!(selected[2].file_size_bytes, 300);
1138 }
1139
1140 fn make_plan_entry(path: &str, size: u64, centroid_b64: Option<String>) -> DataFileEntry {
1141 DataFileEntry {
1142 path: path.to_string(),
1143 record_count: 10,
1144 file_size_bytes: size,
1145 centroid_b64,
1146 radius: None,
1147 hnsw_offset: None,
1148 hnsw_len: None,
1149 vector_column: None,
1150 vector_dim: None,
1151 extra_vector_indexes: vec![],
1152 index_status: IndexStatus::Ready,
1153 index_error: None,
1154 batch_id: None,
1155 embedding_model: None,
1156 partition_value: None,
1157 deletion_vector: None,
1158 first_row_id: None,
1159 column_stats: None,
1160 }
1161 }
1162
1163 #[test]
1164 fn plan_prioritizes_foreign_written_files_regardless_of_size() {
1165 let planner = CompactionPlanner::new(CompactionConfig {
1166 min_files_to_compact: 4, target_file_size_bytes: 1000,
1168 max_files_per_pass: 20,
1169 ..Default::default()
1170 });
1171 let files = vec![
1172 make_plan_entry("native_small.parquet", 500, Some("AAAA".into())),
1175 make_plan_entry("foreign_big.parquet", 200_000_000, None),
1178 ];
1179 let selected = planner.plan(&files);
1180 assert_eq!(
1181 selected.len(),
1182 1,
1183 "foreign file alone should trigger a pass even below min_files_to_compact"
1184 );
1185 assert_eq!(selected[0].path, "foreign_big.parquet");
1186 }
1187
1188 #[test]
1189 fn plan_merges_foreign_and_size_candidates_together() {
1190 let planner = CompactionPlanner::new(CompactionConfig {
1191 min_files_to_compact: 2,
1192 target_file_size_bytes: 1000,
1193 max_files_per_pass: 20,
1194 ..Default::default()
1195 });
1196 let files = vec![
1197 make_plan_entry("native_a.parquet", 300, Some("AAAA".into())),
1198 make_plan_entry("native_b.parquet", 400, Some("AAAA".into())),
1199 make_plan_entry("foreign.parquet", 200_000_000, None),
1200 ];
1201 let selected = planner.plan(&files);
1202 let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1203 assert_eq!(selected.len(), 3, "paths={paths:?}");
1204 assert_eq!(selected[0].path, "foreign.parquet", "paths={paths:?}");
1206 }
1207
1208 #[test]
1209 fn plan_sorts_foreign_files_smallest_first_too() {
1210 let planner = CompactionPlanner::new(CompactionConfig {
1214 min_files_to_compact: 1,
1215 target_file_size_bytes: 1000,
1216 max_files_per_pass: 2,
1217 ..Default::default()
1218 });
1219 let files = vec![
1220 make_plan_entry("foreign_huge.parquet", 500_000_000, None),
1221 make_plan_entry("foreign_small.parquet", 100, None),
1222 make_plan_entry("foreign_medium.parquet", 10_000_000, None),
1223 ];
1224 let selected = planner.plan(&files);
1225 let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1226 assert_eq!(selected.len(), 2, "max_files_per_pass=2, paths={paths:?}");
1227 assert_eq!(
1228 paths,
1229 vec!["foreign_small.parquet", "foreign_medium.parquet"],
1230 "foreign files must be size-sorted before truncation, cheapest first"
1231 );
1232 }
1233
1234 #[tokio::test]
1235 async fn compact_merges_two_files() {
1236 use ailake_core::{VectorMetric, VectorPrecision};
1237 use ailake_store::LocalStore;
1238 use arrow_array::{Int32Array, RecordBatch};
1239 use arrow_schema::{DataType, Field, Schema};
1240 use std::sync::Arc;
1241 use tempfile::TempDir;
1242
1243 let dir = TempDir::new().unwrap();
1244 let store = Arc::new(LocalStore::new(dir.path()));
1245 let policy = VectorStoragePolicy {
1246 column_name: "embedding".into(),
1247 dim: 4,
1248 metric: VectorMetric::Cosine,
1249 precision: VectorPrecision::F16,
1250 pq: None,
1251 keep_raw_for_reranking: true,
1252 pre_normalize: false,
1253 hnsw_m: None,
1254 hnsw_ef_construction: None,
1255 ivf_residual: false,
1256 embedding_model: None,
1257 modality: None,
1258 partition_by: None,
1259 partition_value: None,
1260 partition_column_type: None,
1261 partition_fields: vec![],
1262 };
1263
1264 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1265 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]];
1266 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]];
1267
1268 let batch_a = RecordBatch::try_new(
1269 schema.clone(),
1270 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1271 )
1272 .unwrap();
1273 let batch_b = RecordBatch::try_new(
1274 schema.clone(),
1275 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1276 )
1277 .unwrap();
1278
1279 let writer_a = AilakeFileWriter::new(policy.clone());
1280 let bytes_a = writer_a.write(&batch_a, &embs_a).unwrap();
1281 let writer_b = AilakeFileWriter::new(policy.clone());
1282 let bytes_b = writer_b.write(&batch_b, &embs_b).unwrap();
1283
1284 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1285 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1286
1287 let entries = vec![
1288 DataFileEntry {
1289 path: "data/a.parquet".into(),
1290 record_count: 2,
1291 file_size_bytes: bytes_a.len() as u64,
1292 centroid_b64: None,
1293 radius: None,
1294 hnsw_offset: None,
1295 hnsw_len: None,
1296 vector_column: None,
1297 vector_dim: None,
1298 extra_vector_indexes: vec![],
1299 index_status: IndexStatus::Ready,
1300 index_error: None,
1301 batch_id: None,
1302 embedding_model: None,
1303 partition_value: None,
1304 deletion_vector: None,
1305 first_row_id: None,
1306 column_stats: None,
1307 },
1308 DataFileEntry {
1309 path: "data/b.parquet".into(),
1310 record_count: 2,
1311 file_size_bytes: bytes_b.len() as u64,
1312 centroid_b64: None,
1313 radius: None,
1314 hnsw_offset: None,
1315 hnsw_len: None,
1316 vector_column: None,
1317 vector_dim: None,
1318 extra_vector_indexes: vec![],
1319 index_status: IndexStatus::Ready,
1320 index_error: None,
1321 batch_id: None,
1322 embedding_model: None,
1323 partition_value: None,
1324 deletion_vector: None,
1325 first_row_id: None,
1326 column_stats: None,
1327 },
1328 ];
1329
1330 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1331 let merged = executor
1332 .compact(&entries, "data/merged.parquet")
1333 .await
1334 .unwrap();
1335
1336 assert_eq!(merged.record_count, 4);
1337 assert_eq!(merged.path, "data/merged.parquet");
1338
1339 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1340 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1341 reader.verify_integrity().unwrap();
1342 let (batch, embs) = reader.read_parquet().unwrap();
1343 assert_eq!(batch.num_rows(), 4);
1344 assert_eq!(embs.len(), 4);
1345 }
1346
1347 #[tokio::test]
1352 async fn compact_drops_deletion_vector_masked_rows() {
1353 use ailake_catalog::provider::DeletionVector;
1354 use ailake_core::{VectorMetric, VectorPrecision};
1355 use ailake_store::LocalStore;
1356 use arrow_array::{Int32Array, RecordBatch};
1357 use arrow_schema::{DataType, Field, Schema};
1358 use roaring::RoaringBitmap;
1359 use std::sync::Arc;
1360 use tempfile::TempDir;
1361
1362 let dir = TempDir::new().unwrap();
1363 let store = Arc::new(LocalStore::new(dir.path()));
1364 let policy = VectorStoragePolicy {
1365 column_name: "embedding".into(),
1366 dim: 4,
1367 metric: VectorMetric::Cosine,
1368 precision: VectorPrecision::F16,
1369 pq: None,
1370 keep_raw_for_reranking: true,
1371 pre_normalize: false,
1372 hnsw_m: None,
1373 hnsw_ef_construction: None,
1374 ivf_residual: false,
1375 embedding_model: None,
1376 modality: None,
1377 partition_by: None,
1378 partition_value: None,
1379 partition_column_type: None,
1380 partition_fields: vec![],
1381 };
1382
1383 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1384 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]];
1386 let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
1387
1388 let batch_a = RecordBatch::try_new(
1389 schema.clone(),
1390 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1391 )
1392 .unwrap();
1393 let batch_b =
1394 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![2i32]))])
1395 .unwrap();
1396
1397 let bytes_a = AilakeFileWriter::new(policy.clone())
1398 .write(&batch_a, &embs_a)
1399 .unwrap();
1400 let bytes_b = AilakeFileWriter::new(policy.clone())
1401 .write(&batch_b, &embs_b)
1402 .unwrap();
1403 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1404 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1405
1406 let mut bitmap = RoaringBitmap::new();
1408 bitmap.insert(0);
1409 let (puffin_bytes, offset, length) =
1410 crate::delete::PuffinWriter::write_single_dv(&bitmap, 1).unwrap();
1411 store.put("metadata/dv-1.dvd", puffin_bytes).await.unwrap();
1412
1413 let entry_a = DataFileEntry {
1414 path: "data/a.parquet".into(),
1415 record_count: 2,
1416 file_size_bytes: bytes_a.len() as u64,
1417 centroid_b64: None,
1418 radius: None,
1419 hnsw_offset: None,
1420 hnsw_len: None,
1421 vector_column: None,
1422 vector_dim: None,
1423 extra_vector_indexes: vec![],
1424 index_status: IndexStatus::Ready,
1425 index_error: None,
1426 batch_id: None,
1427 embedding_model: None,
1428 partition_value: None,
1429 deletion_vector: Some(DeletionVector {
1430 path: "metadata/dv-1.dvd".into(),
1431 offset,
1432 length,
1433 cardinality: 1,
1434 }),
1435 first_row_id: None,
1436 column_stats: None,
1437 };
1438 let entry_b = DataFileEntry {
1439 path: "data/b.parquet".into(),
1440 record_count: 1,
1441 file_size_bytes: bytes_b.len() as u64,
1442 centroid_b64: None,
1443 radius: None,
1444 hnsw_offset: None,
1445 hnsw_len: None,
1446 vector_column: None,
1447 vector_dim: None,
1448 extra_vector_indexes: vec![],
1449 index_status: IndexStatus::Ready,
1450 index_error: None,
1451 batch_id: None,
1452 embedding_model: None,
1453 partition_value: None,
1454 deletion_vector: None,
1455 first_row_id: None,
1456 column_stats: None,
1457 };
1458
1459 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1460 let merged = executor
1461 .compact(&[entry_a, entry_b], "data/merged.parquet")
1462 .await
1463 .unwrap();
1464
1465 assert_eq!(merged.record_count, 2);
1467 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1468 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1469 let (batch, embs) = reader.read_parquet().unwrap();
1470 assert_eq!(batch.num_rows(), 2);
1471 assert_eq!(embs.len(), 2);
1472 let ids: Vec<i32> = batch
1473 .column_by_name("id")
1474 .unwrap()
1475 .as_any()
1476 .downcast_ref::<Int32Array>()
1477 .unwrap()
1478 .values()
1479 .to_vec();
1480 assert_eq!(
1481 ids,
1482 vec![1, 2],
1483 "id=0 (deleted) must not survive compaction"
1484 );
1485 }
1486
1487 #[tokio::test]
1492 async fn compact_with_ivf_pq_strategy_does_not_crash_on_verify_integrity() {
1493 use ailake_core::{VectorMetric, VectorPrecision};
1494 use ailake_store::LocalStore;
1495 use arrow_array::{Int32Array, RecordBatch};
1496 use arrow_schema::{DataType, Field, Schema};
1497 use std::sync::Arc;
1498 use tempfile::TempDir;
1499
1500 let dir = TempDir::new().unwrap();
1501 let store = Arc::new(LocalStore::new(dir.path()));
1502 let dim = 8;
1503 let policy = VectorStoragePolicy {
1504 column_name: "embedding".into(),
1505 dim,
1506 metric: VectorMetric::Cosine,
1507 precision: VectorPrecision::F16,
1508 pq: None,
1509 keep_raw_for_reranking: true,
1510 pre_normalize: false,
1511 hnsw_m: None,
1512 hnsw_ef_construction: None,
1513 ivf_residual: false,
1514 embedding_model: None,
1515 modality: None,
1516 partition_by: None,
1517 partition_value: None,
1518 partition_column_type: None,
1519 partition_fields: vec![],
1520 };
1521
1522 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1523 let n_per_file = 30usize;
1524 let make_file = |path: &str, offset: i32| {
1525 let ids: Vec<i32> = (offset..offset + n_per_file as i32).collect();
1526 let embs: Vec<Vec<f32>> = ids
1527 .iter()
1528 .map(|&i| {
1529 (0..dim as i32)
1530 .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1531 .collect()
1532 })
1533 .collect();
1534 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1535 .unwrap();
1536 let bytes = AilakeFileWriter::new(policy.clone())
1537 .write(&batch, &embs)
1538 .unwrap();
1539 (path.to_string(), bytes)
1540 };
1541
1542 let (path_a, bytes_a) = make_file("data/ivfpq_a.parquet", 0);
1543 let (path_b, bytes_b) = make_file("data/ivfpq_b.parquet", n_per_file as i32);
1544 for (path, bytes) in [(&path_a, &bytes_a), (&path_b, &bytes_b)] {
1545 store.put(path, bytes.clone()).await.unwrap();
1546 }
1547
1548 let make_entry = |path: &str, size: u64| DataFileEntry {
1549 path: path.to_string(),
1550 record_count: n_per_file as u64,
1551 file_size_bytes: size,
1552 centroid_b64: None,
1553 radius: None,
1554 hnsw_offset: None,
1555 hnsw_len: None,
1556 vector_column: None,
1557 vector_dim: None,
1558 extra_vector_indexes: vec![],
1559 index_status: IndexStatus::Ready,
1560 index_error: None,
1561 batch_id: None,
1562 embedding_model: None,
1563 partition_value: None,
1564 deletion_vector: None,
1565 first_row_id: None,
1566 column_stats: None,
1567 };
1568 let entries = vec![
1569 make_entry(&path_a, bytes_a.len() as u64),
1570 make_entry(&path_b, bytes_b.len() as u64),
1571 ];
1572
1573 let executor = CompactionExecutor::new(store.clone(), policy.clone())
1574 .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1575 let merged = executor
1576 .compact(&entries, "data/ivfpq_merged.parquet")
1577 .await
1578 .expect("compact() with ForceIvfPq must not fail in verify_integrity()");
1579
1580 assert_eq!(merged.record_count, 2 * n_per_file as u64);
1581
1582 let merged_bytes = store.get("data/ivfpq_merged.parquet").await.unwrap();
1584 let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1585 reader
1586 .verify_integrity()
1587 .expect("verify_integrity() must handle an IVF-PQ-indexed file");
1588 }
1589
1590 #[tokio::test]
1591 async fn compact_incremental_merges_dominant_plus_small() {
1592 use ailake_core::{RowId, VectorMetric, VectorPrecision};
1593 use ailake_store::LocalStore;
1594 use arrow_array::{Int32Array, RecordBatch};
1595 use arrow_schema::{DataType, Field, Schema};
1596 use std::sync::Arc;
1597 use tempfile::TempDir;
1598
1599 let dir = TempDir::new().unwrap();
1600 let store = Arc::new(LocalStore::new(dir.path()));
1601 let policy = VectorStoragePolicy {
1602 column_name: "embedding".into(),
1603 dim: 4,
1604 metric: VectorMetric::Cosine,
1605 precision: VectorPrecision::F16,
1606 pq: None,
1607 keep_raw_for_reranking: true,
1608 pre_normalize: false,
1609 hnsw_m: None,
1610 hnsw_ef_construction: None,
1611 ivf_residual: false,
1612 embedding_model: None,
1613 modality: None,
1614 partition_by: None,
1615 partition_value: None,
1616 partition_column_type: None,
1617 partition_fields: vec![],
1618 };
1619
1620 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1621
1622 let embs_dom: Vec<Vec<f32>> = vec![
1624 vec![1.0, 0.0, 0.0, 0.0],
1625 vec![0.0, 1.0, 0.0, 0.0],
1626 vec![0.0, 0.0, 1.0, 0.0],
1627 vec![0.7, 0.7, 0.0, 0.0],
1628 vec![0.0, 0.7, 0.7, 0.0],
1629 vec![0.0, 0.0, 0.7, 0.7],
1630 ];
1631 let batch_dom = RecordBatch::try_new(
1632 schema.clone(),
1633 vec![Arc::new(Int32Array::from(vec![0i32, 1, 2, 3, 4, 5]))],
1634 )
1635 .unwrap();
1636
1637 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]];
1639 let batch_small = RecordBatch::try_new(
1640 schema.clone(),
1641 vec![Arc::new(Int32Array::from(vec![6i32, 7]))],
1642 )
1643 .unwrap();
1644
1645 let bytes_dom = AilakeFileWriter::new(policy.clone())
1646 .write(&batch_dom, &embs_dom)
1647 .unwrap();
1648 let bytes_small = AilakeFileWriter::new(policy.clone())
1649 .write(&batch_small, &embs_small)
1650 .unwrap();
1651
1652 store
1653 .put("data/dominant.parquet", bytes_dom.clone())
1654 .await
1655 .unwrap();
1656 store
1657 .put("data/small.parquet", bytes_small.clone())
1658 .await
1659 .unwrap();
1660
1661 let entries = vec![
1662 DataFileEntry {
1663 path: "data/dominant.parquet".into(),
1664 record_count: 6,
1665 file_size_bytes: bytes_dom.len() as u64,
1666 centroid_b64: None,
1667 radius: None,
1668 hnsw_offset: None,
1669 hnsw_len: None,
1670 vector_column: None,
1671 vector_dim: None,
1672 extra_vector_indexes: vec![],
1673 index_status: IndexStatus::Ready,
1674 index_error: None,
1675 batch_id: None,
1676 embedding_model: None,
1677 partition_value: None,
1678 deletion_vector: None,
1679 first_row_id: None,
1680 column_stats: None,
1681 },
1682 DataFileEntry {
1683 path: "data/small.parquet".into(),
1684 record_count: 2,
1685 file_size_bytes: bytes_small.len() as u64,
1686 centroid_b64: None,
1687 radius: None,
1688 hnsw_offset: None,
1689 hnsw_len: None,
1690 vector_column: None,
1691 vector_dim: None,
1692 extra_vector_indexes: vec![],
1693 index_status: IndexStatus::Ready,
1694 index_error: None,
1695 batch_id: None,
1696 embedding_model: None,
1697 partition_value: None,
1698 deletion_vector: None,
1699 first_row_id: None,
1700 column_stats: None,
1701 },
1702 ];
1703
1704 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1705 let merged = executor
1706 .compact_incremental(&entries, "data/merged.parquet")
1707 .await
1708 .unwrap();
1709
1710 assert_eq!(merged.record_count, 8);
1712 assert_eq!(merged.path, "data/merged.parquet");
1713
1714 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1716 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1717 reader.verify_integrity().unwrap();
1718
1719 let (batch, embs) = reader.read_parquet().unwrap();
1720 assert_eq!(batch.num_rows(), 8);
1721 assert_eq!(embs.len(), 8);
1722
1723 for f in &embs[..6] {
1725 assert_eq!(f.len(), 4);
1726 }
1727
1728 let hnsw = reader.load_index().unwrap();
1730 assert_eq!(hnsw.node_count(), 8);
1731
1732 let results = hnsw.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
1734 assert_eq!(results[0].0, RowId::new(0));
1735
1736 let results = hnsw.search(&[0.0, 0.0, 0.0, 1.0], 1, 50);
1739 assert_eq!(results[0].0, RowId::new(6));
1740 }
1741
1742 #[tokio::test]
1752 async fn compact_incremental_respects_force_ivf_pq_even_with_dominant_file() {
1753 use ailake_core::{VectorMetric, VectorPrecision};
1754 use ailake_store::LocalStore;
1755 use arrow_array::{Int32Array, RecordBatch};
1756 use arrow_schema::{DataType, Field, Schema};
1757 use std::sync::Arc;
1758 use tempfile::TempDir;
1759
1760 let dir = TempDir::new().unwrap();
1761 let store = Arc::new(LocalStore::new(dir.path()));
1762 let dim = 8;
1763 let policy = VectorStoragePolicy {
1764 column_name: "embedding".into(),
1765 dim,
1766 metric: VectorMetric::Cosine,
1767 precision: VectorPrecision::F16,
1768 pq: None,
1769 keep_raw_for_reranking: true,
1770 pre_normalize: false,
1771 hnsw_m: None,
1772 hnsw_ef_construction: None,
1773 ivf_residual: false,
1774 embedding_model: None,
1775 modality: None,
1776 partition_by: None,
1777 partition_value: None,
1778 partition_column_type: None,
1779 partition_fields: vec![],
1780 };
1781
1782 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1783 let make_file = |path: &str, offset: i32, n: usize| {
1784 let ids: Vec<i32> = (offset..offset + n as i32).collect();
1785 let embs: Vec<Vec<f32>> = ids
1786 .iter()
1787 .map(|&i| {
1788 (0..dim as i32)
1789 .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1790 .collect()
1791 })
1792 .collect();
1793 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1794 .unwrap();
1795 let bytes = AilakeFileWriter::new(policy.clone())
1798 .write(&batch, &embs)
1799 .unwrap();
1800 (path.to_string(), bytes, n as u64)
1801 };
1802
1803 let (path_dom, bytes_dom, n_dom) = make_file("data/dom.parquet", 0, 90);
1806 let (path_small, bytes_small, n_small) = make_file("data/small.parquet", 999, 30);
1807
1808 for (path, bytes) in [(&path_dom, &bytes_dom), (&path_small, &bytes_small)] {
1809 store.put(path, bytes.clone()).await.unwrap();
1810 }
1811
1812 let make_entry = |path: &str, record_count: u64, size: u64| DataFileEntry {
1813 path: path.to_string(),
1814 record_count,
1815 file_size_bytes: size,
1816 centroid_b64: None,
1817 radius: None,
1818 hnsw_offset: None,
1819 hnsw_len: None,
1820 vector_column: None,
1821 vector_dim: None,
1822 extra_vector_indexes: vec![],
1823 index_status: IndexStatus::Ready,
1824 index_error: None,
1825 batch_id: None,
1826 embedding_model: None,
1827 partition_value: None,
1828 deletion_vector: None,
1829 first_row_id: None,
1830 column_stats: None,
1831 };
1832 let entries = vec![
1833 make_entry(&path_dom, n_dom, bytes_dom.len() as u64),
1834 make_entry(&path_small, n_small, bytes_small.len() as u64),
1835 ];
1836
1837 let executor = CompactionExecutor::new(store.clone(), policy.clone())
1838 .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1839 let merged = executor
1840 .compact_incremental(&entries, "data/merged_force_ivfpq.parquet")
1841 .await
1842 .expect("compact_incremental() with ForceIvfPq must fall back to compact() cleanly");
1843
1844 let merged_bytes = store.get("data/merged_force_ivfpq.parquet").await.unwrap();
1845 let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1846 reader.verify_integrity().unwrap();
1847
1848 match reader.load_any_index().unwrap() {
1849 ailake_index::AnyIndex::IvfPq(_) => {}
1850 ailake_index::AnyIndex::Hnsw(_) => panic!(
1851 "ForceIvfPq was silently satisfied with HNSW instead — merged.record_count={}",
1852 merged.record_count
1853 ),
1854 }
1855 }
1856
1857 #[tokio::test]
1858 async fn compact_incremental_falls_back_when_no_dominant() {
1859 use ailake_core::{VectorMetric, VectorPrecision};
1860 use ailake_store::LocalStore;
1861 use arrow_array::{Int32Array, RecordBatch};
1862 use arrow_schema::{DataType, Field, Schema};
1863 use std::sync::Arc;
1864 use tempfile::TempDir;
1865
1866 let dir = TempDir::new().unwrap();
1867 let store = Arc::new(LocalStore::new(dir.path()));
1868 let policy = VectorStoragePolicy {
1869 column_name: "embedding".into(),
1870 dim: 4,
1871 metric: VectorMetric::Cosine,
1872 precision: VectorPrecision::F16,
1873 pq: None,
1874 keep_raw_for_reranking: true,
1875 pre_normalize: false,
1876 hnsw_m: None,
1877 hnsw_ef_construction: None,
1878 ivf_residual: false,
1879 embedding_model: None,
1880 modality: None,
1881 partition_by: None,
1882 partition_value: None,
1883 partition_column_type: None,
1884 partition_fields: vec![],
1885 };
1886
1887 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1888
1889 let make_batch = |ids: Vec<i32>, embs: Vec<Vec<f32>>| {
1891 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1892 .unwrap();
1893 AilakeFileWriter::new(policy.clone())
1894 .write(&batch, &embs)
1895 .unwrap()
1896 };
1897
1898 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]];
1899 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]];
1900 let bytes_a = make_batch(vec![0, 1], embs_a);
1901 let bytes_b = make_batch(vec![2, 3], embs_b);
1902
1903 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1904 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1905
1906 let entries = vec![
1907 DataFileEntry {
1908 path: "data/a.parquet".into(),
1909 record_count: 2,
1910 file_size_bytes: bytes_a.len() as u64,
1911 centroid_b64: None,
1912 radius: None,
1913 hnsw_offset: None,
1914 hnsw_len: None,
1915 vector_column: None,
1916 vector_dim: None,
1917 extra_vector_indexes: vec![],
1918 index_status: IndexStatus::Ready,
1919 index_error: None,
1920 batch_id: None,
1921 embedding_model: None,
1922 partition_value: None,
1923 deletion_vector: None,
1924 first_row_id: None,
1925 column_stats: None,
1926 },
1927 DataFileEntry {
1928 path: "data/b.parquet".into(),
1929 record_count: 2,
1930 file_size_bytes: bytes_b.len() as u64,
1931 centroid_b64: None,
1932 radius: None,
1933 hnsw_offset: None,
1934 hnsw_len: None,
1935 vector_column: None,
1936 vector_dim: None,
1937 extra_vector_indexes: vec![],
1938 index_status: IndexStatus::Ready,
1939 index_error: None,
1940 batch_id: None,
1941 embedding_model: None,
1942 partition_value: None,
1943 deletion_vector: None,
1944 first_row_id: None,
1945 column_stats: None,
1946 },
1947 ];
1948
1949 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1950 let merged = executor
1952 .compact_incremental(&entries, "data/merged.parquet")
1953 .await
1954 .unwrap();
1955
1956 assert_eq!(merged.record_count, 4);
1957
1958 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1959 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1960 reader.verify_integrity().unwrap();
1961 }
1962
1963 #[tokio::test]
1964 async fn compact_deferred_produces_parquet_only_file() {
1965 use ailake_catalog::HadoopCatalog;
1966 use ailake_core::{VectorMetric, VectorPrecision};
1967 use ailake_store::LocalStore;
1968 use arrow_array::{Int32Array, RecordBatch};
1969 use arrow_schema::{DataType, Field, Schema};
1970 use std::sync::Arc;
1971 use tempfile::TempDir;
1972
1973 let dir = TempDir::new().unwrap();
1974 let store = Arc::new(LocalStore::new(dir.path()));
1975 let catalog_dir = TempDir::new().unwrap();
1976 let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
1977 let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
1978 let table = TableIdent {
1979 namespace: "ns".into(),
1980 name: "tbl".into(),
1981 };
1982
1983 let policy = VectorStoragePolicy {
1984 column_name: "embedding".into(),
1985 dim: 4,
1986 metric: VectorMetric::Cosine,
1987 precision: VectorPrecision::F16,
1988 pq: None,
1989 keep_raw_for_reranking: true,
1990 pre_normalize: false,
1991 hnsw_m: None,
1992 hnsw_ef_construction: None,
1993 ivf_residual: false,
1994 embedding_model: None,
1995 modality: None,
1996 partition_by: None,
1997 partition_value: None,
1998 partition_column_type: None,
1999 partition_fields: vec![],
2000 };
2001
2002 use ailake_catalog::TableProperties;
2003 catalog
2004 .create_table(
2005 &table,
2006 &TableProperties {
2007 policy: policy.clone(),
2008 extra: std::collections::HashMap::new(),
2009 format_version: 2,
2010 partition_column_type: None,
2011 },
2012 )
2013 .await
2014 .unwrap();
2015
2016 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2017 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]];
2018 let batch_a = RecordBatch::try_new(
2019 schema.clone(),
2020 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2021 )
2022 .unwrap();
2023 let bytes_a = AilakeFileWriter::new(policy.clone())
2024 .write(&batch_a, &embs_a)
2025 .unwrap();
2026 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
2027
2028 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]];
2029 let batch_b = RecordBatch::try_new(
2030 schema.clone(),
2031 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2032 )
2033 .unwrap();
2034 let bytes_b = AilakeFileWriter::new(policy.clone())
2035 .write(&batch_b, &embs_b)
2036 .unwrap();
2037 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
2038
2039 let entries = vec![
2040 DataFileEntry {
2041 path: "data/a.parquet".into(),
2042 record_count: 2,
2043 file_size_bytes: bytes_a.len() as u64,
2044 centroid_b64: None,
2045 radius: None,
2046 hnsw_offset: None,
2047 hnsw_len: None,
2048 vector_column: None,
2049 vector_dim: None,
2050 extra_vector_indexes: vec![],
2051 index_status: IndexStatus::Ready,
2052 index_error: None,
2053 batch_id: None,
2054 embedding_model: None,
2055 partition_value: None,
2056 deletion_vector: None,
2057 first_row_id: None,
2058 column_stats: None,
2059 },
2060 DataFileEntry {
2061 path: "data/b.parquet".into(),
2062 record_count: 2,
2063 file_size_bytes: bytes_b.len() as u64,
2064 centroid_b64: None,
2065 radius: None,
2066 hnsw_offset: None,
2067 hnsw_len: None,
2068 vector_column: None,
2069 vector_dim: None,
2070 extra_vector_indexes: vec![],
2071 index_status: IndexStatus::Ready,
2072 index_error: None,
2073 batch_id: None,
2074 embedding_model: None,
2075 partition_value: None,
2076 deletion_vector: None,
2077 first_row_id: None,
2078 column_stats: None,
2079 },
2080 ];
2081
2082 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2083 let entry = executor
2084 .compact_deferred(&entries, "data/merged.parquet", catalog.clone(), &table)
2085 .await
2086 .unwrap();
2087
2088 assert_eq!(entry.index_status, IndexStatus::Indexing);
2090 assert_eq!(entry.record_count, 4);
2091
2092 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
2094 let pq_reader = ailake_parquet::ParquetVectorReader::new(merged_bytes, "embedding");
2095 let count = pq_reader.record_count().unwrap();
2096 assert_eq!(count, 4);
2097 }
2098
2099 #[tokio::test]
2107 async fn compact_aggregates_batch_ids_from_sources() {
2108 use ailake_core::{VectorMetric, VectorPrecision};
2109 use ailake_store::LocalStore;
2110 use arrow_array::{Int32Array, RecordBatch};
2111 use arrow_schema::{DataType, Field, Schema};
2112 use std::sync::Arc;
2113 use tempfile::TempDir;
2114
2115 let dir = TempDir::new().unwrap();
2116 let store = Arc::new(LocalStore::new(dir.path()));
2117 let policy = VectorStoragePolicy {
2118 column_name: "embedding".into(),
2119 dim: 4,
2120 metric: VectorMetric::Cosine,
2121 precision: VectorPrecision::F16,
2122 pq: None,
2123 keep_raw_for_reranking: true,
2124 pre_normalize: false,
2125 hnsw_m: None,
2126 hnsw_ef_construction: None,
2127 ivf_residual: false,
2128 embedding_model: None,
2129 modality: None,
2130 partition_by: None,
2131 partition_value: None,
2132 partition_column_type: None,
2133 partition_fields: vec![],
2134 };
2135
2136 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2137 let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0]];
2138 let batch_a =
2139 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![0i32]))])
2140 .unwrap();
2141 let bytes_a = AilakeFileWriter::new(policy.clone())
2142 .write(&batch_a, &embs_a)
2143 .unwrap();
2144 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
2145
2146 let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 1.0, 0.0, 0.0]];
2147 let batch_b =
2148 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1i32]))])
2149 .unwrap();
2150 let bytes_b = AilakeFileWriter::new(policy.clone())
2151 .write(&batch_b, &embs_b)
2152 .unwrap();
2153 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
2154
2155 let embs_c: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
2158 let batch_c =
2159 RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![2i32]))]).unwrap();
2160 let bytes_c = AilakeFileWriter::new(policy.clone())
2161 .write(&batch_c, &embs_c)
2162 .unwrap();
2163 store.put("data/c.parquet", bytes_c.clone()).await.unwrap();
2164
2165 fn entry(path: &str, size: u64, batch_id: Option<&str>) -> DataFileEntry {
2166 DataFileEntry {
2167 path: path.into(),
2168 record_count: 1,
2169 file_size_bytes: size,
2170 centroid_b64: None,
2171 radius: None,
2172 hnsw_offset: None,
2173 hnsw_len: None,
2174 vector_column: None,
2175 vector_dim: None,
2176 extra_vector_indexes: vec![],
2177 index_status: IndexStatus::Ready,
2178 index_error: None,
2179 batch_id: batch_id.map(String::from),
2180 embedding_model: None,
2181 partition_value: None,
2182 deletion_vector: None,
2183 first_row_id: None,
2184 column_stats: None,
2185 }
2186 }
2187 let entries = vec![
2188 entry("data/a.parquet", bytes_a.len() as u64, Some("k-a")),
2189 entry("data/b.parquet", bytes_b.len() as u64, Some("k-b")),
2190 entry("data/c.parquet", bytes_c.len() as u64, None),
2191 ];
2192
2193 let executor = CompactionExecutor::new(store.clone(), policy);
2194 let merged = executor
2195 .compact(&entries, "data/merged.parquet")
2196 .await
2197 .unwrap();
2198
2199 assert_eq!(merged.record_count, 3);
2200 assert_eq!(
2201 merged.batch_ids(),
2202 vec!["k-a".to_string(), "k-b".to_string()],
2203 "merged file must carry every source's idempotency key, none invented"
2204 );
2205 }
2206
2207 #[tokio::test]
2214 async fn run_preserves_untouched_files_outside_compaction_pass() {
2215 use ailake_catalog::HadoopCatalog;
2216 use ailake_core::{VectorMetric, VectorPrecision};
2217 use ailake_store::LocalStore;
2218 use arrow_array::{Int32Array, RecordBatch};
2219 use arrow_schema::{DataType, Field, Schema};
2220 use std::sync::Arc;
2221 use tempfile::TempDir;
2222
2223 let dir = TempDir::new().unwrap();
2224 let store = Arc::new(LocalStore::new(dir.path()));
2225 let catalog_dir = TempDir::new().unwrap();
2226 let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
2227 let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
2228 let table = TableIdent {
2229 namespace: "ns".into(),
2230 name: "tbl".into(),
2231 };
2232
2233 let policy = VectorStoragePolicy {
2234 column_name: "embedding".into(),
2235 dim: 4,
2236 metric: VectorMetric::Cosine,
2237 precision: VectorPrecision::F16,
2238 pq: None,
2239 keep_raw_for_reranking: true,
2240 pre_normalize: false,
2241 hnsw_m: None,
2242 hnsw_ef_construction: None,
2243 ivf_residual: false,
2244 embedding_model: None,
2245 modality: None,
2246 partition_by: None,
2247 partition_value: None,
2248 partition_column_type: None,
2249 partition_fields: vec![],
2250 };
2251
2252 use ailake_catalog::TableProperties;
2253 catalog
2254 .create_table(
2255 &table,
2256 &TableProperties {
2257 policy: policy.clone(),
2258 extra: std::collections::HashMap::new(),
2259 format_version: 2,
2260 partition_column_type: None,
2261 },
2262 )
2263 .await
2264 .unwrap();
2265
2266 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2267
2268 let write_file = |path: &str, ids: Vec<i32>, embs: Vec<Vec<f32>>| {
2270 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
2271 .unwrap();
2272 let bytes = AilakeFileWriter::new(policy.clone())
2273 .write(&batch, &embs)
2274 .unwrap();
2275 (path.to_string(), bytes)
2276 };
2277
2278 let (path_a, bytes_a) = write_file(
2279 "data/small_a.parquet",
2280 vec![0, 1],
2281 vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
2282 );
2283 let (path_b, bytes_b) = write_file(
2284 "data/small_b.parquet",
2285 vec![2, 3],
2286 vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]],
2287 );
2288 let (path_big, bytes_big) = write_file(
2291 "data/big.parquet",
2292 vec![4, 5],
2293 vec![vec![1.0, 1.0, 0.0, 0.0], vec![0.0, 1.0, 1.0, 0.0]],
2294 );
2295
2296 for (path, bytes) in [
2297 (&path_a, &bytes_a),
2298 (&path_b, &bytes_b),
2299 (&path_big, &bytes_big),
2300 ] {
2301 store.put(path, bytes.clone()).await.unwrap();
2302 }
2303
2304 let make_entry = |path: &str, size: u64| DataFileEntry {
2305 path: path.to_string(),
2306 record_count: 2,
2307 file_size_bytes: size,
2308 centroid_b64: Some("AAAA".into()),
2313 radius: None,
2314 hnsw_offset: None,
2315 hnsw_len: None,
2316 vector_column: None,
2317 vector_dim: None,
2318 extra_vector_indexes: vec![],
2319 index_status: IndexStatus::Ready,
2320 index_error: None,
2321 batch_id: None,
2322 embedding_model: None,
2323 partition_value: None,
2324 deletion_vector: None,
2325 first_row_id: None,
2326 column_stats: None,
2327 };
2328
2329 let initial_snap_id = ailake_catalog::new_snapshot_id();
2330 let initial_snapshot = NewSnapshot {
2331 snapshot_id: initial_snap_id,
2332 parent_snapshot_id: None,
2333 files: vec![
2334 make_entry(&path_a, 500),
2335 make_entry(&path_b, 500),
2336 make_entry(&path_big, 200_000_000), ],
2338 operation: SnapshotOperation::Append,
2339 iceberg_schema: None,
2340 extra_properties: std::collections::HashMap::new(),
2341 bloom_filters: vec![],
2342 equality_delete_files: vec![],
2343 };
2344 catalog
2345 .commit_snapshot(&table, initial_snapshot)
2346 .await
2347 .unwrap();
2348
2349 let planner = CompactionPlanner::new(CompactionConfig {
2350 min_files_to_compact: 2,
2351 target_file_size_bytes: 1000,
2352 index_strategy: CompactionIndexStrategy::ForceHnsw,
2353 max_files_per_pass: 20,
2354 });
2355 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2356
2357 let merged = executor
2358 .run(&planner, &table, catalog.clone(), "data")
2359 .await
2360 .unwrap()
2361 .expect("compaction should have run — 2 eligible small files");
2362
2363 let files_after = catalog.list_files(&table, None).await.unwrap();
2364 let paths_after: Vec<&str> = files_after.iter().map(|f| f.path.as_str()).collect();
2365
2366 assert!(
2367 paths_after.contains(&path_big.as_str()),
2368 "BUG: untouched 'big.parquet' vanished after run() — files_after={paths_after:?}"
2369 );
2370 assert!(
2371 paths_after.contains(&merged.path.as_str()),
2372 "merged output file must be present — files_after={paths_after:?}"
2373 );
2374
2375 let meta_dir = catalog_dir.path().join("ns/tbl/metadata");
2382 let latest_metadata = std::fs::read_dir(&meta_dir)
2383 .unwrap()
2384 .filter_map(|e| e.ok())
2385 .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
2386 .max_by_key(|e| e.metadata().unwrap().modified().unwrap())
2387 .expect("metadata.json must exist after commit");
2388 let json: serde_json::Value =
2389 serde_json::from_slice(&std::fs::read(latest_metadata.path()).unwrap()).unwrap();
2390 let last_snapshot = json["snapshots"].as_array().unwrap().last().unwrap();
2391 assert_eq!(
2392 last_snapshot["parent-snapshot-id"].as_i64(),
2393 Some(initial_snap_id),
2394 "compaction's Replace snapshot must chain to the pre-compaction snapshot, not be orphaned"
2395 );
2396 assert!(
2397 !paths_after.contains(&path_a.as_str()) && !paths_after.contains(&path_b.as_str()),
2398 "compacted input files must no longer be listed — files_after={paths_after:?}"
2399 );
2400 assert_eq!(
2401 files_after.len(),
2402 2,
2403 "expected exactly [big.parquet, merged] — files_after={paths_after:?}"
2404 );
2405 }
2406
2407 #[tokio::test]
2413 async fn compact_preserves_rows_from_footerless_file() {
2414 use ailake_core::{VectorMetric, VectorPrecision};
2415 use ailake_store::LocalStore;
2416 use arrow_array::{Int32Array, RecordBatch};
2417 use arrow_schema::{DataType, Field, Schema};
2418 use std::sync::Arc;
2419 use tempfile::TempDir;
2420
2421 let dir = TempDir::new().unwrap();
2422 let store = Arc::new(LocalStore::new(dir.path()));
2423 let policy = VectorStoragePolicy {
2424 column_name: "embedding".into(),
2425 dim: 4,
2426 metric: VectorMetric::Cosine,
2427 precision: VectorPrecision::F16,
2428 pq: None,
2429 keep_raw_for_reranking: true,
2430 pre_normalize: false,
2431 hnsw_m: None,
2432 hnsw_ef_construction: None,
2433 ivf_residual: false,
2434 embedding_model: None,
2435 modality: None,
2436 partition_by: None,
2437 partition_value: None,
2438 partition_column_type: None,
2439 partition_fields: vec![],
2440 };
2441
2442 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2443
2444 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]];
2446 let batch_native = RecordBatch::try_new(
2447 schema.clone(),
2448 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2449 )
2450 .unwrap();
2451 let bytes_native = AilakeFileWriter::new(policy.clone())
2452 .write(&batch_native, &embs_native)
2453 .unwrap();
2454 store
2455 .put("data/native.parquet", bytes_native.clone())
2456 .await
2457 .unwrap();
2458
2459 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]];
2463 let batch_foreign = RecordBatch::try_new(
2464 schema.clone(),
2465 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2466 )
2467 .unwrap();
2468 let bytes_foreign = AilakeFileWriter::new(policy.clone())
2469 .write_parquet_only(&batch_foreign, &embs_foreign)
2470 .unwrap();
2471 store
2472 .put("data/foreign.parquet", bytes_foreign.clone())
2473 .await
2474 .unwrap();
2475
2476 let reader_foreign = AilakeFileReader::new(bytes_foreign.clone(), "embedding", 4);
2477 assert!(
2478 !reader_foreign.is_ailake_file(),
2479 "sanity: write_parquet_only must not embed an AILK footer"
2480 );
2481
2482 let entries = vec![
2483 DataFileEntry {
2484 path: "data/native.parquet".into(),
2485 record_count: 2,
2486 file_size_bytes: bytes_native.len() as u64,
2487 centroid_b64: None,
2488 radius: None,
2489 hnsw_offset: None,
2490 hnsw_len: None,
2491 vector_column: None,
2492 vector_dim: None,
2493 extra_vector_indexes: vec![],
2494 index_status: IndexStatus::Ready,
2495 index_error: None,
2496 batch_id: None,
2497 embedding_model: None,
2498 partition_value: None,
2499 deletion_vector: None,
2500 first_row_id: None,
2501 column_stats: None,
2502 },
2503 DataFileEntry {
2504 path: "data/foreign.parquet".into(),
2505 record_count: 2,
2506 file_size_bytes: bytes_foreign.len() as u64,
2507 centroid_b64: None,
2508 radius: None,
2509 hnsw_offset: None,
2510 hnsw_len: None,
2511 vector_column: None,
2512 vector_dim: None,
2513 extra_vector_indexes: vec![],
2514 index_status: IndexStatus::Ready,
2515 index_error: None,
2516 batch_id: None,
2517 embedding_model: None,
2518 partition_value: None,
2519 deletion_vector: None,
2520 first_row_id: None,
2521 column_stats: None,
2522 },
2523 ];
2524
2525 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2527 let merged = executor
2528 .compact(&entries, "data/merged_full.parquet")
2529 .await
2530 .unwrap();
2531 assert_eq!(
2532 merged.record_count, 4,
2533 "compact() must include all 4 rows — 2 native + 2 from the footerless file"
2534 );
2535 let merged_bytes = store.get("data/merged_full.parquet").await.unwrap();
2536 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
2537 reader.verify_integrity().unwrap();
2538
2539 let merged_inc = executor
2542 .compact_incremental(&entries, "data/merged_inc.parquet")
2543 .await
2544 .unwrap();
2545 assert_eq!(
2546 merged_inc.record_count, 4,
2547 "compact_incremental() must include all 4 rows — 2 native + 2 from the footerless file"
2548 );
2549 }
2550
2551 #[tokio::test]
2554 async fn compact_preserves_search_results() {
2555 use crate::scanner::SearchConfig;
2556 use ailake_catalog::HadoopCatalog;
2557 use ailake_core::{VectorMetric, VectorPrecision};
2558 use ailake_store::LocalStore;
2559 use arrow_array::{Int32Array, RecordBatch};
2560 use arrow_schema::{DataType, Field, Schema};
2561 use std::sync::Arc;
2562 use tempfile::TempDir;
2563 let dir = TempDir::new().unwrap();
2564 let store = Arc::new(LocalStore::new(dir.path()));
2565 let policy = VectorStoragePolicy {
2566 column_name: "embedding".into(),
2567 dim: 4,
2568 metric: VectorMetric::Cosine,
2569 precision: VectorPrecision::F16,
2570 pq: None,
2571 keep_raw_for_reranking: true,
2572 pre_normalize: false,
2573 hnsw_m: None,
2574 hnsw_ef_construction: None,
2575 ivf_residual: false,
2576 embedding_model: None,
2577 modality: None,
2578 partition_by: None,
2579 partition_value: None,
2580 partition_column_type: None,
2581 partition_fields: vec![],
2582 };
2583
2584 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2585
2586 let embs_a: Vec<Vec<f32>> = vec![
2588 vec![1.0, 0.0, 0.0, 0.0],
2589 vec![0.0, 1.0, 0.0, 0.0],
2590 vec![0.0, 0.0, 1.0, 0.0],
2591 ];
2592 let batch_a = RecordBatch::try_new(
2593 schema.clone(),
2594 vec![Arc::new(Int32Array::from(vec![0i32, 1, 2]))],
2595 )
2596 .unwrap();
2597
2598 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]];
2600 let batch_b = RecordBatch::try_new(
2601 schema.clone(),
2602 vec![Arc::new(Int32Array::from(vec![3i32, 4]))],
2603 )
2604 .unwrap();
2605
2606 let bytes_a = AilakeFileWriter::new(policy.clone())
2607 .write(&batch_a, &embs_a)
2608 .unwrap();
2609 let bytes_b = AilakeFileWriter::new(policy.clone())
2610 .write(&batch_b, &embs_b)
2611 .unwrap();
2612
2613 store.put("data/a.parquet", bytes_a).await.unwrap();
2614 store.put("data/b.parquet", bytes_b).await.unwrap();
2615 let file_size_a = store.file_size("data/a.parquet").await.unwrap();
2616 let file_size_b = store.file_size("data/b.parquet").await.unwrap();
2617
2618 let catalog: Arc<dyn CatalogProvider> =
2621 Arc::new(HadoopCatalog::new(store.clone(), "warehouse"));
2622 let table = TableIdent::new("default", "compact_search_test");
2623
2624 catalog
2626 .create_table(
2627 &table,
2628 &ailake_catalog::TableProperties {
2629 policy: policy.clone(),
2630 extra: std::collections::HashMap::new(),
2631 format_version: 2,
2632 partition_column_type: None,
2633 },
2634 )
2635 .await
2636 .unwrap();
2637
2638 let meta = catalog.load_table(&table).await.unwrap();
2640 let parent_snapshot_id = meta.current_snapshot_id;
2641
2642 let entry_a = DataFileEntry {
2643 path: "data/a.parquet".into(),
2644 record_count: 3,
2645 file_size_bytes: file_size_a,
2646 centroid_b64: None,
2647 radius: None,
2648 hnsw_offset: None,
2649 hnsw_len: None,
2650 vector_column: Some("embedding".into()),
2651 vector_dim: Some(4),
2652 extra_vector_indexes: vec![],
2653 index_status: IndexStatus::Ready,
2654 index_error: None,
2655 batch_id: None,
2656 embedding_model: None,
2657 partition_value: None,
2658 deletion_vector: None,
2659 first_row_id: None,
2660 column_stats: None,
2661 };
2662 let entry_b = DataFileEntry {
2663 path: "data/b.parquet".into(),
2664 record_count: 2,
2665 file_size_bytes: file_size_b,
2666 centroid_b64: None,
2667 radius: None,
2668 hnsw_offset: None,
2669 hnsw_len: None,
2670 vector_column: Some("embedding".into()),
2671 vector_dim: Some(4),
2672 extra_vector_indexes: vec![],
2673 index_status: IndexStatus::Ready,
2674 index_error: None,
2675 batch_id: None,
2676 embedding_model: None,
2677 partition_value: None,
2678 deletion_vector: None,
2679 first_row_id: None,
2680 column_stats: None,
2681 };
2682
2683 catalog
2684 .commit_snapshot(
2685 &table,
2686 NewSnapshot {
2687 snapshot_id: 200,
2688 parent_snapshot_id,
2689 files: vec![entry_a, entry_b],
2690 operation: SnapshotOperation::Append,
2691 iceberg_schema: None,
2692 extra_properties: std::collections::HashMap::new(),
2693 bloom_filters: vec![],
2694 equality_delete_files: vec![],
2695 },
2696 )
2697 .await
2698 .unwrap();
2699
2700 let query_before = vec![1.0f32, 0.0, 0.0, 0.0];
2702 let config = SearchConfig {
2703 top_k: 5,
2704 ef_search: 100,
2705 pruning_threshold: f32::INFINITY,
2706 ..Default::default()
2707 };
2708 let before = crate::scanner::search(
2709 &table,
2710 &query_before,
2711 config.clone(),
2712 "embedding",
2713 4,
2714 catalog.clone(),
2715 store.clone(),
2716 )
2717 .await
2718 .unwrap();
2719 assert!(
2720 !before.is_empty(),
2721 "search before compaction must return results"
2722 );
2723 let before_dists: Vec<f32> = before.iter().map(|r| r.distance).collect();
2724
2725 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2727 let files_to_compact = catalog.list_files(&table, None).await.unwrap();
2728 let merged = executor
2729 .compact(&files_to_compact, "data/merged.parquet")
2730 .await
2731 .unwrap();
2732
2733 let meta2 = catalog.load_table(&table).await.unwrap();
2735 let merged_entry = DataFileEntry {
2736 path: merged.path.clone(),
2737 record_count: merged.record_count,
2738 file_size_bytes: merged.file_size_bytes,
2739 centroid_b64: merged.centroid_b64,
2740 radius: merged.radius,
2741 hnsw_offset: merged.hnsw_offset,
2742 hnsw_len: merged.hnsw_len,
2743 vector_column: merged.vector_column,
2744 vector_dim: merged.vector_dim,
2745 extra_vector_indexes: merged.extra_vector_indexes,
2746 index_status: merged.index_status,
2747 index_error: merged.index_error,
2748 batch_id: merged.batch_id,
2749 embedding_model: merged.embedding_model,
2750 partition_value: merged.partition_value,
2751 deletion_vector: merged.deletion_vector,
2752 first_row_id: None,
2753 column_stats: None,
2754 };
2755 catalog
2756 .commit_snapshot(
2757 &table,
2758 NewSnapshot {
2759 snapshot_id: 201,
2760 parent_snapshot_id: meta2.current_snapshot_id,
2761 files: vec![merged_entry],
2762 operation: SnapshotOperation::Replace,
2763 iceberg_schema: None,
2764 extra_properties: std::collections::HashMap::new(),
2765 bloom_filters: vec![],
2766 equality_delete_files: vec![],
2767 },
2768 )
2769 .await
2770 .unwrap();
2771
2772 let after = crate::scanner::search(
2774 &table,
2775 &query_before,
2776 config.clone(),
2777 "embedding",
2778 4,
2779 catalog,
2780 store,
2781 )
2782 .await
2783 .unwrap();
2784 assert!(
2785 !after.is_empty(),
2786 "search after compaction must return results"
2787 );
2788 let after_dists: Vec<f32> = after.iter().map(|r| r.distance).collect();
2789
2790 assert_eq!(
2792 before.len(),
2793 after.len(),
2794 "number of results should match before ({}) and after ({}) compaction",
2795 before.len(),
2796 after.len(),
2797 );
2798 let mut total_diff = 0.0f32;
2803 for (i, (bd, ad)) in before_dists.iter().zip(after_dists.iter()).enumerate() {
2804 let diff = (bd - ad).abs();
2805 total_diff += diff;
2806 assert!(
2807 diff < 0.1,
2808 "distance mismatch at position {i}: before={bd}, after={ad}, diff={diff}"
2809 );
2810 }
2811 let avg_diff = total_diff / before_dists.len() as f32;
2812 assert!(
2813 avg_diff < 0.05,
2814 "average distance error across all results too high: {avg_diff}"
2815 );
2816 }
2817}