Skip to main content

chronix_engine/segment/
reader.rs

1//! Segment reader — reads `.csx` segment files.
2//!
3//! The reader validates the file integrity (magic bytes, `CRC32c` checksum),
4//! parses header/footer/metadata, and provides methods to read individual
5//! columns or entire row groups as Apache Arrow arrays.
6//!
7//! # I/O strategy — memory-mapped prefetch
8//!
9//! The entire segment file is **memory-mapped** via [`memmap2::Mmap`].  No
10//! per-column `seek()`/`read()` system calls are issued; instead, column data
11//! is accessed by slicing the mapped byte range, and the OS page cache
12//! transparently pages data in from disk on first access.
13//!
14//! Two advisory hints accelerate sequential scans:
15//!
16//! 1. **`madvise(MADV_SEQUENTIAL)`** — set once at open time, tells the kernel
17//!    the file will be read front-to-back so it can perform aggressive
18//!    read-ahead (typically 128–256 KiB on Linux).
19//! 2. **`madvise(MADV_WILLNEED)`** — issued per-column in
20//!    [`SegmentReader::read_columns`]: while decoding column *N*, the reader
21//!    hints the kernel to begin paging in column *N+1* (and the same column
22//!    in the next row group for multi-row-group scans).  This overlaps I/O
23//!    with CPU decoding.
24//!
25//! Because the `Mmap` holds the entire file mapping, **there is no separate
26//! read-ahead buffer or prefetch thread**; the OS page cache *is* the buffer.
27//! This is the most efficient approach for columnar segment files that fit in
28//! the page cache — it avoids double-copying into user-space buffers and lets
29//! the kernel coalesce I/O for adjacent columns.
30//!
31//! # Predicate pushdown
32//!
33//! The reader supports multi-level predicate pushdown to minimise the amount
34//! of data that is read and decoded:
35//!
36//! | Level | Mechanism | Method |
37//! |---|---|---|
38//! | **Segment** | Catalog-level time-range + tag-index + bloom filter | `Chronix::prune_segments()` |
39//! | **Row group — time** | `min_ts`/`max_ts` from timestamp column stats | `read_projected_filtered` |
40//! | **Row group — tags** | `value_count == 0` with a non-zero `null_count` (a genuinely all-null block → impossible match) | `read_projected_filtered_with_predicates` |
41//! | **Column** | projection pushdown (only requested columns decoded) | `read_projected` |
42
43use std::fs::File;
44use std::path::{Path, PathBuf};
45// The only bare `Arc` in this file is the key provider, which is gated; every
46// other use is written out as `std::sync::Arc`. Without the same gate the
47// import is unused under `--no-default-features`.
48#[cfg(feature = "field-encryption")]
49use std::sync::Arc;
50
51use arrow::array::{
52    ArrayRef, BooleanArray, Decimal128Array, Float64Array, Int64Array, StringArray, UInt64Array,
53};
54use arrow::buffer::NullBuffer;
55use arrow::datatypes::{Field, Schema};
56use arrow::record_batch::RecordBatch;
57use memmap2::{Advice, Mmap, UncheckedAdvice};
58
59use chronix_encoding::{ColumnDecoder, DecodedColumn, EncodedBlock};
60
61use crate::segment::compression::decompress_block_with_optional_dict;
62use crate::segment::error::{Result, SegmentError};
63use crate::segment::header::{SegmentFooter, SegmentHeader, FOOTER_SIZE, HEADER_SIZE};
64use crate::segment::metadata::{data_types, ColumnBlockMeta, ColumnMeta, SegmentMetadata};
65use crate::segment::stats::ordered_i64_to_f64;
66
67// ── Zone-map field predicates for late-materialisation ─────────────────
68
69/// A numeric comparison operator for zone-map row-group pruning.
70#[derive(Debug, Clone, Copy, PartialEq)]
71pub enum ZoneMapOp {
72    /// Equals.
73    Eq,
74    /// Greater than.
75    Gt,
76    /// Greater than or equal.
77    GtEq,
78    /// Less than.
79    Lt,
80    /// Less than or equal.
81    LtEq,
82}
83
84/// A field predicate that can be evaluated against per-row-group
85/// [`ColumnStats`](crate::segment::stats::ColumnStats) min/max to skip row groups
86/// that cannot possibly match.
87#[derive(Debug, Clone)]
88pub struct FieldPredicate {
89    /// Column name.
90    pub column: String,
91    /// Comparison operator.
92    pub op: ZoneMapOp,
93    /// Predicate value as `f64`.  Integer predicates are cast to `f64`.
94    pub value: f64,
95}
96
97impl FieldPredicate {
98    /// Test whether a row group with the given float min/max statistics
99    /// could possibly contain rows matching this predicate.
100    ///
101    /// A NaN bound makes every comparison false, which would prune the
102    /// group — so an unusable bound is reported as "may match" instead.
103    /// **A pruning step may only ever be wrong in the direction of doing
104    /// more work**; one that can drop rows is a wrong answer, not a slow one.
105    fn may_match_f64(&self, rg_min: f64, rg_max: f64) -> bool {
106        if rg_min.is_nan() || rg_max.is_nan() || rg_min > rg_max {
107            return true;
108        }
109        match self.op {
110            ZoneMapOp::Eq => self.value >= rg_min && self.value <= rg_max,
111            ZoneMapOp::Gt => rg_max > self.value,
112            ZoneMapOp::GtEq => rg_max >= self.value,
113            ZoneMapOp::Lt => rg_min < self.value,
114            ZoneMapOp::LtEq => rg_min <= self.value,
115        }
116    }
117
118    /// Test whether a row group with i64 min/max statistics could
119    /// possibly match.
120    ///
121    /// The predicate value is an `f64`, so it is compared as one: casting
122    /// it to `i64` first truncated the bound, and `n < 1.5` then pruned a
123    /// row group holding `{1, 2, 3}` because `1 < 1` is false. Above 2^53 an
124    /// `i64` bound is not representable as an `f64` at all, and the group is
125    /// kept rather than compared wrongly.
126    fn may_match_i64(&self, rg_min: i64, rg_max: i64) -> bool {
127        if rg_min > rg_max {
128            return true;
129        }
130        const EXACT: i64 = 1 << 53;
131        if rg_min <= -EXACT || rg_max >= EXACT {
132            return true;
133        }
134        #[allow(clippy::cast_precision_loss)] // bounded above by 2^53
135        let (lo, hi) = (rg_min as f64, rg_max as f64);
136        match self.op {
137            // The column is integral, so an equality against a fractional
138            // value cannot match any row in it, whatever the range says.
139            ZoneMapOp::Eq if self.value.fract() != 0.0 => false,
140            ZoneMapOp::Eq => self.value >= lo && self.value <= hi,
141            ZoneMapOp::Gt => hi > self.value,
142            ZoneMapOp::GtEq => hi >= self.value,
143            ZoneMapOp::Lt => lo < self.value,
144            ZoneMapOp::LtEq => lo <= self.value,
145        }
146    }
147
148    /// Test whether a row group of decimal *mantissas* could match.
149    ///
150    /// The zone map holds mantissas; the predicate holds an `f64`. Rather
151    /// than pretend those are the same number, the bounds are widened by a
152    /// whole mantissa unit before the comparison, so a predicate value that
153    /// binary floating point cannot represent exactly — which is most of the
154    /// values anyone writes against a decimal column — can never prune a row
155    /// group that contains it. The rule from
156    /// [`may_match_f64`](Self::may_match_f64) holds here too: wrong in the
157    /// direction of more work, never in the direction of fewer rows.
158    ///
159    /// Above 2⁵³ the `i64 → f64` conversion loses more than the widening
160    /// covers, so the group is kept rather than compared wrongly. Saturated
161    /// bounds — what a mantissa beyond `i64` records — land there by
162    /// construction.
163    fn may_match_decimal(&self, rg_min: i64, rg_max: i64, scale: u8) -> bool {
164        if rg_min > rg_max {
165            return true;
166        }
167        const EXACT: i64 = 1 << 53;
168        if rg_min <= -EXACT || rg_max >= EXACT {
169            return true;
170        }
171        let pow = 10f64.powi(i32::from(scale));
172        if !pow.is_finite() || pow <= 0.0 {
173            return true;
174        }
175        #[allow(clippy::cast_precision_loss)] // bounded above by 2^53
176        let (lo, hi) = ((rg_min as f64 - 1.0) / pow, (rg_max as f64 + 1.0) / pow);
177        self.may_match_f64(lo, hi)
178    }
179}
180
181/// Reads data from a `.csx` segment file.
182///
183/// # I/O model
184///
185/// The segment file is **memory-mapped** ([`Mmap`]) so all column block
186/// reads are zero-copy slices into the OS page cache.  No per-column
187/// `File::read()` calls are issued; instead the kernel pages data in on
188/// first access and the `MADV_WILLNEED` advisory hint triggers
189/// asynchronous prefetch of upcoming blocks while the current block is
190/// being decoded.  This means sequential scans benefit from kernel-level
191/// read-ahead without any user-space prefetch thread or buffer.
192///
193/// # Predicate pushdown
194///
195/// [`read_projected_filtered_with_predicates`](Self::read_projected_filtered_with_predicates)
196/// evaluates predicates at row-group granularity using per-block
197/// [`ColumnStats`](crate::segment::stats::ColumnStats) before touching any
198/// column bytes.  This can skip entire row groups when the query's
199/// time window doesn't overlap a row group's timestamp range, or when
200/// a tag column's `value_count` proves it cannot match an equality
201/// filter.
202pub struct SegmentReader {
203    /// Memory-mapped file data (zero-copy, backed by OS page cache).
204    data: Mmap,
205    /// Parsed segment header.
206    header: SegmentHeader,
207    /// Parsed segment metadata.
208    metadata: SegmentMetadata,
209    /// File path (for error messages).
210    path: PathBuf,
211    /// Optional field-level encryption key provider for decrypting
212    /// encrypted column blocks at read time.
213    #[cfg(feature = "field-encryption")]
214    key_provider: Option<Arc<dyn crate::segment::field_encryption::FieldKeyProvider>>,
215}
216
217impl SegmentReader {
218    /// Open a segment file, validate integrity, and parse metadata.
219    ///
220    /// # Errors
221    ///
222    /// Returns an error if the file doesn't exist, is corrupt, or has an
223    /// invalid checksum.
224    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
225        let path = path.as_ref().to_path_buf();
226        let file = File::open(&path)?;
227        // SAFETY: The file is opened read-only and the mapping lives as long
228        // as `SegmentReader`. External mutation of the underlying file while
229        // mapped is undefined behaviour, but segment files are immutable
230        // after write.
231        #[allow(unsafe_code)] // audited: read-only mmap of an immutable segment file
232        let data = unsafe { Mmap::map(&file)? };
233        // Hint the OS that we will read sequentially through the
234        // mapping.  On Linux this sets MADV_SEQUENTIAL which doubles the
235        // default read-ahead window, reducing page faults on sequential scans.
236        data.advise(Advice::Sequential)?;
237
238        if data.len() < HEADER_SIZE + FOOTER_SIZE {
239            return Err(SegmentError::CorruptFile {
240                detail: format!(
241                    "file too small: {} bytes (need at least {})",
242                    data.len(),
243                    HEADER_SIZE + FOOTER_SIZE
244                ),
245            });
246        }
247
248        // Parse header
249        let header = SegmentHeader::from_bytes(&data[..HEADER_SIZE])?;
250
251        // Parse footer
252        let footer_start = data.len() - FOOTER_SIZE;
253        let footer = SegmentFooter::from_bytes(&data[footer_start..])?;
254
255        // Verify CRC32c checksum (over everything except footer)
256        let computed_checksum = crc32c::crc32c(&data[..footer_start]);
257        if computed_checksum != footer.checksum {
258            return Err(SegmentError::ChecksumMismatch {
259                expected: footer.checksum,
260                actual: computed_checksum,
261            });
262        }
263
264        // Verify independent metadata checksum
265        let metadata_start = footer.metadata_offset as usize;
266        if metadata_start < footer_start {
267            let metadata_bytes = &data[metadata_start..footer_start];
268            let computed_meta_crc = crc32c::crc32c(metadata_bytes);
269            if computed_meta_crc != footer.metadata_checksum {
270                return Err(SegmentError::ChecksumMismatch {
271                    expected: footer.metadata_checksum,
272                    actual: computed_meta_crc,
273                });
274            }
275        }
276
277        // Parse metadata section
278        if metadata_start < HEADER_SIZE {
279            return Err(SegmentError::CorruptFile {
280                detail: format!(
281                    "metadata_offset ({metadata_start}) is before end of header ({HEADER_SIZE})"
282                ),
283            });
284        }
285        if metadata_start > footer_start {
286            return Err(SegmentError::CorruptFile {
287                detail: format!(
288                    "metadata_offset ({metadata_start}) exceeds footer start ({footer_start})"
289                ),
290            });
291        }
292        let metadata = SegmentMetadata::from_bytes(&data[metadata_start..footer_start])?;
293
294        // Validate that header column_count matches the actual metadata.
295        if metadata.columns.len() != header.column_count as usize {
296            return Err(SegmentError::CorruptFile {
297                detail: format!(
298                    "column count mismatch: header says {}, metadata has {}",
299                    header.column_count,
300                    metadata.columns.len()
301                ),
302            });
303        }
304
305        Ok(Self {
306            data,
307            header,
308            metadata,
309            path,
310            #[cfg(feature = "field-encryption")]
311            key_provider: None,
312        })
313    }
314
315    /// Returns the segment header.
316    #[must_use]
317    pub fn header(&self) -> &SegmentHeader {
318        &self.header
319    }
320
321    /// Returns the column metadata.
322    #[must_use]
323    pub fn column_metadata(&self) -> &[crate::segment::metadata::ColumnMeta] {
324        &self.metadata.columns
325    }
326
327    /// Returns the number of row groups.
328    #[must_use]
329    pub fn row_group_count(&self) -> usize {
330        self.metadata.row_group_blocks.len()
331    }
332
333    /// Returns the per-column block metadata for one row group.
334    ///
335    /// Each entry carries the codec that was actually used for that block —
336    /// [`ColumnMeta::default_encoding`](crate::segment::metadata::ColumnMeta::default_encoding)
337    /// is only the writer's initial guess — along with its on-disk length.
338    /// Together they answer "which encoding won, and what did it cost", which
339    /// is what compression regressions and codec benchmarks need to see.
340    ///
341    /// Returns `None` if `row_group` is out of range.
342    #[must_use]
343    pub fn row_group_blocks(&self, row_group: usize) -> Option<&[ColumnBlockMeta]> {
344        self.metadata
345            .row_group_blocks
346            .get(row_group)
347            .map(Vec::as_slice)
348    }
349
350    /// Returns the file path.
351    #[must_use]
352    pub fn path(&self) -> &Path {
353        &self.path
354    }
355
356    /// Set a field-level encryption key provider for decrypting encrypted
357    /// column blocks.  Must be called before reading any encrypted columns.
358    #[cfg(feature = "field-encryption")]
359    pub fn set_key_provider(
360        &mut self,
361        provider: Arc<dyn crate::segment::field_encryption::FieldKeyProvider>,
362    ) {
363        self.key_provider = Some(provider);
364    }
365
366    /// Read specific columns from a specific row group as Arrow arrays.
367    ///
368    /// # Read-ahead / prefetch
369    ///
370    /// When reading column N, the reader issues an `madvise(MADV_WILLNEED)`
371    /// hint for column N+1's byte range so the OS can begin paging it in
372    /// from disk while the current column is being decoded.  Combined with
373    /// the `Sequential` advice set at open time, this reduces I/O stalls
374    /// on sequential scans.
375    ///
376    /// # Errors
377    ///
378    /// Returns an error if the row group index is out of range or a column
379    /// is not found.
380    pub fn read_columns(&self, column_names: &[&str], row_group: usize) -> Result<Vec<ArrayRef>> {
381        if row_group >= self.metadata.row_group_blocks.len() {
382            return Err(SegmentError::RowGroupOutOfRange {
383                index: row_group,
384                max: self.metadata.row_group_blocks.len().saturating_sub(1),
385            });
386        }
387
388        let rg_blocks = &self.metadata.row_group_blocks[row_group];
389
390        // Resolve all column indices and block metas upfront for prefetching.
391        let resolved: Vec<(usize, &ColumnBlockMeta)> = column_names
392            .iter()
393            .map(|&name| {
394                let col_idx = self
395                    .metadata
396                    .columns
397                    .iter()
398                    .position(|c| c.name == name)
399                    .ok_or_else(|| SegmentError::ColumnNotFound {
400                        name: name.to_string(),
401                    })?;
402                let block_meta = rg_blocks
403                    .iter()
404                    .find(|b| b.column_index == col_idx as u16)
405                    .ok_or_else(|| SegmentError::ColumnNotFound {
406                        name: name.to_string(),
407                    })?;
408                Ok((col_idx, block_meta))
409            })
410            .collect::<Result<Vec<_>>>()?;
411
412        let mut result = Vec::with_capacity(column_names.len());
413
414        for (i, &(col_idx, block_meta)) in resolved.iter().enumerate() {
415            // Prefetch the next column block via madvise(WillNeed)
416            // so the OS starts paging it in while we decode the current one.
417            if i + 1 < resolved.len() {
418                let (_, next_block) = resolved[i + 1];
419                self.prefetch_block(next_block);
420            }
421
422            // Also prefetch the same columns in the next row group if
423            // we're doing a sequential multi-row-group scan.
424            if row_group + 1 < self.metadata.row_group_blocks.len() {
425                let next_rg_blocks = &self.metadata.row_group_blocks[row_group + 1];
426                if let Some(next_rg_block) = next_rg_blocks
427                    .iter()
428                    .find(|b| b.column_index == col_idx as u16)
429                {
430                    self.prefetch_block(next_rg_block);
431                }
432            }
433
434            let array = self.read_column_block(block_meta, col_idx)?;
435            result.push(array);
436        }
437
438        Ok(result)
439    }
440
441    /// Issue an `madvise(MADV_WILLNEED)` hint for a column block's byte
442    /// range, asking the OS to begin paging the data into the page cache.
443    ///
444    /// This is a non-blocking kernel hint — the data may already be in the
445    /// page cache (in which case this is a no-op) or the kernel may begin
446    /// an asynchronous read from disk.  Either way, the call returns
447    /// immediately and never blocks the caller.
448    ///
449    /// Errors are silently ignored — prefetch is purely advisory.
450    fn prefetch_block(&self, block_meta: &ColumnBlockMeta) {
451        let start = block_meta.offset as usize;
452        let len = block_meta.length as usize;
453        if start.saturating_add(len) <= self.data.len() {
454            // madvise(MADV_WILLNEED) — non-blocking hint to the kernel.
455            if let Err(e) = self.data.advise_range(Advice::WillNeed, start, len) {
456                tracing::trace!(
457                    offset = start,
458                    length = len,
459                    error = %e,
460                    "prefetch advise failed (non-fatal)"
461                );
462                metrics::counter!("chronix_segment_prefetch_errors_total").increment(1);
463            } else {
464                metrics::counter!("chronix_segment_prefetch_hints_total").increment(1);
465                metrics::counter!("chronix_segment_prefetch_bytes_total").increment(len as u64);
466            }
467        }
468    }
469
470    /// Read all data as a single Arrow `RecordBatch`.
471    ///
472    /// This is a convenience wrapper around [`read_projected`](Self::read_projected)
473    /// that passes every column name.  When only a subset of columns is
474    /// needed, prefer calling [`read_projected`](Self::read_projected) or
475    /// [`read_projected_filtered`](Self::read_projected_filtered) directly
476    /// — they skip I/O and decompression for unrequested columns (
477    /// column pruning).
478    ///
479    /// # Errors
480    ///
481    /// Returns an error if any column data is corrupt.
482    pub fn read_all(&self) -> Result<RecordBatch> {
483        let column_names: Vec<&str> = self
484            .metadata
485            .columns
486            .iter()
487            .map(|c| c.name.as_str())
488            .collect();
489        self.read_projected(&column_names)
490    }
491
492    /// Read all columns from a single row group as a `RecordBatch`.
493    ///
494    /// This is useful for reverse-order scanning (e.g. `last_value`)
495    /// where only one row group at a time needs to be materialised.
496    ///
497    /// # Errors
498    ///
499    /// Returns an error if `row_group` is out of range or column data is corrupt.
500    pub fn read_row_group(&self, row_group: usize) -> Result<RecordBatch> {
501        let column_names: Vec<&str> = self
502            .metadata
503            .columns
504            .iter()
505            .map(|c| c.name.as_str())
506            .collect();
507        let arrays = self.read_columns(&column_names, row_group)?;
508        let fields: Vec<Field> = column_names
509            .iter()
510            .map(|&name| {
511                let col_meta = self
512                    .metadata
513                    .columns
514                    .iter()
515                    .find(|c| c.name == name)
516                    .ok_or_else(|| SegmentError::CorruptFile {
517                        detail: format!("column '{name}' not in segment metadata"),
518                    })?;
519                Ok(col_meta.arrow_field())
520            })
521            .collect::<Result<Vec<_>>>()?;
522        let schema = std::sync::Arc::new(Schema::new(fields));
523        RecordBatch::try_new(schema, arrays).map_err(|e| SegmentError::CorruptFile {
524            detail: format!("failed to build row-group RecordBatch: {e}"),
525        })
526    }
527
528    /// Read selected columns across all row groups as a single `RecordBatch`.
529    ///
530    /// Only the byte ranges for the requested columns are decoded and
531    /// decompressed — all other column data on disk is skipped entirely.
532    /// The `timestamp` column is always included even if not listed.
533    ///
534    /// # Errors
535    ///
536    /// Returns an error if any requested column name is not found in the
537    /// segment metadata, or if column data is corrupt.
538    pub fn read_projected(&self, column_names: &[&str]) -> Result<RecordBatch> {
539        self.read_projected_filtered(column_names, None)
540    }
541
542    /// Read selected columns with optional row-group-level time-range
543    /// predicate pushdown.
544    ///
545    /// When `time_range` is `Some((start_ns, end_ns))`, row groups whose
546    /// timestamp range does not overlap `[start_ns, end_ns]` are skipped
547    /// entirely — no I/O or decompression is performed for them.
548    ///
549    /// Both boundaries are **inclusive**, matching the semantics of
550    /// `chronix_query::plan::TimeRange`.
551    ///
552    /// This is the primary entry point for query-time segment reads where
553    /// the caller already knows the desired time window.
554    pub fn read_projected_filtered(
555        &self,
556        column_names: &[&str],
557        time_range: Option<(i64, i64)>,
558    ) -> Result<RecordBatch> {
559        self.read_projected_filtered_with_predicates(column_names, time_range, &[])
560    }
561
562    /// Read selected columns with row-group-level time-range and tag
563    /// predicate pushdown.
564    ///
565    /// In addition to time-range pruning, this method also pushes tag
566    /// equality predicates down to row-group level. A row group is
567    /// skipped when a tag column's per-block statistics indicate that
568    /// it has zero non-null values (`value_count == 0`), which means it
569    /// cannot possibly match an equality filter on that tag.
570    ///
571    /// # Arguments
572    ///
573    /// - `column_names` — Columns to read (empty = all).
574    /// - `time_range` — Optional inclusive time range `(start, end)`.
575    /// - `tag_predicates` — Tag equality filters as `(column_name, value)` pairs.
576    pub fn read_projected_filtered_with_predicates(
577        &self,
578        column_names: &[&str],
579        time_range: Option<(i64, i64)>,
580        tag_predicates: &[(&str, &str)],
581    ) -> Result<RecordBatch> {
582        self.read_projected_with_zone_maps(column_names, time_range, tag_predicates, &[])
583    }
584
585    /// Read selected columns with full predicate pushdown: time-range,
586    /// tag equality, and zone-map field predicates (late materialisation).
587    ///
588    /// Zone-map field predicates evaluate per-row-group column min/max
589    /// statistics to skip row groups that cannot possibly match numeric
590    /// comparisons.  This avoids decompressing and decoding column data
591    /// for entire row groups.
592    ///
593    /// # Arguments
594    ///
595    /// - `column_names` — Columns to read (empty = all).
596    /// - `time_range` — Optional inclusive time range `(start, end)`.
597    /// - `tag_predicates` — Tag equality filters as `(column_name, value)` pairs.
598    /// - `field_predicates` — Numeric field predicates for zone-map pruning.
599    pub fn read_projected_with_zone_maps(
600        &self,
601        column_names: &[&str],
602        time_range: Option<(i64, i64)>,
603        tag_predicates: &[(&str, &str)],
604        field_predicates: &[FieldPredicate],
605    ) -> Result<RecordBatch> {
606        // If column_names is empty, read ALL columns (equivalent to read_all).
607        let all_col_names: Vec<String> = if column_names.is_empty() {
608            self.metadata
609                .columns
610                .iter()
611                .map(|c| c.name.clone())
612                .collect()
613        } else {
614            Vec::new()
615        };
616
617        // Ensure timestamp is always included
618        let mut names: Vec<&str> = Vec::with_capacity(if column_names.is_empty() {
619            all_col_names.len()
620        } else {
621            column_names.len() + 1
622        });
623        if column_names.is_empty() {
624            // Read all columns
625            for c in &all_col_names {
626                names.push(c.as_str());
627            }
628        } else {
629            if !column_names.contains(&chronix_core::TIME_COLUMN) {
630                names.push(chronix_core::TIME_COLUMN);
631            }
632            for &n in column_names {
633                if !names.contains(&n) {
634                    names.push(n);
635                }
636            }
637        }
638
639        // Find the timestamp column index for row-group pruning.
640        let ts_col_idx = self
641            .metadata
642            .columns
643            .iter()
644            .position(|c| c.name == chronix_core::TIME_COLUMN);
645
646        // Resolve tag predicate column indices for row-group pushdown.
647        let tag_col_indices: Vec<usize> = tag_predicates
648            .iter()
649            .filter_map(|(col_name, _)| {
650                self.metadata
651                    .columns
652                    .iter()
653                    .position(|c| c.name == *col_name)
654            })
655            .collect();
656
657        let mut all_arrays: Vec<Vec<ArrayRef>> = vec![Vec::new(); names.len()];
658
659        // Segment-level bloom filter check.
660        // If any tag predicate value is definitely absent from the
661        // segment's bloom filter, the entire segment cannot match —
662        // return an empty RecordBatch immediately.
663        if !tag_predicates.is_empty() {
664            for &(col_name, value) in tag_predicates {
665                if let Some(col_meta) = self.metadata.columns.iter().find(|c| c.name == col_name) {
666                    if let Some(ref bf) = col_meta.bloom_filter {
667                        if !crate::segment::bloom::bloom_filter_contains(bf, value) {
668                            metrics::counter!("chronix_segment_pruned_by_bloom_filter_total")
669                                .increment(1);
670                            // Build empty RecordBatch with correct schema
671                            let mut empty_arrays: Vec<ArrayRef> = Vec::with_capacity(names.len());
672                            for &name in &names {
673                                let cm = self
674                                    .metadata
675                                    .columns
676                                    .iter()
677                                    .find(|c| c.name == name)
678                                    .ok_or_else(|| SegmentError::CorruptFile {
679                                        detail: format!("column '{name}' not in metadata"),
680                                    })?;
681                                empty_arrays.push(cm.empty_array());
682                            }
683                            let fields: Vec<Field> = names
684                                .iter()
685                                .map(|&n| {
686                                    let cm2 = self
687                                        .metadata
688                                        .columns
689                                        .iter()
690                                        .find(|c| c.name == n)
691                                        .ok_or_else(|| SegmentError::CorruptFile {
692                                            detail: format!("column '{n}' not in metadata"),
693                                        })?;
694                                    Ok(cm2.arrow_field())
695                                })
696                                .collect::<Result<Vec<_>>>()?;
697                            let schema = Schema::new(fields);
698                            return RecordBatch::try_new(std::sync::Arc::new(schema), empty_arrays)
699                                .map_err(|e| SegmentError::CorruptFile {
700                                    detail: format!("failed to create empty RecordBatch: {e}"),
701                                });
702                        }
703                    }
704                }
705            }
706        }
707
708        for rg_idx in 0..self.metadata.row_group_blocks.len() {
709            let rg_blocks = &self.metadata.row_group_blocks[rg_idx];
710
711            // Row-group-level time-range predicate pushdown:
712            // Compare the query window [start_ns, end_ns] against the row
713            // group's timestamp column stats [rg_min, rg_max].  If the
714            // intervals don't overlap, no rows in this row group can match
715            // so we skip it entirely — no I/O, no decompression, no
716            // decoding.
717            if let (Some((start_ns, end_ns)), Some(ts_idx)) = (time_range, ts_col_idx) {
718                if let Some(ts_block) = rg_blocks.iter().find(|b| b.column_index == ts_idx as u16) {
719                    let rg_min = ts_block.stats.min_value;
720                    let rg_max = ts_block.stats.max_value;
721                    // Skip if row group is entirely before or after query range.
722                    // Both boundaries are inclusive: [start_ns, end_ns]
723                    if rg_max < start_ns || rg_min > end_ns {
724                        metrics::counter!("chronix_segment_rg_pruned_by_time_range_total")
725                            .increment(1);
726                        continue;
727                    }
728                }
729            }
730
731            // Row-group-level tag predicate pushdown.
732            //
733            // For each requested tag column, inspect the per-block stats.
734            // If any tag column has `value_count == 0` (i.e. all values in
735            // this row group are NULL for that tag), the row group cannot
736            // contain a matching row — skip it without reading any data.
737            //
738            // This is conservative: a row group with non-zero value_count
739            // may still not match the equality value, but skipping it
740            // would require decoding the block (which is what we're trying
741            // to avoid).  Full value-level filtering is done post-decode
742            // by `filter_batch()`.
743            //
744            // An encrypted column's statistics are suppressed deliberately,
745            // which leaves `value_count` at zero — indistinguishable from
746            // an all-null block unless `null_count` is consulted too. Only
747            // a genuinely all-null block is impossible to match.
748            if !tag_predicates.is_empty() {
749                let mut skip = false;
750                for &tag_idx in &tag_col_indices {
751                    if let Some(tag_block) =
752                        rg_blocks.iter().find(|b| b.column_index == tag_idx as u16)
753                    {
754                        if tag_block.stats.value_count == 0 && !tag_block.stats.says_nothing() {
755                            skip = true;
756                            break;
757                        }
758                    }
759                }
760                if skip {
761                    metrics::counter!("chronix_segment_rg_pruned_by_tag_stats_total").increment(1);
762                    continue;
763                }
764            }
765
766            // Per-row-group bloom filter pruning for tag columns.
767            //
768            // If a tag column has a per-RG bloom filter and the queried value
769            // is definitely not present in this row group, skip it entirely.
770            if !tag_predicates.is_empty() {
771                let mut skip = false;
772                for &(col_name, value) in tag_predicates {
773                    if let Some(col_meta) =
774                        self.metadata.columns.iter().find(|c| c.name == col_name)
775                    {
776                        if let Some(ref rg_blooms) = col_meta.row_group_blooms {
777                            if let Some(bloom_bytes) = rg_blooms.get(rg_idx) {
778                                if !bloom_bytes.is_empty()
779                                    && !crate::segment::bloom::bloom_filter_contains(
780                                        bloom_bytes,
781                                        value,
782                                    )
783                                {
784                                    skip = true;
785                                    break;
786                                }
787                            }
788                        }
789                    }
790                }
791                if skip {
792                    metrics::counter!("chronix_segment_rg_pruned_by_bloom_total").increment(1);
793                    continue;
794                }
795            }
796
797            // Late materialisation — zone-map field predicate pushdown.
798            //
799            // For each numeric field predicate, retrieve the column's
800            // per-row-group min/max from ColumnBlockMeta.stats and check
801            // whether any row in this group could possibly match.  If the
802            // zone-map test rules out all rows, skip the entire row group.
803            if !field_predicates.is_empty() {
804                let mut skip = false;
805                for pred in field_predicates {
806                    if let Some((col_idx, col_meta)) = self
807                        .metadata
808                        .columns
809                        .iter()
810                        .enumerate()
811                        .find(|(_, c)| c.name == pred.column)
812                    {
813                        if let Some(block) =
814                            rg_blocks.iter().find(|b| b.column_index == col_idx as u16)
815                        {
816                            // Statistics that say nothing — an encrypted
817                            // block, whose stats are suppressed so the zone
818                            // map cannot leak its values — must not prune.
819                            let matches = if block.stats.says_nothing() {
820                                true
821                            } else {
822                                match col_meta.data_type {
823                                    data_types::F64 => {
824                                        let rg_min = ordered_i64_to_f64(block.stats.min_value);
825                                        let rg_max = ordered_i64_to_f64(block.stats.max_value);
826                                        pred.may_match_f64(rg_min, rg_max)
827                                    }
828                                    data_types::TIMESTAMP | data_types::I64 => pred.may_match_i64(
829                                        block.stats.min_value,
830                                        block.stats.max_value,
831                                    ),
832                                    data_types::DECIMAL => pred.may_match_decimal(
833                                        block.stats.min_value,
834                                        block.stats.max_value,
835                                        col_meta.decimal_scale.unwrap_or(0),
836                                    ),
837                                    _ => true, // non-numeric → can't prune
838                                }
839                            };
840                            if !matches {
841                                skip = true;
842                                break;
843                            }
844                        }
845                    }
846                }
847                if skip {
848                    metrics::counter!("chronix_segment_rg_pruned_by_zone_map_total").increment(1);
849                    continue;
850                }
851            }
852
853            let arrays = self.read_columns(&names, rg_idx)?;
854            for (col_idx, array) in arrays.into_iter().enumerate() {
855                all_arrays[col_idx].push(array);
856            }
857        }
858
859        // Concatenate row groups per column.
860        // If all row groups were pruned, return an empty RecordBatch
861        // with the correct schema.
862        let all_pruned = all_arrays.iter().all(std::vec::Vec::is_empty);
863
864        let mut final_arrays: Vec<ArrayRef> = Vec::with_capacity(names.len());
865
866        if all_pruned {
867            // Build empty arrays with correct types
868            for &name in &names {
869                let col_meta = self
870                    .metadata
871                    .columns
872                    .iter()
873                    .find(|c| c.name == name)
874                    .ok_or_else(|| SegmentError::CorruptFile {
875                        detail: format!("projected column '{name}' not in segment metadata"),
876                    })?;
877                final_arrays.push(col_meta.empty_array());
878            }
879        } else {
880            for (col_idx, arrays) in all_arrays.iter().enumerate() {
881                let refs: Vec<&dyn arrow::array::Array> =
882                    arrays.iter().map(std::convert::AsRef::as_ref).collect();
883                let concatenated =
884                    arrow::compute::concat(&refs).map_err(|e| SegmentError::CorruptFile {
885                        detail: format!("failed to concatenate column {}: {e}", names[col_idx]),
886                    })?;
887                final_arrays.push(concatenated);
888            }
889        }
890
891        // Build schema for only the projected columns
892        let fields: Vec<Field> = names
893            .iter()
894            .map(|&name| {
895                let col_meta = self
896                    .metadata
897                    .columns
898                    .iter()
899                    .find(|c| c.name == name)
900                    .ok_or_else(|| SegmentError::CorruptFile {
901                        detail: format!("projected column '{name}' not in segment metadata"),
902                    })?;
903                Ok(col_meta.arrow_field())
904            })
905            .collect::<Result<Vec<Field>>>()?;
906
907        let schema = Schema::new(fields);
908        let batch =
909            RecordBatch::try_new(std::sync::Arc::new(schema), final_arrays).map_err(|e| {
910                SegmentError::CorruptFile {
911                    detail: format!("failed to create RecordBatch: {e}"),
912                }
913            })?;
914
915        // All column data has been decoded into owned Arrow arrays.
916        // Hint the kernel to release the mmap pages back to the page cache.
917        // This reduces RSS during parallel scans where completed segments
918        // would otherwise hold pages until SegmentReader is dropped.
919        // SAFETY: We hold a read-only file-backed Mmap and all column data has
920        // already been decoded into owned Arrow arrays above.  No references into
921        // the mapped region remain, so releasing the pages is safe.
922        #[allow(unsafe_code)]
923        let _ = unsafe { self.data.unchecked_advise(UncheckedAdvice::DontNeed) };
924
925        Ok(batch)
926    }
927
928    /// Read and decode a single column block.
929    fn read_column_block(&self, block_meta: &ColumnBlockMeta, col_idx: usize) -> Result<ArrayRef> {
930        let start = block_meta.offset as usize;
931        let end = start
932            .checked_add(block_meta.length as usize)
933            .ok_or_else(|| SegmentError::CorruptFile {
934                detail: format!(
935                    "column block offset+length overflow: offset={start}, length={}",
936                    block_meta.length,
937                ),
938            })?;
939
940        if end > self.data.len() {
941            return Err(SegmentError::CorruptFile {
942                detail: format!(
943                    "column block extends beyond file: offset={start}, \
944                     length={}, file_size={}",
945                    block_meta.length,
946                    self.data.len()
947                ),
948            });
949        }
950
951        let raw_data = &self.data[start..end];
952
953        // Verify per-block CRC32c if present (block_crc == 0 means legacy/no checksum)
954        if block_meta.block_crc != 0 {
955            let actual_crc = crc32c::crc32c(raw_data);
956            if actual_crc != block_meta.block_crc {
957                return Err(SegmentError::CorruptFile {
958                    detail: format!(
959                        "block CRC mismatch for column {}: expected {:#010x}, got {:#010x}",
960                        col_idx, block_meta.block_crc, actual_crc,
961                    ),
962                });
963            }
964        } else {
965            tracing::warn!(
966                column = col_idx,
967                "block CRC is zero, skipping integrity check — this segment was written without per-block CRCs"
968            );
969            metrics::counter!("chronix_segment_crc_skip_total").increment(1);
970        }
971
972        // Field-level decryption: decrypt after CRC check (CRC covers the
973        // encrypted on-disk bytes) but before decompression.
974        #[cfg(feature = "field-encryption")]
975        let raw_data: std::borrow::Cow<'_, [u8]> = if block_meta.encrypted {
976            let col_meta = &self.metadata.columns[col_idx];
977            let key_id = col_meta
978                .key_id
979                .as_deref()
980                .ok_or_else(|| SegmentError::CorruptFile {
981                    detail: format!(
982                        "column {} is encrypted but has no key_id in metadata",
983                        col_meta.name,
984                    ),
985                })?;
986            let provider = self
987                .key_provider
988                .as_ref()
989                .ok_or_else(|| SegmentError::CorruptFile {
990                    detail: format!(
991                        "column {} is encrypted but no key provider is configured",
992                        col_meta.name,
993                    ),
994                })?;
995            let key = provider
996                .get_key(key_id)
997                .ok_or_else(|| SegmentError::CorruptFile {
998                    detail: format!(
999                        "key '{}' not found in key provider for encrypted column {}",
1000                        key_id, col_meta.name,
1001                    ),
1002                })?;
1003            let decrypted = crate::segment::field_encryption::decrypt_block(
1004                raw_data,
1005                &key,
1006                crate::segment::field_encryption::BlockContext {
1007                    column: &col_meta.name,
1008                    segment_created_at: self.header.created_at,
1009                },
1010            )
1011            .map_err(|e| SegmentError::CorruptFile {
1012                detail: format!("failed to decrypt column {}: {e}", col_meta.name,),
1013            })?;
1014            std::borrow::Cow::Owned(decrypted)
1015        } else {
1016            std::borrow::Cow::Borrowed(raw_data)
1017        };
1018
1019        #[cfg(not(feature = "field-encryption"))]
1020        let raw_data: &[u8] = raw_data;
1021
1022        // Decompress if needed — the dictionary is passed for ZSTD_DICT blocks.
1023        let decompressed = if block_meta.compressed {
1024            decompress_block_with_optional_dict(
1025                &raw_data,
1026                self.metadata.zstd_dictionary.as_deref(),
1027            )?
1028        } else {
1029            raw_data.to_vec()
1030        };
1031
1032        // Parse the encoded block (first byte is encoding tag)
1033        let encoded_block = EncodedBlock::from_bytes(&decompressed)?;
1034
1035        // Decode to typed values
1036        let decoded = ColumnDecoder::decode(&encoded_block)?;
1037
1038        // Verify decoded value count matches metadata expectation.
1039        let decoded_len = decoded.len();
1040        let expected = block_meta.value_count as usize;
1041        if decoded_len != expected {
1042            return Err(SegmentError::CorruptFile {
1043                detail: format!(
1044                    "decoded value count mismatch for column {col_idx}: \
1045                     metadata says {expected} values but decoded {decoded_len}",
1046                ),
1047            });
1048        }
1049
1050        // D-NULL: attach the block's validity bitmap so absent values read
1051        // back as SQL NULL rather than as a 0/""/false sentinel.
1052        let nulls = self.read_validity(block_meta)?;
1053
1054        // Convert to Arrow array
1055        let col_meta = &self.metadata.columns[col_idx];
1056        decoded_to_arrow(decoded, col_meta, nulls)
1057    }
1058
1059    /// Load a block's validity bitmap, if it has one (`.csx` v2).
1060    ///
1061    /// Returns `None` for blocks with no nulls, which store no bitmap.
1062    fn read_validity(&self, block_meta: &ColumnBlockMeta) -> Result<Option<NullBuffer>> {
1063        if block_meta.validity_length == 0 {
1064            return Ok(None);
1065        }
1066        let start = block_meta.validity_offset as usize;
1067        let end = start
1068            .checked_add(block_meta.validity_length as usize)
1069            .ok_or_else(|| SegmentError::CorruptFile {
1070                detail: "validity bitmap offset+length overflow".to_string(),
1071            })?;
1072        if end > self.data.len() {
1073            return Err(SegmentError::CorruptFile {
1074                detail: format!(
1075                    "validity bitmap extends beyond file: offset={start}, \
1076                     length={}, file_size={}",
1077                    block_meta.validity_length,
1078                    self.data.len()
1079                ),
1080            });
1081        }
1082        let value_count = block_meta.value_count as usize;
1083        let bytes = &self.data[start..end];
1084        if bytes.len() < crate::segment::validity::bitmap_len(value_count) {
1085            return Err(SegmentError::CorruptFile {
1086                detail: format!(
1087                    "validity bitmap too short: {} bytes for {value_count} values",
1088                    bytes.len()
1089                ),
1090            });
1091        }
1092        Ok(crate::segment::validity::null_buffer_from_bytes(
1093            bytes,
1094            value_count,
1095        ))
1096    }
1097}
1098
1099/// Re-attach a validity bitmap to a freshly decoded Arrow array.
1100///
1101/// The decoder always produces a dense array (nulls were stored as
1102/// sentinels in the encoded payload); this stamps the v2 validity bitmap
1103/// back on so the sentinels become real SQL NULLs.
1104fn apply_nulls<A>(array: A, nulls: Option<NullBuffer>) -> Result<A>
1105where
1106    A: arrow::array::Array + From<arrow::array::ArrayData>,
1107{
1108    let Some(nulls) = nulls else {
1109        return Ok(array);
1110    };
1111    let data = array.into_data();
1112    if nulls.len() != data.len() {
1113        return Err(SegmentError::CorruptFile {
1114            detail: format!(
1115                "validity bitmap length {} does not match column length {}",
1116                nulls.len(),
1117                data.len()
1118            ),
1119        });
1120    }
1121    // SAFETY-equivalent: `into_builder`/`build` revalidates offsets and
1122    // buffers, and the null buffer length is checked above.
1123    let data = data
1124        .into_builder()
1125        .nulls(Some(nulls))
1126        .build()
1127        .map_err(|e| SegmentError::CorruptFile {
1128            detail: format!("failed to attach validity bitmap: {e}"),
1129        })?;
1130    Ok(A::from(data))
1131}
1132
1133/// Convert decoded column values to an Arrow `ArrayRef`.
1134///
1135/// Consumes the decoded column to avoid redundant cloning.
1136fn decoded_to_arrow(
1137    decoded: DecodedColumn,
1138    col_meta: &ColumnMeta,
1139    nulls: Option<NullBuffer>,
1140) -> Result<ArrayRef> {
1141    let data_type = col_meta.data_type;
1142    match (decoded, data_type) {
1143        (DecodedColumn::I64(values), data_types::TIMESTAMP | data_types::I64) => Ok(
1144            std::sync::Arc::new(apply_nulls(Int64Array::from(values), nulls)?),
1145        ),
1146        (DecodedColumn::U64(values), data_types::U64) => Ok(std::sync::Arc::new(apply_nulls(
1147            UInt64Array::from(values),
1148            nulls,
1149        )?)),
1150        (DecodedColumn::F64(values), data_types::F64) => Ok(std::sync::Arc::new(apply_nulls(
1151            Float64Array::from(values),
1152            nulls,
1153        )?)),
1154        (DecodedColumn::Bool(values), data_types::BOOL) => Ok(std::sync::Arc::new(apply_nulls(
1155            BooleanArray::from(values),
1156            nulls,
1157        )?)),
1158        (DecodedColumn::String(values), data_types::STRING) => {
1159            let refs: Vec<&str> = values.iter().map(String::as_str).collect();
1160            Ok(std::sync::Arc::new(apply_nulls(
1161                StringArray::from(refs),
1162                nulls,
1163            )?))
1164        }
1165        (DecodedColumn::Decimal(values), data_types::DECIMAL) => {
1166            // The scale is applied here and nowhere else: the mantissas were
1167            // stored raw, and this is where they become a number again.
1168            let scale = i8::try_from(col_meta.decimal_scale.unwrap_or(0)).map_err(|_| {
1169                SegmentError::CorruptFile {
1170                    detail: format!(
1171                        "decimal column '{}' declares an unrepresentable scale",
1172                        col_meta.name
1173                    ),
1174                }
1175            })?;
1176            let array = Decimal128Array::from(values)
1177                .with_precision_and_scale(chronix_core::DECIMAL_PRECISION, scale)
1178                .map_err(|e| SegmentError::CorruptFile {
1179                    detail: format!(
1180                        "decimal column '{}' does not fit decimal(38, {scale}): {e}",
1181                        col_meta.name
1182                    ),
1183                })?;
1184            Ok(std::sync::Arc::new(apply_nulls(array, nulls)?))
1185        }
1186        (decoded, _) => Err(SegmentError::CorruptFile {
1187            detail: format!(
1188                "type mismatch: decoded variant {:?} does not match expected data_type {data_type}",
1189                std::mem::discriminant(&decoded),
1190            ),
1191        }),
1192    }
1193}
1194
1195#[cfg(test)]
1196mod zone_map_tests {
1197    use super::*;
1198
1199    fn pred(op: ZoneMapOp, value: f64) -> FieldPredicate {
1200        FieldPredicate {
1201            column: "n".into(),
1202            op,
1203            value,
1204        }
1205    }
1206
1207    /// A fractional bound must not be truncated to an integer. `n < 1.5`
1208    /// used to become `n < 1`, which prunes a block holding `{1, 2, 3}`
1209    /// even though 1 matches — a pruning step that loses rows.
1210    #[test]
1211    fn a_fractional_bound_does_not_prune_an_integer_block() {
1212        assert!(pred(ZoneMapOp::Lt, 1.5).may_match_i64(1, 3));
1213        assert!(pred(ZoneMapOp::LtEq, 1.5).may_match_i64(1, 3));
1214        assert!(pred(ZoneMapOp::Gt, 2.5).may_match_i64(1, 3));
1215        assert!(pred(ZoneMapOp::GtEq, 2.5).may_match_i64(1, 3));
1216        assert!(pred(ZoneMapOp::Eq, 2.0).may_match_i64(1, 3));
1217        // And it still prunes what it should.
1218        assert!(!pred(ZoneMapOp::Lt, 1.0).may_match_i64(1, 3));
1219        assert!(!pred(ZoneMapOp::Gt, 3.0).may_match_i64(1, 3));
1220        // An integer column cannot hold 2.5 at all.
1221        assert!(!pred(ZoneMapOp::Eq, 2.5).may_match_i64(1, 3));
1222    }
1223
1224    /// Beyond 2^53 an `i64` bound cannot be compared as an `f64`, so the
1225    /// row group is kept rather than compared wrongly.
1226    #[test]
1227    fn a_bound_beyond_exact_float_range_never_prunes() {
1228        let big = (1_i64 << 53) + 1;
1229        assert!(pred(ZoneMapOp::Eq, big as f64).may_match_i64(big, big));
1230        assert!(pred(ZoneMapOp::Lt, 0.0).may_match_i64(big, big + 10));
1231    }
1232
1233    /// Statistics that say nothing — the sentinel range an encrypted block
1234    /// carries — must read as "unknown", not as an empty range. As an empty
1235    /// range every comparison is false, so a predicate on an encrypted
1236    /// column pruned every row group and the query returned nothing.
1237    #[test]
1238    fn suppressed_statistics_never_prune() {
1239        let empty = crate::segment::stats::ColumnStats::empty();
1240        assert!(empty.says_nothing());
1241        assert!(
1242            pred(ZoneMapOp::Gt, 0.0).may_match_i64(empty.min_value, empty.max_value),
1243            "an inverted sentinel range must not prune"
1244        );
1245        assert!(pred(ZoneMapOp::Eq, 42.0).may_match_f64(f64::NAN, f64::NAN));
1246
1247        // A genuinely all-null block *is* impossible to match, and pruning
1248        // it stays correct.
1249        let mut all_null = crate::segment::stats::ColumnStats::empty();
1250        all_null.record_null();
1251        assert!(!all_null.says_nothing());
1252    }
1253}
1254
1255#[cfg(test)]
1256mod tests {
1257    use super::*;
1258    use crate::segment::writer::{SegmentWriter, SegmentWriterConfig};
1259    use chronix_core::config::{CompressionCodec, FloatEncoding};
1260    use chronix_core::types::{FieldValue, Point, SeriesKey};
1261    use std::collections::BTreeMap;
1262
1263    fn make_point(host: &str, cpu: f64, ts: i64) -> Point {
1264        let tags = BTreeMap::from([("host".to_string(), host.to_string())]);
1265        let key = SeriesKey::new("cpu_usage", tags).unwrap();
1266        let fields = BTreeMap::from([("cpu".to_string(), FieldValue::F64(cpu))]);
1267        Point::new(key, fields, ts).unwrap()
1268    }
1269
1270    #[test]
1271    fn write_and_read_segment() {
1272        let dir = tempfile::tempdir().unwrap();
1273        let path = dir.path().join("test.csx");
1274
1275        // Write
1276        let points: Vec<Point> = (0..100)
1277            .map(|i| make_point("server-1", 50.0 + i as f64, 1_000_000 + i * 10_000))
1278            .collect();
1279
1280        let config = SegmentWriterConfig {
1281            row_group_size: 50,
1282            compress: false,
1283            float_encoding: FloatEncoding::Gorilla,
1284            compression_codec: CompressionCodec::Lz4,
1285            zstd_level: 3,
1286            column_codec_overrides: std::collections::HashMap::new(),
1287            ..Default::default()
1288        };
1289
1290        let mut writer = SegmentWriter::new(&path, config).unwrap();
1291        writer.write_rows(&points).unwrap();
1292        let meta = writer.finalize().unwrap();
1293
1294        assert_eq!(meta.row_count, 100);
1295        assert_eq!(meta.row_group_count, 2); // 100 rows / 50 per group
1296
1297        // Read
1298        let reader = SegmentReader::open(&path).unwrap();
1299        assert_eq!(reader.header().row_count, 100);
1300        assert_eq!(reader.row_group_count(), 2);
1301
1302        // Read specific columns from first row group
1303        let columns = reader
1304            .read_columns(&[chronix_core::TIME_COLUMN, "cpu"], 0)
1305            .unwrap();
1306        assert_eq!(columns.len(), 2);
1307
1308        // Read all data
1309        let batch = reader.read_all().unwrap();
1310        assert_eq!(batch.num_rows(), 100);
1311        assert_eq!(batch.num_columns(), 3); // timestamp, host, cpu
1312    }
1313
1314    #[test]
1315    fn write_and_read_compressed() {
1316        let dir = tempfile::tempdir().unwrap();
1317        let path = dir.path().join("compressed.csx");
1318
1319        let points: Vec<Point> = (0..200)
1320            .map(|i| make_point("server-1", 50.0 + (i % 10) as f64, 1_000_000 + i * 10_000))
1321            .collect();
1322
1323        let config = SegmentWriterConfig {
1324            row_group_size: 100,
1325            compress: true,
1326            float_encoding: FloatEncoding::Gorilla,
1327            compression_codec: CompressionCodec::Lz4,
1328            zstd_level: 3,
1329            column_codec_overrides: std::collections::HashMap::new(),
1330            ..Default::default()
1331        };
1332
1333        let mut writer = SegmentWriter::new(&path, config).unwrap();
1334        writer.write_rows(&points).unwrap();
1335        writer.finalize().unwrap();
1336
1337        let reader = SegmentReader::open(&path).unwrap();
1338        let batch = reader.read_all().unwrap();
1339        assert_eq!(batch.num_rows(), 200);
1340    }
1341
1342    #[test]
1343    fn checksum_validation() {
1344        let dir = tempfile::tempdir().unwrap();
1345        let path = dir.path().join("checksum.csx");
1346
1347        let points = vec![make_point("s1", 1.0, 1000)];
1348        let mut writer = SegmentWriter::new(&path, SegmentWriterConfig::default()).unwrap();
1349        writer.write_rows(&points).unwrap();
1350        writer.finalize().unwrap();
1351
1352        // Corrupt one byte in the middle
1353        let mut data = std::fs::read(&path).unwrap();
1354        let mid = data.len() / 2;
1355        data[mid] ^= 0xFF;
1356        std::fs::write(&path, &data).unwrap();
1357
1358        assert!(matches!(
1359            SegmentReader::open(&path),
1360            Err(SegmentError::ChecksumMismatch { .. })
1361        ));
1362    }
1363
1364    #[test]
1365    fn column_projection() {
1366        let dir = tempfile::tempdir().unwrap();
1367        let path = dir.path().join("projection.csx");
1368
1369        let points: Vec<Point> = (0..50)
1370            .map(|i| make_point("h1", i as f64, i * 1000))
1371            .collect();
1372
1373        let config = SegmentWriterConfig {
1374            row_group_size: 100,
1375            compress: false,
1376            float_encoding: FloatEncoding::Gorilla,
1377            compression_codec: CompressionCodec::Lz4,
1378            zstd_level: 3,
1379            column_codec_overrides: std::collections::HashMap::new(),
1380            ..Default::default()
1381        };
1382
1383        let mut writer = SegmentWriter::new(&path, config).unwrap();
1384        writer.write_rows(&points).unwrap();
1385        writer.finalize().unwrap();
1386
1387        let reader = SegmentReader::open(&path).unwrap();
1388
1389        // Read only the cpu column
1390        let cols = reader.read_columns(&["cpu"], 0).unwrap();
1391        assert_eq!(cols.len(), 1);
1392
1393        let cpu_arr = cols[0].as_any().downcast_ref::<Float64Array>().unwrap();
1394        assert_eq!(cpu_arr.len(), 50);
1395    }
1396
1397    #[test]
1398    fn row_group_out_of_range() {
1399        let dir = tempfile::tempdir().unwrap();
1400        let path = dir.path().join("rg_oor.csx");
1401
1402        let points = vec![make_point("h1", 1.0, 1000)];
1403        let mut writer = SegmentWriter::new(&path, SegmentWriterConfig::default()).unwrap();
1404        writer.write_rows(&points).unwrap();
1405        writer.finalize().unwrap();
1406
1407        let reader = SegmentReader::open(&path).unwrap();
1408        assert!(reader
1409            .read_columns(&[chronix_core::TIME_COLUMN], 999)
1410            .is_err());
1411    }
1412
1413    #[test]
1414    fn column_not_found() {
1415        let dir = tempfile::tempdir().unwrap();
1416        let path = dir.path().join("col_404.csx");
1417
1418        let points = vec![make_point("h1", 1.0, 1000)];
1419        let mut writer = SegmentWriter::new(&path, SegmentWriterConfig::default()).unwrap();
1420        writer.write_rows(&points).unwrap();
1421        writer.finalize().unwrap();
1422
1423        let reader = SegmentReader::open(&path).unwrap();
1424        assert!(reader.read_columns(&["nonexistent"], 0).is_err());
1425    }
1426
1427    #[test]
1428    fn read_projected_selects_subset() {
1429        let dir = tempfile::tempdir().unwrap();
1430        let path = dir.path().join("projected.csx");
1431
1432        // Write a segment with 3 columns: timestamp, host (tag), cpu (field)
1433        let points: Vec<Point> = (0..80)
1434            .map(|i| make_point("server-1", 50.0 + i as f64, 1_000_000 + i * 10_000))
1435            .collect();
1436
1437        let config = SegmentWriterConfig {
1438            row_group_size: 50,
1439            compress: false,
1440            float_encoding: FloatEncoding::Gorilla,
1441            compression_codec: CompressionCodec::Lz4,
1442            zstd_level: 3,
1443            column_codec_overrides: std::collections::HashMap::new(),
1444            ..Default::default()
1445        };
1446
1447        let mut writer = SegmentWriter::new(&path, config).unwrap();
1448        writer.write_rows(&points).unwrap();
1449        writer.finalize().unwrap();
1450
1451        let reader = SegmentReader::open(&path).unwrap();
1452
1453        // Project only "cpu" — timestamp is auto-included
1454        let batch = reader.read_projected(&["cpu"]).unwrap();
1455        assert_eq!(batch.num_rows(), 80);
1456        assert_eq!(batch.num_columns(), 2); // timestamp + cpu
1457        assert!(batch.column_by_name(chronix_core::TIME_COLUMN).is_some());
1458        assert!(batch.column_by_name("cpu").is_some());
1459        assert!(batch.column_by_name("host").is_none());
1460
1461        // Verify values are correct
1462        let cpu_col = batch
1463            .column_by_name("cpu")
1464            .unwrap()
1465            .as_any()
1466            .downcast_ref::<Float64Array>()
1467            .unwrap();
1468        assert!((cpu_col.value(0) - 50.0).abs() < f64::EPSILON);
1469    }
1470
1471    #[test]
1472    fn read_projected_includes_timestamp_implicitly() {
1473        let dir = tempfile::tempdir().unwrap();
1474        let path = dir.path().join("proj_ts.csx");
1475
1476        let points = vec![make_point("h1", 42.0, 9999)];
1477        let mut writer = SegmentWriter::new(&path, SegmentWriterConfig::default()).unwrap();
1478        writer.write_rows(&points).unwrap();
1479        writer.finalize().unwrap();
1480
1481        let reader = SegmentReader::open(&path).unwrap();
1482
1483        // Only request "host" — should still get timestamp
1484        let batch = reader.read_projected(&["host"]).unwrap();
1485        assert_eq!(batch.num_columns(), 2); // timestamp + host
1486        assert!(batch.column_by_name(chronix_core::TIME_COLUMN).is_some());
1487    }
1488
1489    #[test]
1490    fn read_projected_all_columns_matches_read_all() {
1491        let dir = tempfile::tempdir().unwrap();
1492        let path = dir.path().join("proj_all.csx");
1493
1494        let points: Vec<Point> = (0..30)
1495            .map(|i| make_point("h1", i as f64, i * 1000))
1496            .collect();
1497
1498        let mut writer = SegmentWriter::new(&path, SegmentWriterConfig::default()).unwrap();
1499        writer.write_rows(&points).unwrap();
1500        writer.finalize().unwrap();
1501
1502        let reader = SegmentReader::open(&path).unwrap();
1503
1504        let all = reader.read_all().unwrap();
1505        let projected = reader
1506            .read_projected(&[chronix_core::TIME_COLUMN, "host", "cpu"])
1507            .unwrap();
1508
1509        assert_eq!(all.num_rows(), projected.num_rows());
1510        assert_eq!(all.num_columns(), projected.num_columns());
1511    }
1512
1513    #[test]
1514    fn read_projected_unknown_column_errors() {
1515        let dir = tempfile::tempdir().unwrap();
1516        let path = dir.path().join("proj_err.csx");
1517
1518        let points = vec![make_point("h1", 1.0, 1000)];
1519        let mut writer = SegmentWriter::new(&path, SegmentWriterConfig::default()).unwrap();
1520        writer.write_rows(&points).unwrap();
1521        writer.finalize().unwrap();
1522
1523        let reader = SegmentReader::open(&path).unwrap();
1524        assert!(reader.read_projected(&["nonexistent"]).is_err());
1525    }
1526
1527    #[test]
1528    fn read_projected_across_multiple_row_groups() {
1529        let dir = tempfile::tempdir().unwrap();
1530        let path = dir.path().join("proj_multi_rg.csx");
1531
1532        let points: Vec<Point> = (0..200)
1533            .map(|i| make_point("server-1", i as f64, i * 1000))
1534            .collect();
1535
1536        let config = SegmentWriterConfig {
1537            row_group_size: 64,
1538            compress: true,
1539            float_encoding: FloatEncoding::Gorilla,
1540            compression_codec: CompressionCodec::Lz4,
1541            zstd_level: 3,
1542            column_codec_overrides: std::collections::HashMap::new(),
1543            ..Default::default()
1544        };
1545
1546        let mut writer = SegmentWriter::new(&path, config).unwrap();
1547        writer.write_rows(&points).unwrap();
1548        writer.finalize().unwrap();
1549
1550        let reader = SegmentReader::open(&path).unwrap();
1551        assert!(reader.row_group_count() > 1);
1552
1553        let batch = reader.read_projected(&["cpu"]).unwrap();
1554        assert_eq!(batch.num_rows(), 200);
1555        assert_eq!(batch.num_columns(), 2); // timestamp + cpu
1556    }
1557
1558    #[test]
1559    fn multiple_series() {
1560        let dir = tempfile::tempdir().unwrap();
1561        let path = dir.path().join("multi_series.csx");
1562
1563        let mut points = Vec::new();
1564        for i in 0..20 {
1565            points.push(make_point("server-1", i as f64, i * 1000));
1566            points.push(make_point("server-2", i as f64 + 100.0, i * 1000));
1567        }
1568
1569        let config = SegmentWriterConfig {
1570            row_group_size: 100,
1571            compress: false,
1572            float_encoding: FloatEncoding::Gorilla,
1573            compression_codec: CompressionCodec::Lz4,
1574            zstd_level: 3,
1575            column_codec_overrides: std::collections::HashMap::new(),
1576            ..Default::default()
1577        };
1578
1579        let mut writer = SegmentWriter::new(&path, config).unwrap();
1580        writer.write_rows(&points).unwrap();
1581        let meta = writer.finalize().unwrap();
1582
1583        assert_eq!(meta.series_count, 2);
1584        assert_eq!(meta.row_count, 40);
1585
1586        let reader = SegmentReader::open(&path).unwrap();
1587        let batch = reader.read_all().unwrap();
1588        assert_eq!(batch.num_rows(), 40);
1589    }
1590
1591    // ── P1: time-range predicate pushdown tests ──────────────────────
1592
1593    /// Helper: create a segment with `n` points for a single series,
1594    /// timestamps `0, step, 2*step, …` and a small `row_group_size`
1595    /// so we get multiple row groups to prune.
1596    fn write_segment_with_rg_size(path: &std::path::Path, n: usize, step: i64, rg_size: usize) {
1597        let points: Vec<Point> = (0..n)
1598            .map(|i| make_point("server-1", i as f64, i as i64 * step))
1599            .collect();
1600
1601        let config = SegmentWriterConfig {
1602            row_group_size: rg_size,
1603            compress: false,
1604            float_encoding: FloatEncoding::Gorilla,
1605            compression_codec: CompressionCodec::Lz4,
1606            zstd_level: 3,
1607            column_codec_overrides: std::collections::HashMap::new(),
1608            ..Default::default()
1609        };
1610
1611        let mut writer = SegmentWriter::new(path, config).unwrap();
1612        writer.write_rows(&points).unwrap();
1613        writer.finalize().unwrap();
1614    }
1615
1616    #[test]
1617    fn time_range_prunes_row_groups() {
1618        // 30 points, ts = 0, 10_000, …, 290_000.
1619        // row_group_size = 10 → 3 row groups:
1620        //   rg0: ts [0, 90_000]
1621        //   rg1: ts [100_000, 190_000]
1622        //   rg2: ts [200_000, 290_000]
1623        let dir = tempfile::tempdir().unwrap();
1624        let path = dir.path().join("prune.csx");
1625        write_segment_with_rg_size(&path, 30, 10_000, 10);
1626
1627        let reader = SegmentReader::open(&path).unwrap();
1628        assert_eq!(reader.row_group_count(), 3);
1629
1630        // Query only the middle row group's time range.
1631        let batch = reader
1632            .read_projected_filtered(&["cpu"], Some((100_000, 190_000)))
1633            .unwrap();
1634        assert_eq!(batch.num_rows(), 10); // only rg1
1635
1636        // Verify timestamps are in the expected range.
1637        let ts = batch
1638            .column_by_name(chronix_core::TIME_COLUMN)
1639            .unwrap()
1640            .as_any()
1641            .downcast_ref::<Int64Array>()
1642            .unwrap();
1643        for i in 0..ts.len() {
1644            assert!(ts.value(i) >= 100_000);
1645            assert!(ts.value(i) <= 190_000);
1646        }
1647    }
1648
1649    #[test]
1650    fn time_range_inclusive_end_boundary() {
1651        // Same layout: 3 row groups.
1652        //   rg0: ts [0, 90_000]
1653        //   rg1: ts [100_000, 190_000]
1654        //   rg2: ts [200_000, 290_000]
1655        //
1656        // Query [0, 100_000] — end equals rg1.min.
1657        // With inclusive end, rg1 MUST be included (it contains ts=100_000).
1658        let dir = tempfile::tempdir().unwrap();
1659        let path = dir.path().join("boundary.csx");
1660        write_segment_with_rg_size(&path, 30, 10_000, 10);
1661
1662        let reader = SegmentReader::open(&path).unwrap();
1663
1664        let batch = reader
1665            .read_projected_filtered(&["cpu"], Some((0, 100_000)))
1666            .unwrap();
1667        // rg0 (10 rows) + rg1 (10 rows) = 20 rows
1668        // (rg1 is included because rg1.min == 100_000 <= end 100_000)
1669        assert_eq!(batch.num_rows(), 20);
1670    }
1671
1672    #[test]
1673    fn time_range_all_pruned_returns_empty() {
1674        let dir = tempfile::tempdir().unwrap();
1675        let path = dir.path().join("empty.csx");
1676        // ts range [0, 290_000]
1677        write_segment_with_rg_size(&path, 30, 10_000, 10);
1678
1679        let reader = SegmentReader::open(&path).unwrap();
1680
1681        // Query range completely after segment data.
1682        let batch = reader
1683            .read_projected_filtered(&["cpu"], Some((500_000, 600_000)))
1684            .unwrap();
1685        assert_eq!(batch.num_rows(), 0);
1686        // Schema should still have timestamp + cpu.
1687        assert_eq!(batch.num_columns(), 2);
1688        assert!(batch.column_by_name(chronix_core::TIME_COLUMN).is_some());
1689        assert!(batch.column_by_name("cpu").is_some());
1690    }
1691
1692    #[test]
1693    fn time_range_none_reads_all() {
1694        let dir = tempfile::tempdir().unwrap();
1695        let path = dir.path().join("no_filter.csx");
1696        write_segment_with_rg_size(&path, 30, 10_000, 10);
1697
1698        let reader = SegmentReader::open(&path).unwrap();
1699
1700        let batch = reader.read_projected_filtered(&["cpu"], None).unwrap();
1701        assert_eq!(batch.num_rows(), 30);
1702    }
1703
1704    #[test]
1705    fn read_row_group_returns_single_group() {
1706        let dir = tempfile::tempdir().unwrap();
1707        let path = dir.path().join("single_rg.csx");
1708        write_segment_with_rg_size(&path, 30, 10_000, 10);
1709
1710        let reader = SegmentReader::open(&path).unwrap();
1711        assert_eq!(reader.row_group_count(), 3);
1712
1713        let rg1 = reader.read_row_group(1).unwrap();
1714        assert_eq!(rg1.num_rows(), 10);
1715
1716        // Verify timestamps are from the second row group.
1717        let ts = rg1
1718            .column_by_name(chronix_core::TIME_COLUMN)
1719            .unwrap()
1720            .as_any()
1721            .downcast_ref::<Int64Array>()
1722            .unwrap();
1723        assert_eq!(ts.value(0), 100_000);
1724        assert_eq!(ts.value(9), 190_000);
1725    }
1726
1727    #[test]
1728    fn read_row_group_out_of_range_errors() {
1729        let dir = tempfile::tempdir().unwrap();
1730        let path = dir.path().join("rg_oor2.csx");
1731        write_segment_with_rg_size(&path, 10, 1000, 10);
1732
1733        let reader = SegmentReader::open(&path).unwrap();
1734        assert!(reader.read_row_group(5).is_err());
1735    }
1736
1737    // ── Bloom filter integration tests ─────────────────────────
1738
1739    /// Helper: create a segment with multiple hosts (tag values) for
1740    /// bloom filter testing.
1741    fn write_segment_with_tags(path: &std::path::Path, hosts: &[&str], points_per_host: usize) {
1742        let mut points = Vec::new();
1743        for (h_idx, &host) in hosts.iter().enumerate() {
1744            for i in 0..points_per_host {
1745                points.push(make_point(
1746                    host,
1747                    (h_idx * 100 + i) as f64,
1748                    (h_idx * points_per_host + i) as i64 * 1000,
1749                ));
1750            }
1751        }
1752
1753        let config = SegmentWriterConfig {
1754            row_group_size: 64,
1755            compress: false,
1756            float_encoding: FloatEncoding::Gorilla,
1757            compression_codec: CompressionCodec::Lz4,
1758            zstd_level: 3,
1759            column_codec_overrides: std::collections::HashMap::new(),
1760            ..Default::default()
1761        };
1762
1763        let mut writer = SegmentWriter::new(path, config).unwrap();
1764        writer.write_rows(&points).unwrap();
1765        writer.finalize().unwrap();
1766    }
1767
1768    #[test]
1769    fn bloom_filter_present_in_tag_columns() {
1770        let dir = tempfile::tempdir().unwrap();
1771        let path = dir.path().join("bloom_tag.csx");
1772        write_segment_with_tags(&path, &["server-1", "server-2", "server-3"], 10);
1773
1774        let reader = SegmentReader::open(&path).unwrap();
1775
1776        // The "host" column (role=TAG) should have a bloom filter
1777        let host_meta = reader
1778            .column_metadata()
1779            .iter()
1780            .find(|c| c.name == "host")
1781            .unwrap();
1782        assert!(
1783            host_meta.bloom_filter.is_some(),
1784            "tag column should have a bloom filter"
1785        );
1786
1787        // Non-tag columns should NOT have bloom filters
1788        let ts_meta = reader
1789            .column_metadata()
1790            .iter()
1791            .find(|c| c.name == chronix_core::TIME_COLUMN)
1792            .unwrap();
1793        assert!(ts_meta.bloom_filter.is_none());
1794
1795        let cpu_meta = reader
1796            .column_metadata()
1797            .iter()
1798            .find(|c| c.name == "cpu")
1799            .unwrap();
1800        assert!(cpu_meta.bloom_filter.is_none());
1801    }
1802
1803    #[test]
1804    fn bloom_filter_prunes_segment_for_unknown_tag() {
1805        let dir = tempfile::tempdir().unwrap();
1806        let path = dir.path().join("bloom_prune.csx");
1807        write_segment_with_tags(&path, &["server-1", "server-2"], 20);
1808
1809        let reader = SegmentReader::open(&path).unwrap();
1810
1811        // Query with a tag value that does NOT exist in the segment.
1812        // The bloom filter should prune the entire segment, returning
1813        // zero rows.
1814        let batch = reader
1815            .read_projected_filtered_with_predicates(&["cpu"], None, &[("host", "server-999")])
1816            .unwrap();
1817        assert_eq!(batch.num_rows(), 0, "bloom filter should prune unknown tag");
1818    }
1819
1820    #[test]
1821    fn bloom_filter_passes_known_tag() {
1822        let dir = tempfile::tempdir().unwrap();
1823        let path = dir.path().join("bloom_pass.csx");
1824        write_segment_with_tags(&path, &["server-1", "server-2"], 20);
1825
1826        let reader = SegmentReader::open(&path).unwrap();
1827
1828        // Query with a tag value that DOES exist — bloom filter must
1829        // not prune it (no false negatives).
1830        let batch = reader
1831            .read_projected_filtered_with_predicates(
1832                &["cpu", "host"],
1833                None,
1834                &[("host", "server-1")],
1835            )
1836            .unwrap();
1837        // Should return rows (bloom filter passes, all row groups read).
1838        // The batch has all rows (not just server-1) because row-group
1839        // level filtering doesn't do value-level matching.
1840        assert!(
1841            batch.num_rows() > 0,
1842            "bloom filter should not prune known tag"
1843        );
1844    }
1845
1846    #[test]
1847    fn bloom_filter_roundtrip_metadata() {
1848        // Verify that bloom filter data survives the write → read
1849        // round-trip through segment metadata serialization.
1850        let dir = tempfile::tempdir().unwrap();
1851        let path = dir.path().join("bloom_rt.csx");
1852        write_segment_with_tags(&path, &["alpha", "beta", "gamma"], 5);
1853
1854        let reader = SegmentReader::open(&path).unwrap();
1855        let host_meta = reader
1856            .column_metadata()
1857            .iter()
1858            .find(|c| c.name == "host")
1859            .unwrap();
1860
1861        let bf = host_meta.bloom_filter.as_ref().unwrap();
1862
1863        // All written tag values must be present
1864        assert!(crate::segment::bloom::bloom_filter_contains(bf, "alpha"));
1865        assert!(crate::segment::bloom::bloom_filter_contains(bf, "beta"));
1866        assert!(crate::segment::bloom::bloom_filter_contains(bf, "gamma"));
1867
1868        // An unwritten value should (almost certainly) be absent
1869        assert!(!crate::segment::bloom::bloom_filter_contains(bf, "delta"));
1870    }
1871
1872    #[test]
1873    #[cfg(feature = "field-encryption")]
1874    fn field_encryption_roundtrip() {
1875        use crate::segment::field_encryption::{
1876            FieldEncryptionConfig, FieldEncryptionKey, StaticKeyProvider,
1877        };
1878
1879        let dir = tempfile::tempdir().unwrap();
1880        let path = dir.path().join("encrypted.csx");
1881
1882        let key_bytes: [u8; 32] = [0xAB; 32];
1883        let enc_key = FieldEncryptionKey::new("key-1".to_string(), key_bytes);
1884
1885        let mut enc_config = FieldEncryptionConfig::default();
1886        enc_config.encrypt_column("cpu", enc_key);
1887
1888        let points: Vec<Point> = (0..100)
1889            .map(|i| make_point("server-1", 50.0 + i as f64, 1_000_000 + i * 10_000))
1890            .collect();
1891
1892        let config = SegmentWriterConfig {
1893            row_group_size: 50,
1894            compress: true,
1895            field_encryption: enc_config,
1896            ..Default::default()
1897        };
1898
1899        let mut writer = SegmentWriter::new(&path, config).unwrap();
1900        writer.write_rows(&points).unwrap();
1901        let meta = writer.finalize().unwrap();
1902        assert_eq!(meta.row_count, 100);
1903
1904        // Verify column metadata flags
1905        let reader_no_key = SegmentReader::open(&path).unwrap();
1906        let cpu_meta = reader_no_key
1907            .column_metadata()
1908            .iter()
1909            .find(|c| c.name == "cpu")
1910            .unwrap();
1911        assert!(cpu_meta.encrypted);
1912        assert_eq!(cpu_meta.key_id.as_deref(), Some("key-1"));
1913
1914        // Stats should be suppressed for encrypted columns
1915        let ts_meta = reader_no_key
1916            .column_metadata()
1917            .iter()
1918            .find(|c| c.name == chronix_core::TIME_COLUMN)
1919            .unwrap();
1920        assert!(!ts_meta.encrypted);
1921        assert!(ts_meta.stats.min_value != i64::MAX || ts_meta.stats.max_value != i64::MIN);
1922
1923        // Bloom filter should NOT exist for encrypted tag columns — but
1924        // "cpu" is a field, not a tag, so check that "host" still has a bloom.
1925        let host_meta = reader_no_key
1926            .column_metadata()
1927            .iter()
1928            .find(|c| c.name == "host")
1929            .unwrap();
1930        assert!(host_meta.bloom_filter.is_some());
1931
1932        // Read back with a key provider
1933        let mut provider = StaticKeyProvider::new();
1934        provider.add_key("key-1", key_bytes);
1935
1936        let mut reader = SegmentReader::open(&path).unwrap();
1937        reader.set_key_provider(Arc::new(provider));
1938
1939        let batch = reader.read_all().unwrap();
1940        assert_eq!(batch.num_rows(), 100);
1941
1942        let cpu_col = batch
1943            .column_by_name("cpu")
1944            .unwrap()
1945            .as_any()
1946            .downcast_ref::<Float64Array>()
1947            .unwrap();
1948        for i in 0..100 {
1949            assert!(
1950                (cpu_col.value(i) - (50.0 + i as f64)).abs() < 1e-9,
1951                "cpu mismatch at row {i}"
1952            );
1953        }
1954    }
1955
1956    /// The validity bitmap lives in the data region, so the whole-file CRC32c
1957    /// verified at `open()` must cover it — otherwise a flipped bit would
1958    /// silently turn a real value into a NULL (or vice versa) with no error.
1959    #[test]
1960    fn corrupted_validity_bitmap_is_detected() {
1961        let dir = tempfile::tempdir().unwrap();
1962        let path = dir.path().join("nulls.csx");
1963
1964        // Sparse column => at least one block carries a validity bitmap.
1965        let tags = BTreeMap::from([("host".to_string(), "h".to_string())]);
1966        let key = SeriesKey::new("m", tags).unwrap();
1967        let points: Vec<Point> = (0..40i64)
1968            .map(|i| {
1969                let fields = if i % 2 == 0 {
1970                    BTreeMap::from([("v".to_string(), FieldValue::F64(i as f64))])
1971                } else {
1972                    BTreeMap::from([("w".to_string(), FieldValue::F64(1.0))])
1973                };
1974                Point::new(key.clone(), fields, 1_000 + i * 10).unwrap()
1975            })
1976            .collect();
1977
1978        let mut writer = SegmentWriter::new(&path, SegmentWriterConfig::default()).unwrap();
1979        writer.write_rows(&points).unwrap();
1980        writer.finalize().unwrap();
1981
1982        // Locate a bitmap and flip a bit inside it.
1983        let reader = SegmentReader::open(&path).unwrap();
1984        let (off, len) = reader
1985            .metadata
1986            .row_group_blocks
1987            .iter()
1988            .flatten()
1989            .find(|b| b.validity_length > 0)
1990            .map(|b| (b.validity_offset as usize, b.validity_length))
1991            .expect("a sparse column must have written a validity bitmap");
1992        assert!(len > 0);
1993        drop(reader);
1994
1995        let mut bytes = std::fs::read(&path).unwrap();
1996        bytes[off] ^= 0b0000_0001;
1997        std::fs::write(&path, &bytes).unwrap();
1998
1999        assert!(
2000            matches!(
2001                SegmentReader::open(&path),
2002                Err(SegmentError::ChecksumMismatch { .. })
2003            ),
2004            "a corrupted validity bitmap must fail the file checksum"
2005        );
2006    }
2007
2008    /// D-NULL + field encryption: the validity bitmap is stored *outside*
2009    /// the encrypted payload, so a reader without the key must still be
2010    /// blocked on values while a reader with the key sees correct nulls.
2011    #[test]
2012    #[cfg(feature = "field-encryption")]
2013    fn encrypted_column_preserves_nulls() {
2014        use crate::segment::field_encryption::{
2015            FieldEncryptionConfig, FieldEncryptionKey, StaticKeyProvider,
2016        };
2017        use arrow::array::Array;
2018
2019        let dir = tempfile::tempdir().unwrap();
2020        let path = dir.path().join("encrypted_nulls.csx");
2021
2022        let enc_key = FieldEncryptionKey::new("key-1".to_string(), [0xCD; 32]);
2023        let mut enc_config = FieldEncryptionConfig::default();
2024        enc_config.encrypt_column("cpu", enc_key);
2025
2026        // Every other point omits `cpu` entirely.
2027        let tags = BTreeMap::from([("host".to_string(), "server-1".to_string())]);
2028        let key = SeriesKey::new("cpu_usage", tags).unwrap();
2029        let points: Vec<Point> = (0..40i64)
2030            .map(|i| {
2031                let fields = if i % 2 == 0 {
2032                    BTreeMap::from([("cpu".to_string(), FieldValue::F64(i as f64))])
2033                } else {
2034                    BTreeMap::from([("other".to_string(), FieldValue::F64(1.0))])
2035                };
2036                Point::new(key.clone(), fields, 1_000_000 + i * 10_000).unwrap()
2037            })
2038            .collect();
2039
2040        let config = SegmentWriterConfig {
2041            row_group_size: 16,
2042            compress: true,
2043            field_encryption: enc_config,
2044            ..Default::default()
2045        };
2046        let mut writer = SegmentWriter::new(&path, config).unwrap();
2047        writer.write_rows(&points).unwrap();
2048        writer.finalize().unwrap();
2049
2050        let mut provider = StaticKeyProvider::new();
2051        provider.add_key("key-1", [0xCD; 32]);
2052        let mut reader = SegmentReader::open(&path).unwrap();
2053        reader.set_key_provider(Arc::new(provider));
2054        let batch = reader.read_all().unwrap();
2055
2056        let cpu = batch
2057            .column_by_name("cpu")
2058            .expect("cpu column")
2059            .as_any()
2060            .downcast_ref::<arrow::array::Float64Array>()
2061            .expect("f64");
2062        assert_eq!(cpu.null_count(), 20, "half the rows have no cpu value");
2063        // And the surviving values must be intact, not shifted by the bitmap.
2064        let present: Vec<f64> = (0..cpu.len())
2065            .filter(|i| cpu.is_valid(*i))
2066            .map(|i| cpu.value(i))
2067            .collect();
2068        assert_eq!(present.len(), 20);
2069        assert!(
2070            present.iter().all(|v| *v % 2.0 == 0.0),
2071            "values corrupted: {present:?}"
2072        );
2073    }
2074
2075    #[test]
2076    #[cfg(feature = "field-encryption")]
2077    fn field_encryption_wrong_key_fails() {
2078        use crate::segment::field_encryption::{
2079            FieldEncryptionConfig, FieldEncryptionKey, StaticKeyProvider,
2080        };
2081
2082        let dir = tempfile::tempdir().unwrap();
2083        let path = dir.path().join("enc_wrong_key.csx");
2084
2085        let key_bytes: [u8; 32] = [0xCD; 32];
2086        let enc_key = FieldEncryptionKey::new("key-2".to_string(), key_bytes);
2087
2088        let mut enc_config = FieldEncryptionConfig::default();
2089        enc_config.encrypt_column("cpu", enc_key);
2090
2091        let points: Vec<Point> = (0..20)
2092            .map(|i| make_point("server-1", 10.0 + i as f64, 1_000_000 + i * 10_000))
2093            .collect();
2094
2095        let config = SegmentWriterConfig {
2096            row_group_size: 100,
2097            compress: false,
2098            field_encryption: enc_config,
2099            ..Default::default()
2100        };
2101
2102        let mut writer = SegmentWriter::new(&path, config).unwrap();
2103        writer.write_rows(&points).unwrap();
2104        writer.finalize().unwrap();
2105
2106        // Wrong key — decryption should fail
2107        let mut wrong_provider = StaticKeyProvider::new();
2108        wrong_provider.add_key("key-2", [0xFF; 32]);
2109
2110        let mut reader = SegmentReader::open(&path).unwrap();
2111        reader.set_key_provider(Arc::new(wrong_provider));
2112
2113        let result = reader.read_all();
2114        assert!(result.is_err(), "reading with wrong key should fail");
2115    }
2116
2117    #[test]
2118    #[cfg(feature = "field-encryption")]
2119    fn field_encryption_no_provider_fails() {
2120        use crate::segment::field_encryption::{FieldEncryptionConfig, FieldEncryptionKey};
2121
2122        let dir = tempfile::tempdir().unwrap();
2123        let path = dir.path().join("enc_no_provider.csx");
2124
2125        let key_bytes: [u8; 32] = [0xEF; 32];
2126        let enc_key = FieldEncryptionKey::new("key-3".to_string(), key_bytes);
2127
2128        let mut enc_config = FieldEncryptionConfig::default();
2129        enc_config.encrypt_column("cpu", enc_key);
2130
2131        let points: Vec<Point> = (0..10)
2132            .map(|i| make_point("server-1", 1.0, 1_000_000 + i * 10_000))
2133            .collect();
2134
2135        let config = SegmentWriterConfig {
2136            row_group_size: 100,
2137            compress: false,
2138            field_encryption: enc_config,
2139            ..Default::default()
2140        };
2141
2142        let mut writer = SegmentWriter::new(&path, config).unwrap();
2143        writer.write_rows(&points).unwrap();
2144        writer.finalize().unwrap();
2145
2146        // No key provider at all
2147        let reader = SegmentReader::open(&path).unwrap();
2148        let result = reader.read_all();
2149        assert!(
2150            result.is_err(),
2151            "reading encrypted column without key provider should fail"
2152        );
2153    }
2154
2155    // ── Per-row-group bloom filter tests ─────────────────
2156
2157    #[test]
2158    fn row_group_bloom_filters_present_for_tag_columns() {
2159        let dir = tempfile::tempdir().unwrap();
2160        let path = dir.path().join("rg_bloom.csx");
2161
2162        // 30 points across 3 hosts, sorted by time → hosts interleave across
2163        // row groups when row_group_size=10.
2164        write_segment_with_tags(&path, &["alpha", "beta", "gamma"], 10);
2165
2166        let reader = SegmentReader::open(&path).unwrap();
2167        let host_meta = reader
2168            .column_metadata()
2169            .iter()
2170            .find(|c| c.name == "host")
2171            .unwrap();
2172
2173        // Per-RG blooms should exist for tag columns
2174        assert!(
2175            host_meta.row_group_blooms.is_some(),
2176            "tag column should have per-row-group bloom filters"
2177        );
2178
2179        let rg_blooms = host_meta.row_group_blooms.as_ref().unwrap();
2180        assert_eq!(rg_blooms.len(), reader.row_group_count());
2181
2182        // Each bloom should be non-empty
2183        for (i, bloom) in rg_blooms.iter().enumerate() {
2184            assert!(!bloom.is_empty(), "bloom for rg{i} should be non-empty");
2185        }
2186    }
2187
2188    #[test]
2189    fn row_group_bloom_prunes_non_matching_row_groups() {
2190        // Create a segment where host "alpha" is ONLY in the first row group
2191        // and host "beta" is ONLY in the second row group.
2192        //
2193        // Points are written in order (alpha first, beta second) with
2194        // increasing timestamps, so the writer's sort will keep them in
2195        // their respective row groups.
2196        let dir = tempfile::tempdir().unwrap();
2197        let path = dir.path().join("rg_bloom_prune.csx");
2198
2199        let mut points = Vec::new();
2200        // First 10 points: host=alpha, ts=[0, 9000]
2201        for i in 0..10 {
2202            points.push(make_point("alpha", i as f64, i as i64 * 1000));
2203        }
2204        // Next 10 points: host=beta, ts=[10000, 19000]
2205        for i in 0..10 {
2206            points.push(make_point("beta", (100 + i) as f64, (10 + i) as i64 * 1000));
2207        }
2208
2209        let config = SegmentWriterConfig {
2210            row_group_size: 10,
2211            compress: false,
2212            float_encoding: FloatEncoding::Gorilla,
2213            compression_codec: CompressionCodec::Lz4,
2214            zstd_level: 3,
2215            column_codec_overrides: std::collections::HashMap::new(),
2216            ..Default::default()
2217        };
2218
2219        let mut writer = SegmentWriter::new(&path, config).unwrap();
2220        writer.write_rows(&points).unwrap();
2221        writer.finalize().unwrap();
2222
2223        let reader = SegmentReader::open(&path).unwrap();
2224        assert_eq!(reader.row_group_count(), 2);
2225
2226        // Verify per-RG blooms: rg0 should contain "alpha" but NOT "beta"
2227        let host_meta = reader
2228            .column_metadata()
2229            .iter()
2230            .find(|c| c.name == "host")
2231            .unwrap();
2232        let rg_blooms = host_meta.row_group_blooms.as_ref().unwrap();
2233
2234        assert!(crate::segment::bloom::bloom_filter_contains(
2235            &rg_blooms[0],
2236            "alpha"
2237        ));
2238        assert!(!crate::segment::bloom::bloom_filter_contains(
2239            &rg_blooms[0],
2240            "beta"
2241        ));
2242        assert!(crate::segment::bloom::bloom_filter_contains(
2243            &rg_blooms[1],
2244            "beta"
2245        ));
2246        assert!(!crate::segment::bloom::bloom_filter_contains(
2247            &rg_blooms[1],
2248            "alpha"
2249        ));
2250
2251        // Query for "alpha" — should only return rg0 (10 rows)
2252        let batch = reader
2253            .read_projected_filtered_with_predicates(&["cpu", "host"], None, &[("host", "alpha")])
2254            .unwrap();
2255        assert_eq!(
2256            batch.num_rows(),
2257            10,
2258            "only rg0 should be returned for alpha"
2259        );
2260
2261        // Query for "beta" — should only return rg1 (10 rows)
2262        let batch = reader
2263            .read_projected_filtered_with_predicates(&["cpu", "host"], None, &[("host", "beta")])
2264            .unwrap();
2265        assert_eq!(batch.num_rows(), 10, "only rg1 should be returned for beta");
2266
2267        // Query for "gamma" (non-existent) — should return 0 rows
2268        let batch = reader
2269            .read_projected_filtered_with_predicates(&["cpu"], None, &[("host", "gamma")])
2270            .unwrap();
2271        assert_eq!(batch.num_rows(), 0, "non-existent tag should prune all RGs");
2272    }
2273}