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 self.store.put(output_path, file_bytes.clone()).await?;
319
320 let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
322 let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
323 let header = reader.read_header()?;
324 let ailk_start = reader.ailk_offset()?;
325
326 reader.verify_integrity()?;
331
332 let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
335
336 let mut entry = make_data_file_entry(
337 output_path,
338 record_count,
339 file_size,
340 ¢roid,
341 VectorIndexInfo {
342 column: &self.policy.column_name,
343 dim: self.policy.dim,
344 hnsw_offset: ailk_start + header.hnsw_offset,
345 hnsw_len: header.hnsw_len,
346 },
347 );
348 entry.first_row_id = source_first_row_id;
349 Ok(entry)
350 }
351
352 pub async fn compact_incremental(
371 &self,
372 files: &[DataFileEntry],
373 output_path: &str,
374 ) -> AilakeResult<DataFileEntry> {
375 const DOMINANT_RATIO: f64 = 0.40;
376
377 if files.is_empty() {
378 return Err(ailake_core::AilakeError::Catalog(
379 "compact_incremental: no files provided".into(),
380 ));
381 }
382
383 if matches!(self.index_strategy, CompactionIndexStrategy::ForceIvfPq) {
392 debug!(
393 "ailake: compact_incremental — index_strategy=ForceIvfPq, which this method \
394 can never produce (HNSW graph extension only); falling back to full rebuild"
395 );
396 return self.compact(files, output_path).await;
397 }
398
399 let total_rows: u64 = files.iter().map(|f| f.record_count).sum();
401 let dom_idx = files
402 .iter()
403 .enumerate()
404 .max_by_key(|(_, f)| f.record_count)
405 .map(|(i, _)| i)
406 .unwrap_or(0);
407 let dom_rows = files[dom_idx].record_count;
408
409 if files[dom_idx].deletion_vector.is_some() {
416 debug!(
417 "ailake: compact_incremental — dominant file {} has DV-masked rows, \
418 falling back to full rebuild (graph reuse would desync row positions)",
419 files[dom_idx].path
420 );
421 return self.compact(files, output_path).await;
422 }
423
424 if (dom_rows as f64 / total_rows as f64) < DOMINANT_RATIO {
425 debug!(
426 "ailake: compact_incremental — no dominant file ({}/{} rows < {:.0}% threshold), \
427 falling back to full rebuild",
428 dom_rows,
429 total_rows,
430 DOMINANT_RATIO * 100.0
431 );
432 return self.compact(files, output_path).await;
433 }
434
435 let column = self.policy.column_name.clone();
436 let dim = self.policy.dim;
437 let dom_path = files[dom_idx].path.clone();
438
439 let futs: Vec<_> =
448 files
449 .iter()
450 .map(|entry| {
451 let store = self.store.clone();
452 let path = entry.path.clone();
453 let col = column.clone();
454 let is_dom = path == dom_path;
455 let dv = entry.deletion_vector.clone();
456 async move {
457 let bytes: Bytes = store.get(&path).await?;
458 let reader = AilakeFileReader::new(bytes.clone(), &col, dim);
459 let has_index = reader.is_ailake_file();
460 if is_dom && !has_index {
461 debug!(
462 "ailake: compact_incremental — dominant candidate {} has no \
463 AI-Lake index; will fall back to full rebuild if no HNSW to reuse",
464 path
465 );
466 }
467 let (raw_batch, raw_vecs) = reader.read_parquet()?;
468 let (batch, vecs) = if let Some(dv) = dv {
472 let bitmap = crate::dv::load_deletion_vector(&store, &dv).await?;
473 crate::dv::filter_deleted_rows(raw_batch, raw_vecs, &bitmap)?
474 } else {
475 (raw_batch, raw_vecs)
476 };
477 let retained = if is_dom && has_index {
478 Some(bytes)
479 } else {
480 None
481 };
482 Ok::<
483 (RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>),
484 ailake_core::AilakeError,
485 >((batch, vecs, is_dom, retained))
486 }
487 })
488 .collect();
489
490 #[allow(clippy::type_complexity)]
494 let raw: Vec<(RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>)> =
495 try_join_all(futs).await?;
496
497 let mut dom_batch: Option<RecordBatch> = None;
499 let mut dom_vecs: Vec<Vec<f32>> = Vec::new();
500 let mut dom_bytes_found: Option<Bytes> = None;
501 let mut other_batches: Vec<RecordBatch> = Vec::new();
502 let mut other_vecs: Vec<Vec<f32>> = Vec::new();
503
504 for (batch, vecs, is_dom, retained) in raw {
505 if is_dom {
506 dom_batch = Some(batch);
507 dom_vecs = vecs;
508 dom_bytes_found = retained;
509 } else {
510 other_batches.push(batch);
511 other_vecs.extend(vecs);
512 }
513 }
514
515 let (dom_batch, dom_bytes) = match (dom_batch, dom_bytes_found) {
516 (Some(b), Some(byt)) => (b, byt),
517 _ => {
518 debug!(
519 "ailake: compact_incremental — dominant file missing from read results, \
520 falling back to full rebuild"
521 );
522 return self.compact(files, output_path).await;
523 }
524 };
525
526 let dom_reader = AilakeFileReader::new(dom_bytes, &column, dim);
528 let mut hnsw = match dom_reader.load_index() {
529 Ok(idx) => idx,
530 Err(e) => {
531 debug!(
532 "ailake: compact_incremental — cannot load dominant HNSW ({}), \
533 falling back to full rebuild",
534 e
535 );
536 return self.compact(files, output_path).await;
537 }
538 };
539
540 let dom_count = dom_batch.num_rows() as u64;
541
542 for (j, vec) in other_vecs.iter().enumerate() {
545 hnsw.insert_node(RowId::new(dom_count + j as u64), vec.clone());
546 }
547 hnsw.quantize_to_f16();
548
549 let schema: SchemaRef = dom_batch.schema();
551 let mut all_batches = vec![dom_batch];
552 all_batches.extend(other_batches);
553 let merged_batch = concat_batches(schema, &all_batches)?;
554 let record_count = merged_batch.num_rows() as u64;
555
556 let mut all_embeddings = dom_vecs;
557 all_embeddings.extend(other_vecs);
558
559 let writer = {
562 let base = AilakeFileWriter::new(self.policy.clone());
563 if let Some(ref fts_cfg) = self.fts_config {
564 match ailake_fts::merge_fts_blobs(fts_cfg, &merged_batch) {
565 Ok(blob) => base.with_prebuilt_fts_blob(blob),
566 Err(e) => {
567 warn!("ailake: FTS re-index during incremental compaction failed: {e}");
568 base
569 }
570 }
571 } else {
572 base
573 }
574 };
575 let file_bytes = writer.write_with_prebuilt_hnsw(&merged_batch, &all_embeddings, &hnsw)?;
576 let file_size = file_bytes.len() as u64;
577 self.store.put(output_path, file_bytes.clone()).await?;
578
579 let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
580 let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
581 let header = reader.read_header()?;
582 let ailk_start = reader.ailk_offset()?;
583
584 reader.verify_integrity()?;
587
588 let source_first_row_id = files[dom_idx].first_row_id;
592
593 let mut entry = make_data_file_entry(
594 output_path,
595 record_count,
596 file_size,
597 ¢roid,
598 VectorIndexInfo {
599 column: &self.policy.column_name,
600 dim: self.policy.dim,
601 hnsw_offset: ailk_start + header.hnsw_offset,
602 hnsw_len: header.hnsw_len,
603 },
604 );
605 entry.first_row_id = source_first_row_id;
606
607 info!(
608 "ailake: compact_incremental — merged {} files into {} \
609 ({} rows from dominant + {} inserted incrementally)",
610 files.len(),
611 output_path,
612 dom_count,
613 record_count - dom_count
614 );
615
616 Ok(entry)
617 }
618
619 pub async fn compact_deferred(
629 &self,
630 files: &[DataFileEntry],
631 output_path: &str,
632 catalog: Arc<dyn CatalogProvider>,
633 table: &TableIdent,
634 ) -> AilakeResult<DataFileEntry> {
635 if files.is_empty() {
636 return Err(ailake_core::AilakeError::Catalog(
637 "compact_deferred: no files provided".into(),
638 ));
639 }
640
641 let pairs = self.read_files_parallel(files).await?;
644
645 let schema: SchemaRef = pairs[0].0.schema();
646 let (all_batches, all_embeddings): (Vec<_>, Vec<_>) = pairs.into_iter().unzip();
647 let all_embeddings: Vec<Vec<f32>> = all_embeddings.into_iter().flatten().collect();
648
649 let merged_batch = concat_batches(schema, &all_batches)?;
650 let record_count = merged_batch.num_rows() as u64;
651
652 let file_writer = AilakeFileWriter::new(self.policy.clone());
654 let parquet_bytes = file_writer.write_parquet_only(&merged_batch, &all_embeddings)?;
655 let file_size = parquet_bytes.len() as u64;
656 self.store.put(output_path, parquet_bytes).await?;
657
658 let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
660 let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
661 let mut entry = make_data_file_entry_indexing(
662 output_path,
663 record_count,
664 file_size,
665 ¢roid,
666 &self.policy.column_name,
667 self.policy.dim,
668 );
669 entry.first_row_id = source_first_row_id;
670
671 let store = self.store.clone();
673 let policy = self.policy.clone();
674 let table_id = table.clone();
675 let fp = output_path.to_string();
676 tokio::spawn(async move {
677 if let Err(e) = build_and_patch_index(store, catalog, policy, table_id, fp).await {
678 error!(
679 "ailake: compaction deferred HNSW build failed — file indexed as \
680 Parquet-only until next compaction rebuilds the index: {}",
681 e
682 );
683 }
684 });
685
686 Ok(entry)
687 }
688
689 pub async fn run(
692 &self,
693 planner: &CompactionPlanner,
694 table: &TableIdent,
695 catalog: Arc<dyn CatalogProvider>,
696 output_prefix: &str,
697 ) -> AilakeResult<Option<DataFileEntry>> {
698 let all_files = catalog.list_files(table, None).await?;
699 let to_compact = planner.plan(&all_files);
700 if to_compact.is_empty() {
701 return Ok(None);
702 }
703
704 let meta_props = catalog
707 .load_table(table)
708 .await
709 .map(|m| m.properties)
710 .unwrap_or_default();
711 let executor = self.with_effective_fts(&meta_props);
712
713 let ts = std::time::SystemTime::now()
714 .duration_since(std::time::UNIX_EPOCH)
715 .unwrap_or_else(|e| e.duration())
716 .as_millis();
717 let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
718
719 let merged = executor
721 .compact_incremental(&to_compact, &output_path)
722 .await?;
723
724 let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
726 let parent_snapshot_id = catalog
729 .load_table(table)
730 .await
731 .ok()
732 .and_then(|m| m.current_snapshot_id);
733 let snapshot = NewSnapshot {
734 snapshot_id: ailake_catalog::new_snapshot_id(),
735 parent_snapshot_id,
736 files,
737 operation: SnapshotOperation::Replace,
738 iceberg_schema: None,
739 extra_properties: std::collections::HashMap::new(),
740 bloom_filters: vec![],
741 equality_delete_files: vec![],
742 };
743 catalog.commit_snapshot(table, snapshot).await?;
744
745 info!(
746 "ailake: compaction committed — merged {} files into {}",
747 to_compact.len(),
748 output_path
749 );
750
751 if catalog.retires_files_physically() {
752 delete_old_files(&self.store, &to_compact).await;
753 } else {
754 info!(
755 "ailake: compaction — leaving {} retired file(s) in place; catalog backend \
756 manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
757 to_compact.len()
758 );
759 }
760
761 Ok(Some(merged))
762 }
763
764 pub async fn run_deferred(
774 &self,
775 planner: &CompactionPlanner,
776 table: &TableIdent,
777 catalog: Arc<dyn CatalogProvider>,
778 output_prefix: &str,
779 ) -> AilakeResult<Option<DataFileEntry>> {
780 if !catalog.supports_in_place_rewrite() {
785 return Err(ailake_core::AilakeError::Catalog(
786 "deferred compaction is not supported with this catalog backend: the \
787 background index build patches the merged file in place at its committed \
788 path, which this catalog cannot re-register — run a blocking compact"
789 .into(),
790 ));
791 }
792 let all_files = catalog.list_files(table, None).await?;
793 let to_compact = planner.plan(&all_files);
794 if to_compact.is_empty() {
795 return Ok(None);
796 }
797
798 let ts = std::time::SystemTime::now()
799 .duration_since(std::time::UNIX_EPOCH)
800 .unwrap_or_else(|e| e.duration())
801 .as_millis();
802 let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
803
804 let merged = self
805 .compact_deferred(&to_compact, &output_path, catalog.clone(), table)
806 .await?;
807
808 let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
810 let parent_snapshot_id = catalog
813 .load_table(table)
814 .await
815 .ok()
816 .and_then(|m| m.current_snapshot_id);
817 let snapshot = NewSnapshot {
818 snapshot_id: ailake_catalog::new_snapshot_id(),
819 parent_snapshot_id,
820 files,
821 operation: SnapshotOperation::Replace,
822 iceberg_schema: None,
823 extra_properties: std::collections::HashMap::new(),
824 bloom_filters: vec![],
825 equality_delete_files: vec![],
826 };
827 catalog.commit_snapshot(table, snapshot).await?;
828
829 info!(
830 "ailake: compaction committed (deferred) — merged {} files into {} \
831 (index building in background)",
832 to_compact.len(),
833 output_path
834 );
835
836 if catalog.retires_files_physically() {
837 delete_old_files(&self.store, &to_compact).await;
838 } else {
839 info!(
840 "ailake: compaction — leaving {} retired file(s) in place; catalog backend \
841 manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
842 to_compact.len()
843 );
844 }
845
846 Ok(Some(merged))
847 }
848}
849
850async fn build_replace_file_list(
861 catalog: &Arc<dyn CatalogProvider>,
862 table: &TableIdent,
863 to_compact: &[DataFileEntry],
864 merged: DataFileEntry,
865) -> AilakeResult<Vec<DataFileEntry>> {
866 let current_files = catalog.list_files(table, None).await?;
867 let compacted_paths: std::collections::HashSet<&str> =
868 to_compact.iter().map(|f| f.path.as_str()).collect();
869 let mut files: Vec<DataFileEntry> = current_files
870 .into_iter()
871 .filter(|f| !compacted_paths.contains(f.path.as_str()))
872 .collect();
873 files.push(merged);
874 Ok(files)
875}
876
877async fn delete_old_files(store: &Arc<dyn Store>, files: &[DataFileEntry]) {
878 for entry in files {
879 if let Err(e) = store.delete(&entry.path).await {
880 error!(
881 "ailake: compaction cleanup failed — could not delete {}: {} \
882 (orphan file in object store after successful catalog commit; \
883 delete manually to reclaim storage)",
884 entry.path, e
885 );
886 }
887 }
888}
889
890fn concat_batches(schema: SchemaRef, batches: &[RecordBatch]) -> AilakeResult<RecordBatch> {
891 arrow_select::concat::concat_batches(&schema, batches)
892 .map_err(|e| ailake_core::AilakeError::Arrow(e.to_string()))
893}
894
895#[cfg(test)]
896mod tests {
897 use super::*;
898 use ailake_catalog::IndexStatus;
899
900 #[test]
901 fn plan_returns_empty_if_too_few_files() {
902 let planner = CompactionPlanner::new(CompactionConfig {
903 min_files_to_compact: 4,
904 target_file_size_bytes: 1024 * 1024,
905 ..Default::default()
906 });
907 let files: Vec<DataFileEntry> = (0..3)
908 .map(|i| DataFileEntry {
909 path: format!("file-{i}.parquet"),
910 record_count: 10,
911 file_size_bytes: 100,
912 centroid_b64: Some("AAAA".into()),
913 radius: None,
914 hnsw_offset: None,
915 hnsw_len: None,
916 vector_column: None,
917 vector_dim: None,
918 extra_vector_indexes: vec![],
919 index_status: IndexStatus::Ready,
920 index_error: None,
921 batch_id: None,
922 embedding_model: None,
923 partition_value: None,
924 deletion_vector: None,
925 first_row_id: None,
926 })
927 .collect();
928 assert!(planner.plan(&files).is_empty());
929 }
930
931 #[test]
932 fn plan_selects_small_files() {
933 let planner = CompactionPlanner::new(CompactionConfig {
934 min_files_to_compact: 2,
935 target_file_size_bytes: 1000,
936 ..Default::default()
937 });
938 let files = vec![
939 DataFileEntry {
940 path: "small.parquet".into(),
941 record_count: 5,
942 file_size_bytes: 500,
943 centroid_b64: Some("AAAA".into()),
944 radius: None,
945 hnsw_offset: None,
946 hnsw_len: None,
947 vector_column: None,
948 vector_dim: None,
949 extra_vector_indexes: vec![],
950 index_status: IndexStatus::Ready,
951 index_error: None,
952 batch_id: None,
953 embedding_model: None,
954 partition_value: None,
955 deletion_vector: None,
956 first_row_id: None,
957 },
958 DataFileEntry {
959 path: "large.parquet".into(),
960 record_count: 5000,
961 file_size_bytes: 200_000_000,
962 centroid_b64: Some("AAAA".into()),
963 radius: None,
964 hnsw_offset: None,
965 hnsw_len: None,
966 vector_column: None,
967 vector_dim: None,
968 extra_vector_indexes: vec![],
969 index_status: IndexStatus::Ready,
970 index_error: None,
971 batch_id: None,
972 embedding_model: None,
973 partition_value: None,
974 deletion_vector: None,
975 first_row_id: None,
976 },
977 DataFileEntry {
978 path: "also-small.parquet".into(),
979 record_count: 5,
980 file_size_bytes: 800,
981 centroid_b64: Some("AAAA".into()),
982 radius: None,
983 hnsw_offset: None,
984 hnsw_len: None,
985 vector_column: None,
986 vector_dim: None,
987 extra_vector_indexes: vec![],
988 index_status: IndexStatus::Ready,
989 index_error: None,
990 batch_id: None,
991 embedding_model: None,
992 partition_value: None,
993 deletion_vector: None,
994 first_row_id: None,
995 },
996 ];
997 let selected = planner.plan(&files);
998 assert_eq!(selected.len(), 2);
999 assert!(selected.iter().any(|f| f.path == "small.parquet"));
1000 assert!(selected.iter().any(|f| f.path == "also-small.parquet"));
1001 }
1002
1003 #[test]
1004 fn plan_respects_max_files_per_pass() {
1005 let planner = CompactionPlanner::new(CompactionConfig {
1006 min_files_to_compact: 2,
1007 target_file_size_bytes: 1_000_000,
1008 max_files_per_pass: 3,
1009 ..Default::default()
1010 });
1011 let files: Vec<DataFileEntry> = (0..5)
1012 .map(|i| DataFileEntry {
1013 path: format!("f{i}.parquet"),
1014 record_count: 10,
1015 file_size_bytes: 100 + i as u64 * 100,
1016 centroid_b64: Some("AAAA".into()),
1017 radius: None,
1018 hnsw_offset: None,
1019 hnsw_len: None,
1020 vector_column: None,
1021 vector_dim: None,
1022 extra_vector_indexes: vec![],
1023 index_status: IndexStatus::Ready,
1024 index_error: None,
1025 batch_id: None,
1026 embedding_model: None,
1027 partition_value: None,
1028 deletion_vector: None,
1029 first_row_id: None,
1030 })
1031 .collect();
1032 let selected = planner.plan(&files);
1033 assert_eq!(selected.len(), 3);
1034 assert_eq!(selected[0].file_size_bytes, 100);
1035 assert_eq!(selected[1].file_size_bytes, 200);
1036 assert_eq!(selected[2].file_size_bytes, 300);
1037 }
1038
1039 #[test]
1040 fn plan_sorts_smallest_first() {
1041 let planner = CompactionPlanner::new(CompactionConfig {
1042 min_files_to_compact: 2,
1043 target_file_size_bytes: 10_000,
1044 max_files_per_pass: 4,
1045 ..Default::default()
1046 });
1047 let files = vec![
1048 DataFileEntry {
1049 path: "c.parquet".into(),
1050 record_count: 1,
1051 file_size_bytes: 300,
1052 centroid_b64: Some("AAAA".into()),
1053 radius: None,
1054 hnsw_offset: None,
1055 hnsw_len: None,
1056 vector_column: None,
1057 vector_dim: None,
1058 extra_vector_indexes: vec![],
1059 index_status: IndexStatus::Ready,
1060 index_error: None,
1061 batch_id: None,
1062 embedding_model: None,
1063 partition_value: None,
1064 deletion_vector: None,
1065 first_row_id: None,
1066 },
1067 DataFileEntry {
1068 path: "a.parquet".into(),
1069 record_count: 1,
1070 file_size_bytes: 100,
1071 centroid_b64: Some("AAAA".into()),
1072 radius: None,
1073 hnsw_offset: None,
1074 hnsw_len: None,
1075 vector_column: None,
1076 vector_dim: None,
1077 extra_vector_indexes: vec![],
1078 index_status: IndexStatus::Ready,
1079 index_error: None,
1080 batch_id: None,
1081 embedding_model: None,
1082 partition_value: None,
1083 deletion_vector: None,
1084 first_row_id: None,
1085 },
1086 DataFileEntry {
1087 path: "b.parquet".into(),
1088 record_count: 1,
1089 file_size_bytes: 200,
1090 centroid_b64: Some("AAAA".into()),
1091 radius: None,
1092 hnsw_offset: None,
1093 hnsw_len: None,
1094 vector_column: None,
1095 vector_dim: None,
1096 extra_vector_indexes: vec![],
1097 index_status: IndexStatus::Ready,
1098 index_error: None,
1099 batch_id: None,
1100 embedding_model: None,
1101 partition_value: None,
1102 deletion_vector: None,
1103 first_row_id: None,
1104 },
1105 ];
1106 let selected = planner.plan(&files);
1107 assert_eq!(selected[0].file_size_bytes, 100);
1108 assert_eq!(selected[1].file_size_bytes, 200);
1109 assert_eq!(selected[2].file_size_bytes, 300);
1110 }
1111
1112 fn make_plan_entry(path: &str, size: u64, centroid_b64: Option<String>) -> DataFileEntry {
1113 DataFileEntry {
1114 path: path.to_string(),
1115 record_count: 10,
1116 file_size_bytes: size,
1117 centroid_b64,
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 }
1132 }
1133
1134 #[test]
1135 fn plan_prioritizes_foreign_written_files_regardless_of_size() {
1136 let planner = CompactionPlanner::new(CompactionConfig {
1137 min_files_to_compact: 4, target_file_size_bytes: 1000,
1139 max_files_per_pass: 20,
1140 ..Default::default()
1141 });
1142 let files = vec![
1143 make_plan_entry("native_small.parquet", 500, Some("AAAA".into())),
1146 make_plan_entry("foreign_big.parquet", 200_000_000, None),
1149 ];
1150 let selected = planner.plan(&files);
1151 assert_eq!(
1152 selected.len(),
1153 1,
1154 "foreign file alone should trigger a pass even below min_files_to_compact"
1155 );
1156 assert_eq!(selected[0].path, "foreign_big.parquet");
1157 }
1158
1159 #[test]
1160 fn plan_merges_foreign_and_size_candidates_together() {
1161 let planner = CompactionPlanner::new(CompactionConfig {
1162 min_files_to_compact: 2,
1163 target_file_size_bytes: 1000,
1164 max_files_per_pass: 20,
1165 ..Default::default()
1166 });
1167 let files = vec![
1168 make_plan_entry("native_a.parquet", 300, Some("AAAA".into())),
1169 make_plan_entry("native_b.parquet", 400, Some("AAAA".into())),
1170 make_plan_entry("foreign.parquet", 200_000_000, None),
1171 ];
1172 let selected = planner.plan(&files);
1173 let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1174 assert_eq!(selected.len(), 3, "paths={paths:?}");
1175 assert_eq!(selected[0].path, "foreign.parquet", "paths={paths:?}");
1177 }
1178
1179 #[test]
1180 fn plan_sorts_foreign_files_smallest_first_too() {
1181 let planner = CompactionPlanner::new(CompactionConfig {
1185 min_files_to_compact: 1,
1186 target_file_size_bytes: 1000,
1187 max_files_per_pass: 2,
1188 ..Default::default()
1189 });
1190 let files = vec![
1191 make_plan_entry("foreign_huge.parquet", 500_000_000, None),
1192 make_plan_entry("foreign_small.parquet", 100, None),
1193 make_plan_entry("foreign_medium.parquet", 10_000_000, None),
1194 ];
1195 let selected = planner.plan(&files);
1196 let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1197 assert_eq!(selected.len(), 2, "max_files_per_pass=2, paths={paths:?}");
1198 assert_eq!(
1199 paths,
1200 vec!["foreign_small.parquet", "foreign_medium.parquet"],
1201 "foreign files must be size-sorted before truncation, cheapest first"
1202 );
1203 }
1204
1205 #[tokio::test]
1206 async fn compact_merges_two_files() {
1207 use ailake_core::{VectorMetric, VectorPrecision};
1208 use ailake_store::LocalStore;
1209 use arrow_array::{Int32Array, RecordBatch};
1210 use arrow_schema::{DataType, Field, Schema};
1211 use std::sync::Arc;
1212 use tempfile::TempDir;
1213
1214 let dir = TempDir::new().unwrap();
1215 let store = Arc::new(LocalStore::new(dir.path()));
1216 let policy = VectorStoragePolicy {
1217 column_name: "embedding".into(),
1218 dim: 4,
1219 metric: VectorMetric::Cosine,
1220 precision: VectorPrecision::F16,
1221 pq: None,
1222 keep_raw_for_reranking: true,
1223 pre_normalize: false,
1224 hnsw_m: None,
1225 hnsw_ef_construction: None,
1226 ivf_residual: false,
1227 embedding_model: None,
1228 modality: None,
1229 partition_by: None,
1230 partition_value: None,
1231 partition_column_type: None,
1232 partition_fields: vec![],
1233 };
1234
1235 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1236 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]];
1237 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]];
1238
1239 let batch_a = RecordBatch::try_new(
1240 schema.clone(),
1241 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1242 )
1243 .unwrap();
1244 let batch_b = RecordBatch::try_new(
1245 schema.clone(),
1246 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1247 )
1248 .unwrap();
1249
1250 let writer_a = AilakeFileWriter::new(policy.clone());
1251 let bytes_a = writer_a.write(&batch_a, &embs_a).unwrap();
1252 let writer_b = AilakeFileWriter::new(policy.clone());
1253 let bytes_b = writer_b.write(&batch_b, &embs_b).unwrap();
1254
1255 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1256 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1257
1258 let entries = vec![
1259 DataFileEntry {
1260 path: "data/a.parquet".into(),
1261 record_count: 2,
1262 file_size_bytes: bytes_a.len() as u64,
1263 centroid_b64: None,
1264 radius: None,
1265 hnsw_offset: None,
1266 hnsw_len: None,
1267 vector_column: None,
1268 vector_dim: None,
1269 extra_vector_indexes: vec![],
1270 index_status: IndexStatus::Ready,
1271 index_error: None,
1272 batch_id: None,
1273 embedding_model: None,
1274 partition_value: None,
1275 deletion_vector: None,
1276 first_row_id: None,
1277 },
1278 DataFileEntry {
1279 path: "data/b.parquet".into(),
1280 record_count: 2,
1281 file_size_bytes: bytes_b.len() as u64,
1282 centroid_b64: None,
1283 radius: None,
1284 hnsw_offset: None,
1285 hnsw_len: None,
1286 vector_column: None,
1287 vector_dim: None,
1288 extra_vector_indexes: vec![],
1289 index_status: IndexStatus::Ready,
1290 index_error: None,
1291 batch_id: None,
1292 embedding_model: None,
1293 partition_value: None,
1294 deletion_vector: None,
1295 first_row_id: None,
1296 },
1297 ];
1298
1299 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1300 let merged = executor
1301 .compact(&entries, "data/merged.parquet")
1302 .await
1303 .unwrap();
1304
1305 assert_eq!(merged.record_count, 4);
1306 assert_eq!(merged.path, "data/merged.parquet");
1307
1308 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1309 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1310 reader.verify_integrity().unwrap();
1311 let (batch, embs) = reader.read_parquet().unwrap();
1312 assert_eq!(batch.num_rows(), 4);
1313 assert_eq!(embs.len(), 4);
1314 }
1315
1316 #[tokio::test]
1321 async fn compact_drops_deletion_vector_masked_rows() {
1322 use ailake_catalog::provider::DeletionVector;
1323 use ailake_core::{VectorMetric, VectorPrecision};
1324 use ailake_store::LocalStore;
1325 use arrow_array::{Int32Array, RecordBatch};
1326 use arrow_schema::{DataType, Field, Schema};
1327 use roaring::RoaringBitmap;
1328 use std::sync::Arc;
1329 use tempfile::TempDir;
1330
1331 let dir = TempDir::new().unwrap();
1332 let store = Arc::new(LocalStore::new(dir.path()));
1333 let policy = VectorStoragePolicy {
1334 column_name: "embedding".into(),
1335 dim: 4,
1336 metric: VectorMetric::Cosine,
1337 precision: VectorPrecision::F16,
1338 pq: None,
1339 keep_raw_for_reranking: true,
1340 pre_normalize: false,
1341 hnsw_m: None,
1342 hnsw_ef_construction: None,
1343 ivf_residual: false,
1344 embedding_model: None,
1345 modality: None,
1346 partition_by: None,
1347 partition_value: None,
1348 partition_column_type: None,
1349 partition_fields: vec![],
1350 };
1351
1352 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1353 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]];
1355 let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
1356
1357 let batch_a = RecordBatch::try_new(
1358 schema.clone(),
1359 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1360 )
1361 .unwrap();
1362 let batch_b =
1363 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![2i32]))])
1364 .unwrap();
1365
1366 let bytes_a = AilakeFileWriter::new(policy.clone())
1367 .write(&batch_a, &embs_a)
1368 .unwrap();
1369 let bytes_b = AilakeFileWriter::new(policy.clone())
1370 .write(&batch_b, &embs_b)
1371 .unwrap();
1372 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1373 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1374
1375 let mut bitmap = RoaringBitmap::new();
1377 bitmap.insert(0);
1378 let (puffin_bytes, offset, length) =
1379 crate::delete::PuffinWriter::write_single_dv(&bitmap, 1).unwrap();
1380 store.put("metadata/dv-1.dvd", puffin_bytes).await.unwrap();
1381
1382 let entry_a = DataFileEntry {
1383 path: "data/a.parquet".into(),
1384 record_count: 2,
1385 file_size_bytes: bytes_a.len() as u64,
1386 centroid_b64: None,
1387 radius: None,
1388 hnsw_offset: None,
1389 hnsw_len: None,
1390 vector_column: None,
1391 vector_dim: None,
1392 extra_vector_indexes: vec![],
1393 index_status: IndexStatus::Ready,
1394 index_error: None,
1395 batch_id: None,
1396 embedding_model: None,
1397 partition_value: None,
1398 deletion_vector: Some(DeletionVector {
1399 path: "metadata/dv-1.dvd".into(),
1400 offset,
1401 length,
1402 cardinality: 1,
1403 }),
1404 first_row_id: None,
1405 };
1406 let entry_b = DataFileEntry {
1407 path: "data/b.parquet".into(),
1408 record_count: 1,
1409 file_size_bytes: bytes_b.len() as u64,
1410 centroid_b64: None,
1411 radius: None,
1412 hnsw_offset: None,
1413 hnsw_len: None,
1414 vector_column: None,
1415 vector_dim: None,
1416 extra_vector_indexes: vec![],
1417 index_status: IndexStatus::Ready,
1418 index_error: None,
1419 batch_id: None,
1420 embedding_model: None,
1421 partition_value: None,
1422 deletion_vector: None,
1423 first_row_id: None,
1424 };
1425
1426 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1427 let merged = executor
1428 .compact(&[entry_a, entry_b], "data/merged.parquet")
1429 .await
1430 .unwrap();
1431
1432 assert_eq!(merged.record_count, 2);
1434 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1435 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1436 let (batch, embs) = reader.read_parquet().unwrap();
1437 assert_eq!(batch.num_rows(), 2);
1438 assert_eq!(embs.len(), 2);
1439 let ids: Vec<i32> = batch
1440 .column_by_name("id")
1441 .unwrap()
1442 .as_any()
1443 .downcast_ref::<Int32Array>()
1444 .unwrap()
1445 .values()
1446 .to_vec();
1447 assert_eq!(
1448 ids,
1449 vec![1, 2],
1450 "id=0 (deleted) must not survive compaction"
1451 );
1452 }
1453
1454 #[tokio::test]
1459 async fn compact_with_ivf_pq_strategy_does_not_crash_on_verify_integrity() {
1460 use ailake_core::{VectorMetric, VectorPrecision};
1461 use ailake_store::LocalStore;
1462 use arrow_array::{Int32Array, RecordBatch};
1463 use arrow_schema::{DataType, Field, Schema};
1464 use std::sync::Arc;
1465 use tempfile::TempDir;
1466
1467 let dir = TempDir::new().unwrap();
1468 let store = Arc::new(LocalStore::new(dir.path()));
1469 let dim = 8;
1470 let policy = VectorStoragePolicy {
1471 column_name: "embedding".into(),
1472 dim,
1473 metric: VectorMetric::Cosine,
1474 precision: VectorPrecision::F16,
1475 pq: None,
1476 keep_raw_for_reranking: true,
1477 pre_normalize: false,
1478 hnsw_m: None,
1479 hnsw_ef_construction: None,
1480 ivf_residual: false,
1481 embedding_model: None,
1482 modality: None,
1483 partition_by: None,
1484 partition_value: None,
1485 partition_column_type: None,
1486 partition_fields: vec![],
1487 };
1488
1489 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1490 let n_per_file = 30usize;
1491 let make_file = |path: &str, offset: i32| {
1492 let ids: Vec<i32> = (offset..offset + n_per_file as i32).collect();
1493 let embs: Vec<Vec<f32>> = ids
1494 .iter()
1495 .map(|&i| {
1496 (0..dim as i32)
1497 .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1498 .collect()
1499 })
1500 .collect();
1501 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1502 .unwrap();
1503 let bytes = AilakeFileWriter::new(policy.clone())
1504 .write(&batch, &embs)
1505 .unwrap();
1506 (path.to_string(), bytes)
1507 };
1508
1509 let (path_a, bytes_a) = make_file("data/ivfpq_a.parquet", 0);
1510 let (path_b, bytes_b) = make_file("data/ivfpq_b.parquet", n_per_file as i32);
1511 for (path, bytes) in [(&path_a, &bytes_a), (&path_b, &bytes_b)] {
1512 store.put(path, bytes.clone()).await.unwrap();
1513 }
1514
1515 let make_entry = |path: &str, size: u64| DataFileEntry {
1516 path: path.to_string(),
1517 record_count: n_per_file as u64,
1518 file_size_bytes: size,
1519 centroid_b64: None,
1520 radius: None,
1521 hnsw_offset: None,
1522 hnsw_len: None,
1523 vector_column: None,
1524 vector_dim: None,
1525 extra_vector_indexes: vec![],
1526 index_status: IndexStatus::Ready,
1527 index_error: None,
1528 batch_id: None,
1529 embedding_model: None,
1530 partition_value: None,
1531 deletion_vector: None,
1532 first_row_id: None,
1533 };
1534 let entries = vec![
1535 make_entry(&path_a, bytes_a.len() as u64),
1536 make_entry(&path_b, bytes_b.len() as u64),
1537 ];
1538
1539 let executor = CompactionExecutor::new(store.clone(), policy.clone())
1540 .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1541 let merged = executor
1542 .compact(&entries, "data/ivfpq_merged.parquet")
1543 .await
1544 .expect("compact() with ForceIvfPq must not fail in verify_integrity()");
1545
1546 assert_eq!(merged.record_count, 2 * n_per_file as u64);
1547
1548 let merged_bytes = store.get("data/ivfpq_merged.parquet").await.unwrap();
1550 let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1551 reader
1552 .verify_integrity()
1553 .expect("verify_integrity() must handle an IVF-PQ-indexed file");
1554 }
1555
1556 #[tokio::test]
1557 async fn compact_incremental_merges_dominant_plus_small() {
1558 use ailake_core::{RowId, VectorMetric, VectorPrecision};
1559 use ailake_store::LocalStore;
1560 use arrow_array::{Int32Array, RecordBatch};
1561 use arrow_schema::{DataType, Field, Schema};
1562 use std::sync::Arc;
1563 use tempfile::TempDir;
1564
1565 let dir = TempDir::new().unwrap();
1566 let store = Arc::new(LocalStore::new(dir.path()));
1567 let policy = VectorStoragePolicy {
1568 column_name: "embedding".into(),
1569 dim: 4,
1570 metric: VectorMetric::Cosine,
1571 precision: VectorPrecision::F16,
1572 pq: None,
1573 keep_raw_for_reranking: true,
1574 pre_normalize: false,
1575 hnsw_m: None,
1576 hnsw_ef_construction: None,
1577 ivf_residual: false,
1578 embedding_model: None,
1579 modality: None,
1580 partition_by: None,
1581 partition_value: None,
1582 partition_column_type: None,
1583 partition_fields: vec![],
1584 };
1585
1586 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1587
1588 let embs_dom: Vec<Vec<f32>> = vec![
1590 vec![1.0, 0.0, 0.0, 0.0],
1591 vec![0.0, 1.0, 0.0, 0.0],
1592 vec![0.0, 0.0, 1.0, 0.0],
1593 vec![0.7, 0.7, 0.0, 0.0],
1594 vec![0.0, 0.7, 0.7, 0.0],
1595 vec![0.0, 0.0, 0.7, 0.7],
1596 ];
1597 let batch_dom = RecordBatch::try_new(
1598 schema.clone(),
1599 vec![Arc::new(Int32Array::from(vec![0i32, 1, 2, 3, 4, 5]))],
1600 )
1601 .unwrap();
1602
1603 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]];
1605 let batch_small = RecordBatch::try_new(
1606 schema.clone(),
1607 vec![Arc::new(Int32Array::from(vec![6i32, 7]))],
1608 )
1609 .unwrap();
1610
1611 let bytes_dom = AilakeFileWriter::new(policy.clone())
1612 .write(&batch_dom, &embs_dom)
1613 .unwrap();
1614 let bytes_small = AilakeFileWriter::new(policy.clone())
1615 .write(&batch_small, &embs_small)
1616 .unwrap();
1617
1618 store
1619 .put("data/dominant.parquet", bytes_dom.clone())
1620 .await
1621 .unwrap();
1622 store
1623 .put("data/small.parquet", bytes_small.clone())
1624 .await
1625 .unwrap();
1626
1627 let entries = vec![
1628 DataFileEntry {
1629 path: "data/dominant.parquet".into(),
1630 record_count: 6,
1631 file_size_bytes: bytes_dom.len() as u64,
1632 centroid_b64: None,
1633 radius: None,
1634 hnsw_offset: None,
1635 hnsw_len: None,
1636 vector_column: None,
1637 vector_dim: None,
1638 extra_vector_indexes: vec![],
1639 index_status: IndexStatus::Ready,
1640 index_error: None,
1641 batch_id: None,
1642 embedding_model: None,
1643 partition_value: None,
1644 deletion_vector: None,
1645 first_row_id: None,
1646 },
1647 DataFileEntry {
1648 path: "data/small.parquet".into(),
1649 record_count: 2,
1650 file_size_bytes: bytes_small.len() as u64,
1651 centroid_b64: None,
1652 radius: None,
1653 hnsw_offset: None,
1654 hnsw_len: None,
1655 vector_column: None,
1656 vector_dim: None,
1657 extra_vector_indexes: vec![],
1658 index_status: IndexStatus::Ready,
1659 index_error: None,
1660 batch_id: None,
1661 embedding_model: None,
1662 partition_value: None,
1663 deletion_vector: None,
1664 first_row_id: None,
1665 },
1666 ];
1667
1668 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1669 let merged = executor
1670 .compact_incremental(&entries, "data/merged.parquet")
1671 .await
1672 .unwrap();
1673
1674 assert_eq!(merged.record_count, 8);
1676 assert_eq!(merged.path, "data/merged.parquet");
1677
1678 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1680 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1681 reader.verify_integrity().unwrap();
1682
1683 let (batch, embs) = reader.read_parquet().unwrap();
1684 assert_eq!(batch.num_rows(), 8);
1685 assert_eq!(embs.len(), 8);
1686
1687 for f in &embs[..6] {
1689 assert_eq!(f.len(), 4);
1690 }
1691
1692 let hnsw = reader.load_index().unwrap();
1694 assert_eq!(hnsw.node_count(), 8);
1695
1696 let results = hnsw.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
1698 assert_eq!(results[0].0, RowId::new(0));
1699
1700 let results = hnsw.search(&[0.0, 0.0, 0.0, 1.0], 1, 50);
1703 assert_eq!(results[0].0, RowId::new(6));
1704 }
1705
1706 #[tokio::test]
1716 async fn compact_incremental_respects_force_ivf_pq_even_with_dominant_file() {
1717 use ailake_core::{VectorMetric, VectorPrecision};
1718 use ailake_store::LocalStore;
1719 use arrow_array::{Int32Array, RecordBatch};
1720 use arrow_schema::{DataType, Field, Schema};
1721 use std::sync::Arc;
1722 use tempfile::TempDir;
1723
1724 let dir = TempDir::new().unwrap();
1725 let store = Arc::new(LocalStore::new(dir.path()));
1726 let dim = 8;
1727 let policy = VectorStoragePolicy {
1728 column_name: "embedding".into(),
1729 dim,
1730 metric: VectorMetric::Cosine,
1731 precision: VectorPrecision::F16,
1732 pq: None,
1733 keep_raw_for_reranking: true,
1734 pre_normalize: false,
1735 hnsw_m: None,
1736 hnsw_ef_construction: None,
1737 ivf_residual: false,
1738 embedding_model: None,
1739 modality: None,
1740 partition_by: None,
1741 partition_value: None,
1742 partition_column_type: None,
1743 partition_fields: vec![],
1744 };
1745
1746 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1747 let make_file = |path: &str, offset: i32, n: usize| {
1748 let ids: Vec<i32> = (offset..offset + n as i32).collect();
1749 let embs: Vec<Vec<f32>> = ids
1750 .iter()
1751 .map(|&i| {
1752 (0..dim as i32)
1753 .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1754 .collect()
1755 })
1756 .collect();
1757 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1758 .unwrap();
1759 let bytes = AilakeFileWriter::new(policy.clone())
1762 .write(&batch, &embs)
1763 .unwrap();
1764 (path.to_string(), bytes, n as u64)
1765 };
1766
1767 let (path_dom, bytes_dom, n_dom) = make_file("data/dom.parquet", 0, 90);
1770 let (path_small, bytes_small, n_small) = make_file("data/small.parquet", 999, 30);
1771
1772 for (path, bytes) in [(&path_dom, &bytes_dom), (&path_small, &bytes_small)] {
1773 store.put(path, bytes.clone()).await.unwrap();
1774 }
1775
1776 let make_entry = |path: &str, record_count: u64, size: u64| DataFileEntry {
1777 path: path.to_string(),
1778 record_count,
1779 file_size_bytes: size,
1780 centroid_b64: None,
1781 radius: None,
1782 hnsw_offset: None,
1783 hnsw_len: None,
1784 vector_column: None,
1785 vector_dim: None,
1786 extra_vector_indexes: vec![],
1787 index_status: IndexStatus::Ready,
1788 index_error: None,
1789 batch_id: None,
1790 embedding_model: None,
1791 partition_value: None,
1792 deletion_vector: None,
1793 first_row_id: None,
1794 };
1795 let entries = vec![
1796 make_entry(&path_dom, n_dom, bytes_dom.len() as u64),
1797 make_entry(&path_small, n_small, bytes_small.len() as u64),
1798 ];
1799
1800 let executor = CompactionExecutor::new(store.clone(), policy.clone())
1801 .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1802 let merged = executor
1803 .compact_incremental(&entries, "data/merged_force_ivfpq.parquet")
1804 .await
1805 .expect("compact_incremental() with ForceIvfPq must fall back to compact() cleanly");
1806
1807 let merged_bytes = store.get("data/merged_force_ivfpq.parquet").await.unwrap();
1808 let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1809 reader.verify_integrity().unwrap();
1810
1811 match reader.load_any_index().unwrap() {
1812 ailake_index::AnyIndex::IvfPq(_) => {}
1813 ailake_index::AnyIndex::Hnsw(_) => panic!(
1814 "ForceIvfPq was silently satisfied with HNSW instead — merged.record_count={}",
1815 merged.record_count
1816 ),
1817 }
1818 }
1819
1820 #[tokio::test]
1821 async fn compact_incremental_falls_back_when_no_dominant() {
1822 use ailake_core::{VectorMetric, VectorPrecision};
1823 use ailake_store::LocalStore;
1824 use arrow_array::{Int32Array, RecordBatch};
1825 use arrow_schema::{DataType, Field, Schema};
1826 use std::sync::Arc;
1827 use tempfile::TempDir;
1828
1829 let dir = TempDir::new().unwrap();
1830 let store = Arc::new(LocalStore::new(dir.path()));
1831 let policy = VectorStoragePolicy {
1832 column_name: "embedding".into(),
1833 dim: 4,
1834 metric: VectorMetric::Cosine,
1835 precision: VectorPrecision::F16,
1836 pq: None,
1837 keep_raw_for_reranking: true,
1838 pre_normalize: false,
1839 hnsw_m: None,
1840 hnsw_ef_construction: None,
1841 ivf_residual: false,
1842 embedding_model: None,
1843 modality: None,
1844 partition_by: None,
1845 partition_value: None,
1846 partition_column_type: None,
1847 partition_fields: vec![],
1848 };
1849
1850 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1851
1852 let make_batch = |ids: Vec<i32>, embs: Vec<Vec<f32>>| {
1854 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1855 .unwrap();
1856 AilakeFileWriter::new(policy.clone())
1857 .write(&batch, &embs)
1858 .unwrap()
1859 };
1860
1861 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]];
1862 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]];
1863 let bytes_a = make_batch(vec![0, 1], embs_a);
1864 let bytes_b = make_batch(vec![2, 3], embs_b);
1865
1866 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1867 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1868
1869 let entries = vec![
1870 DataFileEntry {
1871 path: "data/a.parquet".into(),
1872 record_count: 2,
1873 file_size_bytes: bytes_a.len() as u64,
1874 centroid_b64: None,
1875 radius: None,
1876 hnsw_offset: None,
1877 hnsw_len: None,
1878 vector_column: None,
1879 vector_dim: None,
1880 extra_vector_indexes: vec![],
1881 index_status: IndexStatus::Ready,
1882 index_error: None,
1883 batch_id: None,
1884 embedding_model: None,
1885 partition_value: None,
1886 deletion_vector: None,
1887 first_row_id: None,
1888 },
1889 DataFileEntry {
1890 path: "data/b.parquet".into(),
1891 record_count: 2,
1892 file_size_bytes: bytes_b.len() as u64,
1893 centroid_b64: None,
1894 radius: None,
1895 hnsw_offset: None,
1896 hnsw_len: None,
1897 vector_column: None,
1898 vector_dim: None,
1899 extra_vector_indexes: vec![],
1900 index_status: IndexStatus::Ready,
1901 index_error: None,
1902 batch_id: None,
1903 embedding_model: None,
1904 partition_value: None,
1905 deletion_vector: None,
1906 first_row_id: None,
1907 },
1908 ];
1909
1910 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1911 let merged = executor
1913 .compact_incremental(&entries, "data/merged.parquet")
1914 .await
1915 .unwrap();
1916
1917 assert_eq!(merged.record_count, 4);
1918
1919 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1920 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1921 reader.verify_integrity().unwrap();
1922 }
1923
1924 #[tokio::test]
1925 async fn compact_deferred_produces_parquet_only_file() {
1926 use ailake_catalog::HadoopCatalog;
1927 use ailake_core::{VectorMetric, VectorPrecision};
1928 use ailake_store::LocalStore;
1929 use arrow_array::{Int32Array, RecordBatch};
1930 use arrow_schema::{DataType, Field, Schema};
1931 use std::sync::Arc;
1932 use tempfile::TempDir;
1933
1934 let dir = TempDir::new().unwrap();
1935 let store = Arc::new(LocalStore::new(dir.path()));
1936 let catalog_dir = TempDir::new().unwrap();
1937 let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
1938 let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
1939 let table = TableIdent {
1940 namespace: "ns".into(),
1941 name: "tbl".into(),
1942 };
1943
1944 let policy = VectorStoragePolicy {
1945 column_name: "embedding".into(),
1946 dim: 4,
1947 metric: VectorMetric::Cosine,
1948 precision: VectorPrecision::F16,
1949 pq: None,
1950 keep_raw_for_reranking: true,
1951 pre_normalize: false,
1952 hnsw_m: None,
1953 hnsw_ef_construction: None,
1954 ivf_residual: false,
1955 embedding_model: None,
1956 modality: None,
1957 partition_by: None,
1958 partition_value: None,
1959 partition_column_type: None,
1960 partition_fields: vec![],
1961 };
1962
1963 use ailake_catalog::TableProperties;
1964 catalog
1965 .create_table(
1966 &table,
1967 &TableProperties {
1968 policy: policy.clone(),
1969 extra: std::collections::HashMap::new(),
1970 format_version: 2,
1971 partition_column_type: None,
1972 },
1973 )
1974 .await
1975 .unwrap();
1976
1977 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1978 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]];
1979 let batch_a = RecordBatch::try_new(
1980 schema.clone(),
1981 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1982 )
1983 .unwrap();
1984 let bytes_a = AilakeFileWriter::new(policy.clone())
1985 .write(&batch_a, &embs_a)
1986 .unwrap();
1987 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1988
1989 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]];
1990 let batch_b = RecordBatch::try_new(
1991 schema.clone(),
1992 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1993 )
1994 .unwrap();
1995 let bytes_b = AilakeFileWriter::new(policy.clone())
1996 .write(&batch_b, &embs_b)
1997 .unwrap();
1998 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1999
2000 let entries = vec![
2001 DataFileEntry {
2002 path: "data/a.parquet".into(),
2003 record_count: 2,
2004 file_size_bytes: bytes_a.len() as u64,
2005 centroid_b64: None,
2006 radius: None,
2007 hnsw_offset: None,
2008 hnsw_len: None,
2009 vector_column: None,
2010 vector_dim: None,
2011 extra_vector_indexes: vec![],
2012 index_status: IndexStatus::Ready,
2013 index_error: None,
2014 batch_id: None,
2015 embedding_model: None,
2016 partition_value: None,
2017 deletion_vector: None,
2018 first_row_id: None,
2019 },
2020 DataFileEntry {
2021 path: "data/b.parquet".into(),
2022 record_count: 2,
2023 file_size_bytes: bytes_b.len() as u64,
2024 centroid_b64: None,
2025 radius: None,
2026 hnsw_offset: None,
2027 hnsw_len: None,
2028 vector_column: None,
2029 vector_dim: None,
2030 extra_vector_indexes: vec![],
2031 index_status: IndexStatus::Ready,
2032 index_error: None,
2033 batch_id: None,
2034 embedding_model: None,
2035 partition_value: None,
2036 deletion_vector: None,
2037 first_row_id: None,
2038 },
2039 ];
2040
2041 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2042 let entry = executor
2043 .compact_deferred(&entries, "data/merged.parquet", catalog.clone(), &table)
2044 .await
2045 .unwrap();
2046
2047 assert_eq!(entry.index_status, IndexStatus::Indexing);
2049 assert_eq!(entry.record_count, 4);
2050
2051 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
2053 let pq_reader = ailake_parquet::ParquetVectorReader::new(merged_bytes, "embedding");
2054 let count = pq_reader.record_count().unwrap();
2055 assert_eq!(count, 4);
2056 }
2057
2058 #[tokio::test]
2065 async fn run_preserves_untouched_files_outside_compaction_pass() {
2066 use ailake_catalog::HadoopCatalog;
2067 use ailake_core::{VectorMetric, VectorPrecision};
2068 use ailake_store::LocalStore;
2069 use arrow_array::{Int32Array, RecordBatch};
2070 use arrow_schema::{DataType, Field, Schema};
2071 use std::sync::Arc;
2072 use tempfile::TempDir;
2073
2074 let dir = TempDir::new().unwrap();
2075 let store = Arc::new(LocalStore::new(dir.path()));
2076 let catalog_dir = TempDir::new().unwrap();
2077 let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
2078 let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
2079 let table = TableIdent {
2080 namespace: "ns".into(),
2081 name: "tbl".into(),
2082 };
2083
2084 let policy = VectorStoragePolicy {
2085 column_name: "embedding".into(),
2086 dim: 4,
2087 metric: VectorMetric::Cosine,
2088 precision: VectorPrecision::F16,
2089 pq: None,
2090 keep_raw_for_reranking: true,
2091 pre_normalize: false,
2092 hnsw_m: None,
2093 hnsw_ef_construction: None,
2094 ivf_residual: false,
2095 embedding_model: None,
2096 modality: None,
2097 partition_by: None,
2098 partition_value: None,
2099 partition_column_type: None,
2100 partition_fields: vec![],
2101 };
2102
2103 use ailake_catalog::TableProperties;
2104 catalog
2105 .create_table(
2106 &table,
2107 &TableProperties {
2108 policy: policy.clone(),
2109 extra: std::collections::HashMap::new(),
2110 format_version: 2,
2111 partition_column_type: None,
2112 },
2113 )
2114 .await
2115 .unwrap();
2116
2117 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2118
2119 let write_file = |path: &str, ids: Vec<i32>, embs: Vec<Vec<f32>>| {
2121 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
2122 .unwrap();
2123 let bytes = AilakeFileWriter::new(policy.clone())
2124 .write(&batch, &embs)
2125 .unwrap();
2126 (path.to_string(), bytes)
2127 };
2128
2129 let (path_a, bytes_a) = write_file(
2130 "data/small_a.parquet",
2131 vec![0, 1],
2132 vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
2133 );
2134 let (path_b, bytes_b) = write_file(
2135 "data/small_b.parquet",
2136 vec![2, 3],
2137 vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]],
2138 );
2139 let (path_big, bytes_big) = write_file(
2142 "data/big.parquet",
2143 vec![4, 5],
2144 vec![vec![1.0, 1.0, 0.0, 0.0], vec![0.0, 1.0, 1.0, 0.0]],
2145 );
2146
2147 for (path, bytes) in [
2148 (&path_a, &bytes_a),
2149 (&path_b, &bytes_b),
2150 (&path_big, &bytes_big),
2151 ] {
2152 store.put(path, bytes.clone()).await.unwrap();
2153 }
2154
2155 let make_entry = |path: &str, size: u64| DataFileEntry {
2156 path: path.to_string(),
2157 record_count: 2,
2158 file_size_bytes: size,
2159 centroid_b64: Some("AAAA".into()),
2164 radius: None,
2165 hnsw_offset: None,
2166 hnsw_len: None,
2167 vector_column: None,
2168 vector_dim: None,
2169 extra_vector_indexes: vec![],
2170 index_status: IndexStatus::Ready,
2171 index_error: None,
2172 batch_id: None,
2173 embedding_model: None,
2174 partition_value: None,
2175 deletion_vector: None,
2176 first_row_id: None,
2177 };
2178
2179 let initial_snap_id = ailake_catalog::new_snapshot_id();
2180 let initial_snapshot = NewSnapshot {
2181 snapshot_id: initial_snap_id,
2182 parent_snapshot_id: None,
2183 files: vec![
2184 make_entry(&path_a, 500),
2185 make_entry(&path_b, 500),
2186 make_entry(&path_big, 200_000_000), ],
2188 operation: SnapshotOperation::Append,
2189 iceberg_schema: None,
2190 extra_properties: std::collections::HashMap::new(),
2191 bloom_filters: vec![],
2192 equality_delete_files: vec![],
2193 };
2194 catalog
2195 .commit_snapshot(&table, initial_snapshot)
2196 .await
2197 .unwrap();
2198
2199 let planner = CompactionPlanner::new(CompactionConfig {
2200 min_files_to_compact: 2,
2201 target_file_size_bytes: 1000,
2202 index_strategy: CompactionIndexStrategy::ForceHnsw,
2203 max_files_per_pass: 20,
2204 });
2205 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2206
2207 let merged = executor
2208 .run(&planner, &table, catalog.clone(), "data")
2209 .await
2210 .unwrap()
2211 .expect("compaction should have run — 2 eligible small files");
2212
2213 let files_after = catalog.list_files(&table, None).await.unwrap();
2214 let paths_after: Vec<&str> = files_after.iter().map(|f| f.path.as_str()).collect();
2215
2216 assert!(
2217 paths_after.contains(&path_big.as_str()),
2218 "BUG: untouched 'big.parquet' vanished after run() — files_after={paths_after:?}"
2219 );
2220 assert!(
2221 paths_after.contains(&merged.path.as_str()),
2222 "merged output file must be present — files_after={paths_after:?}"
2223 );
2224
2225 let meta_dir = catalog_dir.path().join("ns/tbl/metadata");
2232 let latest_metadata = std::fs::read_dir(&meta_dir)
2233 .unwrap()
2234 .filter_map(|e| e.ok())
2235 .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
2236 .max_by_key(|e| e.metadata().unwrap().modified().unwrap())
2237 .expect("metadata.json must exist after commit");
2238 let json: serde_json::Value =
2239 serde_json::from_slice(&std::fs::read(latest_metadata.path()).unwrap()).unwrap();
2240 let last_snapshot = json["snapshots"].as_array().unwrap().last().unwrap();
2241 assert_eq!(
2242 last_snapshot["parent-snapshot-id"].as_i64(),
2243 Some(initial_snap_id),
2244 "compaction's Replace snapshot must chain to the pre-compaction snapshot, not be orphaned"
2245 );
2246 assert!(
2247 !paths_after.contains(&path_a.as_str()) && !paths_after.contains(&path_b.as_str()),
2248 "compacted input files must no longer be listed — files_after={paths_after:?}"
2249 );
2250 assert_eq!(
2251 files_after.len(),
2252 2,
2253 "expected exactly [big.parquet, merged] — files_after={paths_after:?}"
2254 );
2255 }
2256
2257 #[tokio::test]
2263 async fn compact_preserves_rows_from_footerless_file() {
2264 use ailake_core::{VectorMetric, VectorPrecision};
2265 use ailake_store::LocalStore;
2266 use arrow_array::{Int32Array, RecordBatch};
2267 use arrow_schema::{DataType, Field, Schema};
2268 use std::sync::Arc;
2269 use tempfile::TempDir;
2270
2271 let dir = TempDir::new().unwrap();
2272 let store = Arc::new(LocalStore::new(dir.path()));
2273 let policy = VectorStoragePolicy {
2274 column_name: "embedding".into(),
2275 dim: 4,
2276 metric: VectorMetric::Cosine,
2277 precision: VectorPrecision::F16,
2278 pq: None,
2279 keep_raw_for_reranking: true,
2280 pre_normalize: false,
2281 hnsw_m: None,
2282 hnsw_ef_construction: None,
2283 ivf_residual: false,
2284 embedding_model: None,
2285 modality: None,
2286 partition_by: None,
2287 partition_value: None,
2288 partition_column_type: None,
2289 partition_fields: vec![],
2290 };
2291
2292 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2293
2294 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]];
2296 let batch_native = RecordBatch::try_new(
2297 schema.clone(),
2298 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2299 )
2300 .unwrap();
2301 let bytes_native = AilakeFileWriter::new(policy.clone())
2302 .write(&batch_native, &embs_native)
2303 .unwrap();
2304 store
2305 .put("data/native.parquet", bytes_native.clone())
2306 .await
2307 .unwrap();
2308
2309 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]];
2313 let batch_foreign = RecordBatch::try_new(
2314 schema.clone(),
2315 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2316 )
2317 .unwrap();
2318 let bytes_foreign = AilakeFileWriter::new(policy.clone())
2319 .write_parquet_only(&batch_foreign, &embs_foreign)
2320 .unwrap();
2321 store
2322 .put("data/foreign.parquet", bytes_foreign.clone())
2323 .await
2324 .unwrap();
2325
2326 let reader_foreign = AilakeFileReader::new(bytes_foreign.clone(), "embedding", 4);
2327 assert!(
2328 !reader_foreign.is_ailake_file(),
2329 "sanity: write_parquet_only must not embed an AILK footer"
2330 );
2331
2332 let entries = vec![
2333 DataFileEntry {
2334 path: "data/native.parquet".into(),
2335 record_count: 2,
2336 file_size_bytes: bytes_native.len() as u64,
2337 centroid_b64: None,
2338 radius: None,
2339 hnsw_offset: None,
2340 hnsw_len: None,
2341 vector_column: None,
2342 vector_dim: None,
2343 extra_vector_indexes: vec![],
2344 index_status: IndexStatus::Ready,
2345 index_error: None,
2346 batch_id: None,
2347 embedding_model: None,
2348 partition_value: None,
2349 deletion_vector: None,
2350 first_row_id: None,
2351 },
2352 DataFileEntry {
2353 path: "data/foreign.parquet".into(),
2354 record_count: 2,
2355 file_size_bytes: bytes_foreign.len() as u64,
2356 centroid_b64: None,
2357 radius: None,
2358 hnsw_offset: None,
2359 hnsw_len: None,
2360 vector_column: None,
2361 vector_dim: None,
2362 extra_vector_indexes: vec![],
2363 index_status: IndexStatus::Ready,
2364 index_error: None,
2365 batch_id: None,
2366 embedding_model: None,
2367 partition_value: None,
2368 deletion_vector: None,
2369 first_row_id: None,
2370 },
2371 ];
2372
2373 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2375 let merged = executor
2376 .compact(&entries, "data/merged_full.parquet")
2377 .await
2378 .unwrap();
2379 assert_eq!(
2380 merged.record_count, 4,
2381 "compact() must include all 4 rows — 2 native + 2 from the footerless file"
2382 );
2383 let merged_bytes = store.get("data/merged_full.parquet").await.unwrap();
2384 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
2385 reader.verify_integrity().unwrap();
2386
2387 let merged_inc = executor
2390 .compact_incremental(&entries, "data/merged_inc.parquet")
2391 .await
2392 .unwrap();
2393 assert_eq!(
2394 merged_inc.record_count, 4,
2395 "compact_incremental() must include all 4 rows — 2 native + 2 from the footerless file"
2396 );
2397 }
2398}