Skip to main content

ailake_query/
compaction.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2use 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/// Index strategy for the merged file produced by compaction.
21#[derive(Debug, Clone, Default)]
22pub enum CompactionIndexStrategy {
23    /// Detect GPU / CPU cores at compaction time and pick the best index.
24    /// IVF-PQ on GPU/many-core machines; HNSW elsewhere. (default)
25    #[default]
26    Auto,
27    /// Always rebuild with HNSW — highest recall, larger index.
28    ForceHnsw,
29    /// Always rebuild with IVF-PQ — smaller index, better S3 throughput.
30    ///
31    /// Recommended for large compactions (N > 100 000) on CPU-only machines
32    /// where HNSW rebuild cost becomes prohibitive.
33    ForceIvfPq,
34}
35
36#[derive(Debug, Clone)]
37pub struct CompactionConfig {
38    /// Trigger compaction only if at least this many files are eligible.
39    pub min_files_to_compact: usize,
40    /// Target output file size in bytes. Files below this are merged.
41    pub target_file_size_bytes: u64,
42    /// Index algorithm for the merged output file.
43    pub index_strategy: CompactionIndexStrategy,
44    /// Maximum files merged in a single compaction pass.
45    ///
46    /// Candidates are sorted smallest-first; only the first `max_files_per_pass`
47    /// are compacted each run. This bounds peak RAM and HNSW rebuild CPU cost —
48    /// O(N log N) stays proportional to this limit rather than table size.
49    /// Default: 20. Set to `usize::MAX` to compact all eligible files at once.
50    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, // 128 MB
58            index_strategy: CompactionIndexStrategy::Auto,
59            max_files_per_pass: 20,
60        }
61    }
62}
63
64#[derive(Debug, Clone, Copy)]
65pub enum CompactionMode {
66    Full,    // compact all files below target size
67    Partial, // compact the smallest N files
68}
69
70pub struct CompactionPlanner {
71    config: CompactionConfig,
72}
73
74impl CompactionPlanner {
75    pub fn new(config: CompactionConfig) -> Self {
76        Self { config }
77    }
78
79    /// Select files to compact.
80    ///
81    /// Picks files smaller than `target_file_size_bytes`, sorts them smallest-first
82    /// (cheapest to read), and caps the selection at `max_files_per_pass`. This
83    /// tiered approach prevents a single pass from compacting the entire table into
84    /// memory when thousands of small files exist.
85    ///
86    /// Foreign-write detection: every file written by the AI-Lake SDK always carries a
87    /// `centroid_b64` (see `make_data_file_entry`/`make_data_file_entry_indexing` —
88    /// centroid is computed and stored even for `IndexStatus::Indexing` files, before
89    /// the HNSW build itself completes). A manifest entry with no centroid was never
90    /// produced by AI-Lake — almost certainly a generic Iceberg engine (Spark/Trino
91    /// `OPTIMIZE` / `rewrite_data_files`, DuckDB) rewrote the file with no knowledge of
92    /// AI-Lake. Every query against it silently degrades to an O(N) flat scan (see the
93    /// `flat_scan_unexpected` warning in `scanner.rs::search`). These files bypass both
94    /// the size filter and `min_files_to_compact` — a single such file is worth
95    /// repairing on its own; it shouldn't wait for enough small files to accumulate.
96    pub fn plan(&self, files: &[DataFileEntry]) -> Vec<DataFileEntry> {
97        // Single pass: split into foreign (no AI-Lake index) vs. native candidates.
98        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            // Not enough small files to justify a size-based pass on their own, but
124            // foreign files still need repair regardless of batching thresholds.
125            size_candidates.clear();
126        }
127
128        // Sort both smallest-first so each pass handles the cheapest files first —
129        // this bounds peak RAM to max_files_per_pass * avg_selected_file_size. Foreign
130        // files are NOT size-filtered (repair takes priority over the size threshold —
131        // a foreign file may itself be large, e.g. after an external OPTIMIZE), but they
132        // ARE size-sorted so that if there are more foreign files than max_files_per_pass,
133        // this pass still prefers the cheapest ones to merge rather than an arbitrary subset.
134        foreign.sort_unstable_by_key(|f| f.file_size_bytes);
135        size_candidates.sort_unstable_by_key(|f| f.file_size_bytes);
136        // Foreign files take priority; size-based candidates fill the rest of the pass.
137        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/// Executes compaction plans: reads N small files, merges them into a single
152/// AI-Lake file with a rebuilt index, and commits to the catalog.
153///
154/// The index algorithm is chosen via `CompactionIndexStrategy` (default: `Auto`,
155/// which detects GPU / CPU cores at compaction time — the same heuristic used
156/// by `write_batch_auto`).
157///
158/// For large tables use `compact_deferred` / `run_deferred`: the merged Parquet
159/// is persisted immediately and the HNSW build runs in a background Tokio task,
160/// decoupling I/O cost from CPU cost.
161#[derive(Clone)]
162pub struct CompactionExecutor {
163    store: Arc<dyn Store>,
164    policy: VectorStoragePolicy,
165    index_strategy: CompactionIndexStrategy,
166    /// When set, rebuilds and embeds a Tantivy FTS index in the compacted output file.
167    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    /// Override the default (Auto) index strategy for this executor.
181    pub fn with_index_strategy(mut self, strategy: CompactionIndexStrategy) -> Self {
182        self.index_strategy = strategy;
183        self
184    }
185
186    /// Rebuild and embed a Tantivy FTS index in the compacted output file.
187    pub fn with_fts_config(mut self, cfg: ailake_fts::FtsConfig) -> Self {
188        self.fts_config = Some(cfg);
189        self
190    }
191
192    /// Return a clone of this executor whose `fts_config` is filled from table properties
193    /// when the caller did not explicitly call `with_fts_config`.
194    ///
195    /// This is the auto-detect path: if `ailake.fts.enabled=true` is present in the table
196    /// metadata but the operator forgot (or never knew to) set an FTS config, compaction
197    /// would silently drop the FTS index. This method prevents that loss.
198    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    /// Read all input files in parallel, returning ordered (batch, embeddings) pairs.
216    ///
217    /// `read_parquet()` decodes the vector column straight from Parquet and never
218    /// touches the AILK footer, so a file missing its index (external-engine rewrite,
219    /// or still `IndexStatus::Indexing`) still holds real data — it must be read like
220    /// any other input, not dropped. The merged output rebuilds a fresh index for
221    /// every row regardless, so a missing per-file index costs nothing here.
222    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                // Drop DV-masked rows before they're merged into a new physical file —
244                // row positions are about to change, so the old bitmap can't just be
245                // carried forward (see dv::filter_deleted_rows).
246                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    /// Merge `files` into a single new file at `output_path`.
259    ///
260    /// Reads all input files **in parallel** to minimise S3 latency, then
261    /// rebuilds the HNSW / IVF-PQ index synchronously. For very large merges
262    /// (N > 100 000 vectors) prefer `compact_deferred`, which offloads the
263    /// index build to a background Tokio task.
264    ///
265    /// Returns the DataFileEntry for the merged file.
266    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        // `read_files_parallel` always returns exactly one pair per input file (never
278        // filters), so `pairs.len() == files.len() > 0` given the `files.is_empty()` guard
279        // above — no separate emptiness check needed here.
280        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        // Concatenate all row groups into one batch
287        let merged_batch = concat_batches(schema, &all_batches)?;
288        let record_count = merged_batch.num_rows() as u64;
289
290        // Write merged file with adaptive index selection.
291        let writer = {
292            let base = AilakeFileWriter::new(self.policy.clone());
293            let base = match &self.index_strategy {
294                CompactionIndexStrategy::Auto => base.with_auto_index(),
295                CompactionIndexStrategy::ForceHnsw => base,
296                CompactionIndexStrategy::ForceIvfPq => {
297                    let cfg = ailake_index::IvfPqConfig::for_dataset(
298                        self.policy.dim as usize,
299                        all_embeddings.len(),
300                    );
301                    base.with_ivf_pq(cfg)
302                }
303            };
304            if let Some(ref fts_cfg) = self.fts_config {
305                match ailake_fts::merge_fts_blobs(fts_cfg, &merged_batch) {
306                    Ok(blob) => base.with_prebuilt_fts_blob(blob),
307                    Err(e) => {
308                        warn!("ailake: FTS re-index during compaction failed: {e}");
309                        base
310                    }
311                }
312            } else {
313                base
314            }
315        };
316        let file_bytes = writer.write(&merged_batch, &all_embeddings)?;
317        let file_size = file_bytes.len() as u64;
318        let column_stats =
319            ailake_catalog::extract_column_stats(&file_bytes, &[self.policy.column_name.as_str()])
320                .and_then(|m| serde_json::to_string(&m).ok());
321        self.store.put(output_path, file_bytes.clone()).await?;
322
323        // Compute centroid and HNSW offsets for catalog entry
324        let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
325        let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
326        let header = reader.read_header()?;
327        let ailk_start = reader.ailk_offset()?;
328
329        // Positional invariant check: parquet_count == hnsw_node_count == header.record_count.
330        // Catches a mismatched merge before it's committed to the catalog rather than
331        // surfacing as a wrong search result (or the invariant-violated error path in
332        // scanner.rs) later.
333        reader.verify_integrity()?;
334
335        // Preserve row-ID continuity: merged file inherits the minimum first_row_id of
336        // its sources so commit_snapshot doesn't allocate fresh IDs and grow next_row_id.
337        let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
338
339        let mut entry = make_data_file_entry(
340            output_path,
341            record_count,
342            file_size,
343            &centroid,
344            VectorIndexInfo {
345                column: &self.policy.column_name,
346                dim: self.policy.dim,
347                hnsw_offset: ailk_start + header.hnsw_offset,
348                hnsw_len: header.hnsw_len,
349            },
350        );
351        entry.first_row_id = source_first_row_id;
352        // Preserve idempotency keys: a retry of a source write must still see itself
353        // as already-committed after this merge (see `DataFileEntry::merge_batch_ids`
354        // and `write_batch_idempotent`).
355        entry.batch_id = DataFileEntry::merge_batch_ids(files);
356        entry.column_stats = column_stats;
357        Ok(entry)
358    }
359
360    /// Merge `files` into a single new file using incremental HNSW insertion.
361    ///
362    /// Identifies the **dominant file** — the file holding >= 40 % of the total
363    /// row count — loads its existing HNSW graph from the AILK section, then
364    /// calls `HnswIndex::insert_node` for every vector from the remaining files.
365    ///
366    /// **Complexity vs `compact`**:
367    /// - Full rebuild: O(N log N), N = total rows.
368    /// - Incremental (this method): O(N_dom) deserialization + O(N_small × log N_dom).
369    ///   For a 90 / 10 split (N = 1 M, N_dom = 900 k) the speedup is ~7×.
370    ///
371    /// **Fallbacks** (all degrade gracefully to `compact`):
372    /// - No file holds >= 40 % of rows.
373    /// - Dominant file's HNSW cannot be loaded (IVF-PQ, `IndexStatus::Indexing`, corrupt).
374    ///
375    /// **RowId contract**: dominant file's vectors are placed first in the merged
376    /// Parquet (positions 0..N_dom-1); other files follow. The existing RowIds from
377    /// the dominant HNSW remain valid; new nodes receive RowIds N_dom..N-1.
378    pub async fn compact_incremental(
379        &self,
380        files: &[DataFileEntry],
381        output_path: &str,
382    ) -> AilakeResult<DataFileEntry> {
383        const DOMINANT_RATIO: f64 = 0.40;
384
385        if files.is_empty() {
386            return Err(ailake_core::AilakeError::Catalog(
387                "compact_incremental: no files provided".into(),
388            ));
389        }
390
391        // This method only ever produces an HNSW index (dominant-file graph extended via
392        // `insert_node`) — it never builds IVF-PQ. `ForceIvfPq` is an explicit caller
393        // request for a specific index type, so it must never be silently satisfied with
394        // HNSW instead; fall back to `compact()`, which does respect it. `Auto`/`ForceHnsw`
395        // are unaffected: `Auto` treats "the dominant file already has a valid HNSW"
396        // as a legitimate reason to keep reusing HNSW rather than pay for a fresh
397        // hardware-picked rebuild, and `ForceHnsw` is satisfied by this method by
398        // construction either way.
399        if matches!(self.index_strategy, CompactionIndexStrategy::ForceIvfPq) {
400            debug!(
401                "ailake: compact_incremental — index_strategy=ForceIvfPq, which this method \
402                 can never produce (HNSW graph extension only); falling back to full rebuild"
403            );
404            return self.compact(files, output_path).await;
405        }
406
407        // Find the dominant file by record_count.
408        let total_rows: u64 = files.iter().map(|f| f.record_count).sum();
409        let dom_idx = files
410            .iter()
411            .enumerate()
412            .max_by_key(|(_, f)| f.record_count)
413            .map(|(i, _)| i)
414            .unwrap_or(0);
415        let dom_rows = files[dom_idx].record_count;
416
417        // The dominant file's existing HNSW graph is reused as-is (only non-dominant
418        // vectors are inserted into it) — its node IDs are tied to the dominant file's
419        // current row positions. If it has DV-masked rows, filtering them out of the
420        // Parquet data would desync the reused graph's node-to-row mapping (and HNSW
421        // doesn't support cheap node removal). Fall back to a full rebuild, which
422        // filters correctly via `read_files_parallel`.
423        if files[dom_idx].deletion_vector.is_some() {
424            debug!(
425                "ailake: compact_incremental — dominant file {} has DV-masked rows, \
426                 falling back to full rebuild (graph reuse would desync row positions)",
427                files[dom_idx].path
428            );
429            return self.compact(files, output_path).await;
430        }
431
432        if (dom_rows as f64 / total_rows as f64) < DOMINANT_RATIO {
433            debug!(
434                "ailake: compact_incremental — no dominant file ({}/{} rows < {:.0}% threshold), \
435                 falling back to full rebuild",
436                dom_rows,
437                total_rows,
438                DOMINANT_RATIO * 100.0
439            );
440            return self.compact(files, output_path).await;
441        }
442
443        let column = self.policy.column_name.clone();
444        let dim = self.policy.dim;
445        let dom_path = files[dom_idx].path.clone();
446
447        // Read all files in parallel. `read_parquet()` decodes the vector column
448        // directly from Parquet and never touches the AILK footer, so every file's
449        // rows are read regardless of whether it has an AI-Lake index — a file
450        // missing its footer (external-engine rewrite, or still `IndexStatus::Indexing`)
451        // still holds real data and must not be dropped from the merge. Raw bytes are
452        // retained only for the dominant file, and only when it actually has an index
453        // to reuse (needed to load its HNSW without a second round-trip); otherwise the
454        // dominant-file match below falls back to a full rebuild via `compact()`.
455        let futs: Vec<_> =
456            files
457                .iter()
458                .map(|entry| {
459                    let store = self.store.clone();
460                    let path = entry.path.clone();
461                    let col = column.clone();
462                    let is_dom = path == dom_path;
463                    let dv = entry.deletion_vector.clone();
464                    async move {
465                        let bytes: Bytes = store.get(&path).await?;
466                        let reader = AilakeFileReader::new(bytes.clone(), &col, dim);
467                        let has_index = reader.is_ailake_file();
468                        if is_dom && !has_index {
469                            debug!(
470                                "ailake: compact_incremental — dominant candidate {} has no \
471                             AI-Lake index; will fall back to full rebuild if no HNSW to reuse",
472                                path
473                            );
474                        }
475                        let (raw_batch, raw_vecs) = reader.read_parquet()?;
476                        // Non-dominant rows are freshly inserted into the graph below, so
477                        // DV-masked rows must be dropped here (dominant is guaranteed
478                        // DV-free by the caller's earlier fallback check).
479                        let (batch, vecs) = if let Some(dv) = dv {
480                            let bitmap = crate::dv::load_deletion_vector(&store, &dv).await?;
481                            crate::dv::filter_deleted_rows(raw_batch, raw_vecs, &bitmap)?
482                        } else {
483                            (raw_batch, raw_vecs)
484                        };
485                        let retained = if is_dom && has_index {
486                            Some(bytes)
487                        } else {
488                            None
489                        };
490                        Ok::<
491                            (RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>),
492                            ailake_core::AilakeError,
493                        >((batch, vecs, is_dom, retained))
494                    }
495                })
496                .collect();
497
498        // Every future above always resolves to one tuple per input file (or propagates an
499        // Err via `?`) — no per-file filtering, so no emptiness check is needed given the
500        // `files.is_empty()` guard at the top of this function.
501        #[allow(clippy::type_complexity)]
502        let raw: Vec<(RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>)> =
503            try_join_all(futs).await?;
504
505        // Separate dominant from others; dominant goes first in the merged file.
506        let mut dom_batch: Option<RecordBatch> = None;
507        let mut dom_vecs: Vec<Vec<f32>> = Vec::new();
508        let mut dom_bytes_found: Option<Bytes> = None;
509        let mut other_batches: Vec<RecordBatch> = Vec::new();
510        let mut other_vecs: Vec<Vec<f32>> = Vec::new();
511
512        for (batch, vecs, is_dom, retained) in raw {
513            if is_dom {
514                dom_batch = Some(batch);
515                dom_vecs = vecs;
516                dom_bytes_found = retained;
517            } else {
518                other_batches.push(batch);
519                other_vecs.extend(vecs);
520            }
521        }
522
523        let (dom_batch, dom_bytes) = match (dom_batch, dom_bytes_found) {
524            (Some(b), Some(byt)) => (b, byt),
525            _ => {
526                debug!(
527                    "ailake: compact_incremental — dominant file missing from read results, \
528                     falling back to full rebuild"
529                );
530                return self.compact(files, output_path).await;
531            }
532        };
533
534        // Load the dominant file's existing HNSW graph.
535        let dom_reader = AilakeFileReader::new(dom_bytes, &column, dim);
536        let mut hnsw = match dom_reader.load_index() {
537            Ok(idx) => idx,
538            Err(e) => {
539                debug!(
540                    "ailake: compact_incremental — cannot load dominant HNSW ({}), \
541                     falling back to full rebuild",
542                    e
543                );
544                return self.compact(files, output_path).await;
545            }
546        };
547
548        let dom_count = dom_batch.num_rows() as u64;
549
550        // Insert vectors from non-dominant files into the loaded graph.
551        // RowIds are assigned starting at dom_count to match positions in the merged Parquet.
552        for (j, vec) in other_vecs.iter().enumerate() {
553            hnsw.insert_node(RowId::new(dom_count + j as u64), vec.clone());
554        }
555        hnsw.quantize_to_f16();
556
557        // Assemble merged batch (dominant rows first) and all embeddings.
558        let schema: SchemaRef = dom_batch.schema();
559        let mut all_batches = vec![dom_batch];
560        all_batches.extend(other_batches);
561        let merged_batch = concat_batches(schema, &all_batches)?;
562        let record_count = merged_batch.num_rows() as u64;
563
564        let mut all_embeddings = dom_vecs;
565        all_embeddings.extend(other_vecs);
566
567        // Write the merged file using the pre-built index (no rebuild).
568        // Attach FTS blob when configured — data is already in merged_batch so cost is tokenization only.
569        let writer = {
570            let base = AilakeFileWriter::new(self.policy.clone());
571            if let Some(ref fts_cfg) = self.fts_config {
572                match ailake_fts::merge_fts_blobs(fts_cfg, &merged_batch) {
573                    Ok(blob) => base.with_prebuilt_fts_blob(blob),
574                    Err(e) => {
575                        warn!("ailake: FTS re-index during incremental compaction failed: {e}");
576                        base
577                    }
578                }
579            } else {
580                base
581            }
582        };
583        let file_bytes = writer.write_with_prebuilt_hnsw(&merged_batch, &all_embeddings, &hnsw)?;
584        let file_size = file_bytes.len() as u64;
585        let column_stats =
586            ailake_catalog::extract_column_stats(&file_bytes, &[self.policy.column_name.as_str()])
587                .and_then(|m| serde_json::to_string(&m).ok());
588        self.store.put(output_path, file_bytes.clone()).await?;
589
590        let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
591        let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
592        let header = reader.read_header()?;
593        let ailk_start = reader.ailk_offset()?;
594
595        // Positional invariant check — see `compact()` for rationale. Especially relevant
596        // here since the index is grown incrementally (insert_node) rather than rebuilt.
597        reader.verify_integrity()?;
598
599        // Dominant file goes first in the merged output, so the merged file's first
600        // logical row was the dominant file's first row.  Use its first_row_id so
601        // commit_snapshot doesn't grow next_row_id unnecessarily.
602        let source_first_row_id = files[dom_idx].first_row_id;
603
604        let mut entry = make_data_file_entry(
605            output_path,
606            record_count,
607            file_size,
608            &centroid,
609            VectorIndexInfo {
610                column: &self.policy.column_name,
611                dim: self.policy.dim,
612                hnsw_offset: ailk_start + header.hnsw_offset,
613                hnsw_len: header.hnsw_len,
614            },
615        );
616        entry.first_row_id = source_first_row_id;
617        entry.batch_id = DataFileEntry::merge_batch_ids(files);
618        entry.column_stats = column_stats;
619
620        info!(
621            "ailake: compact_incremental — merged {} files into {} \
622             ({} rows from dominant + {} inserted incrementally)",
623            files.len(),
624            output_path,
625            dom_count,
626            record_count - dom_count
627        );
628
629        Ok(entry)
630    }
631
632    /// Merge `files` into a single new file at `output_path`, writing Parquet
633    /// immediately and building the HNSW / IVF-PQ index in a background Tokio task.
634    ///
635    /// The merged file appears in the catalog as `IndexStatus::Indexing` until
636    /// the background task completes; queries fall back to flat scan during that
637    /// window (same behaviour as `write_batch_deferred`).
638    ///
639    /// Returns the `DataFileEntry` with `IndexStatus::Indexing`. The entry
640    /// transitions to `Ready` automatically when the background build finishes.
641    pub async fn compact_deferred(
642        &self,
643        files: &[DataFileEntry],
644        output_path: &str,
645        catalog: Arc<dyn CatalogProvider>,
646        table: &TableIdent,
647    ) -> AilakeResult<DataFileEntry> {
648        if files.is_empty() {
649            return Err(ailake_core::AilakeError::Catalog(
650                "compact_deferred: no files provided".into(),
651            ));
652        }
653
654        // See `compact()` above: read_files_parallel never filters, so no emptiness
655        // check is needed given the `files.is_empty()` guard above.
656        let pairs = self.read_files_parallel(files).await?;
657
658        let schema: SchemaRef = pairs[0].0.schema();
659        let (all_batches, all_embeddings): (Vec<_>, Vec<_>) = pairs.into_iter().unzip();
660        let all_embeddings: Vec<Vec<f32>> = all_embeddings.into_iter().flatten().collect();
661
662        let merged_batch = concat_batches(schema, &all_batches)?;
663        let record_count = merged_batch.num_rows() as u64;
664
665        // Write Parquet-only immediately — fast path, no HNSW build.
666        let file_writer = AilakeFileWriter::new(self.policy.clone());
667        let parquet_bytes = file_writer.write_parquet_only(&merged_batch, &all_embeddings)?;
668        let file_size = parquet_bytes.len() as u64;
669        let column_stats = ailake_catalog::extract_column_stats(
670            &parquet_bytes,
671            &[self.policy.column_name.as_str()],
672        )
673        .and_then(|m| serde_json::to_string(&m).ok());
674        self.store.put(output_path, parquet_bytes).await?;
675
676        // Centroid available for geometric pruning during the build window.
677        let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
678        let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
679        let mut entry = make_data_file_entry_indexing(
680            output_path,
681            record_count,
682            file_size,
683            &centroid,
684            &self.policy.column_name,
685            self.policy.dim,
686        );
687        entry.first_row_id = source_first_row_id;
688        entry.batch_id = DataFileEntry::merge_batch_ids(files);
689        entry.column_stats = column_stats;
690
691        // Spawn background index build; errors are logged, not propagated.
692        let store = self.store.clone();
693        let policy = self.policy.clone();
694        let table_id = table.clone();
695        let fp = output_path.to_string();
696        tokio::spawn(async move {
697            if let Err(e) = build_and_patch_index(store, catalog, policy, table_id, fp).await {
698                error!(
699                    "ailake: compaction deferred HNSW build failed — file indexed as \
700                     Parquet-only until next compaction rebuilds the index: {}",
701                    e
702                );
703            }
704        });
705
706        Ok(entry)
707    }
708
709    /// Full compaction workflow: plan, compact (synchronous HNSW rebuild),
710    /// drop old files from catalog, commit.
711    pub async fn run(
712        &self,
713        planner: &CompactionPlanner,
714        table: &TableIdent,
715        catalog: Arc<dyn CatalogProvider>,
716        output_prefix: &str,
717    ) -> AilakeResult<Option<DataFileEntry>> {
718        let all_files = catalog.list_files(table, None).await?;
719        let to_compact = planner.plan(&all_files);
720        if to_compact.is_empty() {
721            return Ok(None);
722        }
723
724        // Auto-detect FTS from table metadata so compaction never silently drops an FTS index
725        // that was present in the source files. Uses ailake.fts.* properties written at write time.
726        let meta_props = catalog
727            .load_table(table)
728            .await
729            .map(|m| m.properties)
730            .unwrap_or_default();
731        let executor = self.with_effective_fts(&meta_props);
732
733        let ts = std::time::SystemTime::now()
734            .duration_since(std::time::UNIX_EPOCH)
735            .unwrap_or_else(|e| e.duration())
736            .as_millis();
737        let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
738
739        // Use incremental merge when a dominant file exists (falls back to full rebuild automatically).
740        let merged = executor
741            .compact_incremental(&to_compact, &output_path)
742            .await?;
743
744        // Commit: add merged file, remove input files (via Replace snapshot).
745        let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
746        // Fetched fresh, right before the commit — same freshness rationale as the
747        // file-list re-list in `build_replace_file_list` above.
748        let parent_snapshot_id = catalog
749            .load_table(table)
750            .await
751            .ok()
752            .and_then(|m| m.current_snapshot_id);
753        let snapshot = NewSnapshot {
754            snapshot_id: ailake_catalog::new_snapshot_id(),
755            parent_snapshot_id,
756            files,
757            operation: SnapshotOperation::Replace,
758            iceberg_schema: None,
759            extra_properties: std::collections::HashMap::new(),
760            bloom_filters: vec![],
761            equality_delete_files: vec![],
762        };
763        catalog.commit_snapshot(table, snapshot).await?;
764
765        info!(
766            "ailake: compaction committed — merged {} files into {}",
767            to_compact.len(),
768            output_path
769        );
770
771        if catalog.retires_files_physically() {
772            delete_old_files(&self.store, &to_compact).await;
773        } else {
774            info!(
775                "ailake: compaction — leaving {} retired file(s) in place; catalog backend \
776                 manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
777                to_compact.len()
778            );
779        }
780
781        Ok(Some(merged))
782    }
783
784    /// Full compaction workflow with deferred HNSW build: plan, write merged
785    /// Parquet immediately, commit as `Indexing`, spawn background index build.
786    ///
787    /// Use for large tables where inline HNSW rebuild blocks too long.
788    ///
789    /// Note: FTS index is **not** rebuilt in deferred mode — `compact_deferred` writes
790    /// Parquet-only immediately and the background task (`build_and_patch_index`) only
791    /// builds the HNSW/IVF-PQ index. Use `run` (synchronous) when FTS preservation
792    /// on compaction is required.
793    pub async fn run_deferred(
794        &self,
795        planner: &CompactionPlanner,
796        table: &TableIdent,
797        catalog: Arc<dyn CatalogProvider>,
798        output_prefix: &str,
799    ) -> AilakeResult<Option<DataFileEntry>> {
800        // The background index build patches the merged file in place at its
801        // committed path — refuse on backends where committed bytes are
802        // immutable (see TableWriter::ensure_deferred_supported for why a
803        // commit-time guard alone is too late: the physical put comes first).
804        if !catalog.supports_in_place_rewrite() {
805            return Err(ailake_core::AilakeError::Catalog(
806                "deferred compaction is not supported with this catalog backend: the \
807                 background index build patches the merged file in place at its committed \
808                 path, which this catalog cannot re-register — run a blocking compact"
809                    .into(),
810            ));
811        }
812        let all_files = catalog.list_files(table, None).await?;
813        let to_compact = planner.plan(&all_files);
814        if to_compact.is_empty() {
815            return Ok(None);
816        }
817
818        let ts = std::time::SystemTime::now()
819            .duration_since(std::time::UNIX_EPOCH)
820            .unwrap_or_else(|e| e.duration())
821            .as_millis();
822        let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
823
824        let merged = self
825            .compact_deferred(&to_compact, &output_path, catalog.clone(), table)
826            .await?;
827
828        // Commit immediately: merged file in Indexing state replaces input files.
829        let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
830        // Fetched fresh, right before the commit — same freshness rationale as the
831        // file-list re-list in `build_replace_file_list` above.
832        let parent_snapshot_id = catalog
833            .load_table(table)
834            .await
835            .ok()
836            .and_then(|m| m.current_snapshot_id);
837        let snapshot = NewSnapshot {
838            snapshot_id: ailake_catalog::new_snapshot_id(),
839            parent_snapshot_id,
840            files,
841            operation: SnapshotOperation::Replace,
842            iceberg_schema: None,
843            extra_properties: std::collections::HashMap::new(),
844            bloom_filters: vec![],
845            equality_delete_files: vec![],
846        };
847        catalog.commit_snapshot(table, snapshot).await?;
848
849        info!(
850            "ailake: compaction committed (deferred) — merged {} files into {} \
851             (index building in background)",
852            to_compact.len(),
853            output_path
854        );
855
856        if catalog.retires_files_physically() {
857            delete_old_files(&self.store, &to_compact).await;
858        } else {
859            info!(
860                "ailake: compaction — leaving {} retired file(s) in place; catalog backend \
861                 manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
862                to_compact.len()
863            );
864        }
865
866        Ok(Some(merged))
867    }
868}
869
870/// Builds the file list for a post-compaction `Replace` snapshot: the merged output
871/// plus every current file that wasn't part of this compaction pass.
872///
873/// `Replace` does not inherit the previous manifest (see `HadoopCatalog::commit_snapshot`),
874/// so untouched files (above `target_file_size_bytes`, or beyond `max_files_per_pass`) must
875/// be carried forward explicitly or they vanish from the table. Re-lists via `list_files()`
876/// right before this call returns (rather than reusing a pre-merge snapshot) to narrow —
877/// though not eliminate, `HadoopCatalog` has no optimistic-concurrency check on
878/// `commit_snapshot` — the window for a concurrent writer's commit to be silently dropped
879/// by the `Replace` this list feeds into.
880async fn build_replace_file_list(
881    catalog: &Arc<dyn CatalogProvider>,
882    table: &TableIdent,
883    to_compact: &[DataFileEntry],
884    merged: DataFileEntry,
885) -> AilakeResult<Vec<DataFileEntry>> {
886    let current_files = catalog.list_files(table, None).await?;
887    let compacted_paths: std::collections::HashSet<&str> =
888        to_compact.iter().map(|f| f.path.as_str()).collect();
889    let mut files: Vec<DataFileEntry> = current_files
890        .into_iter()
891        .filter(|f| !compacted_paths.contains(f.path.as_str()))
892        .collect();
893    files.push(merged);
894    Ok(files)
895}
896
897async fn delete_old_files(store: &Arc<dyn Store>, files: &[DataFileEntry]) {
898    for entry in files {
899        if let Err(e) = store.delete(&entry.path).await {
900            error!(
901                "ailake: compaction cleanup failed — could not delete {}: {} \
902                 (orphan file in object store after successful catalog commit; \
903                 delete manually to reclaim storage)",
904                entry.path, e
905            );
906        }
907    }
908}
909
910fn concat_batches(schema: SchemaRef, batches: &[RecordBatch]) -> AilakeResult<RecordBatch> {
911    arrow_select::concat::concat_batches(&schema, batches)
912        .map_err(|e| ailake_core::AilakeError::Arrow(e.to_string()))
913}
914
915#[cfg(test)]
916mod tests {
917    use super::*;
918    use ailake_catalog::IndexStatus;
919
920    #[test]
921    fn plan_returns_empty_if_too_few_files() {
922        let planner = CompactionPlanner::new(CompactionConfig {
923            min_files_to_compact: 4,
924            target_file_size_bytes: 1024 * 1024,
925            ..Default::default()
926        });
927        let files: Vec<DataFileEntry> = (0..3)
928            .map(|i| DataFileEntry {
929                path: format!("file-{i}.parquet"),
930                record_count: 10,
931                file_size_bytes: 100,
932                centroid_b64: Some("AAAA".into()),
933                radius: None,
934                hnsw_offset: None,
935                hnsw_len: None,
936                vector_column: None,
937                vector_dim: None,
938                extra_vector_indexes: vec![],
939                index_status: IndexStatus::Ready,
940                index_error: None,
941                batch_id: None,
942                embedding_model: None,
943                partition_value: None,
944                deletion_vector: None,
945                first_row_id: None,
946                column_stats: None,
947                sequence_number: 0,
948            })
949            .collect();
950        assert!(planner.plan(&files).is_empty());
951    }
952
953    #[test]
954    fn plan_selects_small_files() {
955        let planner = CompactionPlanner::new(CompactionConfig {
956            min_files_to_compact: 2,
957            target_file_size_bytes: 1000,
958            ..Default::default()
959        });
960        let files = vec![
961            DataFileEntry {
962                path: "small.parquet".into(),
963                record_count: 5,
964                file_size_bytes: 500,
965                centroid_b64: Some("AAAA".into()),
966                radius: None,
967                hnsw_offset: None,
968                hnsw_len: None,
969                vector_column: None,
970                vector_dim: None,
971                extra_vector_indexes: vec![],
972                index_status: IndexStatus::Ready,
973                index_error: None,
974                batch_id: None,
975                embedding_model: None,
976                partition_value: None,
977                deletion_vector: None,
978                first_row_id: None,
979                column_stats: None,
980                sequence_number: 0,
981            },
982            DataFileEntry {
983                path: "large.parquet".into(),
984                record_count: 5000,
985                file_size_bytes: 200_000_000,
986                centroid_b64: Some("AAAA".into()),
987                radius: None,
988                hnsw_offset: None,
989                hnsw_len: None,
990                vector_column: None,
991                vector_dim: None,
992                extra_vector_indexes: vec![],
993                index_status: IndexStatus::Ready,
994                index_error: None,
995                batch_id: None,
996                embedding_model: None,
997                partition_value: None,
998                deletion_vector: None,
999                first_row_id: None,
1000                column_stats: None,
1001                sequence_number: 0,
1002            },
1003            DataFileEntry {
1004                path: "also-small.parquet".into(),
1005                record_count: 5,
1006                file_size_bytes: 800,
1007                centroid_b64: Some("AAAA".into()),
1008                radius: None,
1009                hnsw_offset: None,
1010                hnsw_len: None,
1011                vector_column: None,
1012                vector_dim: None,
1013                extra_vector_indexes: vec![],
1014                index_status: IndexStatus::Ready,
1015                index_error: None,
1016                batch_id: None,
1017                embedding_model: None,
1018                partition_value: None,
1019                deletion_vector: None,
1020                first_row_id: None,
1021                column_stats: None,
1022                sequence_number: 0,
1023            },
1024        ];
1025        let selected = planner.plan(&files);
1026        assert_eq!(selected.len(), 2);
1027        assert!(selected.iter().any(|f| f.path == "small.parquet"));
1028        assert!(selected.iter().any(|f| f.path == "also-small.parquet"));
1029    }
1030
1031    #[test]
1032    fn plan_respects_max_files_per_pass() {
1033        let planner = CompactionPlanner::new(CompactionConfig {
1034            min_files_to_compact: 2,
1035            target_file_size_bytes: 1_000_000,
1036            max_files_per_pass: 3,
1037            ..Default::default()
1038        });
1039        let files: Vec<DataFileEntry> = (0..5)
1040            .map(|i| DataFileEntry {
1041                path: format!("f{i}.parquet"),
1042                record_count: 10,
1043                file_size_bytes: 100 + i as u64 * 100,
1044                centroid_b64: Some("AAAA".into()),
1045                radius: None,
1046                hnsw_offset: None,
1047                hnsw_len: None,
1048                vector_column: None,
1049                vector_dim: None,
1050                extra_vector_indexes: vec![],
1051                index_status: IndexStatus::Ready,
1052                index_error: None,
1053                batch_id: None,
1054                embedding_model: None,
1055                partition_value: None,
1056                deletion_vector: None,
1057                first_row_id: None,
1058                column_stats: None,
1059                sequence_number: 0,
1060            })
1061            .collect();
1062        let selected = planner.plan(&files);
1063        assert_eq!(selected.len(), 3);
1064        assert_eq!(selected[0].file_size_bytes, 100);
1065        assert_eq!(selected[1].file_size_bytes, 200);
1066        assert_eq!(selected[2].file_size_bytes, 300);
1067    }
1068
1069    #[test]
1070    fn plan_sorts_smallest_first() {
1071        let planner = CompactionPlanner::new(CompactionConfig {
1072            min_files_to_compact: 2,
1073            target_file_size_bytes: 10_000,
1074            max_files_per_pass: 4,
1075            ..Default::default()
1076        });
1077        let files = vec![
1078            DataFileEntry {
1079                path: "c.parquet".into(),
1080                record_count: 1,
1081                file_size_bytes: 300,
1082                centroid_b64: Some("AAAA".into()),
1083                radius: None,
1084                hnsw_offset: None,
1085                hnsw_len: None,
1086                vector_column: None,
1087                vector_dim: None,
1088                extra_vector_indexes: vec![],
1089                index_status: IndexStatus::Ready,
1090                index_error: None,
1091                batch_id: None,
1092                embedding_model: None,
1093                partition_value: None,
1094                deletion_vector: None,
1095                first_row_id: None,
1096                column_stats: None,
1097                sequence_number: 0,
1098            },
1099            DataFileEntry {
1100                path: "a.parquet".into(),
1101                record_count: 1,
1102                file_size_bytes: 100,
1103                centroid_b64: Some("AAAA".into()),
1104                radius: None,
1105                hnsw_offset: None,
1106                hnsw_len: None,
1107                vector_column: None,
1108                vector_dim: None,
1109                extra_vector_indexes: vec![],
1110                index_status: IndexStatus::Ready,
1111                index_error: None,
1112                batch_id: None,
1113                embedding_model: None,
1114                partition_value: None,
1115                deletion_vector: None,
1116                first_row_id: None,
1117                column_stats: None,
1118                sequence_number: 0,
1119            },
1120            DataFileEntry {
1121                path: "b.parquet".into(),
1122                record_count: 1,
1123                file_size_bytes: 200,
1124                centroid_b64: Some("AAAA".into()),
1125                radius: None,
1126                hnsw_offset: None,
1127                hnsw_len: None,
1128                vector_column: None,
1129                vector_dim: None,
1130                extra_vector_indexes: vec![],
1131                index_status: IndexStatus::Ready,
1132                index_error: None,
1133                batch_id: None,
1134                embedding_model: None,
1135                partition_value: None,
1136                deletion_vector: None,
1137                first_row_id: None,
1138                column_stats: None,
1139                sequence_number: 0,
1140            },
1141        ];
1142        let selected = planner.plan(&files);
1143        assert_eq!(selected[0].file_size_bytes, 100);
1144        assert_eq!(selected[1].file_size_bytes, 200);
1145        assert_eq!(selected[2].file_size_bytes, 300);
1146    }
1147
1148    fn make_plan_entry(path: &str, size: u64, centroid_b64: Option<String>) -> DataFileEntry {
1149        DataFileEntry {
1150            path: path.to_string(),
1151            record_count: 10,
1152            file_size_bytes: size,
1153            centroid_b64,
1154            radius: None,
1155            hnsw_offset: None,
1156            hnsw_len: None,
1157            vector_column: None,
1158            vector_dim: None,
1159            extra_vector_indexes: vec![],
1160            index_status: IndexStatus::Ready,
1161            index_error: None,
1162            batch_id: None,
1163            embedding_model: None,
1164            partition_value: None,
1165            deletion_vector: None,
1166            first_row_id: None,
1167            column_stats: None,
1168            sequence_number: 0,
1169        }
1170    }
1171
1172    #[test]
1173    fn plan_prioritizes_foreign_written_files_regardless_of_size() {
1174        let planner = CompactionPlanner::new(CompactionConfig {
1175            min_files_to_compact: 4, // way more than the 1 native small file below
1176            target_file_size_bytes: 1000,
1177            max_files_per_pass: 20,
1178            ..Default::default()
1179        });
1180        let files = vec![
1181            // Native file, below target size, but alone — not enough to hit
1182            // min_files_to_compact on its own.
1183            make_plan_entry("native_small.parquet", 500, Some("AAAA".into())),
1184            // Foreign file — no centroid — large (would never pass the size filter),
1185            // but must still be selected because it has no AI-Lake index at all.
1186            make_plan_entry("foreign_big.parquet", 200_000_000, None),
1187        ];
1188        let selected = planner.plan(&files);
1189        assert_eq!(
1190            selected.len(),
1191            1,
1192            "foreign file alone should trigger a pass even below min_files_to_compact"
1193        );
1194        assert_eq!(selected[0].path, "foreign_big.parquet");
1195    }
1196
1197    #[test]
1198    fn plan_merges_foreign_and_size_candidates_together() {
1199        let planner = CompactionPlanner::new(CompactionConfig {
1200            min_files_to_compact: 2,
1201            target_file_size_bytes: 1000,
1202            max_files_per_pass: 20,
1203            ..Default::default()
1204        });
1205        let files = vec![
1206            make_plan_entry("native_a.parquet", 300, Some("AAAA".into())),
1207            make_plan_entry("native_b.parquet", 400, Some("AAAA".into())),
1208            make_plan_entry("foreign.parquet", 200_000_000, None),
1209        ];
1210        let selected = planner.plan(&files);
1211        let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1212        assert_eq!(selected.len(), 3, "paths={paths:?}");
1213        // Foreign file must be first (repair priority).
1214        assert_eq!(selected[0].path, "foreign.parquet", "paths={paths:?}");
1215    }
1216
1217    #[test]
1218    fn plan_sorts_foreign_files_smallest_first_too() {
1219        // Regression: foreign files used to keep list_files()'s arbitrary order, so a
1220        // pass could select several large foreign files instead of the cheapest ones —
1221        // violating the documented "bounds peak RAM" invariant of max_files_per_pass.
1222        let planner = CompactionPlanner::new(CompactionConfig {
1223            min_files_to_compact: 1,
1224            target_file_size_bytes: 1000,
1225            max_files_per_pass: 2,
1226            ..Default::default()
1227        });
1228        let files = vec![
1229            make_plan_entry("foreign_huge.parquet", 500_000_000, None),
1230            make_plan_entry("foreign_small.parquet", 100, None),
1231            make_plan_entry("foreign_medium.parquet", 10_000_000, None),
1232        ];
1233        let selected = planner.plan(&files);
1234        let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1235        assert_eq!(selected.len(), 2, "max_files_per_pass=2, paths={paths:?}");
1236        assert_eq!(
1237            paths,
1238            vec!["foreign_small.parquet", "foreign_medium.parquet"],
1239            "foreign files must be size-sorted before truncation, cheapest first"
1240        );
1241    }
1242
1243    #[tokio::test]
1244    async fn compact_merges_two_files() {
1245        use ailake_core::{VectorMetric, VectorPrecision};
1246        use ailake_store::LocalStore;
1247        use arrow_array::{Int32Array, RecordBatch};
1248        use arrow_schema::{DataType, Field, Schema};
1249        use std::sync::Arc;
1250        use tempfile::TempDir;
1251
1252        let dir = TempDir::new().unwrap();
1253        let store = Arc::new(LocalStore::new(dir.path()));
1254        let policy = VectorStoragePolicy {
1255            column_name: "embedding".into(),
1256            dim: 4,
1257            metric: VectorMetric::Cosine,
1258            precision: VectorPrecision::F16,
1259            pq: None,
1260            keep_raw_for_reranking: true,
1261            pre_normalize: false,
1262            hnsw_m: None,
1263            hnsw_ef_construction: None,
1264            ivf_residual: false,
1265            embedding_model: None,
1266            modality: None,
1267            partition_by: None,
1268            partition_value: None,
1269            partition_column_type: None,
1270            partition_fields: vec![],
1271        };
1272
1273        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1274        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
1275        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
1276
1277        let batch_a = RecordBatch::try_new(
1278            schema.clone(),
1279            vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1280        )
1281        .unwrap();
1282        let batch_b = RecordBatch::try_new(
1283            schema.clone(),
1284            vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1285        )
1286        .unwrap();
1287
1288        let writer_a = AilakeFileWriter::new(policy.clone());
1289        let bytes_a = writer_a.write(&batch_a, &embs_a).unwrap();
1290        let writer_b = AilakeFileWriter::new(policy.clone());
1291        let bytes_b = writer_b.write(&batch_b, &embs_b).unwrap();
1292
1293        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1294        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1295
1296        let entries = vec![
1297            DataFileEntry {
1298                path: "data/a.parquet".into(),
1299                record_count: 2,
1300                file_size_bytes: bytes_a.len() as u64,
1301                centroid_b64: None,
1302                radius: None,
1303                hnsw_offset: None,
1304                hnsw_len: None,
1305                vector_column: None,
1306                vector_dim: None,
1307                extra_vector_indexes: vec![],
1308                index_status: IndexStatus::Ready,
1309                index_error: None,
1310                batch_id: None,
1311                embedding_model: None,
1312                partition_value: None,
1313                deletion_vector: None,
1314                first_row_id: None,
1315                column_stats: None,
1316                sequence_number: 0,
1317            },
1318            DataFileEntry {
1319                path: "data/b.parquet".into(),
1320                record_count: 2,
1321                file_size_bytes: bytes_b.len() as u64,
1322                centroid_b64: None,
1323                radius: None,
1324                hnsw_offset: None,
1325                hnsw_len: None,
1326                vector_column: None,
1327                vector_dim: None,
1328                extra_vector_indexes: vec![],
1329                index_status: IndexStatus::Ready,
1330                index_error: None,
1331                batch_id: None,
1332                embedding_model: None,
1333                partition_value: None,
1334                deletion_vector: None,
1335                first_row_id: None,
1336                column_stats: None,
1337                sequence_number: 0,
1338            },
1339        ];
1340
1341        let executor = CompactionExecutor::new(store.clone(), policy.clone());
1342        let merged = executor
1343            .compact(&entries, "data/merged.parquet")
1344            .await
1345            .unwrap();
1346
1347        assert_eq!(merged.record_count, 4);
1348        assert_eq!(merged.path, "data/merged.parquet");
1349
1350        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1351        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1352        reader.verify_integrity().unwrap();
1353        let (batch, embs) = reader.read_parquet().unwrap();
1354        assert_eq!(batch.num_rows(), 4);
1355        assert_eq!(embs.len(), 4);
1356    }
1357
1358    /// Regression test: `compact()`/`read_files_parallel` used to read every input file's
1359    /// rows unconditionally, ignoring any `deletion_vector` on the entry — so a row deleted
1360    /// via `delete_rows()` before compaction reappeared in the merged output (new physical
1361    /// file, fresh row positions, no DV carried over — the row was never actually removed).
1362    #[tokio::test]
1363    async fn compact_drops_deletion_vector_masked_rows() {
1364        use ailake_catalog::provider::DeletionVector;
1365        use ailake_core::{VectorMetric, VectorPrecision};
1366        use ailake_store::LocalStore;
1367        use arrow_array::{Int32Array, RecordBatch};
1368        use arrow_schema::{DataType, Field, Schema};
1369        use roaring::RoaringBitmap;
1370        use std::sync::Arc;
1371        use tempfile::TempDir;
1372
1373        let dir = TempDir::new().unwrap();
1374        let store = Arc::new(LocalStore::new(dir.path()));
1375        let policy = VectorStoragePolicy {
1376            column_name: "embedding".into(),
1377            dim: 4,
1378            metric: VectorMetric::Cosine,
1379            precision: VectorPrecision::F16,
1380            pq: None,
1381            keep_raw_for_reranking: true,
1382            pre_normalize: false,
1383            hnsw_m: None,
1384            hnsw_ef_construction: None,
1385            ivf_residual: false,
1386            embedding_model: None,
1387            modality: None,
1388            partition_by: None,
1389            partition_value: None,
1390            partition_column_type: None,
1391            partition_fields: vec![],
1392        };
1393
1394        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1395        // File A: row 0 (id=0) will be marked deleted; row 1 (id=1) survives.
1396        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
1397        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
1398
1399        let batch_a = RecordBatch::try_new(
1400            schema.clone(),
1401            vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1402        )
1403        .unwrap();
1404        let batch_b =
1405            RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![2i32]))])
1406                .unwrap();
1407
1408        let bytes_a = AilakeFileWriter::new(policy.clone())
1409            .write(&batch_a, &embs_a)
1410            .unwrap();
1411        let bytes_b = AilakeFileWriter::new(policy.clone())
1412            .write(&batch_b, &embs_b)
1413            .unwrap();
1414        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1415        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1416
1417        // Puffin DV marking row 0 of file A as deleted.
1418        let mut bitmap = RoaringBitmap::new();
1419        bitmap.insert(0);
1420        let (puffin_bytes, offset, length) =
1421            crate::delete::PuffinWriter::write_single_dv(&bitmap, 1).unwrap();
1422        store.put("metadata/dv-1.dvd", puffin_bytes).await.unwrap();
1423
1424        let entry_a = DataFileEntry {
1425            path: "data/a.parquet".into(),
1426            record_count: 2,
1427            file_size_bytes: bytes_a.len() as u64,
1428            centroid_b64: None,
1429            radius: None,
1430            hnsw_offset: None,
1431            hnsw_len: None,
1432            vector_column: None,
1433            vector_dim: None,
1434            extra_vector_indexes: vec![],
1435            index_status: IndexStatus::Ready,
1436            index_error: None,
1437            batch_id: None,
1438            embedding_model: None,
1439            partition_value: None,
1440            deletion_vector: Some(DeletionVector {
1441                path: "metadata/dv-1.dvd".into(),
1442                offset,
1443                length,
1444                cardinality: 1,
1445            }),
1446            first_row_id: None,
1447            column_stats: None,
1448            sequence_number: 0,
1449        };
1450        let entry_b = DataFileEntry {
1451            path: "data/b.parquet".into(),
1452            record_count: 1,
1453            file_size_bytes: bytes_b.len() as u64,
1454            centroid_b64: None,
1455            radius: None,
1456            hnsw_offset: None,
1457            hnsw_len: None,
1458            vector_column: None,
1459            vector_dim: None,
1460            extra_vector_indexes: vec![],
1461            index_status: IndexStatus::Ready,
1462            index_error: None,
1463            batch_id: None,
1464            embedding_model: None,
1465            partition_value: None,
1466            deletion_vector: None,
1467            first_row_id: None,
1468            column_stats: None,
1469            sequence_number: 0,
1470        };
1471
1472        let executor = CompactionExecutor::new(store.clone(), policy.clone());
1473        let merged = executor
1474            .compact(&[entry_a, entry_b], "data/merged.parquet")
1475            .await
1476            .unwrap();
1477
1478        // 2 rows in A + 1 in B, minus 1 deleted from A = 2 surviving rows.
1479        assert_eq!(merged.record_count, 2);
1480        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1481        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1482        let (batch, embs) = reader.read_parquet().unwrap();
1483        assert_eq!(batch.num_rows(), 2);
1484        assert_eq!(embs.len(), 2);
1485        let ids: Vec<i32> = batch
1486            .column_by_name("id")
1487            .unwrap()
1488            .as_any()
1489            .downcast_ref::<Int32Array>()
1490            .unwrap()
1491            .values()
1492            .to_vec();
1493        assert_eq!(
1494            ids,
1495            vec![1, 2],
1496            "id=0 (deleted) must not survive compaction"
1497        );
1498    }
1499
1500    /// Regression test: `verify_integrity()` used to always call `load_index()`, which
1501    /// unconditionally deserializes as `HnswIndex` regardless of the IVF-PQ flag —
1502    /// so any compaction using `ForceIvfPq` (or `Auto` selecting IVF-PQ) would write a
1503    /// correct merged file and then fail on this call, aborting the whole compaction.
1504    #[tokio::test]
1505    async fn compact_with_ivf_pq_strategy_does_not_crash_on_verify_integrity() {
1506        use ailake_core::{VectorMetric, VectorPrecision};
1507        use ailake_store::LocalStore;
1508        use arrow_array::{Int32Array, RecordBatch};
1509        use arrow_schema::{DataType, Field, Schema};
1510        use std::sync::Arc;
1511        use tempfile::TempDir;
1512
1513        let dir = TempDir::new().unwrap();
1514        let store = Arc::new(LocalStore::new(dir.path()));
1515        let dim = 8;
1516        let policy = VectorStoragePolicy {
1517            column_name: "embedding".into(),
1518            dim,
1519            metric: VectorMetric::Cosine,
1520            precision: VectorPrecision::F16,
1521            pq: None,
1522            keep_raw_for_reranking: true,
1523            pre_normalize: false,
1524            hnsw_m: None,
1525            hnsw_ef_construction: None,
1526            ivf_residual: false,
1527            embedding_model: None,
1528            modality: None,
1529            partition_by: None,
1530            partition_value: None,
1531            partition_column_type: None,
1532            partition_fields: vec![],
1533        };
1534
1535        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1536        let n_per_file = 30usize;
1537        let make_file = |path: &str, offset: i32| {
1538            let ids: Vec<i32> = (offset..offset + n_per_file as i32).collect();
1539            let embs: Vec<Vec<f32>> = ids
1540                .iter()
1541                .map(|&i| {
1542                    (0..dim as i32)
1543                        .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1544                        .collect()
1545                })
1546                .collect();
1547            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1548                .unwrap();
1549            let bytes = AilakeFileWriter::new(policy.clone())
1550                .write(&batch, &embs)
1551                .unwrap();
1552            (path.to_string(), bytes)
1553        };
1554
1555        let (path_a, bytes_a) = make_file("data/ivfpq_a.parquet", 0);
1556        let (path_b, bytes_b) = make_file("data/ivfpq_b.parquet", n_per_file as i32);
1557        for (path, bytes) in [(&path_a, &bytes_a), (&path_b, &bytes_b)] {
1558            store.put(path, bytes.clone()).await.unwrap();
1559        }
1560
1561        let make_entry = |path: &str, size: u64| DataFileEntry {
1562            path: path.to_string(),
1563            record_count: n_per_file as u64,
1564            file_size_bytes: size,
1565            centroid_b64: None,
1566            radius: None,
1567            hnsw_offset: None,
1568            hnsw_len: None,
1569            vector_column: None,
1570            vector_dim: None,
1571            extra_vector_indexes: vec![],
1572            index_status: IndexStatus::Ready,
1573            index_error: None,
1574            batch_id: None,
1575            embedding_model: None,
1576            partition_value: None,
1577            deletion_vector: None,
1578            first_row_id: None,
1579            column_stats: None,
1580            sequence_number: 0,
1581        };
1582        let entries = vec![
1583            make_entry(&path_a, bytes_a.len() as u64),
1584            make_entry(&path_b, bytes_b.len() as u64),
1585        ];
1586
1587        let executor = CompactionExecutor::new(store.clone(), policy.clone())
1588            .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1589        let merged = executor
1590            .compact(&entries, "data/ivfpq_merged.parquet")
1591            .await
1592            .expect("compact() with ForceIvfPq must not fail in verify_integrity()");
1593
1594        assert_eq!(merged.record_count, 2 * n_per_file as u64);
1595
1596        // The bug: this used to panic/error trying to load IVF-PQ bytes as HNSW.
1597        let merged_bytes = store.get("data/ivfpq_merged.parquet").await.unwrap();
1598        let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1599        reader
1600            .verify_integrity()
1601            .expect("verify_integrity() must handle an IVF-PQ-indexed file");
1602    }
1603
1604    #[tokio::test]
1605    async fn compact_incremental_merges_dominant_plus_small() {
1606        use ailake_core::{RowId, VectorMetric, VectorPrecision};
1607        use ailake_store::LocalStore;
1608        use arrow_array::{Int32Array, RecordBatch};
1609        use arrow_schema::{DataType, Field, Schema};
1610        use std::sync::Arc;
1611        use tempfile::TempDir;
1612
1613        let dir = TempDir::new().unwrap();
1614        let store = Arc::new(LocalStore::new(dir.path()));
1615        let policy = VectorStoragePolicy {
1616            column_name: "embedding".into(),
1617            dim: 4,
1618            metric: VectorMetric::Cosine,
1619            precision: VectorPrecision::F16,
1620            pq: None,
1621            keep_raw_for_reranking: true,
1622            pre_normalize: false,
1623            hnsw_m: None,
1624            hnsw_ef_construction: None,
1625            ivf_residual: false,
1626            embedding_model: None,
1627            modality: None,
1628            partition_by: None,
1629            partition_value: None,
1630            partition_column_type: None,
1631            partition_fields: vec![],
1632        };
1633
1634        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1635
1636        // Dominant file: 6 rows (75% of total 8 rows — above 40% threshold).
1637        let embs_dom: Vec<Vec<f32>> = vec![
1638            vec![1.0, 0.0, 0.0, 0.0],
1639            vec![0.0, 1.0, 0.0, 0.0],
1640            vec![0.0, 0.0, 1.0, 0.0],
1641            vec![0.7, 0.7, 0.0, 0.0],
1642            vec![0.0, 0.7, 0.7, 0.0],
1643            vec![0.0, 0.0, 0.7, 0.7],
1644        ];
1645        let batch_dom = RecordBatch::try_new(
1646            schema.clone(),
1647            vec![Arc::new(Int32Array::from(vec![0i32, 1, 2, 3, 4, 5]))],
1648        )
1649        .unwrap();
1650
1651        // Small file: 2 rows.
1652        let embs_small: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 0.0, 1.0], vec![0.5, 0.5, 0.5, 0.5]];
1653        let batch_small = RecordBatch::try_new(
1654            schema.clone(),
1655            vec![Arc::new(Int32Array::from(vec![6i32, 7]))],
1656        )
1657        .unwrap();
1658
1659        let bytes_dom = AilakeFileWriter::new(policy.clone())
1660            .write(&batch_dom, &embs_dom)
1661            .unwrap();
1662        let bytes_small = AilakeFileWriter::new(policy.clone())
1663            .write(&batch_small, &embs_small)
1664            .unwrap();
1665
1666        store
1667            .put("data/dominant.parquet", bytes_dom.clone())
1668            .await
1669            .unwrap();
1670        store
1671            .put("data/small.parquet", bytes_small.clone())
1672            .await
1673            .unwrap();
1674
1675        let entries = vec![
1676            DataFileEntry {
1677                path: "data/dominant.parquet".into(),
1678                record_count: 6,
1679                file_size_bytes: bytes_dom.len() as u64,
1680                centroid_b64: None,
1681                radius: None,
1682                hnsw_offset: None,
1683                hnsw_len: None,
1684                vector_column: None,
1685                vector_dim: None,
1686                extra_vector_indexes: vec![],
1687                index_status: IndexStatus::Ready,
1688                index_error: None,
1689                batch_id: None,
1690                embedding_model: None,
1691                partition_value: None,
1692                deletion_vector: None,
1693                first_row_id: None,
1694                column_stats: None,
1695                sequence_number: 0,
1696            },
1697            DataFileEntry {
1698                path: "data/small.parquet".into(),
1699                record_count: 2,
1700                file_size_bytes: bytes_small.len() as u64,
1701                centroid_b64: None,
1702                radius: None,
1703                hnsw_offset: None,
1704                hnsw_len: None,
1705                vector_column: None,
1706                vector_dim: None,
1707                extra_vector_indexes: vec![],
1708                index_status: IndexStatus::Ready,
1709                index_error: None,
1710                batch_id: None,
1711                embedding_model: None,
1712                partition_value: None,
1713                deletion_vector: None,
1714                first_row_id: None,
1715                column_stats: None,
1716                sequence_number: 0,
1717            },
1718        ];
1719
1720        let executor = CompactionExecutor::new(store.clone(), policy.clone());
1721        let merged = executor
1722            .compact_incremental(&entries, "data/merged.parquet")
1723            .await
1724            .unwrap();
1725
1726        // Structural checks.
1727        assert_eq!(merged.record_count, 8);
1728        assert_eq!(merged.path, "data/merged.parquet");
1729
1730        // Load merged file and verify it's a valid AI-Lake file.
1731        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1732        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1733        reader.verify_integrity().unwrap();
1734
1735        let (batch, embs) = reader.read_parquet().unwrap();
1736        assert_eq!(batch.num_rows(), 8);
1737        assert_eq!(embs.len(), 8);
1738
1739        // Dominant rows must come first (positions 0..5).
1740        for f in &embs[..6] {
1741            assert_eq!(f.len(), 4);
1742        }
1743
1744        // HNSW must be searchable and return the nearest neighbor for a known query.
1745        let hnsw = reader.load_index().unwrap();
1746        assert_eq!(hnsw.node_count(), 8);
1747
1748        // Query [1, 0, 0, 0] → nearest should be RowId 0 (embs_dom[0]).
1749        let results = hnsw.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
1750        assert_eq!(results[0].0, RowId::new(0));
1751
1752        // Query [0, 0, 0, 1] → nearest should be RowId 6 (first row of small file,
1753        // inserted at position 6 in the merged file).
1754        let results = hnsw.search(&[0.0, 0.0, 0.0, 1.0], 1, 50);
1755        assert_eq!(results[0].0, RowId::new(6));
1756    }
1757
1758    /// Regression: `compact_incremental()` only ever produces HNSW (it extends the
1759    /// dominant file's existing graph) — with a dominant file present, it never checked
1760    /// `self.index_strategy` before taking that path, so an explicit `ForceIvfPq` request
1761    /// was silently satisfied with HNSW instead. Reachable from `CompactionExecutor::run()`/
1762    /// `run_deferred()` (used by the CLI `compact` command and every JNI/Spark/Trino/Flink/
1763    /// DuckDB compact call), which always try `compact_incremental()` first — including on
1764    /// a GPU machine where `Auto` would pick IVF-PQ for a fresh build, or where a caller
1765    /// explicitly forces IVF-PQ. Same dominant/small file shape as
1766    /// `compact_incremental_merges_dominant_plus_small`, but with `ForceIvfPq` set.
1767    #[tokio::test]
1768    async fn compact_incremental_respects_force_ivf_pq_even_with_dominant_file() {
1769        use ailake_core::{VectorMetric, VectorPrecision};
1770        use ailake_store::LocalStore;
1771        use arrow_array::{Int32Array, RecordBatch};
1772        use arrow_schema::{DataType, Field, Schema};
1773        use std::sync::Arc;
1774        use tempfile::TempDir;
1775
1776        let dir = TempDir::new().unwrap();
1777        let store = Arc::new(LocalStore::new(dir.path()));
1778        let dim = 8;
1779        let policy = VectorStoragePolicy {
1780            column_name: "embedding".into(),
1781            dim,
1782            metric: VectorMetric::Cosine,
1783            precision: VectorPrecision::F16,
1784            pq: None,
1785            keep_raw_for_reranking: true,
1786            pre_normalize: false,
1787            hnsw_m: None,
1788            hnsw_ef_construction: None,
1789            ivf_residual: false,
1790            embedding_model: None,
1791            modality: None,
1792            partition_by: None,
1793            partition_value: None,
1794            partition_column_type: None,
1795            partition_fields: vec![],
1796        };
1797
1798        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1799        let make_file = |path: &str, offset: i32, n: usize| {
1800            let ids: Vec<i32> = (offset..offset + n as i32).collect();
1801            let embs: Vec<Vec<f32>> = ids
1802                .iter()
1803                .map(|&i| {
1804                    (0..dim as i32)
1805                        .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1806                        .collect()
1807                })
1808                .collect();
1809            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1810                .unwrap();
1811            // Dominant file already has a real, loadable HNSW index — the exact
1812            // condition that used to make compact_incremental() bypass ForceIvfPq.
1813            let bytes = AilakeFileWriter::new(policy.clone())
1814                .write(&batch, &embs)
1815                .unwrap();
1816            (path.to_string(), bytes, n as u64)
1817        };
1818
1819        // Dominant: 90 rows (75% of 120 total — comfortably above the dominant-file
1820        // threshold), small: 30 rows.
1821        let (path_dom, bytes_dom, n_dom) = make_file("data/dom.parquet", 0, 90);
1822        let (path_small, bytes_small, n_small) = make_file("data/small.parquet", 999, 30);
1823
1824        for (path, bytes) in [(&path_dom, &bytes_dom), (&path_small, &bytes_small)] {
1825            store.put(path, bytes.clone()).await.unwrap();
1826        }
1827
1828        let make_entry = |path: &str, record_count: u64, size: u64| DataFileEntry {
1829            path: path.to_string(),
1830            record_count,
1831            file_size_bytes: size,
1832            centroid_b64: None,
1833            radius: None,
1834            hnsw_offset: None,
1835            hnsw_len: None,
1836            vector_column: None,
1837            vector_dim: None,
1838            extra_vector_indexes: vec![],
1839            index_status: IndexStatus::Ready,
1840            index_error: None,
1841            batch_id: None,
1842            embedding_model: None,
1843            partition_value: None,
1844            deletion_vector: None,
1845            first_row_id: None,
1846            column_stats: None,
1847            sequence_number: 0,
1848        };
1849        let entries = vec![
1850            make_entry(&path_dom, n_dom, bytes_dom.len() as u64),
1851            make_entry(&path_small, n_small, bytes_small.len() as u64),
1852        ];
1853
1854        let executor = CompactionExecutor::new(store.clone(), policy.clone())
1855            .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1856        let merged = executor
1857            .compact_incremental(&entries, "data/merged_force_ivfpq.parquet")
1858            .await
1859            .expect("compact_incremental() with ForceIvfPq must fall back to compact() cleanly");
1860
1861        let merged_bytes = store.get("data/merged_force_ivfpq.parquet").await.unwrap();
1862        let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1863        reader.verify_integrity().unwrap();
1864
1865        match reader.load_any_index().unwrap() {
1866            ailake_index::AnyIndex::IvfPq(_) => {}
1867            ailake_index::AnyIndex::Hnsw(_) => panic!(
1868                "ForceIvfPq was silently satisfied with HNSW instead — merged.record_count={}",
1869                merged.record_count
1870            ),
1871        }
1872    }
1873
1874    #[tokio::test]
1875    async fn compact_incremental_falls_back_when_no_dominant() {
1876        use ailake_core::{VectorMetric, VectorPrecision};
1877        use ailake_store::LocalStore;
1878        use arrow_array::{Int32Array, RecordBatch};
1879        use arrow_schema::{DataType, Field, Schema};
1880        use std::sync::Arc;
1881        use tempfile::TempDir;
1882
1883        let dir = TempDir::new().unwrap();
1884        let store = Arc::new(LocalStore::new(dir.path()));
1885        let policy = VectorStoragePolicy {
1886            column_name: "embedding".into(),
1887            dim: 4,
1888            metric: VectorMetric::Cosine,
1889            precision: VectorPrecision::F16,
1890            pq: None,
1891            keep_raw_for_reranking: true,
1892            pre_normalize: false,
1893            hnsw_m: None,
1894            hnsw_ef_construction: None,
1895            ivf_residual: false,
1896            embedding_model: None,
1897            modality: None,
1898            partition_by: None,
1899            partition_value: None,
1900            partition_column_type: None,
1901            partition_fields: vec![],
1902        };
1903
1904        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1905
1906        // Two equal-sized files (50/50 split — no dominant, both below 40% threshold).
1907        let make_batch = |ids: Vec<i32>, embs: Vec<Vec<f32>>| {
1908            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1909                .unwrap();
1910            AilakeFileWriter::new(policy.clone())
1911                .write(&batch, &embs)
1912                .unwrap()
1913        };
1914
1915        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
1916        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
1917        let bytes_a = make_batch(vec![0, 1], embs_a);
1918        let bytes_b = make_batch(vec![2, 3], embs_b);
1919
1920        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1921        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1922
1923        let entries = vec![
1924            DataFileEntry {
1925                path: "data/a.parquet".into(),
1926                record_count: 2,
1927                file_size_bytes: bytes_a.len() as u64,
1928                centroid_b64: None,
1929                radius: None,
1930                hnsw_offset: None,
1931                hnsw_len: None,
1932                vector_column: None,
1933                vector_dim: None,
1934                extra_vector_indexes: vec![],
1935                index_status: IndexStatus::Ready,
1936                index_error: None,
1937                batch_id: None,
1938                embedding_model: None,
1939                partition_value: None,
1940                deletion_vector: None,
1941                first_row_id: None,
1942                column_stats: None,
1943                sequence_number: 0,
1944            },
1945            DataFileEntry {
1946                path: "data/b.parquet".into(),
1947                record_count: 2,
1948                file_size_bytes: bytes_b.len() as u64,
1949                centroid_b64: None,
1950                radius: None,
1951                hnsw_offset: None,
1952                hnsw_len: None,
1953                vector_column: None,
1954                vector_dim: None,
1955                extra_vector_indexes: vec![],
1956                index_status: IndexStatus::Ready,
1957                index_error: None,
1958                batch_id: None,
1959                embedding_model: None,
1960                partition_value: None,
1961                deletion_vector: None,
1962                first_row_id: None,
1963                column_stats: None,
1964                sequence_number: 0,
1965            },
1966        ];
1967
1968        let executor = CompactionExecutor::new(store.clone(), policy.clone());
1969        // Should fall back to full rebuild without error.
1970        let merged = executor
1971            .compact_incremental(&entries, "data/merged.parquet")
1972            .await
1973            .unwrap();
1974
1975        assert_eq!(merged.record_count, 4);
1976
1977        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1978        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1979        reader.verify_integrity().unwrap();
1980    }
1981
1982    #[tokio::test]
1983    async fn compact_deferred_produces_parquet_only_file() {
1984        use ailake_catalog::HadoopCatalog;
1985        use ailake_core::{VectorMetric, VectorPrecision};
1986        use ailake_store::LocalStore;
1987        use arrow_array::{Int32Array, RecordBatch};
1988        use arrow_schema::{DataType, Field, Schema};
1989        use std::sync::Arc;
1990        use tempfile::TempDir;
1991
1992        let dir = TempDir::new().unwrap();
1993        let store = Arc::new(LocalStore::new(dir.path()));
1994        let catalog_dir = TempDir::new().unwrap();
1995        let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
1996        let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
1997        let table = TableIdent {
1998            namespace: "ns".into(),
1999            name: "tbl".into(),
2000        };
2001
2002        let policy = VectorStoragePolicy {
2003            column_name: "embedding".into(),
2004            dim: 4,
2005            metric: VectorMetric::Cosine,
2006            precision: VectorPrecision::F16,
2007            pq: None,
2008            keep_raw_for_reranking: true,
2009            pre_normalize: false,
2010            hnsw_m: None,
2011            hnsw_ef_construction: None,
2012            ivf_residual: false,
2013            embedding_model: None,
2014            modality: None,
2015            partition_by: None,
2016            partition_value: None,
2017            partition_column_type: None,
2018            partition_fields: vec![],
2019        };
2020
2021        use ailake_catalog::TableProperties;
2022        catalog
2023            .create_table(
2024                &table,
2025                &TableProperties {
2026                    policy: policy.clone(),
2027                    extra: std::collections::HashMap::new(),
2028                    format_version: 2,
2029                    partition_column_type: None,
2030                },
2031            )
2032            .await
2033            .unwrap();
2034
2035        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2036        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
2037        let batch_a = RecordBatch::try_new(
2038            schema.clone(),
2039            vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2040        )
2041        .unwrap();
2042        let bytes_a = AilakeFileWriter::new(policy.clone())
2043            .write(&batch_a, &embs_a)
2044            .unwrap();
2045        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
2046
2047        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
2048        let batch_b = RecordBatch::try_new(
2049            schema.clone(),
2050            vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2051        )
2052        .unwrap();
2053        let bytes_b = AilakeFileWriter::new(policy.clone())
2054            .write(&batch_b, &embs_b)
2055            .unwrap();
2056        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
2057
2058        let entries = vec![
2059            DataFileEntry {
2060                path: "data/a.parquet".into(),
2061                record_count: 2,
2062                file_size_bytes: bytes_a.len() as u64,
2063                centroid_b64: None,
2064                radius: None,
2065                hnsw_offset: None,
2066                hnsw_len: None,
2067                vector_column: None,
2068                vector_dim: None,
2069                extra_vector_indexes: vec![],
2070                index_status: IndexStatus::Ready,
2071                index_error: None,
2072                batch_id: None,
2073                embedding_model: None,
2074                partition_value: None,
2075                deletion_vector: None,
2076                first_row_id: None,
2077                column_stats: None,
2078                sequence_number: 0,
2079            },
2080            DataFileEntry {
2081                path: "data/b.parquet".into(),
2082                record_count: 2,
2083                file_size_bytes: bytes_b.len() as u64,
2084                centroid_b64: None,
2085                radius: None,
2086                hnsw_offset: None,
2087                hnsw_len: None,
2088                vector_column: None,
2089                vector_dim: None,
2090                extra_vector_indexes: vec![],
2091                index_status: IndexStatus::Ready,
2092                index_error: None,
2093                batch_id: None,
2094                embedding_model: None,
2095                partition_value: None,
2096                deletion_vector: None,
2097                first_row_id: None,
2098                column_stats: None,
2099                sequence_number: 0,
2100            },
2101        ];
2102
2103        let executor = CompactionExecutor::new(store.clone(), policy.clone());
2104        let entry = executor
2105            .compact_deferred(&entries, "data/merged.parquet", catalog.clone(), &table)
2106            .await
2107            .unwrap();
2108
2109        // Entry is Indexing — HNSW build pending in background
2110        assert_eq!(entry.index_status, IndexStatus::Indexing);
2111        assert_eq!(entry.record_count, 4);
2112
2113        // The written file must be valid Parquet (readable) even without HNSW
2114        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
2115        let pq_reader = ailake_parquet::ParquetVectorReader::new(merged_bytes, "embedding");
2116        let count = pq_reader.record_count().unwrap();
2117        assert_eq!(count, 4);
2118    }
2119
2120    /// Regression test (ADR-018 / CLAUDE.md Fase 5 "Idempotência batch_id sobrevivendo
2121    /// a compaction"): before `DataFileEntry::merge_batch_ids` existed, `compact()`
2122    /// always produced `batch_id: None` on the merged file — a retry of either source
2123    /// write, dispatched after compaction swept up its file, would find no existing
2124    /// entry carrying its key and silently re-insert. All three merge entry-points
2125    /// (`compact`, `compact_incremental`, `compact_deferred`) must aggregate source
2126    /// `batch_id`s; this covers the plain (non-deferred) `compact()` path.
2127    #[tokio::test]
2128    async fn compact_aggregates_batch_ids_from_sources() {
2129        use ailake_core::{VectorMetric, VectorPrecision};
2130        use ailake_store::LocalStore;
2131        use arrow_array::{Int32Array, RecordBatch};
2132        use arrow_schema::{DataType, Field, Schema};
2133        use std::sync::Arc;
2134        use tempfile::TempDir;
2135
2136        let dir = TempDir::new().unwrap();
2137        let store = Arc::new(LocalStore::new(dir.path()));
2138        let policy = VectorStoragePolicy {
2139            column_name: "embedding".into(),
2140            dim: 4,
2141            metric: VectorMetric::Cosine,
2142            precision: VectorPrecision::F16,
2143            pq: None,
2144            keep_raw_for_reranking: true,
2145            pre_normalize: false,
2146            hnsw_m: None,
2147            hnsw_ef_construction: None,
2148            ivf_residual: false,
2149            embedding_model: None,
2150            modality: None,
2151            partition_by: None,
2152            partition_value: None,
2153            partition_column_type: None,
2154            partition_fields: vec![],
2155        };
2156
2157        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2158        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0]];
2159        let batch_a =
2160            RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![0i32]))])
2161                .unwrap();
2162        let bytes_a = AilakeFileWriter::new(policy.clone())
2163            .write(&batch_a, &embs_a)
2164            .unwrap();
2165        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
2166
2167        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 1.0, 0.0, 0.0]];
2168        let batch_b =
2169            RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1i32]))])
2170                .unwrap();
2171        let bytes_b = AilakeFileWriter::new(policy.clone())
2172            .write(&batch_b, &embs_b)
2173            .unwrap();
2174        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
2175
2176        // A file compacted with no batch_id at all (e.g. written via plain
2177        // `write_batch`, never idempotently) — must not inject a spurious key.
2178        let embs_c: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
2179        let batch_c =
2180            RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![2i32]))]).unwrap();
2181        let bytes_c = AilakeFileWriter::new(policy.clone())
2182            .write(&batch_c, &embs_c)
2183            .unwrap();
2184        store.put("data/c.parquet", bytes_c.clone()).await.unwrap();
2185
2186        fn entry(path: &str, size: u64, batch_id: Option<&str>) -> DataFileEntry {
2187            DataFileEntry {
2188                path: path.into(),
2189                record_count: 1,
2190                file_size_bytes: size,
2191                centroid_b64: None,
2192                radius: None,
2193                hnsw_offset: None,
2194                hnsw_len: None,
2195                vector_column: None,
2196                vector_dim: None,
2197                extra_vector_indexes: vec![],
2198                index_status: IndexStatus::Ready,
2199                index_error: None,
2200                batch_id: batch_id.map(String::from),
2201                embedding_model: None,
2202                partition_value: None,
2203                deletion_vector: None,
2204                first_row_id: None,
2205                column_stats: None,
2206                sequence_number: 0,
2207            }
2208        }
2209        let entries = vec![
2210            entry("data/a.parquet", bytes_a.len() as u64, Some("k-a")),
2211            entry("data/b.parquet", bytes_b.len() as u64, Some("k-b")),
2212            entry("data/c.parquet", bytes_c.len() as u64, None),
2213        ];
2214
2215        let executor = CompactionExecutor::new(store.clone(), policy);
2216        let merged = executor
2217            .compact(&entries, "data/merged.parquet")
2218            .await
2219            .unwrap();
2220
2221        assert_eq!(merged.record_count, 3);
2222        assert_eq!(
2223            merged.batch_ids(),
2224            vec!["k-a".to_string(), "k-b".to_string()],
2225            "merged file must carry every source's idempotency key, none invented"
2226        );
2227    }
2228
2229    /// Regression test: `CompactionExecutor::run()` must not drop files that fall
2230    /// outside the compaction pass (too large for `target_file_size_bytes`, or beyond
2231    /// `max_files_per_pass`). `Replace` snapshots don't inherit the previous manifest
2232    /// (see `HadoopCatalog::commit_snapshot`), so `run()` must explicitly carry forward
2233    /// every untouched file alongside the merged output — mirroring the pattern already
2234    /// used by the CLI `compact` command and `MemoryDecayJob::run`.
2235    #[tokio::test]
2236    async fn run_preserves_untouched_files_outside_compaction_pass() {
2237        use ailake_catalog::HadoopCatalog;
2238        use ailake_core::{VectorMetric, VectorPrecision};
2239        use ailake_store::LocalStore;
2240        use arrow_array::{Int32Array, RecordBatch};
2241        use arrow_schema::{DataType, Field, Schema};
2242        use std::sync::Arc;
2243        use tempfile::TempDir;
2244
2245        let dir = TempDir::new().unwrap();
2246        let store = Arc::new(LocalStore::new(dir.path()));
2247        let catalog_dir = TempDir::new().unwrap();
2248        let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
2249        let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
2250        let table = TableIdent {
2251            namespace: "ns".into(),
2252            name: "tbl".into(),
2253        };
2254
2255        let policy = VectorStoragePolicy {
2256            column_name: "embedding".into(),
2257            dim: 4,
2258            metric: VectorMetric::Cosine,
2259            precision: VectorPrecision::F16,
2260            pq: None,
2261            keep_raw_for_reranking: true,
2262            pre_normalize: false,
2263            hnsw_m: None,
2264            hnsw_ef_construction: None,
2265            ivf_residual: false,
2266            embedding_model: None,
2267            modality: None,
2268            partition_by: None,
2269            partition_value: None,
2270            partition_column_type: None,
2271            partition_fields: vec![],
2272        };
2273
2274        use ailake_catalog::TableProperties;
2275        catalog
2276            .create_table(
2277                &table,
2278                &TableProperties {
2279                    policy: policy.clone(),
2280                    extra: std::collections::HashMap::new(),
2281                    format_version: 2,
2282                    partition_column_type: None,
2283                },
2284            )
2285            .await
2286            .unwrap();
2287
2288        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2289
2290        // Two small files — eligible for compaction.
2291        let write_file = |path: &str, ids: Vec<i32>, embs: Vec<Vec<f32>>| {
2292            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
2293                .unwrap();
2294            let bytes = AilakeFileWriter::new(policy.clone())
2295                .write(&batch, &embs)
2296                .unwrap();
2297            (path.to_string(), bytes)
2298        };
2299
2300        let (path_a, bytes_a) = write_file(
2301            "data/small_a.parquet",
2302            vec![0, 1],
2303            vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
2304        );
2305        let (path_b, bytes_b) = write_file(
2306            "data/small_b.parquet",
2307            vec![2, 3],
2308            vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]],
2309        );
2310        // "Big" file — same tiny payload in this test, but its DataFileEntry reports a
2311        // size above target_file_size_bytes so the planner must never select it.
2312        let (path_big, bytes_big) = write_file(
2313            "data/big.parquet",
2314            vec![4, 5],
2315            vec![vec![1.0, 1.0, 0.0, 0.0], vec![0.0, 1.0, 1.0, 0.0]],
2316        );
2317
2318        for (path, bytes) in [
2319            (&path_a, &bytes_a),
2320            (&path_b, &bytes_b),
2321            (&path_big, &bytes_big),
2322        ] {
2323            store.put(path, bytes.clone()).await.unwrap();
2324        }
2325
2326        let make_entry = |path: &str, size: u64| DataFileEntry {
2327            path: path.to_string(),
2328            record_count: 2,
2329            file_size_bytes: size,
2330            // Non-None: all three files here are meant to represent normal,
2331            // already-indexed AI-Lake files — only size should decide eligibility.
2332            // `plan_prioritizes_foreign_written_files_regardless_of_size` covers
2333            // the `centroid_b64: None` (foreign-write) case separately.
2334            centroid_b64: Some("AAAA".into()),
2335            radius: None,
2336            hnsw_offset: None,
2337            hnsw_len: None,
2338            vector_column: None,
2339            vector_dim: None,
2340            extra_vector_indexes: vec![],
2341            index_status: IndexStatus::Ready,
2342            index_error: None,
2343            batch_id: None,
2344            embedding_model: None,
2345            partition_value: None,
2346            deletion_vector: None,
2347            first_row_id: None,
2348            column_stats: None,
2349            sequence_number: 0,
2350        };
2351
2352        let initial_snap_id = ailake_catalog::new_snapshot_id();
2353        let initial_snapshot = NewSnapshot {
2354            snapshot_id: initial_snap_id,
2355            parent_snapshot_id: None,
2356            files: vec![
2357                make_entry(&path_a, 500),
2358                make_entry(&path_b, 500),
2359                make_entry(&path_big, 200_000_000), // far above target_file_size_bytes below
2360            ],
2361            operation: SnapshotOperation::Append,
2362            iceberg_schema: None,
2363            extra_properties: std::collections::HashMap::new(),
2364            bloom_filters: vec![],
2365            equality_delete_files: vec![],
2366        };
2367        catalog
2368            .commit_snapshot(&table, initial_snapshot)
2369            .await
2370            .unwrap();
2371
2372        let planner = CompactionPlanner::new(CompactionConfig {
2373            min_files_to_compact: 2,
2374            target_file_size_bytes: 1000,
2375            index_strategy: CompactionIndexStrategy::ForceHnsw,
2376            max_files_per_pass: 20,
2377        });
2378        let executor = CompactionExecutor::new(store.clone(), policy.clone());
2379
2380        let merged = executor
2381            .run(&planner, &table, catalog.clone(), "data")
2382            .await
2383            .unwrap()
2384            .expect("compaction should have run — 2 eligible small files");
2385
2386        let files_after = catalog.list_files(&table, None).await.unwrap();
2387        let paths_after: Vec<&str> = files_after.iter().map(|f| f.path.as_str()).collect();
2388
2389        assert!(
2390            paths_after.contains(&path_big.as_str()),
2391            "BUG: untouched 'big.parquet' vanished after run() — files_after={paths_after:?}"
2392        );
2393        assert!(
2394            paths_after.contains(&merged.path.as_str()),
2395            "merged output file must be present — files_after={paths_after:?}"
2396        );
2397
2398        // Regression: `run()` used to hardcode `parent_snapshot_id: None` on the
2399        // post-compaction Replace snapshot even though a current snapshot always exists
2400        // here, breaking Iceberg snapshot lineage (`expire_snapshots`/`rollback_to_snapshot`)
2401        // for any compacted table. Read the committed metadata.json directly (no public
2402        // CatalogProvider method exposes snapshot lineage) and confirm the new snapshot's
2403        // parent points at the snapshot committed before this compaction ran.
2404        let meta_dir = catalog_dir.path().join("ns/tbl/metadata");
2405        let latest_metadata = std::fs::read_dir(&meta_dir)
2406            .unwrap()
2407            .filter_map(|e| e.ok())
2408            .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
2409            .max_by_key(|e| e.metadata().unwrap().modified().unwrap())
2410            .expect("metadata.json must exist after commit");
2411        let json: serde_json::Value =
2412            serde_json::from_slice(&std::fs::read(latest_metadata.path()).unwrap()).unwrap();
2413        let last_snapshot = json["snapshots"].as_array().unwrap().last().unwrap();
2414        assert_eq!(
2415            last_snapshot["parent-snapshot-id"].as_i64(),
2416            Some(initial_snap_id),
2417            "compaction's Replace snapshot must chain to the pre-compaction snapshot, not be orphaned"
2418        );
2419        assert!(
2420            !paths_after.contains(&path_a.as_str()) && !paths_after.contains(&path_b.as_str()),
2421            "compacted input files must no longer be listed — files_after={paths_after:?}"
2422        );
2423        assert_eq!(
2424            files_after.len(),
2425            2,
2426            "expected exactly [big.parquet, merged] — files_after={paths_after:?}"
2427        );
2428    }
2429
2430    /// Regression test: a file with no AILK footer (e.g. rewritten by a generic
2431    /// Iceberg engine — Spark/Trino `OPTIMIZE` — with no knowledge of AI-Lake) still
2432    /// holds valid Parquet data. `read_parquet()` decodes the vector column directly
2433    /// from Parquet and never touches the footer, so `compact()`/`compact_incremental()`
2434    /// must include such a file's rows in the merge, not silently drop them.
2435    #[tokio::test]
2436    async fn compact_preserves_rows_from_footerless_file() {
2437        use ailake_core::{VectorMetric, VectorPrecision};
2438        use ailake_store::LocalStore;
2439        use arrow_array::{Int32Array, RecordBatch};
2440        use arrow_schema::{DataType, Field, Schema};
2441        use std::sync::Arc;
2442        use tempfile::TempDir;
2443
2444        let dir = TempDir::new().unwrap();
2445        let store = Arc::new(LocalStore::new(dir.path()));
2446        let policy = VectorStoragePolicy {
2447            column_name: "embedding".into(),
2448            dim: 4,
2449            metric: VectorMetric::Cosine,
2450            precision: VectorPrecision::F16,
2451            pq: None,
2452            keep_raw_for_reranking: true,
2453            pre_normalize: false,
2454            hnsw_m: None,
2455            hnsw_ef_construction: None,
2456            ivf_residual: false,
2457            embedding_model: None,
2458            modality: None,
2459            partition_by: None,
2460            partition_value: None,
2461            partition_column_type: None,
2462            partition_fields: vec![],
2463        };
2464
2465        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2466
2467        // Normal AI-Lake file, written the usual way (has an AILK footer).
2468        let embs_native: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
2469        let batch_native = RecordBatch::try_new(
2470            schema.clone(),
2471            vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2472        )
2473        .unwrap();
2474        let bytes_native = AilakeFileWriter::new(policy.clone())
2475            .write(&batch_native, &embs_native)
2476            .unwrap();
2477        store
2478            .put("data/native.parquet", bytes_native.clone())
2479            .await
2480            .unwrap();
2481
2482        // "Foreign" file — plain Parquet, no AILK footer, same shape a generic
2483        // Iceberg engine's rewrite would produce (`write_parquet_only` is the exact
2484        // primitive `compact_deferred` uses for its Parquet-only fast path).
2485        let embs_foreign: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
2486        let batch_foreign = RecordBatch::try_new(
2487            schema.clone(),
2488            vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2489        )
2490        .unwrap();
2491        let bytes_foreign = AilakeFileWriter::new(policy.clone())
2492            .write_parquet_only(&batch_foreign, &embs_foreign)
2493            .unwrap();
2494        store
2495            .put("data/foreign.parquet", bytes_foreign.clone())
2496            .await
2497            .unwrap();
2498
2499        let reader_foreign = AilakeFileReader::new(bytes_foreign.clone(), "embedding", 4);
2500        assert!(
2501            !reader_foreign.is_ailake_file(),
2502            "sanity: write_parquet_only must not embed an AILK footer"
2503        );
2504
2505        let entries = vec![
2506            DataFileEntry {
2507                path: "data/native.parquet".into(),
2508                record_count: 2,
2509                file_size_bytes: bytes_native.len() as u64,
2510                centroid_b64: None,
2511                radius: None,
2512                hnsw_offset: None,
2513                hnsw_len: None,
2514                vector_column: None,
2515                vector_dim: None,
2516                extra_vector_indexes: vec![],
2517                index_status: IndexStatus::Ready,
2518                index_error: None,
2519                batch_id: None,
2520                embedding_model: None,
2521                partition_value: None,
2522                deletion_vector: None,
2523                first_row_id: None,
2524                column_stats: None,
2525                sequence_number: 0,
2526            },
2527            DataFileEntry {
2528                path: "data/foreign.parquet".into(),
2529                record_count: 2,
2530                file_size_bytes: bytes_foreign.len() as u64,
2531                centroid_b64: None,
2532                radius: None,
2533                hnsw_offset: None,
2534                hnsw_len: None,
2535                vector_column: None,
2536                vector_dim: None,
2537                extra_vector_indexes: vec![],
2538                index_status: IndexStatus::Ready,
2539                index_error: None,
2540                batch_id: None,
2541                embedding_model: None,
2542                partition_value: None,
2543                deletion_vector: None,
2544                first_row_id: None,
2545                column_stats: None,
2546                sequence_number: 0,
2547            },
2548        ];
2549
2550        // Full rebuild path (`compact`).
2551        let executor = CompactionExecutor::new(store.clone(), policy.clone());
2552        let merged = executor
2553            .compact(&entries, "data/merged_full.parquet")
2554            .await
2555            .unwrap();
2556        assert_eq!(
2557            merged.record_count, 4,
2558            "compact() must include all 4 rows — 2 native + 2 from the footerless file"
2559        );
2560        let merged_bytes = store.get("data/merged_full.parquet").await.unwrap();
2561        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
2562        reader.verify_integrity().unwrap();
2563
2564        // Incremental path (`compact_incremental` — native file becomes dominant at 50/50,
2565        // exercising the non-dominant-file read path where the bug lived).
2566        let merged_inc = executor
2567            .compact_incremental(&entries, "data/merged_inc.parquet")
2568            .await
2569            .unwrap();
2570        assert_eq!(
2571            merged_inc.record_count, 4,
2572            "compact_incremental() must include all 4 rows — 2 native + 2 from the footerless file"
2573        );
2574    }
2575
2576    /// End-to-end: search returns the same top-K results before and after compaction.
2577    /// This tests the actual recall — not just verify_integrity().
2578    #[tokio::test]
2579    async fn compact_preserves_search_results() {
2580        use crate::scanner::SearchConfig;
2581        use ailake_catalog::HadoopCatalog;
2582        use ailake_core::{VectorMetric, VectorPrecision};
2583        use ailake_store::LocalStore;
2584        use arrow_array::{Int32Array, RecordBatch};
2585        use arrow_schema::{DataType, Field, Schema};
2586        use std::sync::Arc;
2587        use tempfile::TempDir;
2588        let dir = TempDir::new().unwrap();
2589        let store = Arc::new(LocalStore::new(dir.path()));
2590        let policy = VectorStoragePolicy {
2591            column_name: "embedding".into(),
2592            dim: 4,
2593            metric: VectorMetric::Cosine,
2594            precision: VectorPrecision::F16,
2595            pq: None,
2596            keep_raw_for_reranking: true,
2597            pre_normalize: false,
2598            hnsw_m: None,
2599            hnsw_ef_construction: None,
2600            ivf_residual: false,
2601            embedding_model: None,
2602            modality: None,
2603            partition_by: None,
2604            partition_value: None,
2605            partition_column_type: None,
2606            partition_fields: vec![],
2607        };
2608
2609        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2610
2611        // File A: 3 rows — unit basis vectors along dimensions 0, 1, 2
2612        let embs_a: Vec<Vec<f32>> = vec![
2613            vec![1.0, 0.0, 0.0, 0.0],
2614            vec![0.0, 1.0, 0.0, 0.0],
2615            vec![0.0, 0.0, 1.0, 0.0],
2616        ];
2617        let batch_a = RecordBatch::try_new(
2618            schema.clone(),
2619            vec![Arc::new(Int32Array::from(vec![0i32, 1, 2]))],
2620        )
2621        .unwrap();
2622
2623        // File B: 2 rows — unit basis vectors along dimensions 3 and mixed
2624        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 0.0, 1.0], vec![0.7, 0.7, 0.0, 0.0]];
2625        let batch_b = RecordBatch::try_new(
2626            schema.clone(),
2627            vec![Arc::new(Int32Array::from(vec![3i32, 4]))],
2628        )
2629        .unwrap();
2630
2631        let bytes_a = AilakeFileWriter::new(policy.clone())
2632            .write(&batch_a, &embs_a)
2633            .unwrap();
2634        let bytes_b = AilakeFileWriter::new(policy.clone())
2635            .write(&batch_b, &embs_b)
2636            .unwrap();
2637
2638        store.put("data/a.parquet", bytes_a).await.unwrap();
2639        store.put("data/b.parquet", bytes_b).await.unwrap();
2640        let file_size_a = store.file_size("data/a.parquet").await.unwrap();
2641        let file_size_b = store.file_size("data/b.parquet").await.unwrap();
2642
2643        // We need a CatalogProvider to create/snapshot the table.
2644        // Use the catalog_provider helper from ailake-query.
2645        let catalog: Arc<dyn CatalogProvider> =
2646            Arc::new(HadoopCatalog::new(store.clone(), "warehouse"));
2647        let table = TableIdent::new("default", "compact_search_test");
2648
2649        // Bootstrap the table via create_table, then commit the two files
2650        catalog
2651            .create_table(
2652                &table,
2653                &ailake_catalog::TableProperties {
2654                    policy: policy.clone(),
2655                    extra: std::collections::HashMap::new(),
2656                    format_version: 2,
2657                    partition_column_type: None,
2658                },
2659            )
2660            .await
2661            .unwrap();
2662
2663        // Load current iceberg schema so commit_snapshot can point at it
2664        let meta = catalog.load_table(&table).await.unwrap();
2665        let parent_snapshot_id = meta.current_snapshot_id;
2666
2667        let entry_a = DataFileEntry {
2668            path: "data/a.parquet".into(),
2669            record_count: 3,
2670            file_size_bytes: file_size_a,
2671            centroid_b64: None,
2672            radius: None,
2673            hnsw_offset: None,
2674            hnsw_len: None,
2675            vector_column: Some("embedding".into()),
2676            vector_dim: Some(4),
2677            extra_vector_indexes: vec![],
2678            index_status: IndexStatus::Ready,
2679            index_error: None,
2680            batch_id: None,
2681            embedding_model: None,
2682            partition_value: None,
2683            deletion_vector: None,
2684            first_row_id: None,
2685            column_stats: None,
2686            sequence_number: 0,
2687        };
2688        let entry_b = DataFileEntry {
2689            path: "data/b.parquet".into(),
2690            record_count: 2,
2691            file_size_bytes: file_size_b,
2692            centroid_b64: None,
2693            radius: None,
2694            hnsw_offset: None,
2695            hnsw_len: None,
2696            vector_column: Some("embedding".into()),
2697            vector_dim: Some(4),
2698            extra_vector_indexes: vec![],
2699            index_status: IndexStatus::Ready,
2700            index_error: None,
2701            batch_id: None,
2702            embedding_model: None,
2703            partition_value: None,
2704            deletion_vector: None,
2705            first_row_id: None,
2706            column_stats: None,
2707            sequence_number: 0,
2708        };
2709
2710        catalog
2711            .commit_snapshot(
2712                &table,
2713                NewSnapshot {
2714                    snapshot_id: 200,
2715                    parent_snapshot_id,
2716                    files: vec![entry_a, entry_b],
2717                    operation: SnapshotOperation::Append,
2718                    iceberg_schema: None,
2719                    extra_properties: std::collections::HashMap::new(),
2720                    bloom_filters: vec![],
2721                    equality_delete_files: vec![],
2722                },
2723            )
2724            .await
2725            .unwrap();
2726
2727        // ── Search before compaction ──
2728        let query_before = vec![1.0f32, 0.0, 0.0, 0.0];
2729        let config = SearchConfig {
2730            top_k: 5,
2731            ef_search: 100,
2732            pruning_threshold: f32::INFINITY,
2733            ..Default::default()
2734        };
2735        let before = crate::scanner::search(
2736            &table,
2737            &query_before,
2738            config.clone(),
2739            "embedding",
2740            4,
2741            catalog.clone(),
2742            store.clone(),
2743        )
2744        .await
2745        .unwrap();
2746        assert!(
2747            !before.is_empty(),
2748            "search before compaction must return results"
2749        );
2750        let before_dists: Vec<f32> = before.iter().map(|r| r.distance).collect();
2751
2752        // ── Compact the two files ──
2753        let executor = CompactionExecutor::new(store.clone(), policy.clone());
2754        let files_to_compact = catalog.list_files(&table, None).await.unwrap();
2755        let merged = executor
2756            .compact(&files_to_compact, "data/merged.parquet")
2757            .await
2758            .unwrap();
2759
2760        // Commit the merge: Replace old entries with the merged one
2761        let meta2 = catalog.load_table(&table).await.unwrap();
2762        let merged_entry = DataFileEntry {
2763            path: merged.path.clone(),
2764            record_count: merged.record_count,
2765            file_size_bytes: merged.file_size_bytes,
2766            centroid_b64: merged.centroid_b64,
2767            radius: merged.radius,
2768            hnsw_offset: merged.hnsw_offset,
2769            hnsw_len: merged.hnsw_len,
2770            vector_column: merged.vector_column,
2771            vector_dim: merged.vector_dim,
2772            extra_vector_indexes: merged.extra_vector_indexes,
2773            index_status: merged.index_status,
2774            index_error: merged.index_error,
2775            batch_id: merged.batch_id,
2776            embedding_model: merged.embedding_model,
2777            partition_value: merged.partition_value,
2778            deletion_vector: merged.deletion_vector,
2779            first_row_id: None,
2780            column_stats: None,
2781            sequence_number: 0,
2782        };
2783        catalog
2784            .commit_snapshot(
2785                &table,
2786                NewSnapshot {
2787                    snapshot_id: 201,
2788                    parent_snapshot_id: meta2.current_snapshot_id,
2789                    files: vec![merged_entry],
2790                    operation: SnapshotOperation::Replace,
2791                    iceberg_schema: None,
2792                    extra_properties: std::collections::HashMap::new(),
2793                    bloom_filters: vec![],
2794                    equality_delete_files: vec![],
2795                },
2796            )
2797            .await
2798            .unwrap();
2799
2800        // ── Search after compaction ──
2801        let after = crate::scanner::search(
2802            &table,
2803            &query_before,
2804            config.clone(),
2805            "embedding",
2806            4,
2807            catalog,
2808            store,
2809        )
2810        .await
2811        .unwrap();
2812        assert!(
2813            !after.is_empty(),
2814            "search after compaction must return results"
2815        );
2816        let after_dists: Vec<f32> = after.iter().map(|r| r.distance).collect();
2817
2818        // ── Assert recall parity ──
2819        assert_eq!(
2820            before.len(),
2821            after.len(),
2822            "number of results should match before ({}) and after ({}) compaction",
2823            before.len(),
2824            after.len(),
2825        );
2826        // Row-ids are NOT expected to match — compaction reassigns row positions.
2827        // What matters: distances should be approximately the same for each rank,
2828        // meaning the same items are returned with similar relevance scores.
2829        // Mean Average Distance Error (MADE) should be small.
2830        let mut total_diff = 0.0f32;
2831        for (i, (bd, ad)) in before_dists.iter().zip(after_dists.iter()).enumerate() {
2832            let diff = (bd - ad).abs();
2833            total_diff += diff;
2834            assert!(
2835                diff < 0.1,
2836                "distance mismatch at position {i}: before={bd}, after={ad}, diff={diff}"
2837            );
2838        }
2839        let avg_diff = total_diff / before_dists.len() as f32;
2840        assert!(
2841            avg_diff < 0.05,
2842            "average distance error across all results too high: {avg_diff}"
2843        );
2844    }
2845}