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        self.store.put(output_path, file_bytes.clone()).await?;
319
320        // Compute centroid and HNSW offsets for catalog entry
321        let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
322        let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
323        let header = reader.read_header()?;
324        let ailk_start = reader.ailk_offset()?;
325
326        // Positional invariant check: parquet_count == hnsw_node_count == header.record_count.
327        // Catches a mismatched merge before it's committed to the catalog rather than
328        // surfacing as a wrong search result (or the invariant-violated error path in
329        // scanner.rs) later.
330        reader.verify_integrity()?;
331
332        // Preserve row-ID continuity: merged file inherits the minimum first_row_id of
333        // its sources so commit_snapshot doesn't allocate fresh IDs and grow next_row_id.
334        let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
335
336        let mut entry = make_data_file_entry(
337            output_path,
338            record_count,
339            file_size,
340            &centroid,
341            VectorIndexInfo {
342                column: &self.policy.column_name,
343                dim: self.policy.dim,
344                hnsw_offset: ailk_start + header.hnsw_offset,
345                hnsw_len: header.hnsw_len,
346            },
347        );
348        entry.first_row_id = source_first_row_id;
349        Ok(entry)
350    }
351
352    /// Merge `files` into a single new file using incremental HNSW insertion.
353    ///
354    /// Identifies the **dominant file** — the file holding >= 40 % of the total
355    /// row count — loads its existing HNSW graph from the AILK section, then
356    /// calls `HnswIndex::insert_node` for every vector from the remaining files.
357    ///
358    /// **Complexity vs `compact`**:
359    /// - Full rebuild: O(N log N), N = total rows.
360    /// - Incremental (this method): O(N_dom) deserialization + O(N_small × log N_dom).
361    ///   For a 90 / 10 split (N = 1 M, N_dom = 900 k) the speedup is ~7×.
362    ///
363    /// **Fallbacks** (all degrade gracefully to `compact`):
364    /// - No file holds >= 40 % of rows.
365    /// - Dominant file's HNSW cannot be loaded (IVF-PQ, `IndexStatus::Indexing`, corrupt).
366    ///
367    /// **RowId contract**: dominant file's vectors are placed first in the merged
368    /// Parquet (positions 0..N_dom-1); other files follow. The existing RowIds from
369    /// the dominant HNSW remain valid; new nodes receive RowIds N_dom..N-1.
370    pub async fn compact_incremental(
371        &self,
372        files: &[DataFileEntry],
373        output_path: &str,
374    ) -> AilakeResult<DataFileEntry> {
375        const DOMINANT_RATIO: f64 = 0.40;
376
377        if files.is_empty() {
378            return Err(ailake_core::AilakeError::Catalog(
379                "compact_incremental: no files provided".into(),
380            ));
381        }
382
383        // This method only ever produces an HNSW index (dominant-file graph extended via
384        // `insert_node`) — it never builds IVF-PQ. `ForceIvfPq` is an explicit caller
385        // request for a specific index type, so it must never be silently satisfied with
386        // HNSW instead; fall back to `compact()`, which does respect it. `Auto`/`ForceHnsw`
387        // are unaffected: `Auto` treats "the dominant file already has a valid HNSW"
388        // as a legitimate reason to keep reusing HNSW rather than pay for a fresh
389        // hardware-picked rebuild, and `ForceHnsw` is satisfied by this method by
390        // construction either way.
391        if matches!(self.index_strategy, CompactionIndexStrategy::ForceIvfPq) {
392            debug!(
393                "ailake: compact_incremental — index_strategy=ForceIvfPq, which this method \
394                 can never produce (HNSW graph extension only); falling back to full rebuild"
395            );
396            return self.compact(files, output_path).await;
397        }
398
399        // Find the dominant file by record_count.
400        let total_rows: u64 = files.iter().map(|f| f.record_count).sum();
401        let dom_idx = files
402            .iter()
403            .enumerate()
404            .max_by_key(|(_, f)| f.record_count)
405            .map(|(i, _)| i)
406            .unwrap_or(0);
407        let dom_rows = files[dom_idx].record_count;
408
409        // The dominant file's existing HNSW graph is reused as-is (only non-dominant
410        // vectors are inserted into it) — its node IDs are tied to the dominant file's
411        // current row positions. If it has DV-masked rows, filtering them out of the
412        // Parquet data would desync the reused graph's node-to-row mapping (and HNSW
413        // doesn't support cheap node removal). Fall back to a full rebuild, which
414        // filters correctly via `read_files_parallel`.
415        if files[dom_idx].deletion_vector.is_some() {
416            debug!(
417                "ailake: compact_incremental — dominant file {} has DV-masked rows, \
418                 falling back to full rebuild (graph reuse would desync row positions)",
419                files[dom_idx].path
420            );
421            return self.compact(files, output_path).await;
422        }
423
424        if (dom_rows as f64 / total_rows as f64) < DOMINANT_RATIO {
425            debug!(
426                "ailake: compact_incremental — no dominant file ({}/{} rows < {:.0}% threshold), \
427                 falling back to full rebuild",
428                dom_rows,
429                total_rows,
430                DOMINANT_RATIO * 100.0
431            );
432            return self.compact(files, output_path).await;
433        }
434
435        let column = self.policy.column_name.clone();
436        let dim = self.policy.dim;
437        let dom_path = files[dom_idx].path.clone();
438
439        // Read all files in parallel. `read_parquet()` decodes the vector column
440        // directly from Parquet and never touches the AILK footer, so every file's
441        // rows are read regardless of whether it has an AI-Lake index — a file
442        // missing its footer (external-engine rewrite, or still `IndexStatus::Indexing`)
443        // still holds real data and must not be dropped from the merge. Raw bytes are
444        // retained only for the dominant file, and only when it actually has an index
445        // to reuse (needed to load its HNSW without a second round-trip); otherwise the
446        // dominant-file match below falls back to a full rebuild via `compact()`.
447        let futs: Vec<_> =
448            files
449                .iter()
450                .map(|entry| {
451                    let store = self.store.clone();
452                    let path = entry.path.clone();
453                    let col = column.clone();
454                    let is_dom = path == dom_path;
455                    let dv = entry.deletion_vector.clone();
456                    async move {
457                        let bytes: Bytes = store.get(&path).await?;
458                        let reader = AilakeFileReader::new(bytes.clone(), &col, dim);
459                        let has_index = reader.is_ailake_file();
460                        if is_dom && !has_index {
461                            debug!(
462                                "ailake: compact_incremental — dominant candidate {} has no \
463                             AI-Lake index; will fall back to full rebuild if no HNSW to reuse",
464                                path
465                            );
466                        }
467                        let (raw_batch, raw_vecs) = reader.read_parquet()?;
468                        // Non-dominant rows are freshly inserted into the graph below, so
469                        // DV-masked rows must be dropped here (dominant is guaranteed
470                        // DV-free by the caller's earlier fallback check).
471                        let (batch, vecs) = if let Some(dv) = dv {
472                            let bitmap = crate::dv::load_deletion_vector(&store, &dv).await?;
473                            crate::dv::filter_deleted_rows(raw_batch, raw_vecs, &bitmap)?
474                        } else {
475                            (raw_batch, raw_vecs)
476                        };
477                        let retained = if is_dom && has_index {
478                            Some(bytes)
479                        } else {
480                            None
481                        };
482                        Ok::<
483                            (RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>),
484                            ailake_core::AilakeError,
485                        >((batch, vecs, is_dom, retained))
486                    }
487                })
488                .collect();
489
490        // Every future above always resolves to one tuple per input file (or propagates an
491        // Err via `?`) — no per-file filtering, so no emptiness check is needed given the
492        // `files.is_empty()` guard at the top of this function.
493        #[allow(clippy::type_complexity)]
494        let raw: Vec<(RecordBatch, Vec<Vec<f32>>, bool, Option<Bytes>)> =
495            try_join_all(futs).await?;
496
497        // Separate dominant from others; dominant goes first in the merged file.
498        let mut dom_batch: Option<RecordBatch> = None;
499        let mut dom_vecs: Vec<Vec<f32>> = Vec::new();
500        let mut dom_bytes_found: Option<Bytes> = None;
501        let mut other_batches: Vec<RecordBatch> = Vec::new();
502        let mut other_vecs: Vec<Vec<f32>> = Vec::new();
503
504        for (batch, vecs, is_dom, retained) in raw {
505            if is_dom {
506                dom_batch = Some(batch);
507                dom_vecs = vecs;
508                dom_bytes_found = retained;
509            } else {
510                other_batches.push(batch);
511                other_vecs.extend(vecs);
512            }
513        }
514
515        let (dom_batch, dom_bytes) = match (dom_batch, dom_bytes_found) {
516            (Some(b), Some(byt)) => (b, byt),
517            _ => {
518                debug!(
519                    "ailake: compact_incremental — dominant file missing from read results, \
520                     falling back to full rebuild"
521                );
522                return self.compact(files, output_path).await;
523            }
524        };
525
526        // Load the dominant file's existing HNSW graph.
527        let dom_reader = AilakeFileReader::new(dom_bytes, &column, dim);
528        let mut hnsw = match dom_reader.load_index() {
529            Ok(idx) => idx,
530            Err(e) => {
531                debug!(
532                    "ailake: compact_incremental — cannot load dominant HNSW ({}), \
533                     falling back to full rebuild",
534                    e
535                );
536                return self.compact(files, output_path).await;
537            }
538        };
539
540        let dom_count = dom_batch.num_rows() as u64;
541
542        // Insert vectors from non-dominant files into the loaded graph.
543        // RowIds are assigned starting at dom_count to match positions in the merged Parquet.
544        for (j, vec) in other_vecs.iter().enumerate() {
545            hnsw.insert_node(RowId::new(dom_count + j as u64), vec.clone());
546        }
547        hnsw.quantize_to_f16();
548
549        // Assemble merged batch (dominant rows first) and all embeddings.
550        let schema: SchemaRef = dom_batch.schema();
551        let mut all_batches = vec![dom_batch];
552        all_batches.extend(other_batches);
553        let merged_batch = concat_batches(schema, &all_batches)?;
554        let record_count = merged_batch.num_rows() as u64;
555
556        let mut all_embeddings = dom_vecs;
557        all_embeddings.extend(other_vecs);
558
559        // Write the merged file using the pre-built index (no rebuild).
560        // Attach FTS blob when configured — data is already in merged_batch so cost is tokenization only.
561        let writer = {
562            let base = AilakeFileWriter::new(self.policy.clone());
563            if let Some(ref fts_cfg) = self.fts_config {
564                match ailake_fts::merge_fts_blobs(fts_cfg, &merged_batch) {
565                    Ok(blob) => base.with_prebuilt_fts_blob(blob),
566                    Err(e) => {
567                        warn!("ailake: FTS re-index during incremental compaction failed: {e}");
568                        base
569                    }
570                }
571            } else {
572                base
573            }
574        };
575        let file_bytes = writer.write_with_prebuilt_hnsw(&merged_batch, &all_embeddings, &hnsw)?;
576        let file_size = file_bytes.len() as u64;
577        self.store.put(output_path, file_bytes.clone()).await?;
578
579        let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
580        let reader = AilakeFileReader::new(file_bytes, &self.policy.column_name, self.policy.dim);
581        let header = reader.read_header()?;
582        let ailk_start = reader.ailk_offset()?;
583
584        // Positional invariant check — see `compact()` for rationale. Especially relevant
585        // here since the index is grown incrementally (insert_node) rather than rebuilt.
586        reader.verify_integrity()?;
587
588        // Dominant file goes first in the merged output, so the merged file's first
589        // logical row was the dominant file's first row.  Use its first_row_id so
590        // commit_snapshot doesn't grow next_row_id unnecessarily.
591        let source_first_row_id = files[dom_idx].first_row_id;
592
593        let mut entry = make_data_file_entry(
594            output_path,
595            record_count,
596            file_size,
597            &centroid,
598            VectorIndexInfo {
599                column: &self.policy.column_name,
600                dim: self.policy.dim,
601                hnsw_offset: ailk_start + header.hnsw_offset,
602                hnsw_len: header.hnsw_len,
603            },
604        );
605        entry.first_row_id = source_first_row_id;
606
607        info!(
608            "ailake: compact_incremental — merged {} files into {} \
609             ({} rows from dominant + {} inserted incrementally)",
610            files.len(),
611            output_path,
612            dom_count,
613            record_count - dom_count
614        );
615
616        Ok(entry)
617    }
618
619    /// Merge `files` into a single new file at `output_path`, writing Parquet
620    /// immediately and building the HNSW / IVF-PQ index in a background Tokio task.
621    ///
622    /// The merged file appears in the catalog as `IndexStatus::Indexing` until
623    /// the background task completes; queries fall back to flat scan during that
624    /// window (same behaviour as `write_batch_deferred`).
625    ///
626    /// Returns the `DataFileEntry` with `IndexStatus::Indexing`. The entry
627    /// transitions to `Ready` automatically when the background build finishes.
628    pub async fn compact_deferred(
629        &self,
630        files: &[DataFileEntry],
631        output_path: &str,
632        catalog: Arc<dyn CatalogProvider>,
633        table: &TableIdent,
634    ) -> AilakeResult<DataFileEntry> {
635        if files.is_empty() {
636            return Err(ailake_core::AilakeError::Catalog(
637                "compact_deferred: no files provided".into(),
638            ));
639        }
640
641        // See `compact()` above: read_files_parallel never filters, so no emptiness
642        // check is needed given the `files.is_empty()` guard above.
643        let pairs = self.read_files_parallel(files).await?;
644
645        let schema: SchemaRef = pairs[0].0.schema();
646        let (all_batches, all_embeddings): (Vec<_>, Vec<_>) = pairs.into_iter().unzip();
647        let all_embeddings: Vec<Vec<f32>> = all_embeddings.into_iter().flatten().collect();
648
649        let merged_batch = concat_batches(schema, &all_batches)?;
650        let record_count = merged_batch.num_rows() as u64;
651
652        // Write Parquet-only immediately — fast path, no HNSW build.
653        let file_writer = AilakeFileWriter::new(self.policy.clone());
654        let parquet_bytes = file_writer.write_parquet_only(&merged_batch, &all_embeddings)?;
655        let file_size = parquet_bytes.len() as u64;
656        self.store.put(output_path, parquet_bytes).await?;
657
658        // Centroid available for geometric pruning during the build window.
659        let centroid = compute_centroid_and_radius(&all_embeddings, self.policy.metric);
660        let source_first_row_id = files.iter().filter_map(|f| f.first_row_id).min();
661        let mut entry = make_data_file_entry_indexing(
662            output_path,
663            record_count,
664            file_size,
665            &centroid,
666            &self.policy.column_name,
667            self.policy.dim,
668        );
669        entry.first_row_id = source_first_row_id;
670
671        // Spawn background index build; errors are logged, not propagated.
672        let store = self.store.clone();
673        let policy = self.policy.clone();
674        let table_id = table.clone();
675        let fp = output_path.to_string();
676        tokio::spawn(async move {
677            if let Err(e) = build_and_patch_index(store, catalog, policy, table_id, fp).await {
678                error!(
679                    "ailake: compaction deferred HNSW build failed — file indexed as \
680                     Parquet-only until next compaction rebuilds the index: {}",
681                    e
682                );
683            }
684        });
685
686        Ok(entry)
687    }
688
689    /// Full compaction workflow: plan, compact (synchronous HNSW rebuild),
690    /// drop old files from catalog, commit.
691    pub async fn run(
692        &self,
693        planner: &CompactionPlanner,
694        table: &TableIdent,
695        catalog: Arc<dyn CatalogProvider>,
696        output_prefix: &str,
697    ) -> AilakeResult<Option<DataFileEntry>> {
698        let all_files = catalog.list_files(table, None).await?;
699        let to_compact = planner.plan(&all_files);
700        if to_compact.is_empty() {
701            return Ok(None);
702        }
703
704        // Auto-detect FTS from table metadata so compaction never silently drops an FTS index
705        // that was present in the source files. Uses ailake.fts.* properties written at write time.
706        let meta_props = catalog
707            .load_table(table)
708            .await
709            .map(|m| m.properties)
710            .unwrap_or_default();
711        let executor = self.with_effective_fts(&meta_props);
712
713        let ts = std::time::SystemTime::now()
714            .duration_since(std::time::UNIX_EPOCH)
715            .unwrap_or_else(|e| e.duration())
716            .as_millis();
717        let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
718
719        // Use incremental merge when a dominant file exists (falls back to full rebuild automatically).
720        let merged = executor
721            .compact_incremental(&to_compact, &output_path)
722            .await?;
723
724        // Commit: add merged file, remove input files (via Replace snapshot).
725        let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
726        // Fetched fresh, right before the commit — same freshness rationale as the
727        // file-list re-list in `build_replace_file_list` above.
728        let parent_snapshot_id = catalog
729            .load_table(table)
730            .await
731            .ok()
732            .and_then(|m| m.current_snapshot_id);
733        let snapshot = NewSnapshot {
734            snapshot_id: ailake_catalog::new_snapshot_id(),
735            parent_snapshot_id,
736            files,
737            operation: SnapshotOperation::Replace,
738            iceberg_schema: None,
739            extra_properties: std::collections::HashMap::new(),
740            bloom_filters: vec![],
741            equality_delete_files: vec![],
742        };
743        catalog.commit_snapshot(table, snapshot).await?;
744
745        info!(
746            "ailake: compaction committed — merged {} files into {}",
747            to_compact.len(),
748            output_path
749        );
750
751        if catalog.retires_files_physically() {
752            delete_old_files(&self.store, &to_compact).await;
753        } else {
754            info!(
755                "ailake: compaction — leaving {} retired file(s) in place; catalog backend \
756                 manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
757                to_compact.len()
758            );
759        }
760
761        Ok(Some(merged))
762    }
763
764    /// Full compaction workflow with deferred HNSW build: plan, write merged
765    /// Parquet immediately, commit as `Indexing`, spawn background index build.
766    ///
767    /// Use for large tables where inline HNSW rebuild blocks too long.
768    ///
769    /// Note: FTS index is **not** rebuilt in deferred mode — `compact_deferred` writes
770    /// Parquet-only immediately and the background task (`build_and_patch_index`) only
771    /// builds the HNSW/IVF-PQ index. Use `run` (synchronous) when FTS preservation
772    /// on compaction is required.
773    pub async fn run_deferred(
774        &self,
775        planner: &CompactionPlanner,
776        table: &TableIdent,
777        catalog: Arc<dyn CatalogProvider>,
778        output_prefix: &str,
779    ) -> AilakeResult<Option<DataFileEntry>> {
780        // The background index build patches the merged file in place at its
781        // committed path — refuse on backends where committed bytes are
782        // immutable (see TableWriter::ensure_deferred_supported for why a
783        // commit-time guard alone is too late: the physical put comes first).
784        if !catalog.supports_in_place_rewrite() {
785            return Err(ailake_core::AilakeError::Catalog(
786                "deferred compaction is not supported with this catalog backend: the \
787                 background index build patches the merged file in place at its committed \
788                 path, which this catalog cannot re-register — run a blocking compact"
789                    .into(),
790            ));
791        }
792        let all_files = catalog.list_files(table, None).await?;
793        let to_compact = planner.plan(&all_files);
794        if to_compact.is_empty() {
795            return Ok(None);
796        }
797
798        let ts = std::time::SystemTime::now()
799            .duration_since(std::time::UNIX_EPOCH)
800            .unwrap_or_else(|e| e.duration())
801            .as_millis();
802        let output_path = format!("{output_prefix}/compacted-{ts}.parquet");
803
804        let merged = self
805            .compact_deferred(&to_compact, &output_path, catalog.clone(), table)
806            .await?;
807
808        // Commit immediately: merged file in Indexing state replaces input files.
809        let files = build_replace_file_list(&catalog, table, &to_compact, merged.clone()).await?;
810        // Fetched fresh, right before the commit — same freshness rationale as the
811        // file-list re-list in `build_replace_file_list` above.
812        let parent_snapshot_id = catalog
813            .load_table(table)
814            .await
815            .ok()
816            .and_then(|m| m.current_snapshot_id);
817        let snapshot = NewSnapshot {
818            snapshot_id: ailake_catalog::new_snapshot_id(),
819            parent_snapshot_id,
820            files,
821            operation: SnapshotOperation::Replace,
822            iceberg_schema: None,
823            extra_properties: std::collections::HashMap::new(),
824            bloom_filters: vec![],
825            equality_delete_files: vec![],
826        };
827        catalog.commit_snapshot(table, snapshot).await?;
828
829        info!(
830            "ailake: compaction committed (deferred) — merged {} files into {} \
831             (index building in background)",
832            to_compact.len(),
833            output_path
834        );
835
836        if catalog.retires_files_physically() {
837            delete_old_files(&self.store, &to_compact).await;
838        } else {
839            info!(
840                "ailake: compaction — leaving {} retired file(s) in place; catalog backend \
841                 manages physical reclamation itself (see docs/guides/DUCKLAKE_CATALOG.md)",
842                to_compact.len()
843            );
844        }
845
846        Ok(Some(merged))
847    }
848}
849
850/// Builds the file list for a post-compaction `Replace` snapshot: the merged output
851/// plus every current file that wasn't part of this compaction pass.
852///
853/// `Replace` does not inherit the previous manifest (see `HadoopCatalog::commit_snapshot`),
854/// so untouched files (above `target_file_size_bytes`, or beyond `max_files_per_pass`) must
855/// be carried forward explicitly or they vanish from the table. Re-lists via `list_files()`
856/// right before this call returns (rather than reusing a pre-merge snapshot) to narrow —
857/// though not eliminate, `HadoopCatalog` has no optimistic-concurrency check on
858/// `commit_snapshot` — the window for a concurrent writer's commit to be silently dropped
859/// by the `Replace` this list feeds into.
860async fn build_replace_file_list(
861    catalog: &Arc<dyn CatalogProvider>,
862    table: &TableIdent,
863    to_compact: &[DataFileEntry],
864    merged: DataFileEntry,
865) -> AilakeResult<Vec<DataFileEntry>> {
866    let current_files = catalog.list_files(table, None).await?;
867    let compacted_paths: std::collections::HashSet<&str> =
868        to_compact.iter().map(|f| f.path.as_str()).collect();
869    let mut files: Vec<DataFileEntry> = current_files
870        .into_iter()
871        .filter(|f| !compacted_paths.contains(f.path.as_str()))
872        .collect();
873    files.push(merged);
874    Ok(files)
875}
876
877async fn delete_old_files(store: &Arc<dyn Store>, files: &[DataFileEntry]) {
878    for entry in files {
879        if let Err(e) = store.delete(&entry.path).await {
880            error!(
881                "ailake: compaction cleanup failed — could not delete {}: {} \
882                 (orphan file in object store after successful catalog commit; \
883                 delete manually to reclaim storage)",
884                entry.path, e
885            );
886        }
887    }
888}
889
890fn concat_batches(schema: SchemaRef, batches: &[RecordBatch]) -> AilakeResult<RecordBatch> {
891    arrow_select::concat::concat_batches(&schema, batches)
892        .map_err(|e| ailake_core::AilakeError::Arrow(e.to_string()))
893}
894
895#[cfg(test)]
896mod tests {
897    use super::*;
898    use ailake_catalog::IndexStatus;
899
900    #[test]
901    fn plan_returns_empty_if_too_few_files() {
902        let planner = CompactionPlanner::new(CompactionConfig {
903            min_files_to_compact: 4,
904            target_file_size_bytes: 1024 * 1024,
905            ..Default::default()
906        });
907        let files: Vec<DataFileEntry> = (0..3)
908            .map(|i| DataFileEntry {
909                path: format!("file-{i}.parquet"),
910                record_count: 10,
911                file_size_bytes: 100,
912                centroid_b64: Some("AAAA".into()),
913                radius: None,
914                hnsw_offset: None,
915                hnsw_len: None,
916                vector_column: None,
917                vector_dim: None,
918                extra_vector_indexes: vec![],
919                index_status: IndexStatus::Ready,
920                index_error: None,
921                batch_id: None,
922                embedding_model: None,
923                partition_value: None,
924                deletion_vector: None,
925                first_row_id: None,
926            })
927            .collect();
928        assert!(planner.plan(&files).is_empty());
929    }
930
931    #[test]
932    fn plan_selects_small_files() {
933        let planner = CompactionPlanner::new(CompactionConfig {
934            min_files_to_compact: 2,
935            target_file_size_bytes: 1000,
936            ..Default::default()
937        });
938        let files = vec![
939            DataFileEntry {
940                path: "small.parquet".into(),
941                record_count: 5,
942                file_size_bytes: 500,
943                centroid_b64: Some("AAAA".into()),
944                radius: None,
945                hnsw_offset: None,
946                hnsw_len: None,
947                vector_column: None,
948                vector_dim: None,
949                extra_vector_indexes: vec![],
950                index_status: IndexStatus::Ready,
951                index_error: None,
952                batch_id: None,
953                embedding_model: None,
954                partition_value: None,
955                deletion_vector: None,
956                first_row_id: None,
957            },
958            DataFileEntry {
959                path: "large.parquet".into(),
960                record_count: 5000,
961                file_size_bytes: 200_000_000,
962                centroid_b64: Some("AAAA".into()),
963                radius: None,
964                hnsw_offset: None,
965                hnsw_len: None,
966                vector_column: None,
967                vector_dim: None,
968                extra_vector_indexes: vec![],
969                index_status: IndexStatus::Ready,
970                index_error: None,
971                batch_id: None,
972                embedding_model: None,
973                partition_value: None,
974                deletion_vector: None,
975                first_row_id: None,
976            },
977            DataFileEntry {
978                path: "also-small.parquet".into(),
979                record_count: 5,
980                file_size_bytes: 800,
981                centroid_b64: Some("AAAA".into()),
982                radius: None,
983                hnsw_offset: None,
984                hnsw_len: None,
985                vector_column: None,
986                vector_dim: None,
987                extra_vector_indexes: vec![],
988                index_status: IndexStatus::Ready,
989                index_error: None,
990                batch_id: None,
991                embedding_model: None,
992                partition_value: None,
993                deletion_vector: None,
994                first_row_id: None,
995            },
996        ];
997        let selected = planner.plan(&files);
998        assert_eq!(selected.len(), 2);
999        assert!(selected.iter().any(|f| f.path == "small.parquet"));
1000        assert!(selected.iter().any(|f| f.path == "also-small.parquet"));
1001    }
1002
1003    #[test]
1004    fn plan_respects_max_files_per_pass() {
1005        let planner = CompactionPlanner::new(CompactionConfig {
1006            min_files_to_compact: 2,
1007            target_file_size_bytes: 1_000_000,
1008            max_files_per_pass: 3,
1009            ..Default::default()
1010        });
1011        let files: Vec<DataFileEntry> = (0..5)
1012            .map(|i| DataFileEntry {
1013                path: format!("f{i}.parquet"),
1014                record_count: 10,
1015                file_size_bytes: 100 + i as u64 * 100,
1016                centroid_b64: Some("AAAA".into()),
1017                radius: None,
1018                hnsw_offset: None,
1019                hnsw_len: None,
1020                vector_column: None,
1021                vector_dim: None,
1022                extra_vector_indexes: vec![],
1023                index_status: IndexStatus::Ready,
1024                index_error: None,
1025                batch_id: None,
1026                embedding_model: None,
1027                partition_value: None,
1028                deletion_vector: None,
1029                first_row_id: None,
1030            })
1031            .collect();
1032        let selected = planner.plan(&files);
1033        assert_eq!(selected.len(), 3);
1034        assert_eq!(selected[0].file_size_bytes, 100);
1035        assert_eq!(selected[1].file_size_bytes, 200);
1036        assert_eq!(selected[2].file_size_bytes, 300);
1037    }
1038
1039    #[test]
1040    fn plan_sorts_smallest_first() {
1041        let planner = CompactionPlanner::new(CompactionConfig {
1042            min_files_to_compact: 2,
1043            target_file_size_bytes: 10_000,
1044            max_files_per_pass: 4,
1045            ..Default::default()
1046        });
1047        let files = vec![
1048            DataFileEntry {
1049                path: "c.parquet".into(),
1050                record_count: 1,
1051                file_size_bytes: 300,
1052                centroid_b64: Some("AAAA".into()),
1053                radius: None,
1054                hnsw_offset: None,
1055                hnsw_len: None,
1056                vector_column: None,
1057                vector_dim: None,
1058                extra_vector_indexes: vec![],
1059                index_status: IndexStatus::Ready,
1060                index_error: None,
1061                batch_id: None,
1062                embedding_model: None,
1063                partition_value: None,
1064                deletion_vector: None,
1065                first_row_id: None,
1066            },
1067            DataFileEntry {
1068                path: "a.parquet".into(),
1069                record_count: 1,
1070                file_size_bytes: 100,
1071                centroid_b64: Some("AAAA".into()),
1072                radius: None,
1073                hnsw_offset: None,
1074                hnsw_len: None,
1075                vector_column: None,
1076                vector_dim: None,
1077                extra_vector_indexes: vec![],
1078                index_status: IndexStatus::Ready,
1079                index_error: None,
1080                batch_id: None,
1081                embedding_model: None,
1082                partition_value: None,
1083                deletion_vector: None,
1084                first_row_id: None,
1085            },
1086            DataFileEntry {
1087                path: "b.parquet".into(),
1088                record_count: 1,
1089                file_size_bytes: 200,
1090                centroid_b64: Some("AAAA".into()),
1091                radius: None,
1092                hnsw_offset: None,
1093                hnsw_len: None,
1094                vector_column: None,
1095                vector_dim: None,
1096                extra_vector_indexes: vec![],
1097                index_status: IndexStatus::Ready,
1098                index_error: None,
1099                batch_id: None,
1100                embedding_model: None,
1101                partition_value: None,
1102                deletion_vector: None,
1103                first_row_id: None,
1104            },
1105        ];
1106        let selected = planner.plan(&files);
1107        assert_eq!(selected[0].file_size_bytes, 100);
1108        assert_eq!(selected[1].file_size_bytes, 200);
1109        assert_eq!(selected[2].file_size_bytes, 300);
1110    }
1111
1112    fn make_plan_entry(path: &str, size: u64, centroid_b64: Option<String>) -> DataFileEntry {
1113        DataFileEntry {
1114            path: path.to_string(),
1115            record_count: 10,
1116            file_size_bytes: size,
1117            centroid_b64,
1118            radius: None,
1119            hnsw_offset: None,
1120            hnsw_len: None,
1121            vector_column: None,
1122            vector_dim: None,
1123            extra_vector_indexes: vec![],
1124            index_status: IndexStatus::Ready,
1125            index_error: None,
1126            batch_id: None,
1127            embedding_model: None,
1128            partition_value: None,
1129            deletion_vector: None,
1130            first_row_id: None,
1131        }
1132    }
1133
1134    #[test]
1135    fn plan_prioritizes_foreign_written_files_regardless_of_size() {
1136        let planner = CompactionPlanner::new(CompactionConfig {
1137            min_files_to_compact: 4, // way more than the 1 native small file below
1138            target_file_size_bytes: 1000,
1139            max_files_per_pass: 20,
1140            ..Default::default()
1141        });
1142        let files = vec![
1143            // Native file, below target size, but alone — not enough to hit
1144            // min_files_to_compact on its own.
1145            make_plan_entry("native_small.parquet", 500, Some("AAAA".into())),
1146            // Foreign file — no centroid — large (would never pass the size filter),
1147            // but must still be selected because it has no AI-Lake index at all.
1148            make_plan_entry("foreign_big.parquet", 200_000_000, None),
1149        ];
1150        let selected = planner.plan(&files);
1151        assert_eq!(
1152            selected.len(),
1153            1,
1154            "foreign file alone should trigger a pass even below min_files_to_compact"
1155        );
1156        assert_eq!(selected[0].path, "foreign_big.parquet");
1157    }
1158
1159    #[test]
1160    fn plan_merges_foreign_and_size_candidates_together() {
1161        let planner = CompactionPlanner::new(CompactionConfig {
1162            min_files_to_compact: 2,
1163            target_file_size_bytes: 1000,
1164            max_files_per_pass: 20,
1165            ..Default::default()
1166        });
1167        let files = vec![
1168            make_plan_entry("native_a.parquet", 300, Some("AAAA".into())),
1169            make_plan_entry("native_b.parquet", 400, Some("AAAA".into())),
1170            make_plan_entry("foreign.parquet", 200_000_000, None),
1171        ];
1172        let selected = planner.plan(&files);
1173        let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1174        assert_eq!(selected.len(), 3, "paths={paths:?}");
1175        // Foreign file must be first (repair priority).
1176        assert_eq!(selected[0].path, "foreign.parquet", "paths={paths:?}");
1177    }
1178
1179    #[test]
1180    fn plan_sorts_foreign_files_smallest_first_too() {
1181        // Regression: foreign files used to keep list_files()'s arbitrary order, so a
1182        // pass could select several large foreign files instead of the cheapest ones —
1183        // violating the documented "bounds peak RAM" invariant of max_files_per_pass.
1184        let planner = CompactionPlanner::new(CompactionConfig {
1185            min_files_to_compact: 1,
1186            target_file_size_bytes: 1000,
1187            max_files_per_pass: 2,
1188            ..Default::default()
1189        });
1190        let files = vec![
1191            make_plan_entry("foreign_huge.parquet", 500_000_000, None),
1192            make_plan_entry("foreign_small.parquet", 100, None),
1193            make_plan_entry("foreign_medium.parquet", 10_000_000, None),
1194        ];
1195        let selected = planner.plan(&files);
1196        let paths: Vec<&str> = selected.iter().map(|f| f.path.as_str()).collect();
1197        assert_eq!(selected.len(), 2, "max_files_per_pass=2, paths={paths:?}");
1198        assert_eq!(
1199            paths,
1200            vec!["foreign_small.parquet", "foreign_medium.parquet"],
1201            "foreign files must be size-sorted before truncation, cheapest first"
1202        );
1203    }
1204
1205    #[tokio::test]
1206    async fn compact_merges_two_files() {
1207        use ailake_core::{VectorMetric, VectorPrecision};
1208        use ailake_store::LocalStore;
1209        use arrow_array::{Int32Array, RecordBatch};
1210        use arrow_schema::{DataType, Field, Schema};
1211        use std::sync::Arc;
1212        use tempfile::TempDir;
1213
1214        let dir = TempDir::new().unwrap();
1215        let store = Arc::new(LocalStore::new(dir.path()));
1216        let policy = VectorStoragePolicy {
1217            column_name: "embedding".into(),
1218            dim: 4,
1219            metric: VectorMetric::Cosine,
1220            precision: VectorPrecision::F16,
1221            pq: None,
1222            keep_raw_for_reranking: true,
1223            pre_normalize: false,
1224            hnsw_m: None,
1225            hnsw_ef_construction: None,
1226            ivf_residual: false,
1227            embedding_model: None,
1228            modality: None,
1229            partition_by: None,
1230            partition_value: None,
1231            partition_column_type: None,
1232            partition_fields: vec![],
1233        };
1234
1235        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1236        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
1237        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
1238
1239        let batch_a = RecordBatch::try_new(
1240            schema.clone(),
1241            vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1242        )
1243        .unwrap();
1244        let batch_b = RecordBatch::try_new(
1245            schema.clone(),
1246            vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1247        )
1248        .unwrap();
1249
1250        let writer_a = AilakeFileWriter::new(policy.clone());
1251        let bytes_a = writer_a.write(&batch_a, &embs_a).unwrap();
1252        let writer_b = AilakeFileWriter::new(policy.clone());
1253        let bytes_b = writer_b.write(&batch_b, &embs_b).unwrap();
1254
1255        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1256        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1257
1258        let entries = vec![
1259            DataFileEntry {
1260                path: "data/a.parquet".into(),
1261                record_count: 2,
1262                file_size_bytes: bytes_a.len() as u64,
1263                centroid_b64: None,
1264                radius: None,
1265                hnsw_offset: None,
1266                hnsw_len: None,
1267                vector_column: None,
1268                vector_dim: None,
1269                extra_vector_indexes: vec![],
1270                index_status: IndexStatus::Ready,
1271                index_error: None,
1272                batch_id: None,
1273                embedding_model: None,
1274                partition_value: None,
1275                deletion_vector: None,
1276                first_row_id: None,
1277            },
1278            DataFileEntry {
1279                path: "data/b.parquet".into(),
1280                record_count: 2,
1281                file_size_bytes: bytes_b.len() as u64,
1282                centroid_b64: None,
1283                radius: None,
1284                hnsw_offset: None,
1285                hnsw_len: None,
1286                vector_column: None,
1287                vector_dim: None,
1288                extra_vector_indexes: vec![],
1289                index_status: IndexStatus::Ready,
1290                index_error: None,
1291                batch_id: None,
1292                embedding_model: None,
1293                partition_value: None,
1294                deletion_vector: None,
1295                first_row_id: None,
1296            },
1297        ];
1298
1299        let executor = CompactionExecutor::new(store.clone(), policy.clone());
1300        let merged = executor
1301            .compact(&entries, "data/merged.parquet")
1302            .await
1303            .unwrap();
1304
1305        assert_eq!(merged.record_count, 4);
1306        assert_eq!(merged.path, "data/merged.parquet");
1307
1308        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1309        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1310        reader.verify_integrity().unwrap();
1311        let (batch, embs) = reader.read_parquet().unwrap();
1312        assert_eq!(batch.num_rows(), 4);
1313        assert_eq!(embs.len(), 4);
1314    }
1315
1316    /// Regression test: `compact()`/`read_files_parallel` used to read every input file's
1317    /// rows unconditionally, ignoring any `deletion_vector` on the entry — so a row deleted
1318    /// via `delete_rows()` before compaction reappeared in the merged output (new physical
1319    /// file, fresh row positions, no DV carried over — the row was never actually removed).
1320    #[tokio::test]
1321    async fn compact_drops_deletion_vector_masked_rows() {
1322        use ailake_catalog::provider::DeletionVector;
1323        use ailake_core::{VectorMetric, VectorPrecision};
1324        use ailake_store::LocalStore;
1325        use arrow_array::{Int32Array, RecordBatch};
1326        use arrow_schema::{DataType, Field, Schema};
1327        use roaring::RoaringBitmap;
1328        use std::sync::Arc;
1329        use tempfile::TempDir;
1330
1331        let dir = TempDir::new().unwrap();
1332        let store = Arc::new(LocalStore::new(dir.path()));
1333        let policy = VectorStoragePolicy {
1334            column_name: "embedding".into(),
1335            dim: 4,
1336            metric: VectorMetric::Cosine,
1337            precision: VectorPrecision::F16,
1338            pq: None,
1339            keep_raw_for_reranking: true,
1340            pre_normalize: false,
1341            hnsw_m: None,
1342            hnsw_ef_construction: None,
1343            ivf_residual: false,
1344            embedding_model: None,
1345            modality: None,
1346            partition_by: None,
1347            partition_value: None,
1348            partition_column_type: None,
1349            partition_fields: vec![],
1350        };
1351
1352        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1353        // File A: row 0 (id=0) will be marked deleted; row 1 (id=1) survives.
1354        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
1355        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0]];
1356
1357        let batch_a = RecordBatch::try_new(
1358            schema.clone(),
1359            vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1360        )
1361        .unwrap();
1362        let batch_b =
1363            RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![2i32]))])
1364                .unwrap();
1365
1366        let bytes_a = AilakeFileWriter::new(policy.clone())
1367            .write(&batch_a, &embs_a)
1368            .unwrap();
1369        let bytes_b = AilakeFileWriter::new(policy.clone())
1370            .write(&batch_b, &embs_b)
1371            .unwrap();
1372        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1373        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1374
1375        // Puffin DV marking row 0 of file A as deleted.
1376        let mut bitmap = RoaringBitmap::new();
1377        bitmap.insert(0);
1378        let (puffin_bytes, offset, length) =
1379            crate::delete::PuffinWriter::write_single_dv(&bitmap, 1).unwrap();
1380        store.put("metadata/dv-1.dvd", puffin_bytes).await.unwrap();
1381
1382        let entry_a = DataFileEntry {
1383            path: "data/a.parquet".into(),
1384            record_count: 2,
1385            file_size_bytes: bytes_a.len() as u64,
1386            centroid_b64: None,
1387            radius: None,
1388            hnsw_offset: None,
1389            hnsw_len: None,
1390            vector_column: None,
1391            vector_dim: None,
1392            extra_vector_indexes: vec![],
1393            index_status: IndexStatus::Ready,
1394            index_error: None,
1395            batch_id: None,
1396            embedding_model: None,
1397            partition_value: None,
1398            deletion_vector: Some(DeletionVector {
1399                path: "metadata/dv-1.dvd".into(),
1400                offset,
1401                length,
1402                cardinality: 1,
1403            }),
1404            first_row_id: None,
1405        };
1406        let entry_b = DataFileEntry {
1407            path: "data/b.parquet".into(),
1408            record_count: 1,
1409            file_size_bytes: bytes_b.len() as u64,
1410            centroid_b64: None,
1411            radius: None,
1412            hnsw_offset: None,
1413            hnsw_len: None,
1414            vector_column: None,
1415            vector_dim: None,
1416            extra_vector_indexes: vec![],
1417            index_status: IndexStatus::Ready,
1418            index_error: None,
1419            batch_id: None,
1420            embedding_model: None,
1421            partition_value: None,
1422            deletion_vector: None,
1423            first_row_id: None,
1424        };
1425
1426        let executor = CompactionExecutor::new(store.clone(), policy.clone());
1427        let merged = executor
1428            .compact(&[entry_a, entry_b], "data/merged.parquet")
1429            .await
1430            .unwrap();
1431
1432        // 2 rows in A + 1 in B, minus 1 deleted from A = 2 surviving rows.
1433        assert_eq!(merged.record_count, 2);
1434        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1435        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1436        let (batch, embs) = reader.read_parquet().unwrap();
1437        assert_eq!(batch.num_rows(), 2);
1438        assert_eq!(embs.len(), 2);
1439        let ids: Vec<i32> = batch
1440            .column_by_name("id")
1441            .unwrap()
1442            .as_any()
1443            .downcast_ref::<Int32Array>()
1444            .unwrap()
1445            .values()
1446            .to_vec();
1447        assert_eq!(
1448            ids,
1449            vec![1, 2],
1450            "id=0 (deleted) must not survive compaction"
1451        );
1452    }
1453
1454    /// Regression test: `verify_integrity()` used to always call `load_index()`, which
1455    /// unconditionally deserializes as `HnswIndex` regardless of the IVF-PQ flag —
1456    /// so any compaction using `ForceIvfPq` (or `Auto` selecting IVF-PQ) would write a
1457    /// correct merged file and then fail on this call, aborting the whole compaction.
1458    #[tokio::test]
1459    async fn compact_with_ivf_pq_strategy_does_not_crash_on_verify_integrity() {
1460        use ailake_core::{VectorMetric, VectorPrecision};
1461        use ailake_store::LocalStore;
1462        use arrow_array::{Int32Array, RecordBatch};
1463        use arrow_schema::{DataType, Field, Schema};
1464        use std::sync::Arc;
1465        use tempfile::TempDir;
1466
1467        let dir = TempDir::new().unwrap();
1468        let store = Arc::new(LocalStore::new(dir.path()));
1469        let dim = 8;
1470        let policy = VectorStoragePolicy {
1471            column_name: "embedding".into(),
1472            dim,
1473            metric: VectorMetric::Cosine,
1474            precision: VectorPrecision::F16,
1475            pq: None,
1476            keep_raw_for_reranking: true,
1477            pre_normalize: false,
1478            hnsw_m: None,
1479            hnsw_ef_construction: None,
1480            ivf_residual: false,
1481            embedding_model: None,
1482            modality: None,
1483            partition_by: None,
1484            partition_value: None,
1485            partition_column_type: None,
1486            partition_fields: vec![],
1487        };
1488
1489        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1490        let n_per_file = 30usize;
1491        let make_file = |path: &str, offset: i32| {
1492            let ids: Vec<i32> = (offset..offset + n_per_file as i32).collect();
1493            let embs: Vec<Vec<f32>> = ids
1494                .iter()
1495                .map(|&i| {
1496                    (0..dim as i32)
1497                        .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1498                        .collect()
1499                })
1500                .collect();
1501            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1502                .unwrap();
1503            let bytes = AilakeFileWriter::new(policy.clone())
1504                .write(&batch, &embs)
1505                .unwrap();
1506            (path.to_string(), bytes)
1507        };
1508
1509        let (path_a, bytes_a) = make_file("data/ivfpq_a.parquet", 0);
1510        let (path_b, bytes_b) = make_file("data/ivfpq_b.parquet", n_per_file as i32);
1511        for (path, bytes) in [(&path_a, &bytes_a), (&path_b, &bytes_b)] {
1512            store.put(path, bytes.clone()).await.unwrap();
1513        }
1514
1515        let make_entry = |path: &str, size: u64| DataFileEntry {
1516            path: path.to_string(),
1517            record_count: n_per_file as u64,
1518            file_size_bytes: size,
1519            centroid_b64: None,
1520            radius: None,
1521            hnsw_offset: None,
1522            hnsw_len: None,
1523            vector_column: None,
1524            vector_dim: None,
1525            extra_vector_indexes: vec![],
1526            index_status: IndexStatus::Ready,
1527            index_error: None,
1528            batch_id: None,
1529            embedding_model: None,
1530            partition_value: None,
1531            deletion_vector: None,
1532            first_row_id: None,
1533        };
1534        let entries = vec![
1535            make_entry(&path_a, bytes_a.len() as u64),
1536            make_entry(&path_b, bytes_b.len() as u64),
1537        ];
1538
1539        let executor = CompactionExecutor::new(store.clone(), policy.clone())
1540            .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1541        let merged = executor
1542            .compact(&entries, "data/ivfpq_merged.parquet")
1543            .await
1544            .expect("compact() with ForceIvfPq must not fail in verify_integrity()");
1545
1546        assert_eq!(merged.record_count, 2 * n_per_file as u64);
1547
1548        // The bug: this used to panic/error trying to load IVF-PQ bytes as HNSW.
1549        let merged_bytes = store.get("data/ivfpq_merged.parquet").await.unwrap();
1550        let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1551        reader
1552            .verify_integrity()
1553            .expect("verify_integrity() must handle an IVF-PQ-indexed file");
1554    }
1555
1556    #[tokio::test]
1557    async fn compact_incremental_merges_dominant_plus_small() {
1558        use ailake_core::{RowId, VectorMetric, VectorPrecision};
1559        use ailake_store::LocalStore;
1560        use arrow_array::{Int32Array, RecordBatch};
1561        use arrow_schema::{DataType, Field, Schema};
1562        use std::sync::Arc;
1563        use tempfile::TempDir;
1564
1565        let dir = TempDir::new().unwrap();
1566        let store = Arc::new(LocalStore::new(dir.path()));
1567        let policy = VectorStoragePolicy {
1568            column_name: "embedding".into(),
1569            dim: 4,
1570            metric: VectorMetric::Cosine,
1571            precision: VectorPrecision::F16,
1572            pq: None,
1573            keep_raw_for_reranking: true,
1574            pre_normalize: false,
1575            hnsw_m: None,
1576            hnsw_ef_construction: None,
1577            ivf_residual: false,
1578            embedding_model: None,
1579            modality: None,
1580            partition_by: None,
1581            partition_value: None,
1582            partition_column_type: None,
1583            partition_fields: vec![],
1584        };
1585
1586        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1587
1588        // Dominant file: 6 rows (75% of total 8 rows — above 40% threshold).
1589        let embs_dom: Vec<Vec<f32>> = vec![
1590            vec![1.0, 0.0, 0.0, 0.0],
1591            vec![0.0, 1.0, 0.0, 0.0],
1592            vec![0.0, 0.0, 1.0, 0.0],
1593            vec![0.7, 0.7, 0.0, 0.0],
1594            vec![0.0, 0.7, 0.7, 0.0],
1595            vec![0.0, 0.0, 0.7, 0.7],
1596        ];
1597        let batch_dom = RecordBatch::try_new(
1598            schema.clone(),
1599            vec![Arc::new(Int32Array::from(vec![0i32, 1, 2, 3, 4, 5]))],
1600        )
1601        .unwrap();
1602
1603        // Small file: 2 rows.
1604        let embs_small: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 0.0, 1.0], vec![0.5, 0.5, 0.5, 0.5]];
1605        let batch_small = RecordBatch::try_new(
1606            schema.clone(),
1607            vec![Arc::new(Int32Array::from(vec![6i32, 7]))],
1608        )
1609        .unwrap();
1610
1611        let bytes_dom = AilakeFileWriter::new(policy.clone())
1612            .write(&batch_dom, &embs_dom)
1613            .unwrap();
1614        let bytes_small = AilakeFileWriter::new(policy.clone())
1615            .write(&batch_small, &embs_small)
1616            .unwrap();
1617
1618        store
1619            .put("data/dominant.parquet", bytes_dom.clone())
1620            .await
1621            .unwrap();
1622        store
1623            .put("data/small.parquet", bytes_small.clone())
1624            .await
1625            .unwrap();
1626
1627        let entries = vec![
1628            DataFileEntry {
1629                path: "data/dominant.parquet".into(),
1630                record_count: 6,
1631                file_size_bytes: bytes_dom.len() as u64,
1632                centroid_b64: None,
1633                radius: None,
1634                hnsw_offset: None,
1635                hnsw_len: None,
1636                vector_column: None,
1637                vector_dim: None,
1638                extra_vector_indexes: vec![],
1639                index_status: IndexStatus::Ready,
1640                index_error: None,
1641                batch_id: None,
1642                embedding_model: None,
1643                partition_value: None,
1644                deletion_vector: None,
1645                first_row_id: None,
1646            },
1647            DataFileEntry {
1648                path: "data/small.parquet".into(),
1649                record_count: 2,
1650                file_size_bytes: bytes_small.len() as u64,
1651                centroid_b64: None,
1652                radius: None,
1653                hnsw_offset: None,
1654                hnsw_len: None,
1655                vector_column: None,
1656                vector_dim: None,
1657                extra_vector_indexes: vec![],
1658                index_status: IndexStatus::Ready,
1659                index_error: None,
1660                batch_id: None,
1661                embedding_model: None,
1662                partition_value: None,
1663                deletion_vector: None,
1664                first_row_id: None,
1665            },
1666        ];
1667
1668        let executor = CompactionExecutor::new(store.clone(), policy.clone());
1669        let merged = executor
1670            .compact_incremental(&entries, "data/merged.parquet")
1671            .await
1672            .unwrap();
1673
1674        // Structural checks.
1675        assert_eq!(merged.record_count, 8);
1676        assert_eq!(merged.path, "data/merged.parquet");
1677
1678        // Load merged file and verify it's a valid AI-Lake file.
1679        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1680        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1681        reader.verify_integrity().unwrap();
1682
1683        let (batch, embs) = reader.read_parquet().unwrap();
1684        assert_eq!(batch.num_rows(), 8);
1685        assert_eq!(embs.len(), 8);
1686
1687        // Dominant rows must come first (positions 0..5).
1688        for f in &embs[..6] {
1689            assert_eq!(f.len(), 4);
1690        }
1691
1692        // HNSW must be searchable and return the nearest neighbor for a known query.
1693        let hnsw = reader.load_index().unwrap();
1694        assert_eq!(hnsw.node_count(), 8);
1695
1696        // Query [1, 0, 0, 0] → nearest should be RowId 0 (embs_dom[0]).
1697        let results = hnsw.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
1698        assert_eq!(results[0].0, RowId::new(0));
1699
1700        // Query [0, 0, 0, 1] → nearest should be RowId 6 (first row of small file,
1701        // inserted at position 6 in the merged file).
1702        let results = hnsw.search(&[0.0, 0.0, 0.0, 1.0], 1, 50);
1703        assert_eq!(results[0].0, RowId::new(6));
1704    }
1705
1706    /// Regression: `compact_incremental()` only ever produces HNSW (it extends the
1707    /// dominant file's existing graph) — with a dominant file present, it never checked
1708    /// `self.index_strategy` before taking that path, so an explicit `ForceIvfPq` request
1709    /// was silently satisfied with HNSW instead. Reachable from `CompactionExecutor::run()`/
1710    /// `run_deferred()` (used by the CLI `compact` command and every JNI/Spark/Trino/Flink/
1711    /// DuckDB compact call), which always try `compact_incremental()` first — including on
1712    /// a GPU machine where `Auto` would pick IVF-PQ for a fresh build, or where a caller
1713    /// explicitly forces IVF-PQ. Same dominant/small file shape as
1714    /// `compact_incremental_merges_dominant_plus_small`, but with `ForceIvfPq` set.
1715    #[tokio::test]
1716    async fn compact_incremental_respects_force_ivf_pq_even_with_dominant_file() {
1717        use ailake_core::{VectorMetric, VectorPrecision};
1718        use ailake_store::LocalStore;
1719        use arrow_array::{Int32Array, RecordBatch};
1720        use arrow_schema::{DataType, Field, Schema};
1721        use std::sync::Arc;
1722        use tempfile::TempDir;
1723
1724        let dir = TempDir::new().unwrap();
1725        let store = Arc::new(LocalStore::new(dir.path()));
1726        let dim = 8;
1727        let policy = VectorStoragePolicy {
1728            column_name: "embedding".into(),
1729            dim,
1730            metric: VectorMetric::Cosine,
1731            precision: VectorPrecision::F16,
1732            pq: None,
1733            keep_raw_for_reranking: true,
1734            pre_normalize: false,
1735            hnsw_m: None,
1736            hnsw_ef_construction: None,
1737            ivf_residual: false,
1738            embedding_model: None,
1739            modality: None,
1740            partition_by: None,
1741            partition_value: None,
1742            partition_column_type: None,
1743            partition_fields: vec![],
1744        };
1745
1746        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1747        let make_file = |path: &str, offset: i32, n: usize| {
1748            let ids: Vec<i32> = (offset..offset + n as i32).collect();
1749            let embs: Vec<Vec<f32>> = ids
1750                .iter()
1751                .map(|&i| {
1752                    (0..dim as i32)
1753                        .map(|j| ((i * 31 + j * 7) % 97) as f32 / 97.0)
1754                        .collect()
1755                })
1756                .collect();
1757            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1758                .unwrap();
1759            // Dominant file already has a real, loadable HNSW index — the exact
1760            // condition that used to make compact_incremental() bypass ForceIvfPq.
1761            let bytes = AilakeFileWriter::new(policy.clone())
1762                .write(&batch, &embs)
1763                .unwrap();
1764            (path.to_string(), bytes, n as u64)
1765        };
1766
1767        // Dominant: 90 rows (75% of 120 total — comfortably above the dominant-file
1768        // threshold), small: 30 rows.
1769        let (path_dom, bytes_dom, n_dom) = make_file("data/dom.parquet", 0, 90);
1770        let (path_small, bytes_small, n_small) = make_file("data/small.parquet", 999, 30);
1771
1772        for (path, bytes) in [(&path_dom, &bytes_dom), (&path_small, &bytes_small)] {
1773            store.put(path, bytes.clone()).await.unwrap();
1774        }
1775
1776        let make_entry = |path: &str, record_count: u64, size: u64| DataFileEntry {
1777            path: path.to_string(),
1778            record_count,
1779            file_size_bytes: size,
1780            centroid_b64: None,
1781            radius: None,
1782            hnsw_offset: None,
1783            hnsw_len: None,
1784            vector_column: None,
1785            vector_dim: None,
1786            extra_vector_indexes: vec![],
1787            index_status: IndexStatus::Ready,
1788            index_error: None,
1789            batch_id: None,
1790            embedding_model: None,
1791            partition_value: None,
1792            deletion_vector: None,
1793            first_row_id: None,
1794        };
1795        let entries = vec![
1796            make_entry(&path_dom, n_dom, bytes_dom.len() as u64),
1797            make_entry(&path_small, n_small, bytes_small.len() as u64),
1798        ];
1799
1800        let executor = CompactionExecutor::new(store.clone(), policy.clone())
1801            .with_index_strategy(CompactionIndexStrategy::ForceIvfPq);
1802        let merged = executor
1803            .compact_incremental(&entries, "data/merged_force_ivfpq.parquet")
1804            .await
1805            .expect("compact_incremental() with ForceIvfPq must fall back to compact() cleanly");
1806
1807        let merged_bytes = store.get("data/merged_force_ivfpq.parquet").await.unwrap();
1808        let reader = AilakeFileReader::new(merged_bytes, "embedding", dim);
1809        reader.verify_integrity().unwrap();
1810
1811        match reader.load_any_index().unwrap() {
1812            ailake_index::AnyIndex::IvfPq(_) => {}
1813            ailake_index::AnyIndex::Hnsw(_) => panic!(
1814                "ForceIvfPq was silently satisfied with HNSW instead — merged.record_count={}",
1815                merged.record_count
1816            ),
1817        }
1818    }
1819
1820    #[tokio::test]
1821    async fn compact_incremental_falls_back_when_no_dominant() {
1822        use ailake_core::{VectorMetric, VectorPrecision};
1823        use ailake_store::LocalStore;
1824        use arrow_array::{Int32Array, RecordBatch};
1825        use arrow_schema::{DataType, Field, Schema};
1826        use std::sync::Arc;
1827        use tempfile::TempDir;
1828
1829        let dir = TempDir::new().unwrap();
1830        let store = Arc::new(LocalStore::new(dir.path()));
1831        let policy = VectorStoragePolicy {
1832            column_name: "embedding".into(),
1833            dim: 4,
1834            metric: VectorMetric::Cosine,
1835            precision: VectorPrecision::F16,
1836            pq: None,
1837            keep_raw_for_reranking: true,
1838            pre_normalize: false,
1839            hnsw_m: None,
1840            hnsw_ef_construction: None,
1841            ivf_residual: false,
1842            embedding_model: None,
1843            modality: None,
1844            partition_by: None,
1845            partition_value: None,
1846            partition_column_type: None,
1847            partition_fields: vec![],
1848        };
1849
1850        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1851
1852        // Two equal-sized files (50/50 split — no dominant, both below 40% threshold).
1853        let make_batch = |ids: Vec<i32>, embs: Vec<Vec<f32>>| {
1854            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
1855                .unwrap();
1856            AilakeFileWriter::new(policy.clone())
1857                .write(&batch, &embs)
1858                .unwrap()
1859        };
1860
1861        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
1862        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
1863        let bytes_a = make_batch(vec![0, 1], embs_a);
1864        let bytes_b = make_batch(vec![2, 3], embs_b);
1865
1866        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1867        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1868
1869        let entries = vec![
1870            DataFileEntry {
1871                path: "data/a.parquet".into(),
1872                record_count: 2,
1873                file_size_bytes: bytes_a.len() as u64,
1874                centroid_b64: None,
1875                radius: None,
1876                hnsw_offset: None,
1877                hnsw_len: None,
1878                vector_column: None,
1879                vector_dim: None,
1880                extra_vector_indexes: vec![],
1881                index_status: IndexStatus::Ready,
1882                index_error: None,
1883                batch_id: None,
1884                embedding_model: None,
1885                partition_value: None,
1886                deletion_vector: None,
1887                first_row_id: None,
1888            },
1889            DataFileEntry {
1890                path: "data/b.parquet".into(),
1891                record_count: 2,
1892                file_size_bytes: bytes_b.len() as u64,
1893                centroid_b64: None,
1894                radius: None,
1895                hnsw_offset: None,
1896                hnsw_len: None,
1897                vector_column: None,
1898                vector_dim: None,
1899                extra_vector_indexes: vec![],
1900                index_status: IndexStatus::Ready,
1901                index_error: None,
1902                batch_id: None,
1903                embedding_model: None,
1904                partition_value: None,
1905                deletion_vector: None,
1906                first_row_id: None,
1907            },
1908        ];
1909
1910        let executor = CompactionExecutor::new(store.clone(), policy.clone());
1911        // Should fall back to full rebuild without error.
1912        let merged = executor
1913            .compact_incremental(&entries, "data/merged.parquet")
1914            .await
1915            .unwrap();
1916
1917        assert_eq!(merged.record_count, 4);
1918
1919        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
1920        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
1921        reader.verify_integrity().unwrap();
1922    }
1923
1924    #[tokio::test]
1925    async fn compact_deferred_produces_parquet_only_file() {
1926        use ailake_catalog::HadoopCatalog;
1927        use ailake_core::{VectorMetric, VectorPrecision};
1928        use ailake_store::LocalStore;
1929        use arrow_array::{Int32Array, RecordBatch};
1930        use arrow_schema::{DataType, Field, Schema};
1931        use std::sync::Arc;
1932        use tempfile::TempDir;
1933
1934        let dir = TempDir::new().unwrap();
1935        let store = Arc::new(LocalStore::new(dir.path()));
1936        let catalog_dir = TempDir::new().unwrap();
1937        let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
1938        let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
1939        let table = TableIdent {
1940            namespace: "ns".into(),
1941            name: "tbl".into(),
1942        };
1943
1944        let policy = VectorStoragePolicy {
1945            column_name: "embedding".into(),
1946            dim: 4,
1947            metric: VectorMetric::Cosine,
1948            precision: VectorPrecision::F16,
1949            pq: None,
1950            keep_raw_for_reranking: true,
1951            pre_normalize: false,
1952            hnsw_m: None,
1953            hnsw_ef_construction: None,
1954            ivf_residual: false,
1955            embedding_model: None,
1956            modality: None,
1957            partition_by: None,
1958            partition_value: None,
1959            partition_column_type: None,
1960            partition_fields: vec![],
1961        };
1962
1963        use ailake_catalog::TableProperties;
1964        catalog
1965            .create_table(
1966                &table,
1967                &TableProperties {
1968                    policy: policy.clone(),
1969                    extra: std::collections::HashMap::new(),
1970                    format_version: 2,
1971                    partition_column_type: None,
1972                },
1973            )
1974            .await
1975            .unwrap();
1976
1977        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1978        let embs_a: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
1979        let batch_a = RecordBatch::try_new(
1980            schema.clone(),
1981            vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
1982        )
1983        .unwrap();
1984        let bytes_a = AilakeFileWriter::new(policy.clone())
1985            .write(&batch_a, &embs_a)
1986            .unwrap();
1987        store.put("data/a.parquet", bytes_a.clone()).await.unwrap();
1988
1989        let embs_b: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
1990        let batch_b = RecordBatch::try_new(
1991            schema.clone(),
1992            vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
1993        )
1994        .unwrap();
1995        let bytes_b = AilakeFileWriter::new(policy.clone())
1996            .write(&batch_b, &embs_b)
1997            .unwrap();
1998        store.put("data/b.parquet", bytes_b.clone()).await.unwrap();
1999
2000        let entries = vec![
2001            DataFileEntry {
2002                path: "data/a.parquet".into(),
2003                record_count: 2,
2004                file_size_bytes: bytes_a.len() as u64,
2005                centroid_b64: None,
2006                radius: None,
2007                hnsw_offset: None,
2008                hnsw_len: None,
2009                vector_column: None,
2010                vector_dim: None,
2011                extra_vector_indexes: vec![],
2012                index_status: IndexStatus::Ready,
2013                index_error: None,
2014                batch_id: None,
2015                embedding_model: None,
2016                partition_value: None,
2017                deletion_vector: None,
2018                first_row_id: None,
2019            },
2020            DataFileEntry {
2021                path: "data/b.parquet".into(),
2022                record_count: 2,
2023                file_size_bytes: bytes_b.len() as u64,
2024                centroid_b64: None,
2025                radius: None,
2026                hnsw_offset: None,
2027                hnsw_len: None,
2028                vector_column: None,
2029                vector_dim: None,
2030                extra_vector_indexes: vec![],
2031                index_status: IndexStatus::Ready,
2032                index_error: None,
2033                batch_id: None,
2034                embedding_model: None,
2035                partition_value: None,
2036                deletion_vector: None,
2037                first_row_id: None,
2038            },
2039        ];
2040
2041        let executor = CompactionExecutor::new(store.clone(), policy.clone());
2042        let entry = executor
2043            .compact_deferred(&entries, "data/merged.parquet", catalog.clone(), &table)
2044            .await
2045            .unwrap();
2046
2047        // Entry is Indexing — HNSW build pending in background
2048        assert_eq!(entry.index_status, IndexStatus::Indexing);
2049        assert_eq!(entry.record_count, 4);
2050
2051        // The written file must be valid Parquet (readable) even without HNSW
2052        let merged_bytes = store.get("data/merged.parquet").await.unwrap();
2053        let pq_reader = ailake_parquet::ParquetVectorReader::new(merged_bytes, "embedding");
2054        let count = pq_reader.record_count().unwrap();
2055        assert_eq!(count, 4);
2056    }
2057
2058    /// Regression test: `CompactionExecutor::run()` must not drop files that fall
2059    /// outside the compaction pass (too large for `target_file_size_bytes`, or beyond
2060    /// `max_files_per_pass`). `Replace` snapshots don't inherit the previous manifest
2061    /// (see `HadoopCatalog::commit_snapshot`), so `run()` must explicitly carry forward
2062    /// every untouched file alongside the merged output — mirroring the pattern already
2063    /// used by the CLI `compact` command and `MemoryDecayJob::run`.
2064    #[tokio::test]
2065    async fn run_preserves_untouched_files_outside_compaction_pass() {
2066        use ailake_catalog::HadoopCatalog;
2067        use ailake_core::{VectorMetric, VectorPrecision};
2068        use ailake_store::LocalStore;
2069        use arrow_array::{Int32Array, RecordBatch};
2070        use arrow_schema::{DataType, Field, Schema};
2071        use std::sync::Arc;
2072        use tempfile::TempDir;
2073
2074        let dir = TempDir::new().unwrap();
2075        let store = Arc::new(LocalStore::new(dir.path()));
2076        let catalog_dir = TempDir::new().unwrap();
2077        let catalog_store = Arc::new(LocalStore::new(catalog_dir.path()));
2078        let catalog = Arc::new(HadoopCatalog::new(catalog_store, ""));
2079        let table = TableIdent {
2080            namespace: "ns".into(),
2081            name: "tbl".into(),
2082        };
2083
2084        let policy = VectorStoragePolicy {
2085            column_name: "embedding".into(),
2086            dim: 4,
2087            metric: VectorMetric::Cosine,
2088            precision: VectorPrecision::F16,
2089            pq: None,
2090            keep_raw_for_reranking: true,
2091            pre_normalize: false,
2092            hnsw_m: None,
2093            hnsw_ef_construction: None,
2094            ivf_residual: false,
2095            embedding_model: None,
2096            modality: None,
2097            partition_by: None,
2098            partition_value: None,
2099            partition_column_type: None,
2100            partition_fields: vec![],
2101        };
2102
2103        use ailake_catalog::TableProperties;
2104        catalog
2105            .create_table(
2106                &table,
2107                &TableProperties {
2108                    policy: policy.clone(),
2109                    extra: std::collections::HashMap::new(),
2110                    format_version: 2,
2111                    partition_column_type: None,
2112                },
2113            )
2114            .await
2115            .unwrap();
2116
2117        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2118
2119        // Two small files — eligible for compaction.
2120        let write_file = |path: &str, ids: Vec<i32>, embs: Vec<Vec<f32>>| {
2121            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(ids))])
2122                .unwrap();
2123            let bytes = AilakeFileWriter::new(policy.clone())
2124                .write(&batch, &embs)
2125                .unwrap();
2126            (path.to_string(), bytes)
2127        };
2128
2129        let (path_a, bytes_a) = write_file(
2130            "data/small_a.parquet",
2131            vec![0, 1],
2132            vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
2133        );
2134        let (path_b, bytes_b) = write_file(
2135            "data/small_b.parquet",
2136            vec![2, 3],
2137            vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]],
2138        );
2139        // "Big" file — same tiny payload in this test, but its DataFileEntry reports a
2140        // size above target_file_size_bytes so the planner must never select it.
2141        let (path_big, bytes_big) = write_file(
2142            "data/big.parquet",
2143            vec![4, 5],
2144            vec![vec![1.0, 1.0, 0.0, 0.0], vec![0.0, 1.0, 1.0, 0.0]],
2145        );
2146
2147        for (path, bytes) in [
2148            (&path_a, &bytes_a),
2149            (&path_b, &bytes_b),
2150            (&path_big, &bytes_big),
2151        ] {
2152            store.put(path, bytes.clone()).await.unwrap();
2153        }
2154
2155        let make_entry = |path: &str, size: u64| DataFileEntry {
2156            path: path.to_string(),
2157            record_count: 2,
2158            file_size_bytes: size,
2159            // Non-None: all three files here are meant to represent normal,
2160            // already-indexed AI-Lake files — only size should decide eligibility.
2161            // `plan_prioritizes_foreign_written_files_regardless_of_size` covers
2162            // the `centroid_b64: None` (foreign-write) case separately.
2163            centroid_b64: Some("AAAA".into()),
2164            radius: None,
2165            hnsw_offset: None,
2166            hnsw_len: None,
2167            vector_column: None,
2168            vector_dim: None,
2169            extra_vector_indexes: vec![],
2170            index_status: IndexStatus::Ready,
2171            index_error: None,
2172            batch_id: None,
2173            embedding_model: None,
2174            partition_value: None,
2175            deletion_vector: None,
2176            first_row_id: None,
2177        };
2178
2179        let initial_snap_id = ailake_catalog::new_snapshot_id();
2180        let initial_snapshot = NewSnapshot {
2181            snapshot_id: initial_snap_id,
2182            parent_snapshot_id: None,
2183            files: vec![
2184                make_entry(&path_a, 500),
2185                make_entry(&path_b, 500),
2186                make_entry(&path_big, 200_000_000), // far above target_file_size_bytes below
2187            ],
2188            operation: SnapshotOperation::Append,
2189            iceberg_schema: None,
2190            extra_properties: std::collections::HashMap::new(),
2191            bloom_filters: vec![],
2192            equality_delete_files: vec![],
2193        };
2194        catalog
2195            .commit_snapshot(&table, initial_snapshot)
2196            .await
2197            .unwrap();
2198
2199        let planner = CompactionPlanner::new(CompactionConfig {
2200            min_files_to_compact: 2,
2201            target_file_size_bytes: 1000,
2202            index_strategy: CompactionIndexStrategy::ForceHnsw,
2203            max_files_per_pass: 20,
2204        });
2205        let executor = CompactionExecutor::new(store.clone(), policy.clone());
2206
2207        let merged = executor
2208            .run(&planner, &table, catalog.clone(), "data")
2209            .await
2210            .unwrap()
2211            .expect("compaction should have run — 2 eligible small files");
2212
2213        let files_after = catalog.list_files(&table, None).await.unwrap();
2214        let paths_after: Vec<&str> = files_after.iter().map(|f| f.path.as_str()).collect();
2215
2216        assert!(
2217            paths_after.contains(&path_big.as_str()),
2218            "BUG: untouched 'big.parquet' vanished after run() — files_after={paths_after:?}"
2219        );
2220        assert!(
2221            paths_after.contains(&merged.path.as_str()),
2222            "merged output file must be present — files_after={paths_after:?}"
2223        );
2224
2225        // Regression: `run()` used to hardcode `parent_snapshot_id: None` on the
2226        // post-compaction Replace snapshot even though a current snapshot always exists
2227        // here, breaking Iceberg snapshot lineage (`expire_snapshots`/`rollback_to_snapshot`)
2228        // for any compacted table. Read the committed metadata.json directly (no public
2229        // CatalogProvider method exposes snapshot lineage) and confirm the new snapshot's
2230        // parent points at the snapshot committed before this compaction ran.
2231        let meta_dir = catalog_dir.path().join("ns/tbl/metadata");
2232        let latest_metadata = std::fs::read_dir(&meta_dir)
2233            .unwrap()
2234            .filter_map(|e| e.ok())
2235            .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
2236            .max_by_key(|e| e.metadata().unwrap().modified().unwrap())
2237            .expect("metadata.json must exist after commit");
2238        let json: serde_json::Value =
2239            serde_json::from_slice(&std::fs::read(latest_metadata.path()).unwrap()).unwrap();
2240        let last_snapshot = json["snapshots"].as_array().unwrap().last().unwrap();
2241        assert_eq!(
2242            last_snapshot["parent-snapshot-id"].as_i64(),
2243            Some(initial_snap_id),
2244            "compaction's Replace snapshot must chain to the pre-compaction snapshot, not be orphaned"
2245        );
2246        assert!(
2247            !paths_after.contains(&path_a.as_str()) && !paths_after.contains(&path_b.as_str()),
2248            "compacted input files must no longer be listed — files_after={paths_after:?}"
2249        );
2250        assert_eq!(
2251            files_after.len(),
2252            2,
2253            "expected exactly [big.parquet, merged] — files_after={paths_after:?}"
2254        );
2255    }
2256
2257    /// Regression test: a file with no AILK footer (e.g. rewritten by a generic
2258    /// Iceberg engine — Spark/Trino `OPTIMIZE` — with no knowledge of AI-Lake) still
2259    /// holds valid Parquet data. `read_parquet()` decodes the vector column directly
2260    /// from Parquet and never touches the footer, so `compact()`/`compact_incremental()`
2261    /// must include such a file's rows in the merge, not silently drop them.
2262    #[tokio::test]
2263    async fn compact_preserves_rows_from_footerless_file() {
2264        use ailake_core::{VectorMetric, VectorPrecision};
2265        use ailake_store::LocalStore;
2266        use arrow_array::{Int32Array, RecordBatch};
2267        use arrow_schema::{DataType, Field, Schema};
2268        use std::sync::Arc;
2269        use tempfile::TempDir;
2270
2271        let dir = TempDir::new().unwrap();
2272        let store = Arc::new(LocalStore::new(dir.path()));
2273        let policy = VectorStoragePolicy {
2274            column_name: "embedding".into(),
2275            dim: 4,
2276            metric: VectorMetric::Cosine,
2277            precision: VectorPrecision::F16,
2278            pq: None,
2279            keep_raw_for_reranking: true,
2280            pre_normalize: false,
2281            hnsw_m: None,
2282            hnsw_ef_construction: None,
2283            ivf_residual: false,
2284            embedding_model: None,
2285            modality: None,
2286            partition_by: None,
2287            partition_value: None,
2288            partition_column_type: None,
2289            partition_fields: vec![],
2290        };
2291
2292        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
2293
2294        // Normal AI-Lake file, written the usual way (has an AILK footer).
2295        let embs_native: Vec<Vec<f32>> = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]];
2296        let batch_native = RecordBatch::try_new(
2297            schema.clone(),
2298            vec![Arc::new(Int32Array::from(vec![0i32, 1]))],
2299        )
2300        .unwrap();
2301        let bytes_native = AilakeFileWriter::new(policy.clone())
2302            .write(&batch_native, &embs_native)
2303            .unwrap();
2304        store
2305            .put("data/native.parquet", bytes_native.clone())
2306            .await
2307            .unwrap();
2308
2309        // "Foreign" file — plain Parquet, no AILK footer, same shape a generic
2310        // Iceberg engine's rewrite would produce (`write_parquet_only` is the exact
2311        // primitive `compact_deferred` uses for its Parquet-only fast path).
2312        let embs_foreign: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 1.0, 0.0], vec![0.0, 0.0, 0.0, 1.0]];
2313        let batch_foreign = RecordBatch::try_new(
2314            schema.clone(),
2315            vec![Arc::new(Int32Array::from(vec![2i32, 3]))],
2316        )
2317        .unwrap();
2318        let bytes_foreign = AilakeFileWriter::new(policy.clone())
2319            .write_parquet_only(&batch_foreign, &embs_foreign)
2320            .unwrap();
2321        store
2322            .put("data/foreign.parquet", bytes_foreign.clone())
2323            .await
2324            .unwrap();
2325
2326        let reader_foreign = AilakeFileReader::new(bytes_foreign.clone(), "embedding", 4);
2327        assert!(
2328            !reader_foreign.is_ailake_file(),
2329            "sanity: write_parquet_only must not embed an AILK footer"
2330        );
2331
2332        let entries = vec![
2333            DataFileEntry {
2334                path: "data/native.parquet".into(),
2335                record_count: 2,
2336                file_size_bytes: bytes_native.len() as u64,
2337                centroid_b64: None,
2338                radius: None,
2339                hnsw_offset: None,
2340                hnsw_len: None,
2341                vector_column: None,
2342                vector_dim: None,
2343                extra_vector_indexes: vec![],
2344                index_status: IndexStatus::Ready,
2345                index_error: None,
2346                batch_id: None,
2347                embedding_model: None,
2348                partition_value: None,
2349                deletion_vector: None,
2350                first_row_id: None,
2351            },
2352            DataFileEntry {
2353                path: "data/foreign.parquet".into(),
2354                record_count: 2,
2355                file_size_bytes: bytes_foreign.len() as u64,
2356                centroid_b64: None,
2357                radius: None,
2358                hnsw_offset: None,
2359                hnsw_len: None,
2360                vector_column: None,
2361                vector_dim: None,
2362                extra_vector_indexes: vec![],
2363                index_status: IndexStatus::Ready,
2364                index_error: None,
2365                batch_id: None,
2366                embedding_model: None,
2367                partition_value: None,
2368                deletion_vector: None,
2369                first_row_id: None,
2370            },
2371        ];
2372
2373        // Full rebuild path (`compact`).
2374        let executor = CompactionExecutor::new(store.clone(), policy.clone());
2375        let merged = executor
2376            .compact(&entries, "data/merged_full.parquet")
2377            .await
2378            .unwrap();
2379        assert_eq!(
2380            merged.record_count, 4,
2381            "compact() must include all 4 rows — 2 native + 2 from the footerless file"
2382        );
2383        let merged_bytes = store.get("data/merged_full.parquet").await.unwrap();
2384        let reader = AilakeFileReader::new(merged_bytes, "embedding", 4);
2385        reader.verify_integrity().unwrap();
2386
2387        // Incremental path (`compact_incremental` — native file becomes dominant at 50/50,
2388        // exercising the non-dominant-file read path where the bug lived).
2389        let merged_inc = executor
2390            .compact_incremental(&entries, "data/merged_inc.parquet")
2391            .await
2392            .unwrap();
2393        assert_eq!(
2394            merged_inc.record_count, 4,
2395            "compact_incremental() must include all 4 rows — 2 native + 2 from the footerless file"
2396        );
2397    }
2398}