Skip to main content

chronix_engine/memtable/
memtable.rs

1//! Concurrent memtable backed by a lock-free skip list.
2//!
3//! The [`Memtable`] is the primary in-memory write buffer. It accepts
4//! [`Point`] inserts, stores them in a concurrent skip
5//! list ordered by `(series_key_hash, timestamp)`, and supports range scans
6//! for reads.
7//!
8//! # Threading model
9//!
10//! Multiple writers can insert concurrently without external locking.
11//! The skip list (from [`crossbeam_skiplist`]) provides lock-free concurrent
12//! access. An [`AtomicUsize`] tracks the estimated memory footprint, and
13//! [`AtomicU64`] values track the WAL sequence range for flush coordination.
14
15use std::collections::HashSet;
16use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
17use std::sync::Arc;
18
19use crossbeam_skiplist::SkipMap;
20use dashmap::mapref::entry::Entry;
21use dashmap::DashMap;
22
23use chronix_core::types::{Point, SeriesKey, Timestamp};
24
25use crate::memtable::error::{MemtableError, Result};
26use crate::memtable::interner::StringInterner;
27use crate::memtable::key::{MemtableEntry, MemtableKey, NODE_OVERHEAD};
28
29/// One series' tag set, shared by every point of the series in a memtable.
30pub type TagSet = Arc<[(Arc<str>, Arc<str>)]>;
31
32/// A concurrent in-memory write buffer for time-series data.
33///
34/// Points are stored in a [`crossbeam_skiplist::SkipMap`] keyed by
35/// [`MemtableKey`] `(series_key_hash, timestamp)`. The skip list provides
36/// lock-free concurrent inserts and ordered iteration.
37///
38/// # Deduplication
39///
40/// If a point with the same `(series_key_hash, timestamp)` already exists,
41/// the new value overwrites the old one (last-write-wins).
42pub struct Memtable {
43    /// Lock-free concurrent skip list.
44    data: SkipMap<MemtableKey, MemtableEntry>,
45    /// String interner for deduplicating measurement/tag strings.
46    interner: StringInterner,
47    /// Secondary index: measurement name → set of series hashes.
48    ///
49    /// Enables O(S × log N) `scan_measurement()` instead of O(N) full scans,
50    /// where S is the number of distinct series for the target measurement.
51    ///
52    /// Uses [`DashMap`] with per-shard locking instead of a
53    /// global `RwLock<HashMap>` to eliminate write-lock contention on the
54    /// hot insert path.
55    measurement_index: DashMap<Arc<str>, HashSet<u64>>,
56    /// One interned tag set per series (keyed by canonical form), shared by
57    /// every point of the series in this memtable.
58    series_tags: DashMap<Arc<str>, TagSet>,
59    /// Estimated total memory usage in bytes, excluding `measurement_index`.
60    estimated_size: AtomicUsize,
61    /// Estimated `measurement_index` footprint in bytes.
62    ///
63    /// Accumulated at the two sites that write the index, not recomputed on
64    /// read: the flush controller's capacity check calls `estimated_size()`
65    /// once per point, and walking the `DashMap` there locks every shard.
66    index_size: AtomicUsize,
67    /// Minimum WAL sequence number covered by this memtable.
68    min_wal_seq: AtomicU64,
69    /// Maximum WAL sequence number covered by this memtable.
70    max_wal_seq: AtomicU64,
71    /// Whether this memtable has been frozen (read-only).
72    frozen: AtomicBool,
73}
74
75impl std::fmt::Debug for Memtable {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        f.debug_struct("Memtable")
78            .field("len", &self.data.len())
79            .field(
80                "estimated_size",
81                &self.estimated_size.load(Ordering::Relaxed),
82            )
83            .field("interned_strings", &self.interner.len())
84            .field("min_wal_seq", &self.min_wal_seq())
85            .field("max_wal_seq", &self.max_wal_seq())
86            .field("frozen", &self.frozen.load(Ordering::Relaxed))
87            .finish()
88    }
89}
90
91impl Memtable {
92    /// Create a new empty memtable.
93    #[must_use]
94    pub fn new() -> Self {
95        Self {
96            data: SkipMap::new(),
97            interner: StringInterner::new(),
98            measurement_index: DashMap::new(),
99            series_tags: DashMap::new(),
100            estimated_size: AtomicUsize::new(0),
101            index_size: AtomicUsize::new(0),
102            min_wal_seq: AtomicU64::new(u64::MAX),
103            max_wal_seq: AtomicU64::new(0),
104            frozen: AtomicBool::new(false),
105        }
106    }
107
108    /// Insert a point into the memtable.
109    ///
110    /// The point is stored keyed by `(series_key_hash, timestamp)`. If a
111    /// point with the same key already exists, it is overwritten
112    /// (last-write-wins).
113    ///
114    /// # Errors
115    ///
116    /// Returns [`MemtableError::Frozen`] if the memtable has been frozen.
117    pub fn insert(&self, point: &Point) -> Result<()> {
118        if self.frozen.load(Ordering::Acquire) {
119            return Err(MemtableError::Frozen);
120        }
121
122        let canonical: Arc<str> = self.interner.intern(point.series_key().canonical_form());
123        let key = MemtableKey::new(point.series_key().hash_fnv(), point.timestamp(), canonical);
124
125        let measurement = self.interner.intern(point.series_key().measurement());
126        let measurement_for_index = Arc::clone(&measurement);
127        let tags = self.series_tag_set(&key.series_canonical, point.series_key());
128
129        let entry = MemtableEntry {
130            measurement,
131            tags,
132            fields: point
133                .fields()
134                .iter()
135                .map(|(k, v)| (Arc::clone(k), v.clone()))
136                .collect(),
137        };
138
139        let entry_size =
140            entry.estimated_size() + std::mem::size_of::<MemtableKey>() + NODE_OVERHEAD;
141
142        // If the key already exists, subtract the old entry size to avoid
143        // monotonically inflating the estimate on overwrites.
144        // Use saturating subtraction to prevent underflow when concurrent
145        // overwrites race on the same key (TOCTOU between get and insert).
146        // Crossbeam SkipMap::insert() doesn't return the old value,
147        // so a separate get() is required. The two O(log N) traversals are
148        // inherent to the SkipMap API.
149        if let Some(existing) = self.data.get(&key) {
150            let old_size = existing.value().estimated_size()
151                + std::mem::size_of::<MemtableKey>()
152                + NODE_OVERHEAD;
153            let _ =
154                self.estimated_size
155                    .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
156                        Some(current.saturating_sub(old_size))
157                    });
158        }
159
160        self.data.insert(key, entry);
161        let new_size = self.estimated_size.fetch_add(entry_size, Ordering::Relaxed) + entry_size;
162
163        // Emit memtable memory usage gauge so operators can monitor
164        // memory pressure and flush readiness.
165        metrics::gauge!("chronix_memtable_memory_bytes").set(new_size as f64);
166        metrics::gauge!("chronix_memtable_entries").set(self.data.len() as f64);
167
168        // Update secondary measurement index so scan_measurement() can
169        // use targeted range scans instead of full table iteration.
170        // DashMap entry API — only per-shard lock, no global contention.
171        {
172            let hash = point.series_key().hash_fnv();
173            let delta = self.index_insert(measurement_for_index, hash);
174            if delta > 0 {
175                self.index_size.fetch_add(delta, Ordering::Relaxed);
176            }
177        }
178
179        Ok(())
180    }
181
182    /// Insert a point with an associated WAL sequence number.
183    ///
184    /// Updates the min/max WAL sequence tracking in addition to inserting
185    /// the point.
186    ///
187    /// # Errors
188    ///
189    /// Returns [`MemtableError::Frozen`] if the memtable has been frozen.
190    pub fn insert_with_wal_seq(&self, point: &Point, wal_seq: u64) -> Result<()> {
191        self.insert(point)?;
192        self.update_wal_seq(wal_seq);
193        Ok(())
194    }
195
196    /// Insert a batch of points in one call.
197    ///
198    /// More efficient than calling `insert` in a loop because:
199    /// The frozen check is performed once.
200    /// Strings are pre-interned and the measurement index is updated in
201    ///   a single write-lock acquisition.
202    /// Size adjustments are accumulated locally and flushed with a single
203    ///   `fetch_add`.
204    ///
205    /// Returns the number of points successfully inserted.
206    ///
207    /// # Errors
208    ///
209    /// Returns [`MemtableError::Frozen`] if the memtable has been frozen.
210    pub fn insert_batch(&self, points: &[Point]) -> Result<usize> {
211        if self.frozen.load(Ordering::Acquire) {
212            return Err(MemtableError::Frozen);
213        }
214
215        let mut total_size_delta: i64 = 0;
216        // Collect (measurement, hash) pairs for batch measurement index update.
217        let mut new_index_entries: Vec<(Arc<str>, u64)> = Vec::new();
218
219        // Track previous key to skip overwrite check when
220        // consecutive batch points have different keys (the common case
221        // for append-only time-series ingestion).
222        let mut prev_hash: u64 = 0;
223        let mut prev_ts: i64 = i64::MIN;
224
225        for point in points {
226            let canonical: Arc<str> = self.interner.intern(point.series_key().canonical_form());
227            let key = MemtableKey::new(point.series_key().hash_fnv(), point.timestamp(), canonical);
228
229            let measurement = self.interner.intern(point.series_key().measurement());
230            let tags = self.series_tag_set(&key.series_canonical, point.series_key());
231
232            let entry = MemtableEntry {
233                measurement: Arc::clone(&measurement),
234                tags,
235                fields: point
236                    .fields()
237                    .iter()
238                    .map(|(k, v)| (Arc::clone(k), v.clone()))
239                    .collect(),
240            };
241
242            let entry_size =
243                entry.estimated_size() + std::mem::size_of::<MemtableKey>() + NODE_OVERHEAD;
244
245            // Only check for overwrites when the key could collide
246            // with an existing entry. For strictly increasing timestamps
247            // within the same series or different series, skip the O(log N)
248            // get() entirely — the insert alone is sufficient.
249            let cur_hash = point.series_key().hash_fnv();
250            let cur_ts = point.timestamp();
251            let may_overwrite = cur_hash == prev_hash && cur_ts == prev_ts;
252
253            if may_overwrite {
254                if let Some(existing) = self.data.get(&key) {
255                    let old_size = existing.value().estimated_size()
256                        + std::mem::size_of::<MemtableKey>()
257                        + NODE_OVERHEAD;
258                    total_size_delta -= old_size.min(i64::MAX as usize) as i64;
259                }
260            }
261
262            prev_hash = cur_hash;
263            prev_ts = cur_ts;
264
265            self.data.insert(key, entry);
266            // Saturating cast prevents theoretical i64 overflow.
267            total_size_delta += entry_size.min(i64::MAX as usize) as i64;
268
269            new_index_entries.push((measurement, cur_hash));
270        }
271
272        // Flush accumulated size delta in one atomic operation.
273        if total_size_delta >= 0 {
274            self.estimated_size
275                .fetch_add(total_size_delta as usize, Ordering::Relaxed);
276        } else {
277            let _ =
278                self.estimated_size
279                    .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
280                        Some(current.saturating_sub(total_size_delta.unsigned_abs() as usize))
281                    });
282        }
283
284        // DashMap entry API — per-shard locking, no global contention.
285        let mut index_delta = 0usize;
286        for (measurement, hash) in new_index_entries {
287            index_delta += self.index_insert(measurement, hash);
288        }
289        if index_delta > 0 {
290            self.index_size.fetch_add(index_delta, Ordering::Relaxed);
291        }
292
293        // Emit memtable memory usage gauge after batch insert.
294        metrics::gauge!("chronix_memtable_memory_bytes")
295            .set(self.estimated_size.load(Ordering::Relaxed) as f64);
296        metrics::gauge!("chronix_memtable_entries").set(self.data.len() as f64);
297
298        Ok(points.len())
299    }
300
301    /// Insert a batch of points with an associated WAL sequence number.
302    ///
303    /// Uses the batch to set the WAL sequence range, then delegates to
304    /// `insert_batch`.
305    ///
306    /// # Errors
307    ///
308    /// Returns [`MemtableError::Frozen`] if the memtable has been frozen.
309    pub fn insert_batch_with_wal_seq(&self, points: &[Point], wal_seq: u64) -> Result<usize> {
310        let count = self.insert_batch(points)?;
311        self.update_wal_seq(wal_seq);
312        Ok(count)
313    }
314
315    /// Update the WAL sequence range tracked by this memtable.
316    fn update_wal_seq(&self, seq: u64) {
317        // Update min: keep the smallest
318        let mut current = self.min_wal_seq.load(Ordering::Relaxed);
319        while seq < current {
320            match self.min_wal_seq.compare_exchange_weak(
321                current,
322                seq,
323                Ordering::Release,
324                Ordering::Relaxed,
325            ) {
326                Ok(_) => break,
327                Err(actual) => current = actual,
328            }
329        }
330
331        // Update max: keep the largest
332        let mut current = self.max_wal_seq.load(Ordering::Relaxed);
333        while seq > current {
334            match self.max_wal_seq.compare_exchange_weak(
335                current,
336                seq,
337                Ordering::Release,
338                Ordering::Relaxed,
339            ) {
340                Ok(_) => break,
341                Err(actual) => current = actual,
342            }
343        }
344    }
345
346    /// Scan points for a specific series within a time range.
347    ///
348    /// Returns points sorted by timestamp for the series identified by
349    /// `series_key`. Only points within `[min_ts, max_ts]` are returned.
350    /// Uses canonical form comparison (not hash) to avoid false positives
351    /// from hash collisions.
352    #[must_use]
353    pub fn scan(&self, series_key: &SeriesKey, min_ts: Timestamp, max_ts: Timestamp) -> Vec<Point> {
354        let hash = series_key.hash_fnv();
355        let canonical = series_key.canonical_form();
356        // Use the canonical form for exact range boundaries so the
357        // scan covers only entries for this specific series.
358        let start = MemtableKey::new(hash, min_ts, Arc::from(canonical));
359        let end = MemtableKey::new(hash, max_ts, Arc::from(canonical));
360
361        self.data
362            .range(start..=end)
363            .filter(|entry| &*entry.key().series_canonical == canonical)
364            .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
365            .collect()
366    }
367
368    /// Scan all points across all series within a time range.
369    ///
370    /// Returns all points sorted by `(series_key_hash, timestamp)`.
371    ///
372    /// **Note:** this is intentionally O(N) — it scans all series.
373    /// For measurement-scoped queries (the common case), use
374    /// `scan_measurement` which leverages the measurement index for
375    /// O(S × log N) per-series range scans.
376    #[must_use]
377    pub fn scan_range(&self, min_ts: Timestamp, max_ts: Timestamp) -> Vec<Point> {
378        // Scan the full BTree but filter by timestamp. No sentinel upper
379        // bound needed — just iterate all entries and filter.
380        self.data
381            .iter()
382            .filter(|entry| {
383                let ts = entry.key().timestamp;
384                ts >= min_ts && ts <= max_ts
385            })
386            .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
387            .collect()
388    }
389
390    /// Scan all points in the memtable, sorted by `(series_key_hash,
391    /// timestamp)`.
392    #[must_use]
393    pub fn scan_all(&self) -> Vec<Point> {
394        self.data
395            .iter()
396            .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
397            .collect()
398    }
399
400    /// Scan all points belonging to a specific measurement within a time range.
401    ///
402    /// Uses a secondary measurement index to perform targeted range scans
403    /// per series hash, achieving O(S × log N) complexity instead of O(N)
404    /// full-table iteration, where S is the number of distinct series for
405    /// the measurement.
406    #[must_use]
407    pub fn scan_measurement(
408        &self,
409        measurement: &str,
410        min_ts: Timestamp,
411        max_ts: Timestamp,
412    ) -> Vec<Point> {
413        let hashes = match self.measurement_index.get(measurement) {
414            Some(set) => set.iter().copied().collect::<Vec<_>>(),
415            None => return Vec::new(),
416        };
417
418        let mut results = Vec::new();
419        for hash in hashes {
420            // Build range bounds for this series hash across the timestamp window.
421            let start = MemtableKey::lower_bound(hash, min_ts);
422            let end = MemtableKey::upper_bound(hash, max_ts);
423            results.extend(
424                self.data
425                    .range(start..=end)
426                    .filter(|entry| {
427                        let ts = entry.key().timestamp;
428                        ts >= min_ts && ts <= max_ts && &*entry.value().measurement == measurement
429                    })
430                    .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value())),
431            );
432        }
433        results
434    }
435
436    /// The canonical form of every series this memtable holds.
437    ///
438    /// The cardinality budget is derived from what the database holds, and a
439    /// memtable holds series no segment does yet — so a repair that consulted
440    /// only the segments on disk would release a series written a second ago
441    /// and then count it again on the next write.
442    ///
443    /// `series_tags` is already keyed by canonical form, so this walks the
444    /// map rather than scanning the points.
445    #[must_use]
446    pub fn series_canonical_forms(&self) -> Vec<Arc<str>> {
447        self.series_tags
448            .iter()
449            .map(|e| Arc::clone(e.key()))
450            .collect()
451    }
452
453    /// Hand every `(tag key, tag value)` pair this memtable holds to `f`.
454    ///
455    /// Label discovery has to answer from the memtable as well as from the
456    /// segments: the tag inverted index is built at flush, so a series
457    /// written a second ago is in no segment and therefore in no index, and
458    /// an answer drawn from the index alone omits exactly the labels a fresh
459    /// database has. A visitor rather than a `Vec`, because the caller is
460    /// building a set and the pairs repeat once per series.
461    pub fn visit_tags(&self, mut f: impl FnMut(&str, &str)) {
462        for entry in &self.series_tags {
463            for (key, value) in entry.value().iter() {
464                f(key, value);
465            }
466        }
467    }
468
469    /// Returns the number of entries in the memtable.
470    #[must_use]
471    pub fn len(&self) -> usize {
472        self.data.len()
473    }
474
475    /// Returns `true` if the memtable contains no entries.
476    #[must_use]
477    pub fn is_empty(&self) -> bool {
478        self.data.is_empty()
479    }
480
481    /// Returns the estimated memory usage in bytes.
482    ///
483    /// Heap bytes held by this memtable's string interner.
484    ///
485    /// Reported separately from [`estimated_size`](Self::estimated_size),
486    /// which counts rows and index entries: the interner holds the tag keys,
487    /// tag values and measurement names, so it grows with *cardinality* rather
488    /// than with row count and is the term a high-cardinality workload sees
489    /// first.
490    #[must_use]
491    pub fn interner_bytes(&self) -> usize {
492        self.interner.memory_bytes()
493    }
494
495    /// Includes the `measurement_index` DashMap overhead.
496    #[must_use]
497    pub fn estimated_size(&self) -> usize {
498        self.estimated_size.load(Ordering::Relaxed) + self.index_size.load(Ordering::Relaxed)
499    }
500
501    /// The interned tag set of a series, created on its first point.
502    ///
503    /// A new series charges its canonical form, its tag slice and the map
504    /// entry to the size estimate once; every later point of the series
505    /// shares the `Arc`.
506    fn series_tag_set(
507        &self,
508        canonical: &Arc<str>,
509        series_key: &SeriesKey,
510    ) -> Arc<[(Arc<str>, Arc<str>)]> {
511        if let Some(existing) = self.series_tags.get(canonical) {
512            return Arc::clone(existing.value());
513        }
514        let tags: Arc<[(Arc<str>, Arc<str>)]> = series_key
515            .tags()
516            .iter()
517            .map(|(k, v)| (self.interner.intern(k), self.interner.intern(v)))
518            .collect();
519        let mut charged = 0usize;
520        let stored = self
521            .series_tags
522            .entry(Arc::clone(canonical))
523            .or_insert_with(|| {
524                // Canonical form bytes (interned once), the tag slice, the
525                // strings the slice points at (interned, charged here on
526                // first sight), and the map entry.
527                charged = canonical.len()
528                    + tags.len() * std::mem::size_of::<(Arc<str>, Arc<str>)>()
529                    + tags.iter().map(|(k, v)| k.len() + v.len()).sum::<usize>()
530                    + 2 * std::mem::size_of::<usize>()
531                    + 64;
532                tags
533            })
534            .clone();
535        if charged > 0 {
536            self.estimated_size.fetch_add(charged, Ordering::Relaxed);
537        }
538        stored
539    }
540
541    /// Add `hash` to the index under `measurement`, returning the bytes the
542    /// index grew by.
543    ///
544    /// The accounting matches what a full walk of the map would produce: an
545    /// `Arc<str>` key (pointer pair + string bytes) plus the `HashSet`'s
546    /// allocation (capacity × 8) and its own overhead. Growth is charged as
547    /// the set's capacity actually changes, so an amortised rehash is
548    /// accounted for on the insert that triggers it.
549    fn index_insert(&self, measurement: Arc<str>, hash: u64) -> usize {
550        const SET_OVERHEAD: usize = 56;
551        match self.measurement_index.entry(measurement) {
552            Entry::Occupied(mut occupied) => {
553                let before = occupied.get().capacity();
554                occupied.get_mut().insert(hash);
555                (occupied.get().capacity() - before) * std::mem::size_of::<u64>()
556            }
557            Entry::Vacant(vacant) => {
558                let key_bytes = std::mem::size_of::<usize>() * 2 + vacant.key().len();
559                let mut set = HashSet::new();
560                set.insert(hash);
561                let set_bytes = set.capacity() * std::mem::size_of::<u64>() + SET_OVERHEAD;
562                vacant.insert(set);
563                key_bytes + set_bytes
564            }
565        }
566    }
567
568    /// Returns the minimum WAL sequence number tracked.
569    ///
570    /// Returns `None` if no WAL sequences have been recorded.
571    #[must_use]
572    pub fn min_wal_seq(&self) -> Option<u64> {
573        let val = self.min_wal_seq.load(Ordering::Acquire);
574        if val == u64::MAX {
575            None
576        } else {
577            Some(val)
578        }
579    }
580
581    /// Returns the maximum WAL sequence number tracked.
582    ///
583    /// Returns `None` if no WAL sequences have been recorded.
584    #[must_use]
585    pub fn max_wal_seq(&self) -> Option<u64> {
586        let val = self.max_wal_seq.load(Ordering::Acquire);
587        if val == 0 && self.min_wal_seq.load(Ordering::Acquire) == u64::MAX {
588            None
589        } else {
590            Some(val)
591        }
592    }
593
594    /// Returns whether this memtable has been frozen.
595    #[must_use]
596    pub fn is_frozen(&self) -> bool {
597        self.frozen.load(Ordering::Acquire)
598    }
599
600    /// Freeze this memtable, preventing any further inserts.
601    ///
602    /// After freezing, all `insert` calls will return
603    /// [`MemtableError::Frozen`].
604    pub fn freeze(&self) {
605        self.frozen.store(true, Ordering::Release);
606    }
607
608    /// Collect all points as a `Vec<Point>` for flushing to a segment.
609    #[must_use]
610    pub fn to_points(&self) -> Vec<Point> {
611        self.scan_all()
612    }
613
614    /// Convert a skip list entry back into a [`Point`].
615    ///
616    /// Returns `None` if reconstruction fails (e.g. invalid series key),
617    /// which is logged as a warning.
618    fn entry_to_point(key: &MemtableKey, entry: &MemtableEntry) -> Option<Point> {
619        let tags = entry
620            .tags
621            .iter()
622            .map(|(k, v)| (k.to_string(), v.to_string()))
623            .collect();
624        let series_key = SeriesKey::new(entry.measurement.to_string(), tags)
625            .map_err(|e| {
626                // Emit metric so operators can alert on data loss
627                // during flush, not just rely on log scanning.
628                tracing::error!(
629                    hash = key.series_key_hash,
630                    error = %e,
631                    "DATA LOSS: failed to reconstruct series key during flush"
632                );
633                metrics::counter!("chronix_memtable_dropped_points_total").increment(1);
634                e
635            })
636            .ok()?;
637        Point::new(
638            series_key,
639            entry
640                .fields
641                .iter()
642                .map(|(k, v)| (k.to_string(), v.clone()))
643                .collect(),
644            key.timestamp,
645        )
646        .map_err(|e| {
647            tracing::error!(
648                hash = key.series_key_hash,
649                timestamp = key.timestamp,
650                error = %e,
651                "DATA LOSS: failed to reconstruct point during flush"
652            );
653            metrics::counter!("chronix_memtable_dropped_points_total").increment(1);
654            e
655        })
656        .ok()
657    }
658}
659
660impl Default for Memtable {
661    fn default() -> Self {
662        Self::new()
663    }
664}
665
666// SAFETY: `SkipMap` is `Send + Sync`, and all other fields are atomic types.
667// The `Memtable` is safe to share across threads.
668// (This is automatically derived, but we document it explicitly.)
669
670#[cfg(test)]
671mod tests {
672    use super::*;
673    use std::collections::BTreeMap;
674
675    use chronix_core::FieldValue;
676
677    fn make_point(measurement: &str, host: &str, ts: i64, value: f64) -> Point {
678        let tags: BTreeMap<String, String> = [("host".to_string(), host.to_string())]
679            .into_iter()
680            .collect();
681        let series_key = SeriesKey::new(measurement.to_string(), tags).unwrap();
682        let fields: BTreeMap<String, FieldValue> = [("value".to_string(), FieldValue::F64(value))]
683            .into_iter()
684            .collect();
685        Point::new(series_key, fields, ts).unwrap()
686    }
687
688    #[test]
689    fn insert_and_scan_all() {
690        let mt = Memtable::new();
691        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
692        mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
693        mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
694
695        let points = mt.scan_all();
696        assert_eq!(points.len(), 3);
697        assert_eq!(points[0].timestamp(), 100);
698        assert_eq!(points[1].timestamp(), 200);
699        assert_eq!(points[2].timestamp(), 300);
700    }
701
702    #[test]
703    fn insert_and_scan_series() {
704        let mt = Memtable::new();
705        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
706        mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
707        mt.insert(&make_point("cpu", "b", 150, 9.0)).unwrap();
708
709        let key_a = SeriesKey::new(
710            "cpu".to_string(),
711            [("host".to_string(), "a".to_string())]
712                .into_iter()
713                .collect(),
714        )
715        .unwrap();
716
717        let points = mt.scan(&key_a, 0, 300);
718        assert_eq!(points.len(), 2);
719        assert!(points
720            .iter()
721            .all(|p| p.series_key().tag("host") == Some("a")));
722    }
723
724    #[test]
725    fn scan_time_range() {
726        let mt = Memtable::new();
727        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
728        mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
729        mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
730
731        let key = SeriesKey::new(
732            "cpu".to_string(),
733            [("host".to_string(), "a".to_string())]
734                .into_iter()
735                .collect(),
736        )
737        .unwrap();
738
739        let points = mt.scan(&key, 150, 250);
740        assert_eq!(points.len(), 1);
741        assert_eq!(points[0].timestamp(), 200);
742    }
743
744    #[test]
745    fn last_write_wins() {
746        let mt = Memtable::new();
747        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
748        mt.insert(&make_point("cpu", "a", 100, 99.0)).unwrap();
749
750        let points = mt.scan_all();
751        assert_eq!(points.len(), 1);
752        assert_eq!(points[0].field("value"), Some(&FieldValue::F64(99.0)));
753    }
754
755    #[test]
756    fn empty_scan() {
757        let mt = Memtable::new();
758        let key = SeriesKey::new("cpu".to_string(), BTreeMap::new()).unwrap();
759        let points = mt.scan(&key, 0, 1000);
760        assert!(points.is_empty());
761    }
762
763    /// `estimated_size` tracks the index incrementally; this pins it to the
764    /// full walk it replaced, so the two cannot drift apart silently.
765    #[test]
766    fn index_accounting_matches_a_full_walk() {
767        let mt = Memtable::new();
768        for i in 0..500u64 {
769            let measurement = format!("m{}", i % 7);
770            let series = format!("s{}", i % 53);
771            mt.insert(&make_point(&measurement, &series, i as i64, 1.0))
772                .unwrap();
773        }
774
775        // The batch path has its own index-write site; cover it too.
776        let batch: Vec<_> = (500..800u64)
777            .map(|i| {
778                make_point(
779                    &format!("m{}", i % 11),
780                    &format!("s{}", i % 97),
781                    i as i64,
782                    2.0,
783                )
784            })
785            .collect();
786        mt.insert_batch(&batch).unwrap();
787
788        let walked: usize = mt
789            .measurement_index
790            .iter()
791            .map(|entry| {
792                let key_bytes = std::mem::size_of::<usize>() * 2 + entry.key().len();
793                let set_bytes = entry.value().capacity() * std::mem::size_of::<u64>() + 56;
794                key_bytes + set_bytes
795            })
796            .sum();
797
798        assert_eq!(
799            mt.index_size.load(Ordering::Relaxed),
800            walked,
801            "incremental index accounting drifted from a full walk"
802        );
803    }
804
805    #[test]
806    fn estimated_size_increases() {
807        let mt = Memtable::new();
808        assert_eq!(mt.estimated_size(), 0);
809        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
810        assert!(mt.estimated_size() > 0);
811        let size_after_one = mt.estimated_size();
812        mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
813        assert!(mt.estimated_size() > size_after_one);
814    }
815
816    #[test]
817    fn wal_seq_tracking() {
818        let mt = Memtable::new();
819        assert_eq!(mt.min_wal_seq(), None);
820        assert_eq!(mt.max_wal_seq(), None);
821
822        mt.insert_with_wal_seq(&make_point("cpu", "a", 100, 1.0), 10)
823            .unwrap();
824        assert_eq!(mt.min_wal_seq(), Some(10));
825        assert_eq!(mt.max_wal_seq(), Some(10));
826
827        mt.insert_with_wal_seq(&make_point("cpu", "a", 200, 2.0), 5)
828            .unwrap();
829        assert_eq!(mt.min_wal_seq(), Some(5));
830        assert_eq!(mt.max_wal_seq(), Some(10));
831
832        mt.insert_with_wal_seq(&make_point("cpu", "a", 300, 3.0), 20)
833            .unwrap();
834        assert_eq!(mt.min_wal_seq(), Some(5));
835        assert_eq!(mt.max_wal_seq(), Some(20));
836    }
837
838    #[test]
839    fn freeze_prevents_inserts() {
840        let mt = Memtable::new();
841        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
842        mt.freeze();
843        assert!(mt.is_frozen());
844
845        let result = mt.insert(&make_point("cpu", "a", 200, 2.0));
846        assert!(result.is_err());
847
848        // Existing data is still readable
849        let points = mt.scan_all();
850        assert_eq!(points.len(), 1);
851    }
852
853    #[test]
854    fn len_and_is_empty() {
855        let mt = Memtable::new();
856        assert!(mt.is_empty());
857        assert_eq!(mt.len(), 0);
858
859        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
860        assert!(!mt.is_empty());
861        assert_eq!(mt.len(), 1);
862    }
863
864    #[test]
865    fn multiple_series() {
866        let mt = Memtable::new();
867        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
868        mt.insert(&make_point("cpu", "b", 100, 2.0)).unwrap();
869        mt.insert(&make_point("mem", "a", 100, 3.0)).unwrap();
870
871        assert_eq!(mt.len(), 3);
872
873        let all = mt.scan_all();
874        assert_eq!(all.len(), 3);
875    }
876
877    #[test]
878    fn concurrent_inserts() {
879        use std::sync::Arc;
880        use std::thread;
881
882        let mt = Arc::new(Memtable::new());
883        let mut handles = Vec::new();
884
885        for thread_id in 0..4 {
886            let mt = Arc::clone(&mt);
887            handles.push(thread::spawn(move || {
888                for i in 0..100 {
889                    let ts = i64::from(thread_id * 1000 + i);
890                    let host = format!("host-{thread_id}");
891                    mt.insert(&make_point("cpu", &host, ts, ts as f64)).unwrap();
892                }
893            }));
894        }
895
896        for h in handles {
897            h.join().unwrap();
898        }
899
900        assert_eq!(mt.len(), 400);
901        let all = mt.scan_all();
902        assert_eq!(all.len(), 400);
903
904        // Verify ordering: each series should be ordered by timestamp
905        let mut prev_key: Option<MemtableKey> = None;
906        for p in &all {
907            let key = MemtableKey::new(
908                p.series_key().hash_fnv(),
909                p.timestamp(),
910                Arc::from(p.series_key().canonical_form()),
911            );
912            if let Some(prev) = prev_key {
913                assert!(key >= prev, "entries should be ordered");
914            }
915            prev_key = Some(key);
916        }
917    }
918
919    #[allow(deprecated)]
920    #[test]
921    fn to_points_is_scan_all() {
922        let mt = Memtable::new();
923        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
924        mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
925
926        let scan = mt.scan_all();
927        let points = mt.to_points();
928        assert_eq!(scan.len(), points.len());
929        for (s, d) in scan.iter().zip(points.iter()) {
930            assert_eq!(s.timestamp(), d.timestamp());
931            assert_eq!(s.fields(), d.fields());
932        }
933    }
934
935    #[test]
936    fn scan_range_across_series() {
937        let mt = Memtable::new();
938        mt.insert(&make_point("cpu", "a", 50, 1.0)).unwrap();
939        mt.insert(&make_point("cpu", "a", 100, 2.0)).unwrap();
940        mt.insert(&make_point("cpu", "b", 75, 3.0)).unwrap();
941        mt.insert(&make_point("cpu", "b", 200, 4.0)).unwrap();
942
943        let points = mt.scan_range(60, 150);
944        // Should get: cpu/a@100 and cpu/b@75
945        assert_eq!(points.len(), 2);
946        for p in &points {
947            assert!(
948                p.timestamp() >= 60 && p.timestamp() <= 150,
949                "timestamp {} should be in [60, 150]",
950                p.timestamp()
951            );
952        }
953    }
954
955    #[test]
956    fn scan_measurement_filters_by_name() {
957        let mt = Memtable::new();
958        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
959        mt.insert(&make_point("cpu", "b", 200, 2.0)).unwrap();
960        mt.insert(&make_point("mem", "a", 150, 3.0)).unwrap();
961        mt.insert(&make_point("mem", "a", 300, 4.0)).unwrap();
962        mt.insert(&make_point("disk", "c", 250, 5.0)).unwrap();
963
964        // Scan only "cpu"
965        let cpu = mt.scan_measurement("cpu", 0, i64::MAX);
966        assert_eq!(cpu.len(), 2);
967        assert!(cpu.iter().all(|p| p.series_key().measurement() == "cpu"));
968
969        // Scan only "mem"
970        let mem = mt.scan_measurement("mem", 0, i64::MAX);
971        assert_eq!(mem.len(), 2);
972        assert!(mem.iter().all(|p| p.series_key().measurement() == "mem"));
973
974        // Scan "disk"
975        let disk = mt.scan_measurement("disk", 0, i64::MAX);
976        assert_eq!(disk.len(), 1);
977
978        // Scan non-existent measurement
979        let none = mt.scan_measurement("net", 0, i64::MAX);
980        assert!(none.is_empty());
981    }
982
983    #[test]
984    fn scan_measurement_respects_time_range() {
985        let mt = Memtable::new();
986        mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
987        mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
988        mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
989        mt.insert(&make_point("mem", "a", 200, 4.0)).unwrap();
990
991        let result = mt.scan_measurement("cpu", 150, 250);
992        assert_eq!(result.len(), 1);
993        assert_eq!(result[0].timestamp(), 200);
994    }
995
996    /// Regression test: two distinct series with the same FNV-1a hash
997    /// and timestamp must both be stored and retrievable (no silent
998    /// overwrite).
999    #[test]
1000    fn hash_collision_preserves_both_series() {
1001        let mt = Memtable::new();
1002
1003        // Create two points with the same hash and timestamp but
1004        // different canonical forms by manually constructing keys.
1005        // In production, FNV-1a collisions are rare but possible.
1006        let p1 = make_point("cpu", "host-a", 1000, 42.0);
1007        let p2 = make_point("cpu", "host-b", 1000, 99.0);
1008
1009        mt.insert(&p1).unwrap();
1010        mt.insert(&p2).unwrap();
1011
1012        // Both must be present
1013        assert_eq!(
1014            mt.len(),
1015            2,
1016            "two distinct series at same timestamp must both survive"
1017        );
1018
1019        // Each series should return its own value
1020        let s1 = mt.scan(p1.series_key(), 0, i64::MAX);
1021        assert_eq!(s1.len(), 1);
1022        assert_eq!(
1023            *s1[0].field("value").unwrap(),
1024            chronix_core::FieldValue::F64(42.0)
1025        );
1026
1027        let s2 = mt.scan(p2.series_key(), 0, i64::MAX);
1028        assert_eq!(s2.len(), 1);
1029        assert_eq!(
1030            *s2[0].field("value").unwrap(),
1031            chronix_core::FieldValue::F64(99.0)
1032        );
1033    }
1034
1035    #[test]
1036    fn insert_batch_basic() {
1037        let mt = Memtable::new();
1038        let points = vec![
1039            make_point("cpu", "a", 100, 1.0),
1040            make_point("cpu", "a", 200, 2.0),
1041            make_point("cpu", "b", 300, 3.0),
1042            make_point("mem", "a", 400, 4.0),
1043        ];
1044        let count = mt.insert_batch(&points).unwrap();
1045        assert_eq!(count, 4);
1046
1047        let all = mt.scan_all();
1048        assert_eq!(all.len(), 4);
1049
1050        // Measurement index covers both measurements
1051        let cpu = mt.scan_measurement("cpu", 0, i64::MAX);
1052        assert_eq!(cpu.len(), 3);
1053        let mem = mt.scan_measurement("mem", 0, i64::MAX);
1054        assert_eq!(mem.len(), 1);
1055    }
1056
1057    #[test]
1058    fn insert_batch_overwrite() {
1059        let mt = Memtable::new();
1060        let initial = vec![
1061            make_point("cpu", "a", 100, 1.0),
1062            make_point("cpu", "a", 200, 2.0),
1063        ];
1064        mt.insert_batch(&initial).unwrap();
1065
1066        // Overwrite same timestamps with new values
1067        let overwrite = vec![
1068            make_point("cpu", "a", 100, 10.0),
1069            make_point("cpu", "a", 200, 20.0),
1070        ];
1071        mt.insert_batch(&overwrite).unwrap();
1072
1073        let all = mt.scan_all();
1074        assert_eq!(all.len(), 2);
1075        assert_eq!(
1076            *all[0].field("value").unwrap(),
1077            chronix_core::FieldValue::F64(10.0)
1078        );
1079        assert_eq!(
1080            *all[1].field("value").unwrap(),
1081            chronix_core::FieldValue::F64(20.0)
1082        );
1083    }
1084
1085    #[test]
1086    fn insert_batch_frozen_fails() {
1087        let mt = Memtable::new();
1088        mt.freeze();
1089        let points = vec![make_point("cpu", "a", 100, 1.0)];
1090        assert!(mt.insert_batch(&points).is_err());
1091    }
1092
1093    #[test]
1094    fn insert_batch_with_wal_seq_updates_range() {
1095        let mt = Memtable::new();
1096        let points = vec![
1097            make_point("cpu", "a", 100, 1.0),
1098            make_point("cpu", "a", 200, 2.0),
1099        ];
1100        mt.insert_batch_with_wal_seq(&points, 42).unwrap();
1101        assert_eq!(mt.min_wal_seq(), Some(42));
1102        assert_eq!(mt.max_wal_seq(), Some(42));
1103    }
1104}
1105
1106/// The size estimate against the allocator's number, not against itself.
1107///
1108/// This is its own module because it installs a counting global allocator
1109/// for the whole test binary; that is cheap (two atomics per allocation)
1110/// and it is the only way to ask what a memtable *actually* costs.
1111#[cfg(test)]
1112mod calibration {
1113    use std::alloc::{GlobalAlloc, Layout, System};
1114    use std::collections::BTreeMap;
1115    use std::sync::atomic::{AtomicUsize, Ordering};
1116
1117    use chronix_core::{FieldValue, Point, SeriesKey};
1118
1119    use super::Memtable;
1120
1121    struct Counting;
1122    static LIVE: AtomicUsize = AtomicUsize::new(0);
1123
1124    // Bytes allocated minus freed *on this thread*, so the other tests
1125    // running in parallel in this binary do not show up in the window.
1126    std::thread_local! {
1127        static THREAD_LIVE: std::cell::Cell<isize> = const { std::cell::Cell::new(0) };
1128    }
1129
1130    fn thread_live() -> isize {
1131        THREAD_LIVE.with(std::cell::Cell::get)
1132    }
1133
1134    // SAFETY: every operation is forwarded to the system allocator; the
1135    // counters are the only addition. The thread-local is `const`-initialised
1136    // so touching it never allocates, and `try_with` tolerates teardown.
1137    #[allow(unsafe_code)]
1138    unsafe impl GlobalAlloc for Counting {
1139        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
1140            LIVE.fetch_add(layout.size(), Ordering::Relaxed);
1141            let _ = THREAD_LIVE.try_with(|c| c.set(c.get() + layout.size() as isize));
1142            unsafe { System.alloc(layout) }
1143        }
1144        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
1145            LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
1146            let _ = THREAD_LIVE.try_with(|c| c.set(c.get() - layout.size() as isize));
1147            unsafe { System.dealloc(ptr, layout) }
1148        }
1149    }
1150
1151    #[global_allocator]
1152    static ALLOC: Counting = Counting;
1153
1154    /// The design partner's shape: 50 series, two fields, a few hundred
1155    /// points each. The estimate must land within a quarter of what the
1156    /// allocator says — it used to be eight times below it.
1157    #[test]
1158    fn memtable_estimate_tracks_real_allocation() {
1159        let keys: Vec<SeriesKey> = (0..50)
1160            .map(|m| {
1161                SeriesKey::new(
1162                    "power",
1163                    BTreeMap::from([
1164                        ("meter".to_string(), format!("m{m:02}")),
1165                        ("site".to_string(), "home".to_string()),
1166                    ]),
1167                )
1168                .unwrap()
1169            })
1170            .collect();
1171        // Build the points first so their allocation is outside the window.
1172        let points: Vec<Point> = (0..400i64)
1173            .flat_map(|s| {
1174                keys.iter().map(move |k| {
1175                    Point::new(
1176                        k.clone(),
1177                        BTreeMap::from([
1178                            ("w".to_string(), FieldValue::F64(230.0 + s as f64)),
1179                            ("kwh".to_string(), FieldValue::F64(s as f64 / 1000.0)),
1180                        ]),
1181                        s * 1_000_000_000,
1182                    )
1183                    .unwrap()
1184                })
1185            })
1186            .collect();
1187
1188        let before = thread_live();
1189        let mt = Memtable::new();
1190        for p in &points {
1191            mt.insert(p).unwrap();
1192        }
1193        let measured = usize::try_from(thread_live() - before).unwrap_or(0);
1194        let estimated = mt.estimated_size();
1195        let ratio = estimated as f64 / measured as f64;
1196        assert!(
1197            (0.75..=1.25).contains(&ratio),
1198            "estimate {estimated} B vs allocator {measured} B for {} points: ratio {ratio:.2} \
1199             ({:.0} B/point measured, {:.0} B/point estimated)",
1200            points.len(),
1201            measured as f64 / points.len() as f64,
1202            estimated as f64 / points.len() as f64,
1203        );
1204        drop(mt);
1205    }
1206}
1207
1208/// One measurement's rows, ready for the segment writer.
1209#[derive(Debug)]
1210pub struct MeasurementBatch {
1211    /// The measurement.
1212    pub measurement: String,
1213    /// Its rows: `timestamp`, then the tag columns, then the field columns,
1214    /// each group sorted by name; rows in skip-list order (series-major,
1215    /// then time).
1216    pub batch: arrow::record_batch::RecordBatch,
1217    /// Which of the columns are tags.
1218    pub tag_columns: Vec<String>,
1219}
1220
1221/// Per-column Arrow builder for one field.
1222enum FieldBuilder {
1223    F64(arrow::array::Float64Builder),
1224    I64(arrow::array::Int64Builder),
1225    U64(arrow::array::UInt64Builder),
1226    Bool(arrow::array::BooleanBuilder),
1227    Str(arrow::array::StringBuilder),
1228    /// An exact decimal column, with the scale the first value carried.
1229    ///
1230    /// The scale is fixed for the column by the write path, which rescales
1231    /// every value to the schema's declared scale before it reaches the
1232    /// memtable — so the first value's scale is the column's scale.
1233    Decimal(arrow::array::Decimal128Builder, u8),
1234}
1235
1236impl FieldBuilder {
1237    fn for_value(v: &chronix_core::FieldValue, rows_before: usize) -> Self {
1238        use chronix_core::FieldValue as V;
1239        let mut b = match v {
1240            V::F64(_) => Self::F64(arrow::array::Float64Builder::new()),
1241            V::I64(_) => Self::I64(arrow::array::Int64Builder::new()),
1242            V::U64(_) => Self::U64(arrow::array::UInt64Builder::new()),
1243            V::Bool(_) => Self::Bool(arrow::array::BooleanBuilder::new()),
1244            V::String(_) => Self::Str(arrow::array::StringBuilder::new()),
1245            // The Arrow type is set on the builder rather than on the
1246            // finished array: `with_precision_and_scale` re-validates every
1247            // value and can fail, and there is no honest thing to return
1248            // from a failure here. The scale comes from a validated
1249            // `Decimal`, so it is inside Arrow's range by construction.
1250            V::Decimal(d) => Self::Decimal(
1251                arrow::array::Decimal128Builder::new().with_data_type(
1252                    arrow::datatypes::DataType::Decimal128(
1253                        chronix_core::DECIMAL_PRECISION,
1254                        d.scale() as i8,
1255                    ),
1256                ),
1257                d.scale(),
1258            ),
1259        };
1260        for _ in 0..rows_before {
1261            b.append_null();
1262        }
1263        b
1264    }
1265
1266    fn append_null(&mut self) {
1267        match self {
1268            Self::F64(b) => b.append_null(),
1269            Self::I64(b) => b.append_null(),
1270            Self::U64(b) => b.append_null(),
1271            Self::Bool(b) => b.append_null(),
1272            Self::Str(b) => b.append_null(),
1273            Self::Decimal(b, _) => b.append_null(),
1274        }
1275    }
1276
1277    /// Append a value; a value of another type than the column's is a null,
1278    /// because the schema registry has already refused it at write time
1279    /// and this is the last line of defence, not the first.
1280    fn append(&mut self, v: &chronix_core::FieldValue) {
1281        use chronix_core::FieldValue as V;
1282        match (self, v) {
1283            (Self::F64(b), V::F64(x)) => b.append_value(*x),
1284            (Self::I64(b), V::I64(x)) => b.append_value(*x),
1285            (Self::U64(b), V::U64(x)) => b.append_value(*x),
1286            (Self::Bool(b), V::Bool(x)) => b.append_value(*x),
1287            (Self::Str(b), V::String(x)) => b.append_value(x),
1288            // A value at another scale is rescaled rather than dropped: the
1289            // write path has already normalised the batch, so reaching this
1290            // means the value came from somewhere that had not — and a
1291            // lossless widening is still the right answer.
1292            (Self::Decimal(b, scale), V::Decimal(x)) => match x.rescale(*scale) {
1293                Ok(v) => b.append_value(v.mantissa()),
1294                Err(_) => b.append_null(),
1295            },
1296            (this, _) => this.append_null(),
1297        }
1298    }
1299
1300    fn finish(self) -> (arrow::datatypes::DataType, arrow::array::ArrayRef) {
1301        use arrow::datatypes::DataType;
1302        match self {
1303            Self::F64(mut b) => (DataType::Float64, Arc::new(b.finish())),
1304            Self::I64(mut b) => (DataType::Int64, Arc::new(b.finish())),
1305            Self::U64(mut b) => (DataType::UInt64, Arc::new(b.finish())),
1306            Self::Bool(mut b) => (DataType::Boolean, Arc::new(b.finish())),
1307            Self::Str(mut b) => (DataType::Utf8, Arc::new(b.finish())),
1308            Self::Decimal(mut b, scale) => {
1309                let dt = DataType::Decimal128(chronix_core::DECIMAL_PRECISION, scale as i8);
1310                (dt, Arc::new(b.finish()))
1311            }
1312        }
1313    }
1314}
1315
1316/// Builders for one measurement, growing columns as the entries reveal them.
1317struct MeasurementBuilder {
1318    rows: usize,
1319    timestamps: arrow::array::Int64Builder,
1320    tags: std::collections::BTreeMap<Arc<str>, arrow::array::StringBuilder>,
1321    fields: std::collections::BTreeMap<Arc<str>, FieldBuilder>,
1322}
1323
1324impl MeasurementBuilder {
1325    fn new() -> Self {
1326        Self {
1327            rows: 0,
1328            timestamps: arrow::array::Int64Builder::new(),
1329            tags: std::collections::BTreeMap::new(),
1330            fields: std::collections::BTreeMap::new(),
1331        }
1332    }
1333
1334    fn push(&mut self, key: &MemtableKey, entry: &MemtableEntry) {
1335        self.timestamps.append_value(key.timestamp);
1336        // Tags: every known column gets a value or a null; a new column is
1337        // back-filled with nulls for the rows before it.
1338        for (k, v) in entry.tags.iter() {
1339            let rows = self.rows;
1340            self.tags.entry(Arc::clone(k)).or_insert_with(|| {
1341                let mut b = arrow::array::StringBuilder::new();
1342                for _ in 0..rows {
1343                    b.append_null();
1344                }
1345                b
1346            });
1347            let _ = v;
1348        }
1349        for (name, b) in &mut self.tags {
1350            match entry.tags.iter().find(|(k, _)| k == name) {
1351                Some((_, v)) => b.append_value(v),
1352                None => b.append_null(),
1353            }
1354        }
1355        for (k, v) in entry.fields.iter() {
1356            let rows = self.rows;
1357            self.fields
1358                .entry(Arc::clone(k))
1359                .or_insert_with(|| FieldBuilder::for_value(v, rows));
1360        }
1361        for (name, b) in &mut self.fields {
1362            match entry.fields.iter().find(|(k, _)| k == name) {
1363                Some((_, v)) => b.append(v),
1364                None => b.append_null(),
1365            }
1366        }
1367        self.rows += 1;
1368    }
1369
1370    fn finish(mut self, measurement: &str) -> Option<MeasurementBatch> {
1371        use arrow::datatypes::{DataType, Field, Schema};
1372        if self.rows == 0 {
1373            return None;
1374        }
1375        use crate::segment::metadata::roles;
1376        // A tag and a string field are the same Arrow type, so the schema has
1377        // to carry the role or every consumer guesses at it.
1378        let mut fields = vec![
1379            Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false)
1380                .with_metadata(roles::arrow_metadata(roles::TIMESTAMP)),
1381        ];
1382        let mut columns: Vec<arrow::array::ArrayRef> = vec![Arc::new(self.timestamps.finish())];
1383        let tag_columns: Vec<String> = self.tags.keys().map(ToString::to_string).collect();
1384        for (name, mut b) in self.tags {
1385            fields.push(
1386                Field::new(name.as_ref(), DataType::Utf8, true)
1387                    .with_metadata(roles::arrow_metadata(roles::TAG)),
1388            );
1389            columns.push(Arc::new(b.finish()));
1390        }
1391        for (name, b) in self.fields {
1392            let (dt, arr) = b.finish();
1393            fields.push(
1394                Field::new(name.as_ref(), dt, true)
1395                    .with_metadata(roles::arrow_metadata(roles::FIELD)),
1396            );
1397            columns.push(arr);
1398        }
1399        let batch =
1400            arrow::record_batch::RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
1401                .ok()?;
1402        Some(MeasurementBatch {
1403            measurement: measurement.to_string(),
1404            batch,
1405            tag_columns,
1406        })
1407    }
1408}
1409
1410impl Memtable {
1411    /// Every row, as one Arrow batch per measurement, built straight from
1412    /// the skip-list entries.
1413    ///
1414    /// This is what a flush hands the segment writer. It replaced
1415    /// `to_points()` on that path: materialising a full memtable as
1416    /// `Point`s — two `BTreeMap`s each — cost a kilobyte a row and made the
1417    /// flush of an 8 MB memtable peak at ten times that. A batch is the
1418    /// columns themselves, and the writer encodes it without another copy.
1419    #[must_use]
1420    pub fn to_record_batches(&self) -> Vec<MeasurementBatch> {
1421        let mut builders: std::collections::BTreeMap<Arc<str>, MeasurementBuilder> =
1422            std::collections::BTreeMap::new();
1423        for entry in self.data.iter() {
1424            let value = entry.value();
1425            builders
1426                .entry(Arc::clone(&value.measurement))
1427                .or_insert_with(MeasurementBuilder::new)
1428                .push(entry.key(), value);
1429        }
1430        builders
1431            .into_iter()
1432            .filter_map(|(m, b)| b.finish(&m))
1433            .collect()
1434    }
1435}