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 delete_old_files(&self.store, &to_compact).await;
752
753 Ok(Some(merged))
754 }
755
756 pub async fn run_deferred(
766 &self,
767 planner: &CompactionPlanner,
768 table: &TableIdent,
769 catalog: Arc<dyn CatalogProvider>,
770 output_prefix: &str,
771 ) -> AilakeResult<Option<DataFileEntry>> {
772 let all_files = catalog.list_files(table, None).await?;
773 let to_compact = planner.plan(&all_files);
774 if to_compact.is_empty() {
775 return Ok(None);
776 }
777
778 let ts = std::time::SystemTime::now()
779 .duration_since(std::time::UNIX_EPOCH)
780 .unwrap_or_else(|e| e.duration())
781 .as_millis();
782 let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
783
784 let merged = self
785 .compact_deferred(&to_compact, &output_path, catalog.clone(), table)
786 .await?;
787
788 let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
790 let parent_snapshot_id = catalog
793 .load_table(table)
794 .await
795 .ok()
796 .and_then(|m| m.current_snapshot_id);
797 let snapshot = NewSnapshot {
798 snapshot_id: ailake_catalog::new_snapshot_id(),
799 parent_snapshot_id,
800 files,
801 operation: SnapshotOperation::Replace,
802 iceberg_schema: None,
803 extra_properties: std::collections::HashMap::new(),
804 bloom_filters: vec![],
805 equality_delete_files: vec![],
806 };
807 catalog.commit_snapshot(table, snapshot).await?;
808
809 info!(
810 "ailake: compaction committed (deferred) — merged {} files into {} \
811 (index building in background)",
812 to_compact.len(),
813 output_path
814 );
815
816 delete_old_files(&self.store, &to_compact).await;
817
818 Ok(Some(merged))
819 }
820}
821
822async fn build_replace_file_list(
833 catalog: &Arc<dyn CatalogProvider>,
834 table: &TableIdent,
835 to_compact: &[DataFileEntry],
836 merged: DataFileEntry,
837) -> AilakeResult<Vec<DataFileEntry>> {
838 let current_files = catalog.list_files(table, None).await?;
839 let compacted_paths: std::collections::HashSet<&str> =
840 to_compact.iter().map(|f| f.path.as_str()).collect();
841 let mut files: Vec<DataFileEntry> = current_files
842 .into_iter()
843 .filter(|f| !compacted_paths.contains(f.path.as_str()))
844 .collect();
845 files.push(merged);
846 Ok(files)
847}
848
849async fn delete_old_files(store: &Arc<dyn Store>, files: &[DataFileEntry]) {
850 for entry in files {
851 if let Err(e) = store.delete(&entry.path).await {
852 error!(
853 "ailake: compaction cleanup failed — could not delete {}: {} \
854 (orphan file in object store after successful catalog commit; \
855 delete manually to reclaim storage)",
856 entry.path, e
857 );
858 }
859 }
860}
861
862fn concat_batches(schema: SchemaRef, batches: &[RecordBatch]) -> AilakeResult<RecordBatch> {
863 arrow_select::concat::concat_batches(&schema, batches)
864 .map_err(|e| ailake_core::AilakeError::Arrow(e.to_string()))
865}
866
867#[cfg(test)]
868mod tests {
869 use super::*;
870 use ailake_catalog::IndexStatus;
871
872 #[test]
873 fn plan_returns_empty_if_too_few_files() {
874 let planner = CompactionPlanner::new(CompactionConfig {
875 min_files_to_compact: 4,
876 target_file_size_bytes: 1024 * 1024,
877 ..Default::default()
878 });
879 let files: Vec<DataFileEntry> = (0..3)
880 .map(|i| DataFileEntry {
881 path: format!("file-{i}.parquet"),
882 record_count: 10,
883 file_size_bytes: 100,
884 centroid_b64: Some("AAAA".into()),
885 radius: None,
886 hnsw_offset: None,
887 hnsw_len: None,
888 vector_column: None,
889 vector_dim: None,
890 extra_vector_indexes: vec![],
891 index_status: IndexStatus::Ready,
892 index_error: None,
893 batch_id: None,
894 embedding_model: None,
895 partition_value: None,
896 deletion_vector: None,
897 first_row_id: None,
898 })
899 .collect();
900 assert!(planner.plan(&files).is_empty());
901 }
902
903 #[test]
904 fn plan_selects_small_files() {
905 let planner = CompactionPlanner::new(CompactionConfig {
906 min_files_to_compact: 2,
907 target_file_size_bytes: 1000,
908 ..Default::default()
909 });
910 let files = vec![
911 DataFileEntry {
912 path: "small.parquet".into(),
913 record_count: 5,
914 file_size_bytes: 500,
915 centroid_b64: Some("AAAA".into()),
916 radius: None,
917 hnsw_offset: None,
918 hnsw_len: None,
919 vector_column: None,
920 vector_dim: None,
921 extra_vector_indexes: vec![],
922 index_status: IndexStatus::Ready,
923 index_error: None,
924 batch_id: None,
925 embedding_model: None,
926 partition_value: None,
927 deletion_vector: None,
928 first_row_id: None,
929 },
930 DataFileEntry {
931 path: "large.parquet".into(),
932 record_count: 5000,
933 file_size_bytes: 200_000_000,
934 centroid_b64: Some("AAAA".into()),
935 radius: None,
936 hnsw_offset: None,
937 hnsw_len: None,
938 vector_column: None,
939 vector_dim: None,
940 extra_vector_indexes: vec![],
941 index_status: IndexStatus::Ready,
942 index_error: None,
943 batch_id: None,
944 embedding_model: None,
945 partition_value: None,
946 deletion_vector: None,
947 first_row_id: None,
948 },
949 DataFileEntry {
950 path: "also-small.parquet".into(),
951 record_count: 5,
952 file_size_bytes: 800,
953 centroid_b64: Some("AAAA".into()),
954 radius: None,
955 hnsw_offset: None,
956 hnsw_len: None,
957 vector_column: None,
958 vector_dim: None,
959 extra_vector_indexes: vec![],
960 index_status: IndexStatus::Ready,
961 index_error: None,
962 batch_id: None,
963 embedding_model: None,
964 partition_value: None,
965 deletion_vector: None,
966 first_row_id: None,
967 },
968 ];
969 let selected = planner.plan(&files);
970 assert_eq!(selected.len(), 2);
971 assert!(selected.iter().any(|f| f.path == "small.parquet"));
972 assert!(selected.iter().any(|f| f.path == "also-small.parquet"));
973 }
974
975 #[test]
976 fn plan_respects_max_files_per_pass() {
977 let planner = CompactionPlanner::new(CompactionConfig {
978 min_files_to_compact: 2,
979 target_file_size_bytes: 1_000_000,
980 max_files_per_pass: 3,
981 ..Default::default()
982 });
983 let files: Vec<DataFileEntry> = (0..5)
984 .map(|i| DataFileEntry {
985 path: format!("f{i}.parquet"),
986 record_count: 10,
987 file_size_bytes: 100 + i as u64 * 100,
988 centroid_b64: Some("AAAA".into()),
989 radius: None,
990 hnsw_offset: None,
991 hnsw_len: None,
992 vector_column: None,
993 vector_dim: None,
994 extra_vector_indexes: vec![],
995 index_status: IndexStatus::Ready,
996 index_error: None,
997 batch_id: None,
998 embedding_model: None,
999 partition_value: None,
1000 deletion_vector: None,
1001 first_row_id: None,
1002 })
1003 .collect();
1004 let selected = planner.plan(&files);
1005 assert_eq!(selected.len(), 3);
1006 assert_eq!(selected[0].file_size_bytes, 100);
1007 assert_eq!(selected[1].file_size_bytes, 200);
1008 assert_eq!(selected[2].file_size_bytes, 300);
1009 }
1010
1011 #[test]
1012 fn plan_sorts_smallest_first() {
1013 let planner = CompactionPlanner::new(CompactionConfig {
1014 min_files_to_compact: 2,
1015 target_file_size_bytes: 10_000,
1016 max_files_per_pass: 4,
1017 ..Default::default()
1018 });
1019 let files = vec![
1020 DataFileEntry {
1021 path: "c.parquet".into(),
1022 record_count: 1,
1023 file_size_bytes: 300,
1024 centroid_b64: Some("AAAA".into()),
1025 radius: None,
1026 hnsw_offset: None,
1027 hnsw_len: None,
1028 vector_column: None,
1029 vector_dim: None,
1030 extra_vector_indexes: vec![],
1031 index_status: IndexStatus::Ready,
1032 index_error: None,
1033 batch_id: None,
1034 embedding_model: None,
1035 partition_value: None,
1036 deletion_vector: None,
1037 first_row_id: None,
1038 },
1039 DataFileEntry {
1040 path: "a.parquet".into(),
1041 record_count: 1,
1042 file_size_bytes: 100,
1043 centroid_b64: Some("AAAA".into()),
1044 radius: None,
1045 hnsw_offset: None,
1046 hnsw_len: None,
1047 vector_column: None,
1048 vector_dim: None,
1049 extra_vector_indexes: vec![],
1050 index_status: IndexStatus::Ready,
1051 index_error: None,
1052 batch_id: None,
1053 embedding_model: None,
1054 partition_value: None,
1055 deletion_vector: None,
1056 first_row_id: None,
1057 },
1058 DataFileEntry {
1059 path: "b.parquet".into(),
1060 record_count: 1,
1061 file_size_bytes: 200,
1062 centroid_b64: Some("AAAA".into()),
1063 radius: None,
1064 hnsw_offset: None,
1065 hnsw_len: None,
1066 vector_column: None,
1067 vector_dim: None,
1068 extra_vector_indexes: vec![],
1069 index_status: IndexStatus::Ready,
1070 index_error: None,
1071 batch_id: None,
1072 embedding_model: None,
1073 partition_value: None,
1074 deletion_vector: None,
1075 first_row_id: None,
1076 },
1077 ];
1078 let selected = planner.plan(&files);
1079 assert_eq!(selected[0].file_size_bytes, 100);
1080 assert_eq!(selected[1].file_size_bytes, 200);
1081 assert_eq!(selected[2].file_size_bytes, 300);
1082 }
1083
1084 fn make_plan_entry(path: &str, size: u64, centroid_b64: Option<String>) -> DataFileEntry {
1085 DataFileEntry {
1086 path: path.to_string(),
1087 record_count: 10,
1088 file_size_bytes: size,
1089 centroid_b64,
1090 radius: None,
1091 hnsw_offset: None,
1092 hnsw_len: None,
1093 vector_column: None,
1094 vector_dim: None,
1095 extra_vector_indexes: vec![],
1096 index_status: IndexStatus::Ready,
1097 index_error: None,
1098 batch_id: None,
1099 embedding_model: None,
1100 partition_value: None,
1101 deletion_vector: None,
1102 first_row_id: None,
1103 }
1104 }
1105
1106 #[test]
1107 fn plan_prioritizes_foreign_written_files_regardless_of_size() {
1108 let planner = CompactionPlanner::new(CompactionConfig {
1109 min_files_to_compact: 4, target_file_size_bytes: 1000,
1111 max_files_per_pass: 20,
1112 ..Default::default()
1113 });
1114 let files = vec![
1115 make_plan_entry("native_small.parquet", 500, Some("AAAA".into())),
1118 make_plan_entry("foreign_big.parquet", 200_000_000, None),
1121 ];
1122 let selected = planner.plan(&files);
1123 assert_eq!(
1124 selected.len(),
1125 1,
1126 "foreign file alone should trigger a pass even below min_files_to_compact"
1127 );
1128 assert_eq!(selected[0].path, "foreign_big.parquet");
1129 }
1130
1131 #[test]
1132 fn plan_merges_foreign_and_size_candidates_together() {
1133 let planner = CompactionPlanner::new(CompactionConfig {
1134 min_files_to_compact: 2,
1135 target_file_size_bytes: 1000,
1136 max_files_per_pass: 20,
1137 ..Default::default()
1138 });
1139 let files = vec![
1140 make_plan_entry("native_a.parquet", 300, Some("AAAA".into())),
1141 make_plan_entry("native_b.parquet", 400, Some("AAAA".into())),
1142 make_plan_entry("foreign.parquet", 200_000_000, None),
1143 ];
1144 let selected = planner.plan(&files);
1145 let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1146 assert_eq!(selected.len(), 3, "paths={paths:?}");
1147 assert_eq!(selected[0].path, "foreign.parquet", "paths={paths:?}");
1149 }
1150
1151 #[test]
1152 fn plan_sorts_foreign_files_smallest_first_too() {
1153 let planner = CompactionPlanner::new(CompactionConfig {
1157 min_files_to_compact: 1,
1158 target_file_size_bytes: 1000,
1159 max_files_per_pass: 2,
1160 ..Default::default()
1161 });
1162 let files = vec![
1163 make_plan_entry("foreign_huge.parquet", 500_000_000, None),
1164 make_plan_entry("foreign_small.parquet", 100, None),
1165 make_plan_entry("foreign_medium.parquet", 10_000_000, None),
1166 ];
1167 let selected = planner.plan(&files);
1168 let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1169 assert_eq!(selected.len(), 2, "max_files_per_pass=2, paths={paths:?}");
1170 assert_eq!(
1171 paths,
1172 vec!["foreign_small.parquet", "foreign_medium.parquet"],
1173 "foreign files must be size-sorted before truncation, cheapest first"
1174 );
1175 }
1176
1177 #[tokio::test]
1178 async fn compact_merges_two_files() {
1179 use ailake_core::{VectorMetric, VectorPrecision};
1180 use ailake_store::LocalStore;
1181 use arrow_array::{Int32Array, RecordBatch};
1182 use arrow_schema::{DataType, Field, Schema};
1183 use std::sync::Arc;
1184 use tempfile::TempDir;
1185
1186 let dir = TempDir::new().unwrap();
1187 let store = Arc::new(LocalStore::new(dir.path()));
1188 let policy = VectorStoragePolicy {
1189 column_name: "embedding".into(),
1190 dim: 4,
1191 metric: VectorMetric::Cosine,
1192 precision: VectorPrecision::F16,
1193 pq: None,
1194 keep_raw_for_reranking: true,
1195 pre_normalize: false,
1196 hnsw_m: None,
1197 hnsw_ef_construction: None,
1198 ivf_residual: false,
1199 embedding_model: None,
1200 modality: None,
1201 partition_by: None,
1202 partition_value: None,
1203 partition_column_type: None,
1204 partition_fields: vec![],
1205 };
1206
1207 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1208 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]];
1209 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]];
1210
1211 let batch_a = RecordBatch::try_new(
1212 schema.clone(),
1213 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1214 )
1215 .unwrap();
1216 let batch_b = RecordBatch::try_new(
1217 schema.clone(),
1218 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1219 )
1220 .unwrap();
1221
1222 let writer_a = AilakeFileWriter::new(policy.clone());
1223 let bytes_a = writer_a.write(&batch_a, &embs_a).unwrap();
1224 let writer_b = AilakeFileWriter::new(policy.clone());
1225 let bytes_b = writer_b.write(&batch_b, &embs_b).unwrap();
1226
1227 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1228 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1229
1230 let entries = vec![
1231 DataFileEntry {
1232 path: "data/a.parquet".into(),
1233 record_count: 2,
1234 file_size_bytes: bytes_a.len() as u64,
1235 centroid_b64: None,
1236 radius: None,
1237 hnsw_offset: None,
1238 hnsw_len: None,
1239 vector_column: None,
1240 vector_dim: None,
1241 extra_vector_indexes: vec![],
1242 index_status: IndexStatus::Ready,
1243 index_error: None,
1244 batch_id: None,
1245 embedding_model: None,
1246 partition_value: None,
1247 deletion_vector: None,
1248 first_row_id: None,
1249 },
1250 DataFileEntry {
1251 path: "data/b.parquet".into(),
1252 record_count: 2,
1253 file_size_bytes: bytes_b.len() as u64,
1254 centroid_b64: None,
1255 radius: None,
1256 hnsw_offset: None,
1257 hnsw_len: None,
1258 vector_column: None,
1259 vector_dim: None,
1260 extra_vector_indexes: vec![],
1261 index_status: IndexStatus::Ready,
1262 index_error: None,
1263 batch_id: None,
1264 embedding_model: None,
1265 partition_value: None,
1266 deletion_vector: None,
1267 first_row_id: None,
1268 },
1269 ];
1270
1271 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1272 let merged = executor
1273 .compact(&entries, "data/merged.parquet")
1274 .await
1275 .unwrap();
1276
1277 assert_eq!(merged.record_count, 4);
1278 assert_eq!(merged.path, "data/merged.parquet");
1279
1280 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1281 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1282 reader.verify_integrity().unwrap();
1283 let (batch, embs) = reader.read_parquet().unwrap();
1284 assert_eq!(batch.num_rows(), 4);
1285 assert_eq!(embs.len(), 4);
1286 }
1287
1288 #[tokio::test]
1293 async fn compact_drops_deletion_vector_masked_rows() {
1294 use ailake_catalog::provider::DeletionVector;
1295 use ailake_core::{VectorMetric, VectorPrecision};
1296 use ailake_store::LocalStore;
1297 use arrow_array::{Int32Array, RecordBatch};
1298 use arrow_schema::{DataType, Field, Schema};
1299 use roaring::RoaringBitmap;
1300 use std::sync::Arc;
1301 use tempfile::TempDir;
1302
1303 let dir = TempDir::new().unwrap();
1304 let store = Arc::new(LocalStore::new(dir.path()));
1305 let policy = VectorStoragePolicy {
1306 column_name: "embedding".into(),
1307 dim: 4,
1308 metric: VectorMetric::Cosine,
1309 precision: VectorPrecision::F16,
1310 pq: None,
1311 keep_raw_for_reranking: true,
1312 pre_normalize: false,
1313 hnsw_m: None,
1314 hnsw_ef_construction: None,
1315 ivf_residual: false,
1316 embedding_model: None,
1317 modality: None,
1318 partition_by: None,
1319 partition_value: None,
1320 partition_column_type: None,
1321 partition_fields: vec![],
1322 };
1323
1324 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1325 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]];
1327 let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
1328
1329 let batch_a = RecordBatch::try_new(
1330 schema.clone(),
1331 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1332 )
1333 .unwrap();
1334 let batch_b =
1335 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![2i32]))])
1336 .unwrap();
1337
1338 let bytes_a = AilakeFileWriter::new(policy.clone())
1339 .write(&batch_a, &embs_a)
1340 .unwrap();
1341 let bytes_b = AilakeFileWriter::new(policy.clone())
1342 .write(&batch_b, &embs_b)
1343 .unwrap();
1344 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1345 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1346
1347 let mut bitmap = RoaringBitmap::new();
1349 bitmap.insert(0);
1350 let (puffin_bytes, offset, length) =
1351 crate::delete::PuffinWriter::write_single_dv(&bitmap, 1).unwrap();
1352 store.put("metadata/dv-1.dvd", puffin_bytes).await.unwrap();
1353
1354 let entry_a = DataFileEntry {
1355 path: "data/a.parquet".into(),
1356 record_count: 2,
1357 file_size_bytes: bytes_a.len() as u64,
1358 centroid_b64: None,
1359 radius: None,
1360 hnsw_offset: None,
1361 hnsw_len: None,
1362 vector_column: None,
1363 vector_dim: None,
1364 extra_vector_indexes: vec![],
1365 index_status: IndexStatus::Ready,
1366 index_error: None,
1367 batch_id: None,
1368 embedding_model: None,
1369 partition_value: None,
1370 deletion_vector: Some(DeletionVector {
1371 path: "metadata/dv-1.dvd".into(),
1372 offset,
1373 length,
1374 cardinality: 1,
1375 }),
1376 first_row_id: None,
1377 };
1378 let entry_b = DataFileEntry {
1379 path: "data/b.parquet".into(),
1380 record_count: 1,
1381 file_size_bytes: bytes_b.len() as u64,
1382 centroid_b64: None,
1383 radius: None,
1384 hnsw_offset: None,
1385 hnsw_len: None,
1386 vector_column: None,
1387 vector_dim: None,
1388 extra_vector_indexes: vec![],
1389 index_status: IndexStatus::Ready,
1390 index_error: None,
1391 batch_id: None,
1392 embedding_model: None,
1393 partition_value: None,
1394 deletion_vector: None,
1395 first_row_id: None,
1396 };
1397
1398 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1399 let merged = executor
1400 .compact(&[entry_a, entry_b], "data/merged.parquet")
1401 .await
1402 .unwrap();
1403
1404 assert_eq!(merged.record_count, 2);
1406 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1407 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1408 let (batch, embs) = reader.read_parquet().unwrap();
1409 assert_eq!(batch.num_rows(), 2);
1410 assert_eq!(embs.len(), 2);
1411 let ids: Vec<i32> = batch
1412 .column_by_name("id")
1413 .unwrap()
1414 .as_any()
1415 .downcast_ref::<Int32Array>()
1416 .unwrap()
1417 .values()
1418 .to_vec();
1419 assert_eq!(
1420 ids,
1421 vec![1, 2],
1422 "id=0 (deleted) must not survive compaction"
1423 );
1424 }
1425
1426 #[tokio::test]
1431 async fn compact_with_ivf_pq_strategy_does_not_crash_on_verify_integrity() {
1432 use ailake_core::{VectorMetric, VectorPrecision};
1433 use ailake_store::LocalStore;
1434 use arrow_array::{Int32Array, RecordBatch};
1435 use arrow_schema::{DataType, Field, Schema};
1436 use std::sync::Arc;
1437 use tempfile::TempDir;
1438
1439 let dir = TempDir::new().unwrap();
1440 let store = Arc::new(LocalStore::new(dir.path()));
1441 let dim = 8;
1442 let policy = VectorStoragePolicy {
1443 column_name: "embedding".into(),
1444 dim,
1445 metric: VectorMetric::Cosine,
1446 precision: VectorPrecision::F16,
1447 pq: None,
1448 keep_raw_for_reranking: true,
1449 pre_normalize: false,
1450 hnsw_m: None,
1451 hnsw_ef_construction: None,
1452 ivf_residual: false,
1453 embedding_model: None,
1454 modality: None,
1455 partition_by: None,
1456 partition_value: None,
1457 partition_column_type: None,
1458 partition_fields: vec![],
1459 };
1460
1461 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1462 let n_per_file = 30usize;
1463 let make_file = |path: &str, offset: i32| {
1464 let ids: Vec<i32> = (offset..offset + n_per_file as i32).collect();
1465 let embs: Vec<Vec<f32>> = ids
1466 .iter()
1467 .map(|&i| {
1468 (0..dim as i32)
1469 .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1470 .collect()
1471 })
1472 .collect();
1473 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1474 .unwrap();
1475 let bytes = AilakeFileWriter::new(policy.clone())
1476 .write(&batch, &embs)
1477 .unwrap();
1478 (path.to_string(), bytes)
1479 };
1480
1481 let (path_a, bytes_a) = make_file("data/ivfpq_a.parquet", 0);
1482 let (path_b, bytes_b) = make_file("data/ivfpq_b.parquet", n_per_file as i32);
1483 for (path, bytes) in [(&path_a, &bytes_a), (&path_b, &bytes_b)] {
1484 store.put(path, bytes.clone()).await.unwrap();
1485 }
1486
1487 let make_entry = |path: &str, size: u64| DataFileEntry {
1488 path: path.to_string(),
1489 record_count: n_per_file as u64,
1490 file_size_bytes: size,
1491 centroid_b64: None,
1492 radius: None,
1493 hnsw_offset: None,
1494 hnsw_len: None,
1495 vector_column: None,
1496 vector_dim: None,
1497 extra_vector_indexes: vec![],
1498 index_status: IndexStatus::Ready,
1499 index_error: None,
1500 batch_id: None,
1501 embedding_model: None,
1502 partition_value: None,
1503 deletion_vector: None,
1504 first_row_id: None,
1505 };
1506 let entries = vec![
1507 make_entry(&path_a, bytes_a.len() as u64),
1508 make_entry(&path_b, bytes_b.len() as u64),
1509 ];
1510
1511 let executor = CompactionExecutor::new(store.clone(), policy.clone())
1512 .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1513 let merged = executor
1514 .compact(&entries, "data/ivfpq_merged.parquet")
1515 .await
1516 .expect("compact() with ForceIvfPq must not fail in verify_integrity()");
1517
1518 assert_eq!(merged.record_count, 2 * n_per_file as u64);
1519
1520 let merged_bytes = store.get("data/ivfpq_merged.parquet").await.unwrap();
1522 let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1523 reader
1524 .verify_integrity()
1525 .expect("verify_integrity() must handle an IVF-PQ-indexed file");
1526 }
1527
1528 #[tokio::test]
1529 async fn compact_incremental_merges_dominant_plus_small() {
1530 use ailake_core::{RowId, VectorMetric, VectorPrecision};
1531 use ailake_store::LocalStore;
1532 use arrow_array::{Int32Array, RecordBatch};
1533 use arrow_schema::{DataType, Field, Schema};
1534 use std::sync::Arc;
1535 use tempfile::TempDir;
1536
1537 let dir = TempDir::new().unwrap();
1538 let store = Arc::new(LocalStore::new(dir.path()));
1539 let policy = VectorStoragePolicy {
1540 column_name: "embedding".into(),
1541 dim: 4,
1542 metric: VectorMetric::Cosine,
1543 precision: VectorPrecision::F16,
1544 pq: None,
1545 keep_raw_for_reranking: true,
1546 pre_normalize: false,
1547 hnsw_m: None,
1548 hnsw_ef_construction: None,
1549 ivf_residual: false,
1550 embedding_model: None,
1551 modality: None,
1552 partition_by: None,
1553 partition_value: None,
1554 partition_column_type: None,
1555 partition_fields: vec![],
1556 };
1557
1558 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1559
1560 let embs_dom: Vec<Vec<f32>> = vec![
1562 vec![1.0, 0.0, 0.0, 0.0],
1563 vec![0.0, 1.0, 0.0, 0.0],
1564 vec![0.0, 0.0, 1.0, 0.0],
1565 vec![0.7, 0.7, 0.0, 0.0],
1566 vec![0.0, 0.7, 0.7, 0.0],
1567 vec![0.0, 0.0, 0.7, 0.7],
1568 ];
1569 let batch_dom = RecordBatch::try_new(
1570 schema.clone(),
1571 vec![Arc::new(Int32Array::from(vec![0i32, 1, 2, 3, 4, 5]))],
1572 )
1573 .unwrap();
1574
1575 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]];
1577 let batch_small = RecordBatch::try_new(
1578 schema.clone(),
1579 vec![Arc::new(Int32Array::from(vec![6i32, 7]))],
1580 )
1581 .unwrap();
1582
1583 let bytes_dom = AilakeFileWriter::new(policy.clone())
1584 .write(&batch_dom, &embs_dom)
1585 .unwrap();
1586 let bytes_small = AilakeFileWriter::new(policy.clone())
1587 .write(&batch_small, &embs_small)
1588 .unwrap();
1589
1590 store
1591 .put("data/dominant.parquet", bytes_dom.clone())
1592 .await
1593 .unwrap();
1594 store
1595 .put("data/small.parquet", bytes_small.clone())
1596 .await
1597 .unwrap();
1598
1599 let entries = vec![
1600 DataFileEntry {
1601 path: "data/dominant.parquet".into(),
1602 record_count: 6,
1603 file_size_bytes: bytes_dom.len() as u64,
1604 centroid_b64: None,
1605 radius: None,
1606 hnsw_offset: None,
1607 hnsw_len: None,
1608 vector_column: None,
1609 vector_dim: None,
1610 extra_vector_indexes: vec![],
1611 index_status: IndexStatus::Ready,
1612 index_error: None,
1613 batch_id: None,
1614 embedding_model: None,
1615 partition_value: None,
1616 deletion_vector: None,
1617 first_row_id: None,
1618 },
1619 DataFileEntry {
1620 path: "data/small.parquet".into(),
1621 record_count: 2,
1622 file_size_bytes: bytes_small.len() as u64,
1623 centroid_b64: None,
1624 radius: None,
1625 hnsw_offset: None,
1626 hnsw_len: None,
1627 vector_column: None,
1628 vector_dim: None,
1629 extra_vector_indexes: vec![],
1630 index_status: IndexStatus::Ready,
1631 index_error: None,
1632 batch_id: None,
1633 embedding_model: None,
1634 partition_value: None,
1635 deletion_vector: None,
1636 first_row_id: None,
1637 },
1638 ];
1639
1640 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1641 let merged = executor
1642 .compact_incremental(&entries, "data/merged.parquet")
1643 .await
1644 .unwrap();
1645
1646 assert_eq!(merged.record_count, 8);
1648 assert_eq!(merged.path, "data/merged.parquet");
1649
1650 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1652 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1653 reader.verify_integrity().unwrap();
1654
1655 let (batch, embs) = reader.read_parquet().unwrap();
1656 assert_eq!(batch.num_rows(), 8);
1657 assert_eq!(embs.len(), 8);
1658
1659 for f in &embs[..6] {
1661 assert_eq!(f.len(), 4);
1662 }
1663
1664 let hnsw = reader.load_index().unwrap();
1666 assert_eq!(hnsw.node_count(), 8);
1667
1668 let results = hnsw.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
1670 assert_eq!(results[0].0, RowId::new(0));
1671
1672 let results = hnsw.search(&[0.0, 0.0, 0.0, 1.0], 1, 50);
1675 assert_eq!(results[0].0, RowId::new(6));
1676 }
1677
1678 #[tokio::test]
1688 async fn compact_incremental_respects_force_ivf_pq_even_with_dominant_file() {
1689 use ailake_core::{VectorMetric, VectorPrecision};
1690 use ailake_store::LocalStore;
1691 use arrow_array::{Int32Array, RecordBatch};
1692 use arrow_schema::{DataType, Field, Schema};
1693 use std::sync::Arc;
1694 use tempfile::TempDir;
1695
1696 let dir = TempDir::new().unwrap();
1697 let store = Arc::new(LocalStore::new(dir.path()));
1698 let dim = 8;
1699 let policy = VectorStoragePolicy {
1700 column_name: "embedding".into(),
1701 dim,
1702 metric: VectorMetric::Cosine,
1703 precision: VectorPrecision::F16,
1704 pq: None,
1705 keep_raw_for_reranking: true,
1706 pre_normalize: false,
1707 hnsw_m: None,
1708 hnsw_ef_construction: None,
1709 ivf_residual: false,
1710 embedding_model: None,
1711 modality: None,
1712 partition_by: None,
1713 partition_value: None,
1714 partition_column_type: None,
1715 partition_fields: vec![],
1716 };
1717
1718 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1719 let make_file = |path: &str, offset: i32, n: usize| {
1720 let ids: Vec<i32> = (offset..offset + n as i32).collect();
1721 let embs: Vec<Vec<f32>> = ids
1722 .iter()
1723 .map(|&i| {
1724 (0..dim as i32)
1725 .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1726 .collect()
1727 })
1728 .collect();
1729 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1730 .unwrap();
1731 let bytes = AilakeFileWriter::new(policy.clone())
1734 .write(&batch, &embs)
1735 .unwrap();
1736 (path.to_string(), bytes, n as u64)
1737 };
1738
1739 let (path_dom, bytes_dom, n_dom) = make_file("data/dom.parquet", 0, 90);
1742 let (path_small, bytes_small, n_small) = make_file("data/small.parquet", 999, 30);
1743
1744 for (path, bytes) in [(&path_dom, &bytes_dom), (&path_small, &bytes_small)] {
1745 store.put(path, bytes.clone()).await.unwrap();
1746 }
1747
1748 let make_entry = |path: &str, record_count: u64, size: u64| DataFileEntry {
1749 path: path.to_string(),
1750 record_count,
1751 file_size_bytes: size,
1752 centroid_b64: None,
1753 radius: None,
1754 hnsw_offset: None,
1755 hnsw_len: None,
1756 vector_column: None,
1757 vector_dim: None,
1758 extra_vector_indexes: vec![],
1759 index_status: IndexStatus::Ready,
1760 index_error: None,
1761 batch_id: None,
1762 embedding_model: None,
1763 partition_value: None,
1764 deletion_vector: None,
1765 first_row_id: None,
1766 };
1767 let entries = vec![
1768 make_entry(&path_dom, n_dom, bytes_dom.len() as u64),
1769 make_entry(&path_small, n_small, bytes_small.len() as u64),
1770 ];
1771
1772 let executor = CompactionExecutor::new(store.clone(), policy.clone())
1773 .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1774 let merged = executor
1775 .compact_incremental(&entries, "data/merged_force_ivfpq.parquet")
1776 .await
1777 .expect("compact_incremental() with ForceIvfPq must fall back to compact() cleanly");
1778
1779 let merged_bytes = store.get("data/merged_force_ivfpq.parquet").await.unwrap();
1780 let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1781 reader.verify_integrity().unwrap();
1782
1783 match reader.load_any_index().unwrap() {
1784 ailake_index::AnyIndex::IvfPq(_) => {}
1785 ailake_index::AnyIndex::Hnsw(_) => panic!(
1786 "ForceIvfPq was silently satisfied with HNSW instead — merged.record_count={}",
1787 merged.record_count
1788 ),
1789 }
1790 }
1791
1792 #[tokio::test]
1793 async fn compact_incremental_falls_back_when_no_dominant() {
1794 use ailake_core::{VectorMetric, VectorPrecision};
1795 use ailake_store::LocalStore;
1796 use arrow_array::{Int32Array, RecordBatch};
1797 use arrow_schema::{DataType, Field, Schema};
1798 use std::sync::Arc;
1799 use tempfile::TempDir;
1800
1801 let dir = TempDir::new().unwrap();
1802 let store = Arc::new(LocalStore::new(dir.path()));
1803 let policy = VectorStoragePolicy {
1804 column_name: "embedding".into(),
1805 dim: 4,
1806 metric: VectorMetric::Cosine,
1807 precision: VectorPrecision::F16,
1808 pq: None,
1809 keep_raw_for_reranking: true,
1810 pre_normalize: false,
1811 hnsw_m: None,
1812 hnsw_ef_construction: None,
1813 ivf_residual: false,
1814 embedding_model: None,
1815 modality: None,
1816 partition_by: None,
1817 partition_value: None,
1818 partition_column_type: None,
1819 partition_fields: vec![],
1820 };
1821
1822 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1823
1824 let make_batch = |ids: Vec<i32>, embs: Vec<Vec<f32>>| {
1826 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1827 .unwrap();
1828 AilakeFileWriter::new(policy.clone())
1829 .write(&batch, &embs)
1830 .unwrap()
1831 };
1832
1833 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]];
1834 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]];
1835 let bytes_a = make_batch(vec![0, 1], embs_a);
1836 let bytes_b = make_batch(vec![2, 3], embs_b);
1837
1838 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1839 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1840
1841 let entries = vec![
1842 DataFileEntry {
1843 path: "data/a.parquet".into(),
1844 record_count: 2,
1845 file_size_bytes: bytes_a.len() as u64,
1846 centroid_b64: None,
1847 radius: None,
1848 hnsw_offset: None,
1849 hnsw_len: None,
1850 vector_column: None,
1851 vector_dim: None,
1852 extra_vector_indexes: vec![],
1853 index_status: IndexStatus::Ready,
1854 index_error: None,
1855 batch_id: None,
1856 embedding_model: None,
1857 partition_value: None,
1858 deletion_vector: None,
1859 first_row_id: None,
1860 },
1861 DataFileEntry {
1862 path: "data/b.parquet".into(),
1863 record_count: 2,
1864 file_size_bytes: bytes_b.len() as u64,
1865 centroid_b64: None,
1866 radius: None,
1867 hnsw_offset: None,
1868 hnsw_len: None,
1869 vector_column: None,
1870 vector_dim: None,
1871 extra_vector_indexes: vec![],
1872 index_status: IndexStatus::Ready,
1873 index_error: None,
1874 batch_id: None,
1875 embedding_model: None,
1876 partition_value: None,
1877 deletion_vector: None,
1878 first_row_id: None,
1879 },
1880 ];
1881
1882 let executor = CompactionExecutor::new(store.clone(), policy.clone());
1883 let merged = executor
1885 .compact_incremental(&entries, "data/merged.parquet")
1886 .await
1887 .unwrap();
1888
1889 assert_eq!(merged.record_count, 4);
1890
1891 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1892 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1893 reader.verify_integrity().unwrap();
1894 }
1895
1896 #[tokio::test]
1897 async fn compact_deferred_produces_parquet_only_file() {
1898 use ailake_catalog::HadoopCatalog;
1899 use ailake_core::{VectorMetric, VectorPrecision};
1900 use ailake_store::LocalStore;
1901 use arrow_array::{Int32Array, RecordBatch};
1902 use arrow_schema::{DataType, Field, Schema};
1903 use std::sync::Arc;
1904 use tempfile::TempDir;
1905
1906 let dir = TempDir::new().unwrap();
1907 let store = Arc::new(LocalStore::new(dir.path()));
1908 let catalog_dir = TempDir::new().unwrap();
1909 let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
1910 let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
1911 let table = TableIdent {
1912 namespace: "ns".into(),
1913 name: "tbl".into(),
1914 };
1915
1916 let policy = VectorStoragePolicy {
1917 column_name: "embedding".into(),
1918 dim: 4,
1919 metric: VectorMetric::Cosine,
1920 precision: VectorPrecision::F16,
1921 pq: None,
1922 keep_raw_for_reranking: true,
1923 pre_normalize: false,
1924 hnsw_m: None,
1925 hnsw_ef_construction: None,
1926 ivf_residual: false,
1927 embedding_model: None,
1928 modality: None,
1929 partition_by: None,
1930 partition_value: None,
1931 partition_column_type: None,
1932 partition_fields: vec![],
1933 };
1934
1935 use ailake_catalog::TableProperties;
1936 catalog
1937 .create_table(
1938 &table,
1939 &TableProperties {
1940 policy: policy.clone(),
1941 extra: std::collections::HashMap::new(),
1942 format_version: 2,
1943 partition_column_type: None,
1944 },
1945 )
1946 .await
1947 .unwrap();
1948
1949 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1950 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]];
1951 let batch_a = RecordBatch::try_new(
1952 schema.clone(),
1953 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1954 )
1955 .unwrap();
1956 let bytes_a = AilakeFileWriter::new(policy.clone())
1957 .write(&batch_a, &embs_a)
1958 .unwrap();
1959 store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1960
1961 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]];
1962 let batch_b = RecordBatch::try_new(
1963 schema.clone(),
1964 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1965 )
1966 .unwrap();
1967 let bytes_b = AilakeFileWriter::new(policy.clone())
1968 .write(&batch_b, &embs_b)
1969 .unwrap();
1970 store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1971
1972 let entries = vec![
1973 DataFileEntry {
1974 path: "data/a.parquet".into(),
1975 record_count: 2,
1976 file_size_bytes: bytes_a.len() as u64,
1977 centroid_b64: None,
1978 radius: None,
1979 hnsw_offset: None,
1980 hnsw_len: None,
1981 vector_column: None,
1982 vector_dim: None,
1983 extra_vector_indexes: vec![],
1984 index_status: IndexStatus::Ready,
1985 index_error: None,
1986 batch_id: None,
1987 embedding_model: None,
1988 partition_value: None,
1989 deletion_vector: None,
1990 first_row_id: None,
1991 },
1992 DataFileEntry {
1993 path: "data/b.parquet".into(),
1994 record_count: 2,
1995 file_size_bytes: bytes_b.len() as u64,
1996 centroid_b64: None,
1997 radius: None,
1998 hnsw_offset: None,
1999 hnsw_len: None,
2000 vector_column: None,
2001 vector_dim: None,
2002 extra_vector_indexes: vec![],
2003 index_status: IndexStatus::Ready,
2004 index_error: None,
2005 batch_id: None,
2006 embedding_model: None,
2007 partition_value: None,
2008 deletion_vector: None,
2009 first_row_id: None,
2010 },
2011 ];
2012
2013 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2014 let entry = executor
2015 .compact_deferred(&entries, "data/merged.parquet", catalog.clone(), &table)
2016 .await
2017 .unwrap();
2018
2019 assert_eq!(entry.index_status, IndexStatus::Indexing);
2021 assert_eq!(entry.record_count, 4);
2022
2023 let merged_bytes = store.get("data/merged.parquet").await.unwrap();
2025 let pq_reader = ailake_parquet::ParquetVectorReader::new(merged_bytes, "embedding");
2026 let count = pq_reader.record_count().unwrap();
2027 assert_eq!(count, 4);
2028 }
2029
2030 #[tokio::test]
2037 async fn run_preserves_untouched_files_outside_compaction_pass() {
2038 use ailake_catalog::HadoopCatalog;
2039 use ailake_core::{VectorMetric, VectorPrecision};
2040 use ailake_store::LocalStore;
2041 use arrow_array::{Int32Array, RecordBatch};
2042 use arrow_schema::{DataType, Field, Schema};
2043 use std::sync::Arc;
2044 use tempfile::TempDir;
2045
2046 let dir = TempDir::new().unwrap();
2047 let store = Arc::new(LocalStore::new(dir.path()));
2048 let catalog_dir = TempDir::new().unwrap();
2049 let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
2050 let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
2051 let table = TableIdent {
2052 namespace: "ns".into(),
2053 name: "tbl".into(),
2054 };
2055
2056 let policy = VectorStoragePolicy {
2057 column_name: "embedding".into(),
2058 dim: 4,
2059 metric: VectorMetric::Cosine,
2060 precision: VectorPrecision::F16,
2061 pq: None,
2062 keep_raw_for_reranking: true,
2063 pre_normalize: false,
2064 hnsw_m: None,
2065 hnsw_ef_construction: None,
2066 ivf_residual: false,
2067 embedding_model: None,
2068 modality: None,
2069 partition_by: None,
2070 partition_value: None,
2071 partition_column_type: None,
2072 partition_fields: vec![],
2073 };
2074
2075 use ailake_catalog::TableProperties;
2076 catalog
2077 .create_table(
2078 &table,
2079 &TableProperties {
2080 policy: policy.clone(),
2081 extra: std::collections::HashMap::new(),
2082 format_version: 2,
2083 partition_column_type: None,
2084 },
2085 )
2086 .await
2087 .unwrap();
2088
2089 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2090
2091 let write_file = |path: &str, ids: Vec<i32>, embs: Vec<Vec<f32>>| {
2093 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
2094 .unwrap();
2095 let bytes = AilakeFileWriter::new(policy.clone())
2096 .write(&batch, &embs)
2097 .unwrap();
2098 (path.to_string(), bytes)
2099 };
2100
2101 let (path_a, bytes_a) = write_file(
2102 "data/small_a.parquet",
2103 vec![0, 1],
2104 vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
2105 );
2106 let (path_b, bytes_b) = write_file(
2107 "data/small_b.parquet",
2108 vec![2, 3],
2109 vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]],
2110 );
2111 let (path_big, bytes_big) = write_file(
2114 "data/big.parquet",
2115 vec![4, 5],
2116 vec![vec![1.0, 1.0, 0.0, 0.0], vec![0.0, 1.0, 1.0, 0.0]],
2117 );
2118
2119 for (path, bytes) in [
2120 (&path_a, &bytes_a),
2121 (&path_b, &bytes_b),
2122 (&path_big, &bytes_big),
2123 ] {
2124 store.put(path, bytes.clone()).await.unwrap();
2125 }
2126
2127 let make_entry = |path: &str, size: u64| DataFileEntry {
2128 path: path.to_string(),
2129 record_count: 2,
2130 file_size_bytes: size,
2131 centroid_b64: Some("AAAA".into()),
2136 radius: None,
2137 hnsw_offset: None,
2138 hnsw_len: None,
2139 vector_column: None,
2140 vector_dim: None,
2141 extra_vector_indexes: vec![],
2142 index_status: IndexStatus::Ready,
2143 index_error: None,
2144 batch_id: None,
2145 embedding_model: None,
2146 partition_value: None,
2147 deletion_vector: None,
2148 first_row_id: None,
2149 };
2150
2151 let initial_snap_id = ailake_catalog::new_snapshot_id();
2152 let initial_snapshot = NewSnapshot {
2153 snapshot_id: initial_snap_id,
2154 parent_snapshot_id: None,
2155 files: vec![
2156 make_entry(&path_a, 500),
2157 make_entry(&path_b, 500),
2158 make_entry(&path_big, 200_000_000), ],
2160 operation: SnapshotOperation::Append,
2161 iceberg_schema: None,
2162 extra_properties: std::collections::HashMap::new(),
2163 bloom_filters: vec![],
2164 equality_delete_files: vec![],
2165 };
2166 catalog
2167 .commit_snapshot(&table, initial_snapshot)
2168 .await
2169 .unwrap();
2170
2171 let planner = CompactionPlanner::new(CompactionConfig {
2172 min_files_to_compact: 2,
2173 target_file_size_bytes: 1000,
2174 index_strategy: CompactionIndexStrategy::ForceHnsw,
2175 max_files_per_pass: 20,
2176 });
2177 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2178
2179 let merged = executor
2180 .run(&planner, &table, catalog.clone(), "data")
2181 .await
2182 .unwrap()
2183 .expect("compaction should have run — 2 eligible small files");
2184
2185 let files_after = catalog.list_files(&table, None).await.unwrap();
2186 let paths_after: Vec<&str> = files_after.iter().map(|f| f.path.as_str()).collect();
2187
2188 assert!(
2189 paths_after.contains(&path_big.as_str()),
2190 "BUG: untouched 'big.parquet' vanished after run() — files_after={paths_after:?}"
2191 );
2192 assert!(
2193 paths_after.contains(&merged.path.as_str()),
2194 "merged output file must be present — files_after={paths_after:?}"
2195 );
2196
2197 let meta_dir = catalog_dir.path().join("ns/tbl/metadata");
2204 let latest_metadata = std::fs::read_dir(&meta_dir)
2205 .unwrap()
2206 .filter_map(|e| e.ok())
2207 .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
2208 .max_by_key(|e| e.metadata().unwrap().modified().unwrap())
2209 .expect("metadata.json must exist after commit");
2210 let json: serde_json::Value =
2211 serde_json::from_slice(&std::fs::read(latest_metadata.path()).unwrap()).unwrap();
2212 let last_snapshot = json["snapshots"].as_array().unwrap().last().unwrap();
2213 assert_eq!(
2214 last_snapshot["parent-snapshot-id"].as_i64(),
2215 Some(initial_snap_id),
2216 "compaction's Replace snapshot must chain to the pre-compaction snapshot, not be orphaned"
2217 );
2218 assert!(
2219 !paths_after.contains(&path_a.as_str()) && !paths_after.contains(&path_b.as_str()),
2220 "compacted input files must no longer be listed — files_after={paths_after:?}"
2221 );
2222 assert_eq!(
2223 files_after.len(),
2224 2,
2225 "expected exactly [big.parquet, merged] — files_after={paths_after:?}"
2226 );
2227 }
2228
2229 #[tokio::test]
2235 async fn compact_preserves_rows_from_footerless_file() {
2236 use ailake_core::{VectorMetric, VectorPrecision};
2237 use ailake_store::LocalStore;
2238 use arrow_array::{Int32Array, RecordBatch};
2239 use arrow_schema::{DataType, Field, Schema};
2240 use std::sync::Arc;
2241 use tempfile::TempDir;
2242
2243 let dir = TempDir::new().unwrap();
2244 let store = Arc::new(LocalStore::new(dir.path()));
2245 let policy = VectorStoragePolicy {
2246 column_name: "embedding".into(),
2247 dim: 4,
2248 metric: VectorMetric::Cosine,
2249 precision: VectorPrecision::F16,
2250 pq: None,
2251 keep_raw_for_reranking: true,
2252 pre_normalize: false,
2253 hnsw_m: None,
2254 hnsw_ef_construction: None,
2255 ivf_residual: false,
2256 embedding_model: None,
2257 modality: None,
2258 partition_by: None,
2259 partition_value: None,
2260 partition_column_type: None,
2261 partition_fields: vec![],
2262 };
2263
2264 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2265
2266 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]];
2268 let batch_native = RecordBatch::try_new(
2269 schema.clone(),
2270 vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2271 )
2272 .unwrap();
2273 let bytes_native = AilakeFileWriter::new(policy.clone())
2274 .write(&batch_native, &embs_native)
2275 .unwrap();
2276 store
2277 .put("data/native.parquet", bytes_native.clone())
2278 .await
2279 .unwrap();
2280
2281 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]];
2285 let batch_foreign = RecordBatch::try_new(
2286 schema.clone(),
2287 vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2288 )
2289 .unwrap();
2290 let bytes_foreign = AilakeFileWriter::new(policy.clone())
2291 .write_parquet_only(&batch_foreign, &embs_foreign)
2292 .unwrap();
2293 store
2294 .put("data/foreign.parquet", bytes_foreign.clone())
2295 .await
2296 .unwrap();
2297
2298 let reader_foreign = AilakeFileReader::new(bytes_foreign.clone(), "embedding", 4);
2299 assert!(
2300 !reader_foreign.is_ailake_file(),
2301 "sanity: write_parquet_only must not embed an AILK footer"
2302 );
2303
2304 let entries = vec![
2305 DataFileEntry {
2306 path: "data/native.parquet".into(),
2307 record_count: 2,
2308 file_size_bytes: bytes_native.len() as u64,
2309 centroid_b64: None,
2310 radius: None,
2311 hnsw_offset: None,
2312 hnsw_len: None,
2313 vector_column: None,
2314 vector_dim: None,
2315 extra_vector_indexes: vec![],
2316 index_status: IndexStatus::Ready,
2317 index_error: None,
2318 batch_id: None,
2319 embedding_model: None,
2320 partition_value: None,
2321 deletion_vector: None,
2322 first_row_id: None,
2323 },
2324 DataFileEntry {
2325 path: "data/foreign.parquet".into(),
2326 record_count: 2,
2327 file_size_bytes: bytes_foreign.len() as u64,
2328 centroid_b64: None,
2329 radius: None,
2330 hnsw_offset: None,
2331 hnsw_len: None,
2332 vector_column: None,
2333 vector_dim: None,
2334 extra_vector_indexes: vec![],
2335 index_status: IndexStatus::Ready,
2336 index_error: None,
2337 batch_id: None,
2338 embedding_model: None,
2339 partition_value: None,
2340 deletion_vector: None,
2341 first_row_id: None,
2342 },
2343 ];
2344
2345 let executor = CompactionExecutor::new(store.clone(), policy.clone());
2347 let merged = executor
2348 .compact(&entries, "data/merged_full.parquet")
2349 .await
2350 .unwrap();
2351 assert_eq!(
2352 merged.record_count, 4,
2353 "compact() must include all 4 rows — 2 native + 2 from the footerless file"
2354 );
2355 let merged_bytes = store.get("data/merged_full.parquet").await.unwrap();
2356 let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
2357 reader.verify_integrity().unwrap();
2358
2359 let merged_inc = executor
2362 .compact_incremental(&entries, "data/merged_inc.parquet")
2363 .await
2364 .unwrap();
2365 assert_eq!(
2366 merged_inc.record_count, 4,
2367 "compact_incremental() must include all 4 rows — 2 native + 2 from the footerless file"
2368 );
2369 }
2370}