Skip to main content

chronix_engine/segment/
writer.rs

1//! Segment writer — builds `.csx` segment files.
2//!
3//! The writer accepts rows (as [`Point`]s), sorts them by series key and
4//! timestamp, splits them into row groups, encodes each column with
5//! `chronix-encoding`, optionally compresses with LZ4, and writes
6//! the segment file using **streaming I/O with incremental CRC-32c**.
7//!
8//! # Streaming Write Architecture
9//!
10//! Row groups are encoded one at a time and flushed directly to a temp
11//! file via `BufWriter`. An incremental CRC-32c checksum (`crc32c_append`)
12//! is updated as each chunk is written. Peak memory is bounded by
13//! `O(row_group_size)` — typically 1-10 MiB — rather than `O(segment_size)`.
14//!
15//! The write path uses temp-file → `fsync` → `rename` → dir-fsync for
16//! crash-safe atomic visibility.
17//!
18//! # Zstd dictionary compression
19//!
20//! When `zstd_dict_training` is enabled and the compression codec is `Zstd`,
21//! the writer collects encoded column blocks from the first row group as
22//! training samples, trains a dictionary via `zstd::dict::from_samples`,
23//! and compresses subsequent row groups with the trained dictionary.
24//! The dictionary is embedded in the segment metadata for the reader.
25//!
26//! This improves compression ratio by 20–40% on small blocks (< 64 KiB)
27//! where column values have repeating structure (e.g. string tags, enum fields).
28//!
29//! # Usage
30//!
31//! ```no_run
32//! # use chronix_engine::segment::writer::{SegmentWriter, SegmentWriterConfig};
33//! let config = SegmentWriterConfig::default();
34//! let mut writer = SegmentWriter::new("/tmp/segment.csx", config).unwrap();
35//! // writer.write_rows(&points).unwrap();
36//! // let meta = writer.finalize().unwrap();
37//! ```
38
39use std::collections::{BTreeMap, BTreeSet, HashMap};
40use std::io::{BufWriter, Write};
41use std::path::{Path, PathBuf};
42
43use chronix_core::config::{CompressionCodec, FloatEncoding};
44use chronix_core::types::{FieldValue, Point};
45use chronix_encoding::{ColumnEncoder, EncodedBlock};
46
47use arrow::array::{
48    Array, BooleanArray, Decimal128Array, Float64Array, Int64Array, StringArray, UInt64Array,
49};
50use arrow::datatypes::DataType;
51use arrow::record_batch::RecordBatch;
52
53use crate::segment::compression::{
54    compress_block_with_codec, compress_block_zstd_dict, should_skip_compression,
55    train_zstd_dictionary,
56};
57use crate::segment::error::{Result, SegmentError};
58use crate::segment::header::{SegmentFooter, SegmentHeader, VERSION};
59use crate::segment::metadata::{ColumnBlockMeta, ColumnMeta, SegmentMetadata, data_types, roles};
60use crate::segment::stats::ColumnStats;
61use crate::segment::validity::ValidityBuilder;
62
63/// Output of a row-group write pass:
64/// (column metas, per-row-group block metas, rows written, optional trained zstd dictionary).
65type RowGroupWriteOutput = (
66    Vec<ColumnMeta>,
67    Vec<Vec<ColumnBlockMeta>>,
68    u64,
69    Option<Vec<u8>>,
70);
71
72/// Default maximum rows per row group.
73pub const DEFAULT_ROW_GROUP_SIZE: usize = 65_536;
74
75/// Default target segment file size (256 MiB).
76pub const DEFAULT_TARGET_SEGMENT_SIZE_BYTES: usize = 256 * 1024 * 1024;
77
78/// Configuration for the segment writer.
79#[derive(Debug, Clone)]
80pub struct SegmentWriterConfig {
81    /// Maximum rows per row group.
82    pub row_group_size: usize,
83    /// Maximum rows per row group for memory-limiting purposes.
84    ///
85    /// When set, the writer will split row groups to ensure that no single
86    /// row group exceeds this number of rows. This bounds peak memory usage
87    /// during encoding: only `max_row_group_size` rows worth of encoded
88    /// column buffers are held in memory at any time.
89    ///
90    /// Defaults to [`DEFAULT_ROW_GROUP_SIZE`] (65 536 rows). Reducing this
91    /// trades write throughput (more row groups, more metadata) for lower
92    /// peak memory consumption.
93    pub max_row_group_size: usize,
94    /// Whether to apply compression after encoding.
95    pub compress: bool,
96    /// Float encoding strategy.
97    pub float_encoding: FloatEncoding,
98    /// Compression codec (LZ4, Zstd, or None).
99    pub compression_codec: CompressionCodec,
100    /// Zstd compression level (1–22). Only used when `compression_codec` is Zstd.
101    pub zstd_level: i32,
102    /// Per-column codec overrides. Keys are column names; values override
103    /// the default `compression_codec` for that column.
104    pub column_codec_overrides: HashMap<String, CompressionCodec>,
105    /// Advisory maximum segment file size in bytes (default: 256 MiB).
106    ///
107    /// When the estimated segment size exceeds this threshold, the flush
108    /// logic in the database layer should split the data into multiple
109    /// segments. This keeps individual segment files manageable for
110    /// compaction, object-store upload, and cache eviction.
111    ///
112    /// **Note:** This is an advisory limit — the writer itself produces a
113    /// single segment file. Size-based splitting is enforced by the flush
114    /// layer (`chronix-memtable::flush::flush_frozen`), which estimates
115    /// the compressed row size and splits large measurement batches into
116    /// multiple `SegmentWriter` instances.
117    pub target_segment_size_bytes: usize,
118    /// Bloom filter false-positive rate (default: 0.01 = 1%).
119    ///
120    /// Lower values improve query pruning at the cost of larger bloom
121    /// filter bitmaps.  Recommended values:
122    /// Low cardinality (<100 unique): 0.001
123    /// Default workloads: 0.01
124    /// High cardinality (>100K unique): 0.05
125    pub bloom_fpr: f64,
126    /// Enable Zstd dictionary training.
127    ///
128    /// When `true` **and** `compression_codec` is `Zstd`, the writer
129    /// collects encoded column blocks from the first row group as
130    /// training samples, trains a dictionary via `zstd::dict::from_samples`,
131    /// and compresses subsequent row groups with the trained dictionary.
132    /// The dictionary is embedded in the segment metadata for the reader.
133    ///
134    /// Most beneficial when row groups are small (< 64 KiB) and values
135    /// contain repeating structure (e.g. string tags, enum fields).
136    pub zstd_dict_training: bool,
137    /// Per-column field-level encryption configuration.
138    ///
139    /// Only columns listed here are encrypted with AES-256-GCM.
140    /// Encrypted columns have their statistics (min/max/sum/distinct)
141    /// zeroed to prevent information leakage through zone maps.
142    #[cfg(feature = "field-encryption")]
143    pub field_encryption: crate::segment::field_encryption::FieldEncryptionConfig,
144}
145
146impl Default for SegmentWriterConfig {
147    fn default() -> Self {
148        Self {
149            row_group_size: DEFAULT_ROW_GROUP_SIZE,
150            max_row_group_size: DEFAULT_ROW_GROUP_SIZE,
151            compress: true,
152            float_encoding: FloatEncoding::Chimp,
153            compression_codec: CompressionCodec::Lz4,
154            zstd_level: 3,
155            column_codec_overrides: HashMap::new(),
156            target_segment_size_bytes: DEFAULT_TARGET_SEGMENT_SIZE_BYTES,
157            bloom_fpr: 0.01,
158            zstd_dict_training: false,
159            #[cfg(feature = "field-encryption")]
160            field_encryption: crate::segment::field_encryption::FieldEncryptionConfig::default(),
161        }
162    }
163}
164
165/// Metadata returned after a segment is finalized.
166#[derive(Debug, Clone, PartialEq)]
167pub struct SegmentMeta {
168    /// Path to the segment file.
169    pub path: PathBuf,
170    /// Minimum timestamp in the segment.
171    pub min_timestamp: i64,
172    /// Maximum timestamp in the segment.
173    pub max_timestamp: i64,
174    /// Total number of rows.
175    pub row_count: u64,
176    /// Number of unique series.
177    pub series_count: u32,
178    /// File size in bytes.
179    pub byte_size: u64,
180    /// Number of row groups.
181    pub row_group_count: u32,
182    /// Number of columns.
183    pub column_count: u16,
184    /// Total uncompressed payload bytes (before LZ4/Zstd), for compression ratio tracking.
185    pub uncompressed_bytes: u64,
186    /// Segment header (available from the write path, avoids re-reading the file).
187    pub header: crate::segment::header::SegmentHeader,
188    /// Per-column metadata with stats (available from the write path).
189    pub column_metas: Vec<crate::segment::metadata::ColumnMeta>,
190    /// The distinct series the segment holds, in series-major order — the
191    /// writer knows them, so nothing has to decode the segment to learn them.
192    pub series_keys: Vec<chronix_core::SeriesKey>,
193}
194
195/// Builds a `.csx` segment file from time-series data points.
196///
197/// Points are accumulated across multiple [`write_rows`](Self::write_rows)
198/// calls and encoded during [`finalize`](Self::finalize).
199pub struct SegmentWriter {
200    path: PathBuf,
201    config: SegmentWriterConfig,
202    /// Accumulated points from all `write_rows` calls.
203    points: Vec<Point>,
204    finalized: bool,
205}
206
207impl SegmentWriter {
208    /// Create a new segment writer that will write to `path`.
209    ///
210    /// # Errors
211    ///
212    /// Returns an I/O error if the directory cannot be determined.
213    pub fn new(path: impl AsRef<Path>, config: SegmentWriterConfig) -> Result<Self> {
214        Ok(Self {
215            path: path.as_ref().to_path_buf(),
216            config,
217            points: Vec::new(),
218            finalized: false,
219        })
220    }
221
222    /// Write a batch of rows to the segment.
223    ///
224    /// Can be called multiple times — all rows are accumulated and encoded
225    /// together during [`finalize`](Self::finalize).
226    ///
227    /// # Errors
228    ///
229    /// Returns [`SegmentError::AlreadyFinalized`] if the segment has been
230    /// finalized.
231    pub fn write_rows(&mut self, points: &[Point]) -> Result<()> {
232        if self.finalized {
233            return Err(SegmentError::AlreadyFinalized);
234        }
235
236        self.points.extend_from_slice(points);
237        Ok(())
238    }
239
240    /// Returns the number of points accumulated so far.
241    #[must_use]
242    pub fn buffered_point_count(&self) -> usize {
243        self.points.len()
244    }
245
246    /// Finalize the segment: sort, encode, and stream to disk with
247    /// incremental CRC-32c.
248    ///
249    /// # Streaming Write Architecture
250    ///
251    /// Row groups are encoded one at a time and written directly to a
252    /// temp file via `BufWriter<File>`. An incremental CRC-32c checksum
253    /// is updated as each chunk is written. Peak memory is bounded by
254    /// `O(row_group_size)` — typically 1-10 MiB — not `O(segment_size)`.
255    ///
256    /// The write path uses temp-file → `fsync` → `rename` → dir-fsync
257    /// for crash-safe atomic visibility.
258    ///
259    /// # Errors
260    ///
261    /// Returns an error if the segment is empty or already finalized.
262    pub fn finalize(&mut self) -> Result<SegmentMeta> {
263        if self.finalized {
264            return Err(SegmentError::AlreadyFinalized);
265        }
266
267        if self.points.is_empty() {
268            return Err(SegmentError::EmptySegment);
269        }
270
271        let points = &self.points;
272
273        // Discover the column schema from the data. The schema already has
274        // tag columns sorted by ascending cardinality from discover_schema().
275        let schema = discover_schema(points);
276
277        // Extract the cardinality-ordered tag names directly from the
278        // schema — discover_schema() already sorted tags by cardinality,
279        // so we avoid recomputing distinct values.
280        let tag_cardinality_order: Vec<&str> = schema
281            .iter()
282            .filter(|c| c.role == roles::TAG)
283            .map(|c| c.name.as_str())
284            .collect();
285
286        // Sort rows by (measurement, lowest-cardinality tag first … , timestamp).
287        // Uses a borrow-based comparator to avoid per-row String allocations.
288        let mut sorted_indices: Vec<usize> = (0..points.len()).collect();
289        sorted_indices.sort_unstable_by(|&a, &b| {
290            let pa = &points[a];
291            let pb = &points[b];
292            // Compare measurement
293            let cmp = pa
294                .series_key()
295                .measurement()
296                .cmp(pb.series_key().measurement());
297            if cmp != std::cmp::Ordering::Equal {
298                return cmp;
299            }
300            // Compare tags in cardinality order
301            for tag_name in &tag_cardinality_order {
302                let va = pa.series_key().tag(tag_name).unwrap_or("");
303                let vb = pb.series_key().tag(tag_name).unwrap_or("");
304                let cmp = va.cmp(vb);
305                if cmp != std::cmp::Ordering::Equal {
306                    return cmp;
307                }
308            }
309            // Finally compare timestamps
310            pa.timestamp().cmp(&pb.timestamp())
311        });
312
313        // Compute global stats
314        let (min_ts, max_ts, series_keys) = compute_global_stats(points, &sorted_indices);
315        let series_count = u32::try_from(series_keys.len()).unwrap_or(u32::MAX);
316
317        let total_rows = sorted_indices.len();
318
319        let column_count = u16::try_from(schema.len()).map_err(|_| SegmentError::CorruptFile {
320            detail: format!("column count {} exceeds u16::MAX", schema.len()),
321        })?;
322
323        // Build header
324        let header = SegmentHeader {
325            version: VERSION,
326            flags: 0,
327            created_at: std::time::SystemTime::now()
328                .duration_since(std::time::UNIX_EPOCH)
329                .map(|d| i64::try_from(d.as_nanos()).unwrap_or(i64::MAX))
330                .unwrap_or(0),
331            min_timestamp: min_ts,
332            max_timestamp: max_ts,
333            row_count: total_rows as u64,
334            column_count,
335            series_count,
336            compression: match self.config.compression_codec {
337                CompressionCodec::Lz4 => 0,
338                CompressionCodec::Zstd => 1,
339                CompressionCodec::None => 2,
340            },
341            sort_order: 1,
342        };
343
344        // ── Streaming write path ───────────────────────────────────
345        if let Some(parent) = self.path.parent() {
346            std::fs::create_dir_all(parent)?;
347        }
348
349        let temp_path = {
350            let mut p = self.path.clone().into_os_string();
351            p.push(".tmp");
352            PathBuf::from(p)
353        };
354
355        let result = (|| -> Result<(SegmentMeta, u64)> {
356            let file = std::fs::File::create(&temp_path)?;
357            let mut writer = BufWriter::new(file);
358            let mut crc: u32 = 0;
359            let mut pos: u64 = 0;
360
361            // Write header
362            let header_bytes = header.to_bytes();
363            writer.write_all(&header_bytes)?;
364            crc = crc32c::crc32c_append(crc, &header_bytes);
365            pos += header_bytes.len() as u64;
366
367            // Encode and stream row groups one at a time
368            let (column_metas, row_group_blocks, uncompressed_bytes, zstd_dictionary) =
369                encode_row_groups_streaming(
370                    &schema,
371                    &sorted_indices,
372                    points,
373                    &self.config,
374                    header.created_at,
375                    &mut writer,
376                    &mut crc,
377                    &mut pos,
378                )?;
379
380            let metadata = SegmentMetadata {
381                columns: column_metas,
382                row_group_blocks,
383                zstd_dictionary,
384            };
385
386            // Write metadata
387            let metadata_offset = pos;
388            let metadata_bytes = metadata.to_bytes()?;
389            let metadata_checksum = crc32c::crc32c(&metadata_bytes);
390            writer.write_all(&metadata_bytes)?;
391            crc = crc32c::crc32c_append(crc, &metadata_bytes);
392            pos += metadata_bytes.len() as u64;
393
394            // Write footer with accumulated CRC
395            let footer = SegmentFooter {
396                metadata_offset,
397                row_group_count: metadata.row_group_blocks.len() as u32,
398                checksum: crc,
399                metadata_checksum,
400            };
401            let footer_bytes = footer.to_bytes();
402            writer.write_all(&footer_bytes)?;
403            pos += footer_bytes.len() as u64;
404
405            // Flush and fsync
406            writer.flush()?;
407            writer.get_ref().sync_all()?;
408            drop(writer);
409
410            // Atomic rename
411            std::fs::rename(&temp_path, &self.path)?;
412
413            Ok((
414                SegmentMeta {
415                    path: self.path.clone(),
416                    min_timestamp: header.min_timestamp,
417                    max_timestamp: header.max_timestamp,
418                    row_count: header.row_count,
419                    series_count: header.series_count,
420                    byte_size: pos,
421                    row_group_count: footer.row_group_count,
422                    column_count: header.column_count,
423                    uncompressed_bytes,
424                    header,
425                    column_metas: metadata.columns,
426                    series_keys,
427                },
428                pos,
429            ))
430        })();
431
432        match result {
433            Ok((meta, _)) => {
434                self.finalized = true;
435
436                // Fsync parent directory for rename durability.
437                // Failure here means the rename may not survive a power
438                // loss, so propagate the error to the caller.
439                if let Some(parent) = self.path.parent() {
440                    let dir = std::fs::File::open(parent).map_err(|e| {
441                        SegmentError::Io(std::io::Error::new(
442                            e.kind(),
443                            format!("failed to open parent dir for fsync: {e}"),
444                        ))
445                    })?;
446                    dir.sync_all().map_err(|e| {
447                        SegmentError::Io(std::io::Error::new(
448                            e.kind(),
449                            format!("parent directory fsync failed: {e}"),
450                        ))
451                    })?;
452                }
453
454                Ok(meta)
455            }
456            Err(e) => {
457                let _ = std::fs::remove_file(&temp_path);
458                Err(e)
459            }
460        }
461    }
462
463    /// Finalize from a **pre-sorted, pre-deduped** `RecordBatch`.
464    ///
465    /// This is the zero-copy write path used by the compaction executor.
466    /// Because the caller has already sorted and deduped the data, this
467    /// method skips schema discovery, sorting, and — most importantly —
468    /// avoids the `batch_to_points` round-trip that would allocate a
469    /// `BTreeMap` per row.
470    ///
471    /// Column data is extracted directly from Arrow arrays and passed
472    /// to the column encoders. Uses streaming I/O with incremental
473    /// CRC-32c.
474    ///
475    /// # Errors
476    ///
477    /// Returns an error if the batch is empty, the writer was already
478    /// finalized, or encoding/I/O fails.
479    pub fn finalize_batch(
480        &mut self,
481        batch: &RecordBatch,
482        measurement: &str,
483        tag_columns: &[String],
484    ) -> Result<SegmentMeta> {
485        if self.finalized {
486            return Err(SegmentError::AlreadyFinalized);
487        }
488        if batch.num_rows() == 0 {
489            return Err(SegmentError::EmptySegment);
490        }
491
492        let total_rows = batch.num_rows();
493        let tag_set: std::collections::HashSet<&str> =
494            tag_columns.iter().map(String::as_str).collect();
495
496        // Build schema from Arrow schema + tag metadata
497        let schema = discover_schema_from_batch(batch, &tag_set)?;
498
499        // Compute global stats directly from Arrow arrays
500        let ts_col = batch
501            .column_by_name(chronix_core::TIME_COLUMN)
502            .and_then(|c| c.as_any().downcast_ref::<Int64Array>())
503            .ok_or_else(|| SegmentError::CorruptFile {
504                detail: "missing or non-Int64 timestamp column".into(),
505            })?;
506
507        let mut min_ts = i64::MAX;
508        let mut max_ts = i64::MIN;
509        for i in 0..ts_col.len() {
510            if ts_col.is_null(i) {
511                return Err(SegmentError::CorruptFile {
512                    detail: format!("null timestamp at row {i}"),
513                });
514            }
515            let t = ts_col.value(i);
516            min_ts = min_ts.min(t);
517            max_ts = max_ts.max(t);
518        }
519
520        let series_keys = series_keys_from_batch(batch, measurement, tag_columns);
521        let series_count = u32::try_from(series_keys.len()).unwrap_or(u32::MAX);
522
523        let column_count = u16::try_from(schema.len()).map_err(|_| SegmentError::CorruptFile {
524            detail: format!("column count {} exceeds u16::MAX", schema.len()),
525        })?;
526
527        let header = SegmentHeader {
528            version: VERSION,
529            flags: 0,
530            created_at: std::time::SystemTime::now()
531                .duration_since(std::time::UNIX_EPOCH)
532                .map(|d| i64::try_from(d.as_nanos()).unwrap_or(i64::MAX))
533                .unwrap_or(0),
534            min_timestamp: min_ts,
535            max_timestamp: max_ts,
536            row_count: total_rows as u64,
537            column_count,
538            series_count,
539            compression: match self.config.compression_codec {
540                CompressionCodec::Lz4 => 0,
541                CompressionCodec::Zstd => 1,
542                CompressionCodec::None => 2,
543            },
544            sort_order: 1,
545        };
546
547        // ── Streaming write path ───────────────────────────────────
548        if let Some(parent) = self.path.parent() {
549            std::fs::create_dir_all(parent)?;
550        }
551
552        let temp_path = {
553            let mut p = self.path.clone().into_os_string();
554            p.push(".tmp");
555            PathBuf::from(p)
556        };
557
558        let result = (|| -> Result<(SegmentMeta, u64)> {
559            let file = std::fs::File::create(&temp_path)?;
560            let mut writer = BufWriter::new(file);
561            let mut crc: u32 = 0;
562            let mut pos: u64 = 0;
563
564            // Write header
565            let header_bytes = header.to_bytes();
566            writer.write_all(&header_bytes)?;
567            crc = crc32c::crc32c_append(crc, &header_bytes);
568            pos += header_bytes.len() as u64;
569
570            // Encode and stream row groups one at a time
571            let (column_metas, row_group_blocks, uncompressed_bytes, zstd_dictionary) =
572                encode_row_groups_from_batch_streaming(
573                    &schema,
574                    batch,
575                    &self.config,
576                    header.created_at,
577                    &mut writer,
578                    &mut crc,
579                    &mut pos,
580                )?;
581
582            let metadata = SegmentMetadata {
583                columns: column_metas,
584                row_group_blocks,
585                zstd_dictionary,
586            };
587
588            // Write metadata
589            let metadata_offset = pos;
590            let metadata_bytes = metadata.to_bytes()?;
591            let metadata_checksum = crc32c::crc32c(&metadata_bytes);
592            writer.write_all(&metadata_bytes)?;
593            crc = crc32c::crc32c_append(crc, &metadata_bytes);
594            pos += metadata_bytes.len() as u64;
595
596            // Write footer with accumulated CRC
597            let footer = SegmentFooter {
598                metadata_offset,
599                row_group_count: metadata.row_group_blocks.len() as u32,
600                checksum: crc,
601                metadata_checksum,
602            };
603            let footer_bytes = footer.to_bytes();
604            writer.write_all(&footer_bytes)?;
605            pos += footer_bytes.len() as u64;
606
607            // Flush and fsync
608            writer.flush()?;
609            writer.get_ref().sync_all()?;
610            drop(writer);
611
612            // Atomic rename
613            std::fs::rename(&temp_path, &self.path)?;
614
615            Ok((
616                SegmentMeta {
617                    path: self.path.clone(),
618                    min_timestamp: min_ts,
619                    max_timestamp: max_ts,
620                    row_count: total_rows as u64,
621                    series_count,
622                    byte_size: pos,
623                    row_group_count: footer.row_group_count,
624                    column_count,
625                    uncompressed_bytes,
626                    header,
627                    column_metas: metadata.columns,
628                    series_keys,
629                },
630                pos,
631            ))
632        })();
633
634        match result {
635            Ok((meta, _)) => {
636                self.finalized = true;
637
638                // Fsync parent directory for rename durability.
639                if let Some(parent) = self.path.parent() {
640                    let dir = std::fs::File::open(parent).map_err(|e| {
641                        SegmentError::Io(std::io::Error::new(
642                            e.kind(),
643                            format!("failed to open parent dir for fsync: {e}"),
644                        ))
645                    })?;
646                    dir.sync_all().map_err(|e| {
647                        SegmentError::Io(std::io::Error::new(
648                            e.kind(),
649                            format!("parent directory fsync failed: {e}"),
650                        ))
651                    })?;
652                }
653
654                Ok(meta)
655            }
656            Err(e) => {
657                let _ = std::fs::remove_file(&temp_path);
658                Err(e)
659            }
660        }
661    }
662
663    /// Finalize a segment from pre-chunked `RecordBatch` slices.
664    ///
665    /// Unlike [`finalize_batch`](Self::finalize_batch), this avoids
666    /// `concat_batches` by iterating over the chunks directly. Each chunk
667    /// is encoded as one or more row groups and streamed to disk, bounding
668    /// peak memory to `O(chunk_size × columns)`.
669    pub fn finalize_batches(
670        &mut self,
671        batches: &[RecordBatch],
672        measurement: &str,
673        tag_columns: &[String],
674    ) -> Result<SegmentMeta> {
675        if self.finalized {
676            return Err(SegmentError::AlreadyFinalized);
677        }
678        if batches.is_empty() {
679            return Err(SegmentError::EmptySegment);
680        }
681
682        // Concatenate for now — this is the fallback. The streaming
683        // architecture benefit is that the compaction executor no longer
684        // needs to concat first, saving one full copy. We can concat
685        // here since the chunks are already small (row_group_size each).
686        let schema = batches[0].schema();
687        let batch = arrow::compute::concat_batches(&schema, batches).map_err(|e| {
688            SegmentError::CorruptFile {
689                detail: format!("concat batches failed: {e}"),
690            }
691        })?;
692        self.finalize_batch(&batch, measurement, tag_columns)
693    }
694}
695
696// ── Row Group Encoding ─────────────────────────────────────────────────
697
698/// Encode all row groups from sorted point data, streaming each encoded
699/// column block directly to `writer` instead of buffering in memory.
700///
701/// The file position (`pos`) and incremental CRC (`crc`) are updated in-place
702/// as each block is written. Peak memory is bounded by `O(row_group_size)`
703/// instead of `O(segment_size)`.
704fn encode_row_groups_streaming<W: Write>(
705    schema: &[ColumnDef],
706    sorted_indices: &[usize],
707    points: &[Point],
708    config: &SegmentWriterConfig,
709    // Only the encryption path reads it: it is half of the AEAD nonce's
710    // associated data, so a block cannot be replayed into another segment.
711    #[cfg_attr(not(feature = "field-encryption"), allow(unused_variables))] segment_created_at: i64,
712    writer: &mut W,
713    crc: &mut u32,
714    pos: &mut u64,
715) -> Result<RowGroupWriteOutput> {
716    let total_rows = sorted_indices.len();
717
718    let mut column_metas: Vec<ColumnMeta> = schema
719        .iter()
720        .map(|col| -> Result<ColumnMeta> {
721            #[cfg(feature = "field-encryption")]
722            let (encrypted, key_id) = if let Some(k) = config.field_encryption.key_for(&col.name) {
723                // Only a **field** can honestly be encrypted. A tag's value
724                // is part of the series key, and the series key is written
725                // in plaintext right beside this segment — in its `.series`
726                // sidecar, in the inverted tag index and in the bloom filter
727                // — so an "encrypted" tag column would publish the value it
728                // claims to hide, a foot away. The timestamp cannot be
729                // encrypted either: every scan ranges on it.
730                if col.role != crate::segment::metadata::roles::FIELD {
731                    return Err(SegmentError::InvalidConfiguration {
732                        detail: format!(
733                            "column '{}' is declared encrypted but is a {} — only a field can \
734                             be encrypted, because a tag's value is also written in plaintext \
735                             to the segment's series sidecar, the inverted tag index and its \
736                             bloom filter, and the time column is what every scan ranges on",
737                            col.name,
738                            if col.role == crate::segment::metadata::roles::TAG {
739                                "tag"
740                            } else {
741                                "timestamp"
742                            }
743                        ),
744                    });
745                }
746                (true, Some(k.key_id.clone()))
747            } else {
748                (false, None)
749            };
750            #[cfg(not(feature = "field-encryption"))]
751            let (encrypted, key_id) = (false, None);
752
753            Ok(ColumnMeta {
754                name: col.name.clone(),
755                data_type: col.data_type,
756                role: col.role,
757                default_encoding: col.default_encoding,
758                stats: ColumnStats::empty(),
759                bloom_filter: None,
760                encrypted,
761                key_id,
762                decimal_scale: col.decimal_scale,
763                row_group_blooms: None,
764            })
765        })
766        .collect::<Result<Vec<_>>>()?;
767
768    let mut row_group_blocks: Vec<Vec<ColumnBlockMeta>> = Vec::new();
769    let mut total_uncompressed: u64 = 0;
770
771    // Zstd dictionary training state.
772    let use_dict =
773        config.zstd_dict_training && matches!(config.compression_codec, CompressionCodec::Zstd);
774    let mut dict_samples: Vec<Vec<u8>> = Vec::new();
775    let mut trained_dict: Option<Vec<u8>> = None;
776    let mut rg_index: usize = 0;
777
778    #[allow(clippy::explicit_counter_loop)]
779    // rg_index also names files/metadata, not just iterations
780    for chunk_start in (0..total_rows).step_by(config.row_group_size) {
781        let chunk_end = (chunk_start + config.row_group_size).min(total_rows);
782        let chunk_indices = &sorted_indices[chunk_start..chunk_end];
783        let chunk_size = chunk_end - chunk_start;
784
785        let mut rg_blocks = Vec::with_capacity(schema.len());
786
787        for (col_idx, col) in schema.iter().enumerate() {
788            let EncodedColumn {
789                block,
790                stats,
791                validity,
792            } = encode_column(col, chunk_indices, points, config.float_encoding)?;
793
794            let raw_encoded = block.to_bytes();
795            let raw_size = chunk_size * 8;
796            total_uncompressed += raw_encoded.len() as u64;
797
798            // Collect training samples from the first row group.
799            if use_dict && rg_index == 0 {
800                dict_samples.push(raw_encoded.clone());
801            }
802
803            let codec = config
804                .column_codec_overrides
805                .get(&col.name)
806                .copied()
807                .unwrap_or(config.compression_codec);
808            let (final_data, compressed) =
809                if config.compress && !should_skip_compression(raw_encoded.len(), raw_size) {
810                    // Use dictionary compression when available.
811                    if let Some(ref dict) = trained_dict {
812                        if matches!(codec, CompressionCodec::Zstd) {
813                            let compressed_data =
814                                compress_block_zstd_dict(&raw_encoded, config.zstd_level, dict)?;
815                            if compressed_data.len() < raw_encoded.len() {
816                                (compressed_data, true)
817                            } else {
818                                (raw_encoded, false)
819                            }
820                        } else {
821                            let compressed_data = compress_block_with_codec(&raw_encoded, codec)?;
822                            if compressed_data.len() < raw_encoded.len() {
823                                (compressed_data, true)
824                            } else {
825                                (raw_encoded, false)
826                            }
827                        }
828                    } else {
829                        let compressed_data = compress_block_with_codec(&raw_encoded, codec)?;
830                        if compressed_data.len() < raw_encoded.len() {
831                            (compressed_data, true)
832                        } else {
833                            (raw_encoded, false)
834                        }
835                    }
836                } else {
837                    (raw_encoded, false)
838                };
839
840            // Field-level encryption: encrypt after encode+compress.
841            #[cfg(feature = "field-encryption")]
842            let (final_data, block_encrypted) =
843                if let Some(enc_key) = config.field_encryption.key_for(&col.name) {
844                    let encrypted = crate::segment::field_encryption::encrypt_block(
845                        &final_data,
846                        enc_key.key_bytes(),
847                        crate::segment::field_encryption::BlockContext {
848                            column: &col.name,
849                            segment_created_at,
850                        },
851                    )?;
852                    (encrypted, true)
853                } else {
854                    (final_data, false)
855                };
856            #[cfg(not(feature = "field-encryption"))]
857            let block_encrypted = false;
858
859            let offset = *pos;
860            writer.write_all(&final_data)?;
861            *crc = crc32c::crc32c_append(*crc, &final_data);
862            *pos += final_data.len() as u64;
863
864            let block_crc = crc32c::crc32c(&final_data);
865
866            // Suppress statistics for encrypted columns to prevent
867            // information leakage through zone-map predicates.
868            let final_stats = if block_encrypted {
869                ColumnStats::empty()
870            } else {
871                stats
872            };
873
874            merge_stats(&mut column_metas[col_idx].stats, &final_stats);
875
876            // D-NULL: the validity bitmap follows the block payload. It is
877            // stored uncompressed and unencrypted — it reveals only which
878            // rows have values, never what they are — so the reader can turn
879            // it straight into an Arrow NullBuffer.
880            let (validity_offset, validity_length) =
881                write_validity(writer, validity.as_deref(), pos, crc)?;
882
883            rg_blocks.push(ColumnBlockMeta {
884                column_index: col_idx as u16,
885                encoding: block.encoding,
886                compressed,
887                offset,
888                length: u32::try_from(final_data.len()).map_err(|_| SegmentError::CorruptFile {
889                    detail: format!("block length {} exceeds u32::MAX", final_data.len()),
890                })?,
891                value_count: u32::try_from(chunk_size).map_err(|_| SegmentError::CorruptFile {
892                    detail: format!("value_count {chunk_size} exceeds u32::MAX"),
893                })?,
894                block_crc,
895                encrypted: block_encrypted,
896                validity_offset,
897                validity_length,
898                stats: final_stats,
899            });
900        }
901
902        row_group_blocks.push(rg_blocks);
903
904        // Build per-row-group bloom filters for tag columns.
905        for (col_idx, col) in schema.iter().enumerate() {
906            if col.role == roles::TAG && !column_metas[col_idx].encrypted {
907                let mut distinct: std::collections::HashSet<&str> =
908                    std::collections::HashSet::new();
909                for &idx in chunk_indices {
910                    let val = points[idx].series_key().tag(&col.name).unwrap_or("");
911                    if !val.is_empty() {
912                        distinct.insert(val);
913                    }
914                }
915                let bloom = if distinct.is_empty() {
916                    Vec::new()
917                } else {
918                    let vals: Vec<&str> = distinct.into_iter().collect();
919                    crate::segment::bloom::bloom_filter_build(&vals, config.bloom_fpr)
920                };
921                column_metas[col_idx]
922                    .row_group_blooms
923                    .get_or_insert_with(Vec::new)
924                    .push(bloom);
925            }
926        }
927
928        // Train dictionary after the first row group.
929        if use_dict && rg_index == 0 && !dict_samples.is_empty() {
930            let sample_refs: Vec<&[u8]> =
931                dict_samples.iter().map(std::vec::Vec::as_slice).collect();
932            if let Ok(Some(dict)) = train_zstd_dictionary(&sample_refs) {
933                tracing::debug!(
934                    dict_size = dict.len(),
935                    samples = dict_samples.len(),
936                    "trained Zstd dictionary from first row group"
937                );
938                trained_dict = Some(dict);
939            }
940            dict_samples.clear();
941        }
942        rg_index += 1;
943    }
944
945    // Build bloom filters for tag columns.
946    // Skip bloom filter for encrypted columns to prevent information leakage.
947    for (col_idx, col) in schema.iter().enumerate() {
948        if col.role == roles::TAG && !column_metas[col_idx].encrypted {
949            let mut distinct: std::collections::HashSet<&str> = std::collections::HashSet::new();
950            for &idx in sorted_indices {
951                let val = points[idx].series_key().tag(&col.name).unwrap_or("");
952                if !val.is_empty() {
953                    distinct.insert(val);
954                }
955            }
956            if !distinct.is_empty() {
957                let values: Vec<&str> = distinct.into_iter().collect();
958                column_metas[col_idx].bloom_filter = Some(
959                    crate::segment::bloom::bloom_filter_build(&values, config.bloom_fpr),
960                );
961            }
962        }
963    }
964
965    Ok((
966        column_metas,
967        row_group_blocks,
968        total_uncompressed,
969        trained_dict,
970    ))
971}
972
973// ── Schema Discovery ───────────────────────────────────────────────────
974
975/// Discovered column definition from point data.
976#[derive(Debug, Clone)]
977struct ColumnDef {
978    name: String,
979    data_type: u8,
980    role: u8,
981    default_encoding: u8,
982    /// Digits after the decimal point, for a decimal column; `None` otherwise.
983    decimal_scale: Option<u8>,
984}
985
986/// Discover the column schema by scanning all points.
987///
988/// Tag columns are ordered by ascending cardinality (fewest distinct
989/// values first) so that the physical column layout is optimal for
990/// run-length and dictionary encoding.
991fn discover_schema(points: &[Point]) -> Vec<ColumnDef> {
992    let mut columns: Vec<ColumnDef> = Vec::new();
993
994    // Timestamp column is always first
995    columns.push(ColumnDef {
996        name: chronix_core::TIME_COLUMN.to_string(),
997        data_type: data_types::TIMESTAMP,
998        role: roles::TIMESTAMP,
999        default_encoding: chronix_encoding::EncodingType::DeltaOfDelta.tag(),
1000        decimal_scale: None,
1001    });
1002
1003    // Discover tag columns and compute cardinality
1004    let mut tag_names = BTreeSet::new();
1005    for p in points {
1006        for key in p.series_key().tag_keys() {
1007            tag_names.insert(key.to_string());
1008        }
1009    }
1010
1011    // Sort tags by cardinality (lowest first) for optimal compression
1012    let mut tag_with_card: Vec<(String, usize)> = tag_names
1013        .into_iter()
1014        .map(|name: String| {
1015            let distinct: std::collections::HashSet<&str> = points
1016                .iter()
1017                .filter_map(|p| p.series_key().tag(&name))
1018                .collect();
1019            (name, distinct.len())
1020        })
1021        .collect();
1022    tag_with_card.sort_by_key(|(_, card)| *card);
1023
1024    for (name, _) in tag_with_card {
1025        columns.push(ColumnDef {
1026            name,
1027            data_type: data_types::STRING,
1028            role: roles::TAG,
1029            default_encoding: chronix_encoding::EncodingType::Dictionary.tag(),
1030            decimal_scale: None,
1031        });
1032    }
1033
1034    // Discover field columns.
1035    //
1036    // A decimal field also carries its scale, and the *widest* scale in the
1037    // batch wins rather than whichever point happened to come first. The
1038    // write path has already normalised every value to the column's declared
1039    // scale, so in practice they all agree; taking the maximum means a
1040    // segment written from points that had not been through that path — a
1041    // test, a tool, a future caller — still produces one scale for the
1042    // column instead of silently truncating to the first one seen.
1043    let mut field_defs: BTreeMap<String, (u8, Option<u8>)> = BTreeMap::new();
1044    for p in points {
1045        for (key, val) in p.fields() {
1046            let entry = field_defs.entry(key.to_string()).or_insert_with(|| {
1047                let dt = match val {
1048                    FieldValue::F64(_) => data_types::F64,
1049                    FieldValue::I64(_) => data_types::I64,
1050                    FieldValue::U64(_) => data_types::U64,
1051                    FieldValue::Bool(_) => data_types::BOOL,
1052                    FieldValue::String(_) => data_types::STRING,
1053                    FieldValue::Decimal(_) => data_types::DECIMAL,
1054                    FieldValue::Histogram(_) => data_types::HISTOGRAM,
1055                };
1056                (dt, None)
1057            });
1058            if let FieldValue::Decimal(d) = val
1059                && entry.0 == data_types::DECIMAL
1060            {
1061                entry.1 = Some(entry.1.unwrap_or(0).max(d.scale()));
1062            }
1063        }
1064    }
1065    for (name, (dt, scale)) in field_defs {
1066        let enc = match dt {
1067            data_types::F64 => chronix_encoding::EncodingType::Chimp.tag(),
1068            data_types::I64 => chronix_encoding::EncodingType::IntegerI64.tag(),
1069            data_types::U64 => chronix_encoding::EncodingType::IntegerU64.tag(),
1070            data_types::BOOL => chronix_encoding::EncodingType::Bitmap.tag(),
1071            data_types::DECIMAL => chronix_encoding::EncodingType::DecimalI128.tag(),
1072            _ => chronix_encoding::EncodingType::Dictionary.tag(),
1073        };
1074        columns.push(ColumnDef {
1075            name,
1076            data_type: dt,
1077            role: roles::FIELD,
1078            default_encoding: enc,
1079            decimal_scale: scale,
1080        });
1081    }
1082
1083    columns
1084}
1085
1086/// Append a block's validity bitmap to the segment file.
1087///
1088/// Returns the `(offset, length)` pair recorded in
1089/// [`ColumnBlockMeta::validity_offset`] / `validity_length`. A `None` bitmap
1090/// (a block with no nulls) writes nothing and yields `(0, 0)`.
1091fn write_validity<W: Write>(
1092    writer: &mut W,
1093    validity: Option<&[u8]>,
1094    pos: &mut u64,
1095    crc: &mut u32,
1096) -> Result<(u64, u32)> {
1097    let Some(bytes) = validity else {
1098        return Ok((0, 0));
1099    };
1100    let offset = *pos;
1101    writer.write_all(bytes)?;
1102    *crc = crc32c::crc32c_append(*crc, bytes);
1103    *pos += bytes.len() as u64;
1104    let length = u32::try_from(bytes.len()).map_err(|_| SegmentError::CorruptFile {
1105        detail: format!("validity bitmap length {} exceeds u32::MAX", bytes.len()),
1106    })?;
1107    Ok((offset, length))
1108}
1109
1110/// An encoded column block together with its statistics and, when the block
1111/// contains nulls, its validity bitmap (`.csx` v2, D-NULL).
1112struct EncodedColumn {
1113    block: EncodedBlock,
1114    stats: ColumnStats,
1115    /// `None` when every row in the block holds a real value.
1116    validity: Option<Vec<u8>>,
1117}
1118
1119impl EncodedColumn {
1120    fn new(block: EncodedBlock, stats: ColumnStats, validity: ValidityBuilder) -> Self {
1121        Self {
1122            block,
1123            stats,
1124            validity: validity.finish(),
1125        }
1126    }
1127
1128    /// Wrap a column that cannot contain nulls (timestamps, tags).
1129    fn dense((block, stats): (EncodedBlock, ColumnStats)) -> Self {
1130        Self {
1131            block,
1132            stats,
1133            validity: None,
1134        }
1135    }
1136}
1137
1138/// Encode a single column for a chunk of rows.
1139fn encode_column(
1140    col: &ColumnDef,
1141    indices: &[usize],
1142    points: &[Point],
1143    float_encoding: FloatEncoding,
1144) -> Result<EncodedColumn> {
1145    let mut stats = ColumnStats::empty();
1146    // D-NULL: record which rows actually carry a value so the reader can
1147    // distinguish an absent field from a stored zero.
1148    let mut validity = ValidityBuilder::with_capacity(indices.len());
1149
1150    match (col.role, col.data_type) {
1151        // Timestamps are never null; tags are materialised as "" by design
1152        // (an absent tag is not part of the series key).
1153        (roles::TIMESTAMP, _) => encode_timestamp_column(indices, points).map(EncodedColumn::dense),
1154        (roles::TAG, _) => encode_tag_column(col, indices, points).map(EncodedColumn::dense),
1155        (_, data_types::F64) => {
1156            let mut values = Vec::with_capacity(indices.len());
1157            for &idx in indices {
1158                if let Some(FieldValue::F64(v)) = points[idx].field(&col.name) {
1159                    values.push(*v);
1160                    stats.update_f64(*v);
1161                    validity.push(true);
1162                } else {
1163                    values.push(0.0);
1164                    stats.record_null();
1165                    validity.push(false);
1166                }
1167            }
1168            let block = ColumnEncoder::encode_f64(&values, float_encoding)?;
1169            Ok(EncodedColumn::new(block, stats, validity))
1170        }
1171        (_, data_types::I64) => {
1172            let mut values = Vec::with_capacity(indices.len());
1173            for &idx in indices {
1174                if let Some(FieldValue::I64(v)) = points[idx].field(&col.name) {
1175                    values.push(*v);
1176                    stats.update_i64(*v);
1177                    validity.push(true);
1178                } else {
1179                    values.push(0);
1180                    stats.record_null();
1181                    validity.push(false);
1182                }
1183            }
1184            let block = ColumnEncoder::encode_i64(&values)?;
1185            Ok(EncodedColumn::new(block, stats, validity))
1186        }
1187        (_, data_types::U64) => {
1188            let mut values = Vec::with_capacity(indices.len());
1189            for &idx in indices {
1190                if let Some(FieldValue::U64(v)) = points[idx].field(&col.name) {
1191                    values.push(*v);
1192                    stats.update_u64(*v);
1193                    validity.push(true);
1194                } else {
1195                    values.push(0);
1196                    stats.record_null();
1197                    validity.push(false);
1198                }
1199            }
1200            let block = ColumnEncoder::encode_u64(&values)?;
1201            Ok(EncodedColumn::new(block, stats, validity))
1202        }
1203        (_, data_types::BOOL) => {
1204            let mut values = Vec::with_capacity(indices.len());
1205            for &idx in indices {
1206                if let Some(FieldValue::Bool(v)) = points[idx].field(&col.name) {
1207                    values.push(*v);
1208                    stats.update_bool(*v);
1209                    validity.push(true);
1210                } else {
1211                    values.push(false);
1212                    stats.record_null();
1213                    validity.push(false);
1214                }
1215            }
1216            let block = ColumnEncoder::encode_bool(&values)?;
1217            Ok(EncodedColumn::new(block, stats, validity))
1218        }
1219        (_, data_types::DECIMAL) => {
1220            // The scale is the column's, recorded once in its metadata; what
1221            // is encoded is the mantissa column. A value that arrived at a
1222            // different scale is widened losslessly rather than dropped — the
1223            // write path normalises, so this only catches a caller that
1224            // bypassed it.
1225            let scale = col.decimal_scale.unwrap_or(0);
1226            let mut values = Vec::with_capacity(indices.len());
1227            for &idx in indices {
1228                match points[idx].field(&col.name) {
1229                    Some(FieldValue::Decimal(d)) => match d.rescale(scale) {
1230                        Ok(v) => {
1231                            values.push(v.mantissa());
1232                            stats.update_decimal(v.mantissa());
1233                            validity.push(true);
1234                        }
1235                        Err(e) => {
1236                            return Err(SegmentError::CorruptFile {
1237                                detail: format!(
1238                                    "decimal field '{}' at scale {} cannot be stored in a \
1239                                     scale-{scale} column: {e}",
1240                                    col.name,
1241                                    d.scale()
1242                                ),
1243                            });
1244                        }
1245                    },
1246                    _ => {
1247                        values.push(0);
1248                        stats.record_null();
1249                        validity.push(false);
1250                    }
1251                }
1252            }
1253            let block = ColumnEncoder::encode_decimal(&values)?;
1254            Ok(EncodedColumn::new(block, stats, validity))
1255        }
1256        (_, data_types::STRING) => {
1257            // Collect owned strings, then borrow as &str for encoding
1258            let mut owned: Vec<String> = Vec::with_capacity(indices.len());
1259            for &idx in indices {
1260                if let Some(FieldValue::String(v)) = points[idx].field(&col.name) {
1261                    owned.push(v.clone());
1262                    validity.push(true);
1263                } else {
1264                    owned.push(String::new());
1265                    stats.record_null();
1266                    validity.push(false);
1267                }
1268            }
1269            let values: Vec<&str> = owned.iter().map(String::as_str).collect();
1270            let distinct: std::collections::HashSet<&str> = values.iter().copied().collect();
1271            stats.distinct_count = distinct.len() as u32;
1272            // value_count should reflect non-null entries (null_count already tracked).
1273            stats.value_count = (values.len() as u64).saturating_sub(stats.null_count);
1274            let block = ColumnEncoder::encode_string(&values)?;
1275            Ok(EncodedColumn::new(block, stats, validity))
1276        }
1277        (_, data_types::HISTOGRAM) => {
1278            // The value is already encoded — `postcard`, by the histogram
1279            // itself — so this path frames the blobs and records presence.
1280            // No statistics: min/max of a distribution is not a number, and a
1281            // zone map over one would be a comparison nothing can make.
1282            let mut owned: Vec<Vec<u8>> = Vec::with_capacity(indices.len());
1283            for &idx in indices {
1284                match points[idx].field(&col.name) {
1285                    Some(FieldValue::Histogram(h)) => match postcard::to_allocvec(h.as_ref()) {
1286                        Ok(bytes) => {
1287                            owned.push(bytes);
1288                            validity.push(true);
1289                        }
1290                        Err(e) => {
1291                            return Err(SegmentError::CorruptFile {
1292                                detail: format!(
1293                                    "histogram column '{}' could not be encoded: {e}",
1294                                    col.name
1295                                ),
1296                            });
1297                        }
1298                    },
1299                    _ => {
1300                        owned.push(Vec::new());
1301                        stats.record_null();
1302                        validity.push(false);
1303                    }
1304                }
1305            }
1306            stats.value_count = (owned.len() as u64).saturating_sub(stats.null_count);
1307            let refs: Vec<&[u8]> = owned.iter().map(Vec::as_slice).collect();
1308            let block = chronix_encoding::EncodedBlock {
1309                encoding: chronix_encoding::EncodingType::Bytes,
1310                payload: chronix_encoding::PlainEncoder::encode_bytes(&refs)?,
1311            };
1312            Ok(EncodedColumn::new(block, stats, validity))
1313        }
1314        _ => Err(SegmentError::CorruptFile {
1315            detail: format!("unknown column type: {}", col.data_type),
1316        }),
1317    }
1318}
1319
1320/// Compute min/max timestamps and unique series count from sorted indices.
1321///
1322/// Since the indices are already sorted by `(canonical_form, timestamp)`,
1323/// we count distinct series by checking when the canonical form changes.
1324fn compute_global_stats(
1325    points: &[Point],
1326    sorted_indices: &[usize],
1327) -> (i64, i64, Vec<chronix_core::SeriesKey>) {
1328    let mut min_ts = i64::MAX;
1329    let mut max_ts = i64::MIN;
1330    let mut series_keys: Vec<chronix_core::SeriesKey> = Vec::new();
1331
1332    for &idx in sorted_indices {
1333        let p = &points[idx];
1334        let ts = p.timestamp();
1335        min_ts = min_ts.min(ts);
1336        max_ts = max_ts.max(ts);
1337        // Series-major order: a new series starts wherever the canonical form
1338        // changes. Comparing the form rather than the hash means a hash
1339        // collision between adjacent series is counted as two, not one.
1340        let is_new = series_keys
1341            .last()
1342            .is_none_or(|prev| prev.canonical_form() != p.series_key().canonical_form());
1343        if is_new {
1344            series_keys.push(p.series_key().clone());
1345        }
1346    }
1347    (min_ts, max_ts, series_keys)
1348}
1349
1350/// Encode a timestamp column from a set of points at the given indices.
1351fn encode_timestamp_column(
1352    indices: &[usize],
1353    points: &[Point],
1354) -> Result<(EncodedBlock, ColumnStats)> {
1355    let mut stats = ColumnStats::empty();
1356    let mut timestamps = Vec::with_capacity(indices.len());
1357    for &idx in indices {
1358        let ts = points[idx].timestamp();
1359        timestamps.push(ts);
1360        stats.update_i64(ts);
1361    }
1362    let block = ColumnEncoder::encode_timestamps(&timestamps)?;
1363    Ok((block, stats))
1364}
1365
1366/// Encode a tag column by looking up each tag value from the series key.
1367fn encode_tag_column(
1368    col: &ColumnDef,
1369    indices: &[usize],
1370    points: &[Point],
1371) -> Result<(EncodedBlock, ColumnStats)> {
1372    let mut stats = ColumnStats::empty();
1373    let mut values = Vec::with_capacity(indices.len());
1374    for &idx in indices {
1375        let val = points[idx].series_key().tag(&col.name).unwrap_or("");
1376        values.push(val);
1377    }
1378    let distinct: std::collections::HashSet<&str> = values.iter().copied().collect();
1379    stats.distinct_count = distinct.len() as u32;
1380    stats.value_count = values.len() as u64;
1381    let block = ColumnEncoder::encode_string(&values)?;
1382    Ok((block, stats))
1383}
1384
1385/// Merge source stats into destination.
1386///
1387/// **Note:** `distinct_count` uses `max()` when merging across row groups —
1388/// a lower bound on the true segment cardinality, and a tighter estimate than
1389/// summing. The segment-level per-column `distinct_count` computed over all
1390/// rows is the exact figure; this merge only applies to the row-group rollup.
1391fn merge_stats(dest: &mut ColumnStats, src: &ColumnStats) {
1392    dest.min_value = dest.min_value.min(src.min_value);
1393    dest.max_value = dest.max_value.max(src.max_value);
1394    // Also merge u64 min/max — previously these were silently
1395    // dropped, breaking predicate pushdown for u64 columns.
1396    dest.min_value_u64 = dest.min_value_u64.min(src.min_value_u64);
1397    dest.max_value_u64 = dest.max_value_u64.max(src.max_value_u64);
1398    dest.null_count += src.null_count;
1399    dest.value_count += src.value_count;
1400    dest.sum += src.sum;
1401    dest.sum_i128 += src.sum_i128;
1402    // Use max() instead of saturating_add for a tighter bound.
1403    // Row groups may share values, so summing grossly over-counts.
1404    dest.distinct_count = dest.distinct_count.max(src.distinct_count);
1405}
1406
1407// ── Batch-based encoding (zero-copy compaction path) ───────────────────
1408
1409/// Discover column schema from an Arrow `RecordBatch`.
1410///
1411/// Uses `tag_columns` to distinguish tags from fields (since Arrow
1412/// schema alone doesn't carry Chronix role information).
1413fn discover_schema_from_batch(
1414    batch: &RecordBatch,
1415    tag_columns: &std::collections::HashSet<&str>,
1416) -> Result<Vec<ColumnDef>> {
1417    let mut columns = Vec::new();
1418
1419    // Timestamp first
1420    columns.push(ColumnDef {
1421        name: chronix_core::TIME_COLUMN.to_string(),
1422        data_type: data_types::TIMESTAMP,
1423        role: roles::TIMESTAMP,
1424        default_encoding: chronix_encoding::EncodingType::DeltaOfDelta.tag(),
1425        decimal_scale: None,
1426    });
1427
1428    // Tags (sorted for consistency)
1429    let mut tags: Vec<&str> = tag_columns.iter().copied().collect();
1430    tags.sort_unstable();
1431    for tag in tags {
1432        if batch.column_by_name(tag).is_some() {
1433            columns.push(ColumnDef {
1434                name: tag.to_string(),
1435                data_type: data_types::STRING,
1436                role: roles::TAG,
1437                default_encoding: chronix_encoding::EncodingType::Dictionary.tag(),
1438                decimal_scale: None,
1439            });
1440        }
1441    }
1442
1443    // Fields (everything else except timestamp and tags)
1444    for field in batch.schema().fields() {
1445        let name = field.name().as_str();
1446        if name == chronix_core::TIME_COLUMN || tag_columns.contains(name) {
1447            continue;
1448        }
1449        let (dt, enc) = match field.data_type() {
1450            DataType::Float64 => (data_types::F64, chronix_encoding::EncodingType::Chimp.tag()),
1451            DataType::Int64 => (
1452                data_types::I64,
1453                chronix_encoding::EncodingType::IntegerI64.tag(),
1454            ),
1455            DataType::UInt64 => (
1456                data_types::U64,
1457                chronix_encoding::EncodingType::IntegerU64.tag(),
1458            ),
1459            DataType::Boolean => (
1460                data_types::BOOL,
1461                chronix_encoding::EncodingType::Bitmap.tag(),
1462            ),
1463            DataType::Utf8 | DataType::LargeUtf8 => (
1464                data_types::STRING,
1465                chronix_encoding::EncodingType::Dictionary.tag(),
1466            ),
1467            DataType::Decimal128(_, _) => (
1468                data_types::DECIMAL,
1469                chronix_encoding::EncodingType::DecimalI128.tag(),
1470            ),
1471            // A native histogram: a composite value carried as its own
1472            // encoding. There is nothing numeric to zone-map or delta, so
1473            // the framing is a length prefix and the block compression
1474            // underneath does the rest.
1475            DataType::Binary | DataType::LargeBinary => (
1476                data_types::HISTOGRAM,
1477                chronix_encoding::EncodingType::Bytes.tag(),
1478            ),
1479            other => {
1480                return Err(SegmentError::CorruptFile {
1481                    detail: format!(
1482                        "unsupported Arrow data type {:?} for field '{}'",
1483                        other, name
1484                    ),
1485                });
1486            }
1487        };
1488        // The Arrow scale is authoritative on this path: the batch came out
1489        // of a segment or a memtable, both of which already fixed it.
1490        let decimal_scale = match field.data_type() {
1491            DataType::Decimal128(_, scale) => {
1492                Some(u8::try_from(*scale).map_err(|_| SegmentError::CorruptFile {
1493                    detail: format!("negative decimal scale {scale} for field '{name}'"),
1494                })?)
1495            }
1496            _ => None,
1497        };
1498        columns.push(ColumnDef {
1499            name: name.to_string(),
1500            data_type: dt,
1501            role: roles::FIELD,
1502            default_encoding: enc,
1503            decimal_scale,
1504        });
1505    }
1506
1507    Ok(columns)
1508}
1509
1510/// Count distinct series in a RecordBatch by canonical series key form.
1511///
1512/// Uses the full canonical form (not just hash) to avoid undercounting
1513/// in the extremely rare case of FNV hash collisions.
1514fn series_keys_from_batch(
1515    batch: &RecordBatch,
1516    measurement: &str,
1517    tag_columns: &[String],
1518) -> Vec<chronix_core::SeriesKey> {
1519    let mut seen = std::collections::HashSet::new();
1520    let mut keys = Vec::new();
1521    let num_rows = batch.num_rows();
1522
1523    let tag_arrays: Vec<Option<&StringArray>> = tag_columns
1524        .iter()
1525        .map(|name| {
1526            batch
1527                .column_by_name(name)
1528                .and_then(|c| c.as_any().downcast_ref::<StringArray>())
1529        })
1530        .collect();
1531
1532    for row in 0..num_rows {
1533        let mut tags = BTreeMap::new();
1534        for (i, tag_name) in tag_columns.iter().enumerate() {
1535            if let Some(arr) = tag_arrays[i]
1536                && !arr.is_null(row)
1537            {
1538                tags.insert(tag_name.clone(), arr.value(row).to_string());
1539            }
1540        }
1541        if let Ok(key) = chronix_core::SeriesKey::new(measurement, tags)
1542            && seen.insert(key.canonical_form().to_string())
1543        {
1544            keys.push(key);
1545        }
1546    }
1547    keys
1548}
1549
1550/// Encode all row groups from a pre-sorted `RecordBatch`, streaming each
1551/// encoded column block directly to `writer` instead of buffering in memory.
1552///
1553/// This is the batch-based counterpart of [`encode_row_groups_streaming`].
1554fn encode_row_groups_from_batch_streaming<W: Write>(
1555    schema: &[ColumnDef],
1556    batch: &RecordBatch,
1557    config: &SegmentWriterConfig,
1558    // Only the encryption path reads it: it is half of the AEAD nonce's
1559    // associated data, so a block cannot be replayed into another segment.
1560    #[cfg_attr(not(feature = "field-encryption"), allow(unused_variables))] segment_created_at: i64,
1561    writer: &mut W,
1562    crc: &mut u32,
1563    pos: &mut u64,
1564) -> Result<RowGroupWriteOutput> {
1565    let total_rows = batch.num_rows();
1566
1567    let mut column_metas: Vec<ColumnMeta> = schema
1568        .iter()
1569        .map(|col| -> Result<ColumnMeta> {
1570            #[cfg(feature = "field-encryption")]
1571            let (encrypted, key_id) = if let Some(k) = config.field_encryption.key_for(&col.name) {
1572                // Only a **field** can honestly be encrypted. A tag's value
1573                // is part of the series key, and the series key is written
1574                // in plaintext right beside this segment — in its `.series`
1575                // sidecar, in the inverted tag index and in the bloom filter
1576                // — so an "encrypted" tag column would publish the value it
1577                // claims to hide, a foot away. The timestamp cannot be
1578                // encrypted either: every scan ranges on it.
1579                if col.role != crate::segment::metadata::roles::FIELD {
1580                    return Err(SegmentError::InvalidConfiguration {
1581                        detail: format!(
1582                            "column '{}' is declared encrypted but is a {} — only a field can \
1583                             be encrypted, because a tag's value is also written in plaintext \
1584                             to the segment's series sidecar, the inverted tag index and its \
1585                             bloom filter, and the time column is what every scan ranges on",
1586                            col.name,
1587                            if col.role == crate::segment::metadata::roles::TAG {
1588                                "tag"
1589                            } else {
1590                                "timestamp"
1591                            }
1592                        ),
1593                    });
1594                }
1595                (true, Some(k.key_id.clone()))
1596            } else {
1597                (false, None)
1598            };
1599            #[cfg(not(feature = "field-encryption"))]
1600            let (encrypted, key_id) = (false, None);
1601
1602            Ok(ColumnMeta {
1603                name: col.name.clone(),
1604                data_type: col.data_type,
1605                role: col.role,
1606                default_encoding: col.default_encoding,
1607                stats: ColumnStats::empty(),
1608                bloom_filter: None,
1609                encrypted,
1610                key_id,
1611                decimal_scale: col.decimal_scale,
1612                row_group_blooms: None,
1613            })
1614        })
1615        .collect::<Result<Vec<_>>>()?;
1616
1617    let mut row_group_blocks: Vec<Vec<ColumnBlockMeta>> = Vec::new();
1618    let mut total_uncompressed: u64 = 0;
1619
1620    // Zstd dictionary training state.
1621    let use_dict =
1622        config.zstd_dict_training && matches!(config.compression_codec, CompressionCodec::Zstd);
1623    let mut dict_samples: Vec<Vec<u8>> = Vec::new();
1624    let mut trained_dict: Option<Vec<u8>> = None;
1625    let mut rg_index: usize = 0;
1626
1627    #[allow(clippy::explicit_counter_loop)]
1628    // rg_index also names files/metadata, not just iterations
1629    for chunk_start in (0..total_rows).step_by(config.row_group_size) {
1630        let chunk_end = (chunk_start + config.row_group_size).min(total_rows);
1631        let chunk_size = chunk_end - chunk_start;
1632
1633        let mut rg_blocks = Vec::with_capacity(schema.len());
1634
1635        for (col_idx, col) in schema.iter().enumerate() {
1636            let EncodedColumn {
1637                block,
1638                stats,
1639                validity,
1640            } = encode_column_from_batch(
1641                col,
1642                batch,
1643                chunk_start,
1644                chunk_end,
1645                config.float_encoding,
1646            )?;
1647
1648            let raw_encoded = block.to_bytes();
1649            let raw_size = chunk_size * 8;
1650            total_uncompressed += raw_encoded.len() as u64;
1651
1652            // Collect training samples from the first row group.
1653            if use_dict && rg_index == 0 {
1654                dict_samples.push(raw_encoded.clone());
1655            }
1656
1657            let codec = config
1658                .column_codec_overrides
1659                .get(&col.name)
1660                .copied()
1661                .unwrap_or(config.compression_codec);
1662            let (final_data, compressed) =
1663                if config.compress && !should_skip_compression(raw_encoded.len(), raw_size) {
1664                    // Use dictionary compression when available.
1665                    if let Some(ref dict) = trained_dict {
1666                        if matches!(codec, CompressionCodec::Zstd) {
1667                            let compressed_data =
1668                                compress_block_zstd_dict(&raw_encoded, config.zstd_level, dict)?;
1669                            if compressed_data.len() < raw_encoded.len() {
1670                                (compressed_data, true)
1671                            } else {
1672                                (raw_encoded, false)
1673                            }
1674                        } else {
1675                            let compressed_data = compress_block_with_codec(&raw_encoded, codec)?;
1676                            if compressed_data.len() < raw_encoded.len() {
1677                                (compressed_data, true)
1678                            } else {
1679                                (raw_encoded, false)
1680                            }
1681                        }
1682                    } else {
1683                        let compressed_data = compress_block_with_codec(&raw_encoded, codec)?;
1684                        if compressed_data.len() < raw_encoded.len() {
1685                            (compressed_data, true)
1686                        } else {
1687                            (raw_encoded, false)
1688                        }
1689                    }
1690                } else {
1691                    (raw_encoded, false)
1692                };
1693
1694            // Field-level encryption: encrypt after encode+compress.
1695            #[cfg(feature = "field-encryption")]
1696            let (final_data, block_encrypted) =
1697                if let Some(enc_key) = config.field_encryption.key_for(&col.name) {
1698                    let encrypted = crate::segment::field_encryption::encrypt_block(
1699                        &final_data,
1700                        enc_key.key_bytes(),
1701                        crate::segment::field_encryption::BlockContext {
1702                            column: &col.name,
1703                            segment_created_at,
1704                        },
1705                    )?;
1706                    (encrypted, true)
1707                } else {
1708                    (final_data, false)
1709                };
1710            #[cfg(not(feature = "field-encryption"))]
1711            let block_encrypted = false;
1712
1713            let offset = *pos;
1714            writer.write_all(&final_data)?;
1715            *crc = crc32c::crc32c_append(*crc, &final_data);
1716            *pos += final_data.len() as u64;
1717
1718            let block_crc = crc32c::crc32c(&final_data);
1719
1720            // Suppress statistics for encrypted columns to prevent
1721            // information leakage through zone-map predicates.
1722            let final_stats = if block_encrypted {
1723                ColumnStats::empty()
1724            } else {
1725                stats
1726            };
1727
1728            merge_stats(&mut column_metas[col_idx].stats, &final_stats);
1729
1730            // D-NULL: validity bitmap follows the block payload.
1731            let (validity_offset, validity_length) =
1732                write_validity(writer, validity.as_deref(), pos, crc)?;
1733
1734            rg_blocks.push(ColumnBlockMeta {
1735                column_index: col_idx as u16,
1736                encoding: block.encoding,
1737                compressed,
1738                offset,
1739                length: u32::try_from(final_data.len()).map_err(|_| SegmentError::CorruptFile {
1740                    detail: format!("block length {} exceeds u32::MAX", final_data.len()),
1741                })?,
1742                value_count: u32::try_from(chunk_size).map_err(|_| SegmentError::CorruptFile {
1743                    detail: format!("value_count {chunk_size} exceeds u32::MAX"),
1744                })?,
1745                block_crc,
1746                encrypted: block_encrypted,
1747                validity_offset,
1748                validity_length,
1749                stats: final_stats,
1750            });
1751        }
1752
1753        row_group_blocks.push(rg_blocks);
1754
1755        // Build per-row-group bloom filters for tag columns (batch path).
1756        for (col_idx, col) in schema.iter().enumerate() {
1757            if col.role == roles::TAG
1758                && !column_metas[col_idx].encrypted
1759                && let Some(arr) = batch
1760                    .column_by_name(&col.name)
1761                    .and_then(|c| c.as_any().downcast_ref::<StringArray>())
1762            {
1763                let mut distinct: std::collections::HashSet<&str> =
1764                    std::collections::HashSet::new();
1765                for i in chunk_start..chunk_end {
1766                    if !arr.is_null(i) {
1767                        let val = arr.value(i);
1768                        if !val.is_empty() {
1769                            distinct.insert(val);
1770                        }
1771                    }
1772                }
1773                let bloom = if distinct.is_empty() {
1774                    Vec::new()
1775                } else {
1776                    let vals: Vec<&str> = distinct.into_iter().collect();
1777                    crate::segment::bloom::bloom_filter_build(&vals, config.bloom_fpr)
1778                };
1779                column_metas[col_idx]
1780                    .row_group_blooms
1781                    .get_or_insert_with(Vec::new)
1782                    .push(bloom);
1783            }
1784        }
1785
1786        // Train dictionary after the first row group.
1787        if use_dict && rg_index == 0 && !dict_samples.is_empty() {
1788            let sample_refs: Vec<&[u8]> =
1789                dict_samples.iter().map(std::vec::Vec::as_slice).collect();
1790            if let Ok(Some(dict)) = train_zstd_dictionary(&sample_refs) {
1791                tracing::debug!(
1792                    dict_size = dict.len(),
1793                    samples = dict_samples.len(),
1794                    "trained Zstd dictionary from first row group (batch path)"
1795                );
1796                trained_dict = Some(dict);
1797            }
1798            dict_samples.clear();
1799        }
1800        rg_index += 1;
1801    }
1802
1803    // Build bloom filters for tag columns from Arrow arrays.
1804    // Skip bloom filter for encrypted columns to prevent information leakage.
1805    for (col_idx, col) in schema.iter().enumerate() {
1806        if col.role == roles::TAG
1807            && !column_metas[col_idx].encrypted
1808            && let Some(arr) = batch
1809                .column_by_name(&col.name)
1810                .and_then(|c| c.as_any().downcast_ref::<StringArray>())
1811        {
1812            let mut distinct: std::collections::HashSet<&str> = std::collections::HashSet::new();
1813            for i in 0..arr.len() {
1814                if !arr.is_null(i) {
1815                    let val = arr.value(i);
1816                    if !val.is_empty() {
1817                        distinct.insert(val);
1818                    }
1819                }
1820            }
1821            if !distinct.is_empty() {
1822                let values: Vec<&str> = distinct.into_iter().collect();
1823                column_metas[col_idx].bloom_filter = Some(
1824                    crate::segment::bloom::bloom_filter_build(&values, config.bloom_fpr),
1825                );
1826            }
1827        }
1828    }
1829
1830    Ok((
1831        column_metas,
1832        row_group_blocks,
1833        total_uncompressed,
1834        trained_dict,
1835    ))
1836}
1837
1838/// Encode a single column from a RecordBatch slice [start..end).
1839fn encode_column_from_batch(
1840    col: &ColumnDef,
1841    batch: &RecordBatch,
1842    start: usize,
1843    end: usize,
1844    float_encoding: FloatEncoding,
1845) -> Result<EncodedColumn> {
1846    let mut stats = ColumnStats::empty();
1847    // D-NULL: track per-row presence alongside the encoded values.
1848    let mut validity = ValidityBuilder::with_capacity(end.saturating_sub(start));
1849
1850    match (col.role, col.data_type) {
1851        (roles::TIMESTAMP, _) => {
1852            let arr = batch
1853                .column_by_name(chronix_core::TIME_COLUMN)
1854                .and_then(|c| c.as_any().downcast_ref::<Int64Array>())
1855                .ok_or_else(|| SegmentError::CorruptFile {
1856                    detail: "missing timestamp".into(),
1857                })?;
1858            let mut timestamps = Vec::with_capacity(end - start);
1859            for i in start..end {
1860                let t = arr.value(i);
1861                timestamps.push(t);
1862                stats.update_i64(t);
1863            }
1864            let block = ColumnEncoder::encode_timestamps(&timestamps)?;
1865            // Timestamps are never null.
1866            Ok(EncodedColumn::dense((block, stats)))
1867        }
1868        (roles::TAG, _) | (_, data_types::STRING) => {
1869            let arr = batch
1870                .column_by_name(&col.name)
1871                .and_then(|c| c.as_any().downcast_ref::<StringArray>())
1872                .ok_or_else(|| SegmentError::CorruptFile {
1873                    detail: format!("missing string column '{}'", col.name),
1874                })?;
1875            let mut values = Vec::with_capacity(end - start);
1876            for i in start..end {
1877                if arr.is_null(i) {
1878                    values.push("");
1879                    stats.record_null();
1880                    validity.push(false);
1881                } else {
1882                    values.push(arr.value(i));
1883                    validity.push(true);
1884                }
1885            }
1886            let distinct: std::collections::HashSet<&str> = values.iter().copied().collect();
1887            stats.distinct_count = distinct.len() as u32;
1888            stats.value_count = (values.len() as u64).saturating_sub(stats.null_count);
1889            let block = ColumnEncoder::encode_string(&values)?;
1890            Ok(EncodedColumn::new(block, stats, validity))
1891        }
1892        (_, data_types::F64) => {
1893            let arr = batch
1894                .column_by_name(&col.name)
1895                .and_then(|c| c.as_any().downcast_ref::<Float64Array>())
1896                .ok_or_else(|| SegmentError::CorruptFile {
1897                    detail: format!("missing f64 column '{}'", col.name),
1898                })?;
1899            let mut values = Vec::with_capacity(end - start);
1900            for i in start..end {
1901                if arr.is_null(i) {
1902                    values.push(0.0);
1903                    stats.record_null();
1904                    validity.push(false);
1905                } else {
1906                    let v = arr.value(i);
1907                    values.push(v);
1908                    stats.update_f64(v);
1909                    validity.push(true);
1910                }
1911            }
1912            let block = ColumnEncoder::encode_f64(&values, float_encoding)?;
1913            Ok(EncodedColumn::new(block, stats, validity))
1914        }
1915        (_, data_types::I64) => {
1916            let arr = batch
1917                .column_by_name(&col.name)
1918                .and_then(|c| c.as_any().downcast_ref::<Int64Array>())
1919                .ok_or_else(|| SegmentError::CorruptFile {
1920                    detail: format!("missing i64 column '{}'", col.name),
1921                })?;
1922            let mut values = Vec::with_capacity(end - start);
1923            for i in start..end {
1924                if arr.is_null(i) {
1925                    values.push(0);
1926                    stats.record_null();
1927                    validity.push(false);
1928                } else {
1929                    let v = arr.value(i);
1930                    values.push(v);
1931                    stats.update_i64(v);
1932                    validity.push(true);
1933                }
1934            }
1935            let block = ColumnEncoder::encode_i64(&values)?;
1936            Ok(EncodedColumn::new(block, stats, validity))
1937        }
1938        (_, data_types::U64) => {
1939            let arr = batch
1940                .column_by_name(&col.name)
1941                .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
1942                .ok_or_else(|| SegmentError::CorruptFile {
1943                    detail: format!("missing u64 column '{}'", col.name),
1944                })?;
1945            let mut values = Vec::with_capacity(end - start);
1946            for i in start..end {
1947                if arr.is_null(i) {
1948                    values.push(0);
1949                    stats.record_null();
1950                    validity.push(false);
1951                } else {
1952                    let v = arr.value(i);
1953                    values.push(v);
1954                    stats.update_u64(v);
1955                    validity.push(true);
1956                }
1957            }
1958            let block = ColumnEncoder::encode_u64(&values)?;
1959            Ok(EncodedColumn::new(block, stats, validity))
1960        }
1961        (_, data_types::DECIMAL) => {
1962            let arr = batch
1963                .column_by_name(&col.name)
1964                .and_then(|c| c.as_any().downcast_ref::<Decimal128Array>())
1965                .ok_or_else(|| SegmentError::CorruptFile {
1966                    detail: format!("missing decimal column '{}'", col.name),
1967                })?;
1968            // The column's scale and the array's must agree: the mantissas
1969            // are stored raw, so writing a scale-2 array into a scale-4
1970            // column would multiply every value by a hundred.
1971            let column_scale = col.decimal_scale.unwrap_or(0);
1972            let array_scale = u8::try_from(arr.scale()).map_err(|_| SegmentError::CorruptFile {
1973                detail: format!("negative decimal scale for column '{}'", col.name),
1974            })?;
1975            if array_scale != column_scale {
1976                return Err(SegmentError::CorruptFile {
1977                    detail: format!(
1978                        "decimal column '{}' is scale {column_scale} but the batch is scale \
1979                         {array_scale}",
1980                        col.name
1981                    ),
1982                });
1983            }
1984            let mut values = Vec::with_capacity(end - start);
1985            for i in start..end {
1986                if arr.is_null(i) {
1987                    values.push(0);
1988                    stats.record_null();
1989                    validity.push(false);
1990                } else {
1991                    let v = arr.value(i);
1992                    values.push(v);
1993                    stats.update_decimal(v);
1994                    validity.push(true);
1995                }
1996            }
1997            let block = ColumnEncoder::encode_decimal(&values)?;
1998            Ok(EncodedColumn::new(block, stats, validity))
1999        }
2000        (_, data_types::BOOL) => {
2001            let arr = batch
2002                .column_by_name(&col.name)
2003                .and_then(|c| c.as_any().downcast_ref::<BooleanArray>())
2004                .ok_or_else(|| SegmentError::CorruptFile {
2005                    detail: format!("missing bool column '{}'", col.name),
2006                })?;
2007            let mut values = Vec::with_capacity(end - start);
2008            for i in start..end {
2009                if arr.is_null(i) {
2010                    values.push(false);
2011                    stats.record_null();
2012                    validity.push(false);
2013                } else {
2014                    let v = arr.value(i);
2015                    values.push(v);
2016                    stats.update_bool(v);
2017                    validity.push(true);
2018                }
2019            }
2020            let block = ColumnEncoder::encode_bool(&values)?;
2021            Ok(EncodedColumn::new(block, stats, validity))
2022        }
2023        (_, data_types::HISTOGRAM) => {
2024            // The batch path — what a flush and a compaction take. The values
2025            // arrive already encoded, because a histogram column is `Binary`
2026            // in Arrow and the memtable put the `postcard` bytes there.
2027            let arr = batch
2028                .column_by_name(&col.name)
2029                .and_then(|c| c.as_any().downcast_ref::<arrow::array::BinaryArray>())
2030                .ok_or_else(|| SegmentError::CorruptFile {
2031                    detail: format!("missing histogram column '{}'", col.name),
2032                })?;
2033            let mut owned: Vec<&[u8]> = Vec::with_capacity(end - start);
2034            for i in start..end {
2035                if arr.is_null(i) {
2036                    owned.push(&[]);
2037                    stats.record_null();
2038                    validity.push(false);
2039                } else {
2040                    owned.push(arr.value(i));
2041                    validity.push(true);
2042                }
2043            }
2044            stats.value_count = (owned.len() as u64).saturating_sub(stats.null_count);
2045            let block = chronix_encoding::EncodedBlock {
2046                encoding: chronix_encoding::EncodingType::Bytes,
2047                payload: chronix_encoding::PlainEncoder::encode_bytes(&owned)?,
2048            };
2049            Ok(EncodedColumn::new(block, stats, validity))
2050        }
2051        _ => Err(SegmentError::CorruptFile {
2052            detail: format!("unknown column type: {}", col.data_type),
2053        }),
2054    }
2055}
2056
2057#[cfg(test)]
2058mod tests {
2059    use super::*;
2060    use crate::segment::reader::SegmentReader;
2061    use chronix_core::config::{CompressionCodec, FloatEncoding};
2062    use chronix_core::types::{FieldValue, Point, SeriesKey};
2063    use std::collections::{BTreeMap, HashMap};
2064
2065    fn make_point(host: &str, cpu: f64, ts: i64) -> Point {
2066        let tags = BTreeMap::from([("host".to_string(), host.to_string())]);
2067        let key = SeriesKey::new("cpu_usage", tags).unwrap();
2068        let fields = BTreeMap::from([("cpu".to_string(), FieldValue::F64(cpu))]);
2069        Point::new(key, fields, ts).unwrap()
2070    }
2071
2072    #[test]
2073    fn per_column_codec_mixed_lz4_zstd() {
2074        let dir = tempfile::tempdir().unwrap();
2075        let path = dir.path().join("mixed_codec.csx");
2076
2077        let points: Vec<Point> = (0..200)
2078            .map(|i| make_point("server-1", 50.0 + i as f64, 1_000_000 + i * 10_000))
2079            .collect();
2080
2081        // Override the "cpu" column to use Zstd while default is Lz4
2082        let mut overrides = HashMap::new();
2083        overrides.insert("cpu".to_string(), CompressionCodec::Zstd);
2084
2085        let config = SegmentWriterConfig {
2086            row_group_size: 100,
2087            compress: true,
2088            float_encoding: FloatEncoding::Gorilla,
2089            compression_codec: CompressionCodec::Lz4,
2090            zstd_level: 3,
2091            column_codec_overrides: overrides,
2092            ..Default::default()
2093        };
2094
2095        let mut writer = SegmentWriter::new(&path, config).unwrap();
2096        writer.write_rows(&points).unwrap();
2097        let meta = writer.finalize().unwrap();
2098
2099        assert_eq!(meta.row_count, 200);
2100
2101        // Read back and verify data integrity
2102        let reader = SegmentReader::open(&path).unwrap();
2103        let batch = reader.read_all().unwrap();
2104        assert_eq!(batch.num_rows(), 200);
2105
2106        // Verify cpu values
2107        let cpu_col = batch
2108            .column_by_name("cpu")
2109            .unwrap()
2110            .as_any()
2111            .downcast_ref::<arrow::array::Float64Array>()
2112            .unwrap();
2113        for i in 0..200 {
2114            assert!(
2115                (cpu_col.value(i) - (50.0 + i as f64)).abs() < 1e-9,
2116                "cpu mismatch at row {i}"
2117            );
2118        }
2119    }
2120
2121    #[test]
2122    fn p09_zstd_dict_training_roundtrip() {
2123        let dir = tempfile::tempdir().unwrap();
2124        let path = dir.path().join("dict_trained.csx");
2125
2126        // Generate enough data for 3+ row groups so dictionary training
2127        // kicks in on the first and is used on subsequent groups.
2128        let points: Vec<Point> = (0..600)
2129            .map(|i| {
2130                let host = format!("host-{}", i % 5);
2131                make_point(&host, 10.0 + (i as f64) * 0.1, 1_000_000 + i * 10_000)
2132            })
2133            .collect();
2134
2135        let config = SegmentWriterConfig {
2136            row_group_size: 100,
2137            compress: true,
2138            compression_codec: CompressionCodec::Zstd,
2139            zstd_level: 3,
2140            zstd_dict_training: true,
2141            ..Default::default()
2142        };
2143
2144        let mut writer = SegmentWriter::new(&path, config).unwrap();
2145        writer.write_rows(&points).unwrap();
2146        let meta = writer.finalize().unwrap();
2147
2148        assert_eq!(meta.row_count, 600);
2149        assert!(meta.row_group_count >= 3);
2150
2151        // Verify the dictionary was trained and embedded
2152        let reader = SegmentReader::open(&path).unwrap();
2153        let batch = reader.read_all().unwrap();
2154        assert_eq!(batch.num_rows(), 600);
2155
2156        // Verify data integrity
2157        let cpu_col = batch
2158            .column_by_name("cpu")
2159            .unwrap()
2160            .as_any()
2161            .downcast_ref::<arrow::array::Float64Array>()
2162            .unwrap();
2163        assert_eq!(cpu_col.len(), 600);
2164    }
2165
2166    /// Segment metadata must stay proportional to the data, not fixed.
2167    ///
2168    /// Every column used to carry a full-precision HyperLogLog sketch
2169    /// (16 KiB) plus an equi-depth histogram, unconditionally — roughly 49 KiB
2170    /// of fixed overhead per segment that nothing ever read: the DataFusion
2171    /// statistics provider works from the catalog's exact `distinct_count`.
2172    /// On an hourly-sharded gateway that dwarfed the data and burned flash for
2173    /// nothing, so the sketches were removed. This pins the result.
2174    #[test]
2175    fn segment_metadata_overhead_stays_proportional_to_data() {
2176        let dir = tempfile::tempdir().unwrap();
2177        let path = dir.path().join("overhead.csx");
2178
2179        let points: Vec<Point> = (0..500)
2180            .map(|i| {
2181                let key = SeriesKey::new(
2182                    "m".to_string(),
2183                    std::collections::BTreeMap::from([(
2184                        "host".to_string(),
2185                        format!("h{}", i % 10),
2186                    )]),
2187                )
2188                .unwrap();
2189                Point::new(
2190                    key,
2191                    std::collections::BTreeMap::from([(
2192                        "cpu".to_string(),
2193                        FieldValue::F64(f64::from(i % 100) / 10.0),
2194                    )]),
2195                    i64::from(i) * 1_000_000,
2196                )
2197                .unwrap()
2198            })
2199            .collect();
2200
2201        let mut writer = SegmentWriter::new(&path, SegmentWriterConfig::default()).unwrap();
2202        writer.write_rows(&points).unwrap();
2203        writer.finalize().unwrap();
2204
2205        let file_size = std::fs::metadata(&path).unwrap().len();
2206        let block_bytes: u64 = {
2207            let reader = SegmentReader::open(&path).unwrap();
2208            (0..reader.row_group_count())
2209                .filter_map(|rg| reader.row_group_blocks(rg))
2210                .flatten()
2211                .map(|b| u64::from(b.length))
2212                .sum()
2213        };
2214        let overhead = file_size - block_bytes;
2215
2216        assert!(
2217            overhead < 8 * 1024,
2218            "segment carries {overhead} bytes of non-block overhead for {block_bytes} bytes of \
2219             data — fixed-size metadata has crept back in"
2220        );
2221    }
2222}