1use 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
63type RowGroupWriteOutput = (
66 Vec<ColumnMeta>,
67 Vec<Vec<ColumnBlockMeta>>,
68 u64,
69 Option<Vec<u8>>,
70);
71
72pub const DEFAULT_ROW_GROUP_SIZE: usize = 65_536;
74
75pub const DEFAULT_TARGET_SEGMENT_SIZE_BYTES: usize = 256 * 1024 * 1024;
77
78#[derive(Debug, Clone)]
80pub struct SegmentWriterConfig {
81 pub row_group_size: usize,
83 pub max_row_group_size: usize,
94 pub compress: bool,
96 pub float_encoding: FloatEncoding,
98 pub compression_codec: CompressionCodec,
100 pub zstd_level: i32,
102 pub column_codec_overrides: HashMap<String, CompressionCodec>,
105 pub target_segment_size_bytes: usize,
118 pub bloom_fpr: f64,
126 pub zstd_dict_training: bool,
137 #[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#[derive(Debug, Clone, PartialEq)]
167pub struct SegmentMeta {
168 pub path: PathBuf,
170 pub min_timestamp: i64,
172 pub max_timestamp: i64,
174 pub row_count: u64,
176 pub series_count: u32,
178 pub byte_size: u64,
180 pub row_group_count: u32,
182 pub column_count: u16,
184 pub uncompressed_bytes: u64,
186 pub header: crate::segment::header::SegmentHeader,
188 pub column_metas: Vec<crate::segment::metadata::ColumnMeta>,
190 pub series_keys: Vec<chronix_core::SeriesKey>,
193}
194
195pub struct SegmentWriter {
200 path: PathBuf,
201 config: SegmentWriterConfig,
202 points: Vec<Point>,
204 finalized: bool,
205}
206
207impl SegmentWriter {
208 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 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 #[must_use]
242 pub fn buffered_point_count(&self) -> usize {
243 self.points.len()
244 }
245
246 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 let schema = discover_schema(points);
276
277 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 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 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 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 pa.timestamp().cmp(&pb.timestamp())
311 });
312
313 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 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 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 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 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 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 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 writer.flush()?;
407 writer.get_ref().sync_all()?;
408 drop(writer);
409
410 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 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 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 let schema = discover_schema_from_batch(batch, &tag_set)?;
498
499 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 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 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 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 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 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 writer.flush()?;
609 writer.get_ref().sync_all()?;
610 drop(writer);
611
612 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 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 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 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
696fn encode_row_groups_streaming<W: Write>(
705 schema: &[ColumnDef],
706 sorted_indices: &[usize],
707 points: &[Point],
708 config: &SegmentWriterConfig,
709 #[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 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 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 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 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 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 #[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 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 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 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 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 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#[derive(Debug, Clone)]
977struct ColumnDef {
978 name: String,
979 data_type: u8,
980 role: u8,
981 default_encoding: u8,
982 decimal_scale: Option<u8>,
984}
985
986fn discover_schema(points: &[Point]) -> Vec<ColumnDef> {
992 let mut columns: Vec<ColumnDef> = Vec::new();
993
994 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 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 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 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
1086fn 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
1110struct EncodedColumn {
1113 block: EncodedBlock,
1114 stats: ColumnStats,
1115 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 fn dense((block, stats): (EncodedBlock, ColumnStats)) -> Self {
1130 Self {
1131 block,
1132 stats,
1133 validity: None,
1134 }
1135 }
1136}
1137
1138fn 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 let mut validity = ValidityBuilder::with_capacity(indices.len());
1149
1150 match (col.role, col.data_type) {
1151 (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 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 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 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 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
1320fn 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 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
1350fn 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(×tamps)?;
1363 Ok((block, stats))
1364}
1365
1366fn 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
1385fn 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 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 dest.distinct_count = dest.distinct_count.max(src.distinct_count);
1405}
1406
1407fn 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 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 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 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 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 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
1510fn 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
1550fn encode_row_groups_from_batch_streaming<W: Write>(
1555 schema: &[ColumnDef],
1556 batch: &RecordBatch,
1557 config: &SegmentWriterConfig,
1558 #[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 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 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 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 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 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 #[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 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 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 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 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 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
1838fn 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 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(×tamps)?;
1865 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 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 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 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 let reader = SegmentReader::open(&path).unwrap();
2103 let batch = reader.read_all().unwrap();
2104 assert_eq!(batch.num_rows(), 200);
2105
2106 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 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 let reader = SegmentReader::open(&path).unwrap();
2153 let batch = reader.read_all().unwrap();
2154 assert_eq!(batch.num_rows(), 600);
2155
2156 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 #[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}