1use std::collections::HashMap;
7use std::fmt;
8
9use crate::error::{IoError, Result};
10
11pub const COLUMNAR_MAGIC: &[u8; 8] = b"SCIRCOL\x01";
13
14pub const FORMAT_VERSION: u32 = 2;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24#[repr(u8)]
25pub enum ColumnTypeTag {
26 Float64 = 0,
28 Int64 = 1,
30 Str = 2,
32 Bool = 3,
34}
35
36impl TryFrom<u8> for ColumnTypeTag {
37 type Error = IoError;
38
39 fn try_from(value: u8) -> std::result::Result<Self, Self::Error> {
40 match value {
41 0 => Ok(ColumnTypeTag::Float64),
42 1 => Ok(ColumnTypeTag::Int64),
43 2 => Ok(ColumnTypeTag::Str),
44 3 => Ok(ColumnTypeTag::Bool),
45 _ => Err(IoError::FormatError(format!(
46 "Unknown column type tag: {}",
47 value
48 ))),
49 }
50 }
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55#[repr(u8)]
56pub enum EncodingType {
57 Plain = 0,
59 Rle = 1,
61 Dictionary = 2,
63 Delta = 3,
65}
66
67impl TryFrom<u8> for EncodingType {
68 type Error = IoError;
69
70 fn try_from(value: u8) -> std::result::Result<Self, Self::Error> {
71 match value {
72 0 => Ok(EncodingType::Plain),
73 1 => Ok(EncodingType::Rle),
74 2 => Ok(EncodingType::Dictionary),
75 3 => Ok(EncodingType::Delta),
76 _ => Err(IoError::FormatError(format!(
77 "Unknown encoding type: {}",
78 value
79 ))),
80 }
81 }
82}
83
84#[derive(Debug, Clone)]
86pub enum ColumnData {
87 Float64(Vec<f64>),
89 Int64(Vec<i64>),
91 Str(Vec<String>),
93 Bool(Vec<bool>),
95}
96
97impl ColumnData {
98 pub fn len(&self) -> usize {
100 match self {
101 ColumnData::Float64(v) => v.len(),
102 ColumnData::Int64(v) => v.len(),
103 ColumnData::Str(v) => v.len(),
104 ColumnData::Bool(v) => v.len(),
105 }
106 }
107
108 pub fn is_empty(&self) -> bool {
110 self.len() == 0
111 }
112
113 pub fn type_tag(&self) -> ColumnTypeTag {
115 match self {
116 ColumnData::Float64(_) => ColumnTypeTag::Float64,
117 ColumnData::Int64(_) => ColumnTypeTag::Int64,
118 ColumnData::Str(_) => ColumnTypeTag::Str,
119 ColumnData::Bool(_) => ColumnTypeTag::Bool,
120 }
121 }
122
123 pub fn as_f64(&self) -> Result<&[f64]> {
125 match self {
126 ColumnData::Float64(v) => Ok(v),
127 _ => Err(IoError::ConversionError(format!(
128 "Column is {:?}, not Float64",
129 self.type_tag()
130 ))),
131 }
132 }
133
134 pub fn as_i64(&self) -> Result<&[i64]> {
136 match self {
137 ColumnData::Int64(v) => Ok(v),
138 _ => Err(IoError::ConversionError(format!(
139 "Column is {:?}, not Int64",
140 self.type_tag()
141 ))),
142 }
143 }
144
145 pub fn as_str(&self) -> Result<&[String]> {
147 match self {
148 ColumnData::Str(v) => Ok(v),
149 _ => Err(IoError::ConversionError(format!(
150 "Column is {:?}, not Str",
151 self.type_tag()
152 ))),
153 }
154 }
155
156 pub fn as_bool(&self) -> Result<&[bool]> {
158 match self {
159 ColumnData::Bool(v) => Ok(v),
160 _ => Err(IoError::ConversionError(format!(
161 "Column is {:?}, not Bool",
162 self.type_tag()
163 ))),
164 }
165 }
166
167 pub fn best_encoding(&self) -> EncodingType {
169 match self {
170 ColumnData::Float64(v) => {
171 if is_sorted_f64(v) {
172 EncodingType::Delta
173 } else if has_runs_f64(v) {
174 EncodingType::Rle
175 } else {
176 EncodingType::Plain
177 }
178 }
179 ColumnData::Int64(v) => {
180 if is_sorted_i64(v) {
181 EncodingType::Delta
182 } else if has_runs_i64(v) {
183 EncodingType::Rle
184 } else {
185 EncodingType::Plain
186 }
187 }
188 ColumnData::Str(v) => {
189 let unique_count = count_unique_strings(v);
190 if unique_count < v.len() / 2 {
191 EncodingType::Dictionary
192 } else if has_runs_str(v) {
193 EncodingType::Rle
194 } else {
195 EncodingType::Plain
196 }
197 }
198 ColumnData::Bool(v) => {
199 if has_runs_bool(v) {
200 EncodingType::Rle
201 } else {
202 EncodingType::Plain
203 }
204 }
205 }
206 }
207}
208
209impl fmt::Display for ColumnData {
210 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
211 match self {
212 ColumnData::Float64(v) => write!(f, "Float64[{}]", v.len()),
213 ColumnData::Int64(v) => write!(f, "Int64[{}]", v.len()),
214 ColumnData::Str(v) => write!(f, "Str[{}]", v.len()),
215 ColumnData::Bool(v) => write!(f, "Bool[{}]", v.len()),
216 }
217 }
218}
219
220#[derive(Debug, Clone)]
222pub struct Column {
223 pub name: String,
225 pub data: ColumnData,
227 pub null_mask: Option<Vec<bool>>,
238}
239
240impl Column {
241 pub fn float64(name: impl Into<String>, data: Vec<f64>) -> Self {
243 Column {
244 name: name.into(),
245 data: ColumnData::Float64(data),
246 null_mask: None,
247 }
248 }
249
250 pub fn int64(name: impl Into<String>, data: Vec<i64>) -> Self {
252 Column {
253 name: name.into(),
254 data: ColumnData::Int64(data),
255 null_mask: None,
256 }
257 }
258
259 pub fn string(name: impl Into<String>, data: Vec<String>) -> Self {
261 Column {
262 name: name.into(),
263 data: ColumnData::Str(data),
264 null_mask: None,
265 }
266 }
267
268 pub fn boolean(name: impl Into<String>, data: Vec<bool>) -> Self {
270 Column {
271 name: name.into(),
272 data: ColumnData::Bool(data),
273 null_mask: None,
274 }
275 }
276
277 pub fn float64_with_nulls(
286 name: impl Into<String>,
287 data: Vec<f64>,
288 null_mask: Vec<bool>,
289 ) -> Result<Self> {
290 Self::with_nulls(name, ColumnData::Float64(data), null_mask)
291 }
292
293 pub fn int64_with_nulls(
296 name: impl Into<String>,
297 data: Vec<i64>,
298 null_mask: Vec<bool>,
299 ) -> Result<Self> {
300 Self::with_nulls(name, ColumnData::Int64(data), null_mask)
301 }
302
303 pub fn string_with_nulls(
306 name: impl Into<String>,
307 data: Vec<String>,
308 null_mask: Vec<bool>,
309 ) -> Result<Self> {
310 Self::with_nulls(name, ColumnData::Str(data), null_mask)
311 }
312
313 pub fn boolean_with_nulls(
316 name: impl Into<String>,
317 data: Vec<bool>,
318 null_mask: Vec<bool>,
319 ) -> Result<Self> {
320 Self::with_nulls(name, ColumnData::Bool(data), null_mask)
321 }
322
323 fn with_nulls(name: impl Into<String>, data: ColumnData, null_mask: Vec<bool>) -> Result<Self> {
324 if null_mask.len() != data.len() {
325 return Err(IoError::FormatError(format!(
326 "null_mask has {} entries, expected {} (one per row)",
327 null_mask.len(),
328 data.len()
329 )));
330 }
331 Ok(Column {
332 name: name.into(),
333 data,
334 null_mask: Some(null_mask),
335 })
336 }
337
338 pub fn len(&self) -> usize {
340 self.data.len()
341 }
342
343 pub fn is_empty(&self) -> bool {
345 self.data.is_empty()
346 }
347
348 pub fn is_null(&self, row: usize) -> bool {
351 self.null_mask
352 .as_ref()
353 .and_then(|mask| mask.get(row))
354 .copied()
355 .unwrap_or(false)
356 }
357
358 pub fn null_count(&self) -> usize {
361 self.null_mask
362 .as_ref()
363 .map(|mask| mask.iter().filter(|&&is_null| is_null).count())
364 .unwrap_or(0)
365 }
366}
367
368#[derive(Debug, Clone)]
370pub struct ColumnarTable {
371 columns: Vec<Column>,
373 index: HashMap<String, usize>,
375}
376
377impl ColumnarTable {
378 pub fn new() -> Self {
380 ColumnarTable {
381 columns: Vec::new(),
382 index: HashMap::new(),
383 }
384 }
385
386 pub fn from_columns(columns: Vec<Column>) -> Result<Self> {
388 if !columns.is_empty() {
390 let expected_len = columns[0].len();
391 for col in &columns[1..] {
392 if col.len() != expected_len {
393 return Err(IoError::FormatError(format!(
394 "Column '{}' has {} rows, expected {}",
395 col.name,
396 col.len(),
397 expected_len
398 )));
399 }
400 }
401 }
402
403 let mut index = HashMap::new();
404 for (i, col) in columns.iter().enumerate() {
405 if index.contains_key(&col.name) {
406 return Err(IoError::FormatError(format!(
407 "Duplicate column name: '{}'",
408 col.name
409 )));
410 }
411 index.insert(col.name.clone(), i);
412 }
413
414 Ok(ColumnarTable { columns, index })
415 }
416
417 pub fn add_column(&mut self, column: Column) -> Result<()> {
419 if !self.columns.is_empty() && column.len() != self.num_rows() {
420 return Err(IoError::FormatError(format!(
421 "Column '{}' has {} rows, expected {}",
422 column.name,
423 column.len(),
424 self.num_rows()
425 )));
426 }
427 if self.index.contains_key(&column.name) {
428 return Err(IoError::FormatError(format!(
429 "Duplicate column name: '{}'",
430 column.name
431 )));
432 }
433 let idx = self.columns.len();
434 self.index.insert(column.name.clone(), idx);
435 self.columns.push(column);
436 Ok(())
437 }
438
439 pub fn num_rows(&self) -> usize {
441 self.columns.first().map(|c| c.len()).unwrap_or(0)
442 }
443
444 pub fn num_columns(&self) -> usize {
446 self.columns.len()
447 }
448
449 pub fn column_names(&self) -> Vec<&str> {
451 self.columns.iter().map(|c| c.name.as_str()).collect()
452 }
453
454 pub fn column(&self, name: &str) -> Result<&Column> {
456 self.index
457 .get(name)
458 .map(|&idx| &self.columns[idx])
459 .ok_or_else(|| IoError::NotFound(format!("Column '{}' not found", name)))
460 }
461
462 pub fn column_by_index(&self, idx: usize) -> Result<&Column> {
464 self.columns
465 .get(idx)
466 .ok_or_else(|| IoError::NotFound(format!("Column index {} out of range", idx)))
467 }
468
469 pub fn columns(&self) -> &[Column] {
471 &self.columns
472 }
473
474 pub fn get_f64(&self, name: &str) -> Result<&[f64]> {
476 self.column(name)?.data.as_f64()
477 }
478
479 pub fn get_i64(&self, name: &str) -> Result<&[i64]> {
481 self.column(name)?.data.as_i64()
482 }
483
484 pub fn get_str(&self, name: &str) -> Result<&[String]> {
486 self.column(name)?.data.as_str()
487 }
488
489 pub fn get_bool(&self, name: &str) -> Result<&[bool]> {
491 self.column(name)?.data.as_bool()
492 }
493}
494
495impl Default for ColumnarTable {
496 fn default() -> Self {
497 Self::new()
498 }
499}
500
501fn is_sorted_f64(data: &[f64]) -> bool {
504 if data.len() < 2 {
505 return true;
506 }
507 data.windows(2).all(|w| w[0] <= w[1])
508}
509
510fn is_sorted_i64(data: &[i64]) -> bool {
511 if data.len() < 2 {
512 return true;
513 }
514 data.windows(2).all(|w| w[0] <= w[1])
515}
516
517fn has_runs_f64(data: &[f64]) -> bool {
518 if data.len() < 4 {
519 return false;
520 }
521 let mut run_count = 0;
522 let mut i = 0;
523 while i < data.len() {
524 let val = data[i];
525 let mut run_len = 1;
526 while i + run_len < data.len() && data[i + run_len] == val {
527 run_len += 1;
528 }
529 if run_len > 1 {
530 run_count += 1;
531 }
532 i += run_len;
533 }
534 run_count * 5 >= data.len()
536}
537
538fn has_runs_i64(data: &[i64]) -> bool {
539 if data.len() < 4 {
540 return false;
541 }
542 let mut run_count = 0;
543 let mut i = 0;
544 while i < data.len() {
545 let val = data[i];
546 let mut run_len = 1;
547 while i + run_len < data.len() && data[i + run_len] == val {
548 run_len += 1;
549 }
550 if run_len > 1 {
551 run_count += 1;
552 }
553 i += run_len;
554 }
555 run_count * 5 >= data.len()
556}
557
558fn has_runs_str(data: &[String]) -> bool {
559 if data.len() < 4 {
560 return false;
561 }
562 let mut run_count = 0;
563 let mut i = 0;
564 while i < data.len() {
565 let val = &data[i];
566 let mut run_len = 1;
567 while i + run_len < data.len() && &data[i + run_len] == val {
568 run_len += 1;
569 }
570 if run_len > 1 {
571 run_count += 1;
572 }
573 i += run_len;
574 }
575 run_count * 5 >= data.len()
576}
577
578fn has_runs_bool(data: &[bool]) -> bool {
579 if data.len() < 4 {
580 return false;
581 }
582 let mut run_count = 0;
583 let mut i = 0;
584 while i < data.len() {
585 let val = data[i];
586 let mut run_len = 1;
587 while i + run_len < data.len() && data[i + run_len] == val {
588 run_len += 1;
589 }
590 if run_len > 1 {
591 run_count += 1;
592 }
593 i += run_len;
594 }
595 run_count * 5 >= data.len()
596}
597
598fn count_unique_strings(data: &[String]) -> usize {
599 let mut seen = std::collections::HashSet::new();
600 for s in data {
601 seen.insert(s.as_str());
602 }
603 seen.len()
604}