1use crate::arrow_io::utf8_value_at;
4use crate::ndarrow_bridge::{f32_view, f64_view};
5use crate::sigfig::round_to_significant_figures;
6use crate::table::{BatchChunk, BatchColumn, ChunkTable, ColGraph, ColumnVec};
7use ndarray::ArrayView1;
8use std::collections::HashMap;
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum BinType {
12 Numerical,
13 Categorical,
14}
15
16#[derive(Clone, Debug, Hash, PartialEq, Eq)]
17pub enum ValueKey {
18 Num(u32),
19 Str(String),
20}
21
22impl ValueKey {
23 pub fn from_f32(x: f32) -> Self {
24 ValueKey::Num(f32_to_bits(x))
25 }
26
27 fn as_num_bits(&self) -> Option<u32> {
28 match self {
29 ValueKey::Num(b) => Some(*b),
30 _ => None,
31 }
32 }
33}
34
35fn f32_to_bits(x: f32) -> u32 {
36 if x == 0.0 {
37 0.0f32.to_bits()
38 } else {
39 x.to_bits()
40 }
41}
42
43fn bits_to_f32(b: u32) -> f32 {
44 f32::from_bits(b)
45}
46
47#[derive(Clone, Debug)]
48pub struct NumericalBin {
49 pub bottom: f32,
50 pub top: f32,
51 pub label: String,
52}
53
54#[derive(Clone, Debug)]
55pub struct ColumnPreprocess {
56 pub bin_type: BinType,
57 pub min_v: f32,
58 pub max_v: f32,
59 pub values: HashMap<ValueKey, u64>,
60 pub bins: HashMap<ValueKey, u64>,
61 pub bin_labels: Vec<String>,
62 pub numerical_bins: Vec<NumericalBin>,
63}
64
65impl ColumnPreprocess {
66 fn new_numerical(min_v: f32, max_v: f32) -> Self {
67 Self {
68 bin_type: BinType::Numerical,
69 min_v,
70 max_v,
71 values: HashMap::new(),
72 bins: HashMap::new(),
73 bin_labels: Vec::new(),
74 numerical_bins: Vec::new(),
75 }
76 }
77
78 fn new_categorical() -> Self {
79 Self {
80 bin_type: BinType::Categorical,
81 min_v: 0.0,
82 max_v: 0.0,
83 values: HashMap::new(),
84 bins: HashMap::new(),
85 bin_labels: Vec::new(),
86 numerical_bins: Vec::new(),
87 }
88 }
89}
90
91#[derive(Clone, Debug, PartialEq, Eq)]
92pub struct BinDepth {
93 pub main: usize,
94 pub per_column: HashMap<usize, usize>,
95}
96
97impl BinDepth {
98 pub fn new(main: usize) -> Self {
99 Self {
100 main,
101 per_column: HashMap::new(),
102 }
103 }
104
105 fn depth_for(&self, col: usize) -> usize {
106 self.per_column.get(&col).copied().unwrap_or(self.main)
107 }
108}
109
110#[derive(Clone, Debug)]
111pub struct PreprocessStream {
112 pub num_chunks: u64,
113 pub col_graph: ColGraph,
114 pub cols: Vec<usize>,
115 pub preprocess_map: HashMap<usize, ColumnPreprocess>,
116 finished: bool,
117}
118
119impl PreprocessStream {
120 pub fn new(col_graph: ColGraph) -> Self {
121 let cols = col_graph.active_indices();
122 Self {
123 num_chunks: 0,
124 col_graph,
125 cols,
126 preprocess_map: HashMap::new(),
127 finished: false,
128 }
129 }
130
131 pub fn preprocess(&mut self, chunk: &ChunkTable) -> Result<(), String> {
132 if self.finished {
133 return Err("preprocess called after finish_map".into());
134 }
135 chunk.validate()?;
136 if self.num_chunks == 0 {
137 self.initialize_preprocess_map(chunk)?;
138 } else {
139 self.update_preprocess_map(chunk)?;
140 }
141 self.num_chunks += 1;
142 Ok(())
143 }
144
145 pub fn preprocess_batch(&mut self, chunk: &BatchChunk) -> Result<(), String> {
146 if self.finished {
147 return Err("preprocess called after finish_map".into());
148 }
149 chunk.validate()?;
150 if self.num_chunks == 0 {
151 self.initialize_preprocess_map_batch(chunk)?;
152 } else {
153 self.update_preprocess_map_batch(chunk)?;
154 }
155 self.num_chunks += 1;
156 Ok(())
157 }
158
159 fn column_is_numeric(col: &ColumnVec) -> bool {
160 matches!(col, ColumnVec::F32(_) | ColumnVec::F32Array(_))
161 }
162
163 fn initialize_preprocess_map(&mut self, chunk: &ChunkTable) -> Result<(), String> {
164 self.preprocess_map.clear();
165 for &col in &self.cols {
166 let c = chunk
167 .cols
168 .get(col)
169 .ok_or_else(|| format!("missing column index {col}"))?;
170 if Self::column_is_numeric(c) {
171 let arr = Self::col_as_f32(c)?;
172 let arr = round_to_significant_figures(ArrayView1::from(arr.as_slice()), 4);
173 let arr = nan_to_zero_f32(arr.to_vec());
174 let min_v = arr.iter().cloned().fold(f32::INFINITY, f32::min);
175 let max_v = arr.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
176 let mut cp = ColumnPreprocess::new_numerical(min_v, max_v);
177 merge_uniques_numerical(&mut cp.values, &arr);
178 self.preprocess_map.insert(col, cp);
179 } else {
180 let mut s = Self::col_as_utf8(c)?;
181 for t in &mut s {
182 if t == "nan" || t.eq_ignore_ascii_case("nan") {
183 *t = "empty".to_string();
184 }
185 }
186 let mut cp = ColumnPreprocess::new_categorical();
187 merge_uniques_categorical(&mut cp.values, &s);
188 self.preprocess_map.insert(col, cp);
189 }
190 }
191 Ok(())
192 }
193
194 fn update_preprocess_map(&mut self, chunk: &ChunkTable) -> Result<(), String> {
195 for &col in &self.cols {
196 let c = chunk
197 .cols
198 .get(col)
199 .ok_or_else(|| format!("missing column index {col}"))?;
200 let cp = self
201 .preprocess_map
202 .get_mut(&col)
203 .ok_or_else(|| format!("preprocess_map missing col {col}"))?;
204 match cp.bin_type {
205 BinType::Numerical => {
206 let arr = Self::col_as_f32(c)?;
207 let arr = round_to_significant_figures(ArrayView1::from(arr.as_slice()), 4);
208 let arr = nan_to_zero_f32(arr.to_vec());
209 let min_v = arr.iter().cloned().fold(f32::INFINITY, f32::min);
210 let max_v = arr.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
211 if min_v < cp.min_v {
212 cp.min_v = min_v;
213 }
214 if max_v > cp.max_v {
215 cp.max_v = max_v;
216 }
217 merge_uniques_numerical(&mut cp.values, &arr);
218 }
219 BinType::Categorical => {
220 let mut s = Self::col_as_utf8(c)?;
221 for t in &mut s {
222 if t == "nan" || t.eq_ignore_ascii_case("nan") {
223 *t = "empty".to_string();
224 }
225 }
226 merge_uniques_categorical(&mut cp.values, &s);
227 }
228 }
229 }
230 Ok(())
231 }
232
233 fn initialize_preprocess_map_batch(&mut self, chunk: &BatchChunk) -> Result<(), String> {
234 self.preprocess_map.clear();
235 for &col in &self.cols {
236 let c = chunk
237 .cols
238 .get(col)
239 .ok_or_else(|| format!("missing column index {col}"))?;
240 if c.is_numeric() {
241 let arr = Self::batch_numeric_processed(c)?;
242 let min_v = arr.iter().cloned().fold(f32::INFINITY, f32::min);
243 let max_v = arr.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
244 let mut cp = ColumnPreprocess::new_numerical(min_v, max_v);
245 merge_uniques_numerical(&mut cp.values, &arr);
246 self.preprocess_map.insert(col, cp);
247 } else {
248 let mut cp = ColumnPreprocess::new_categorical();
249 merge_uniques_categorical_arrow(&mut cp.values, c)?;
250 self.preprocess_map.insert(col, cp);
251 }
252 }
253 Ok(())
254 }
255
256 fn update_preprocess_map_batch(&mut self, chunk: &BatchChunk) -> Result<(), String> {
257 for &col in &self.cols {
258 let c = chunk
259 .cols
260 .get(col)
261 .ok_or_else(|| format!("missing column index {col}"))?;
262 let cp = self
263 .preprocess_map
264 .get_mut(&col)
265 .ok_or_else(|| format!("preprocess_map missing col {col}"))?;
266 match cp.bin_type {
267 BinType::Numerical => {
268 let arr = Self::batch_numeric_processed(c)?;
269 let min_v = arr.iter().cloned().fold(f32::INFINITY, f32::min);
270 let max_v = arr.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
271 if min_v < cp.min_v {
272 cp.min_v = min_v;
273 }
274 if max_v > cp.max_v {
275 cp.max_v = max_v;
276 }
277 merge_uniques_numerical(&mut cp.values, &arr);
278 }
279 BinType::Categorical => {
280 merge_uniques_categorical_arrow(&mut cp.values, c)?;
281 }
282 }
283 }
284 Ok(())
285 }
286
287 fn batch_numeric_processed(col: &BatchColumn) -> Result<Vec<f32>, String> {
288 match col {
289 BatchColumn::F32(a) => {
290 let view = f32_view(a)?;
291 Ok(nan_to_zero_f32(
292 round_to_significant_figures(view, 4).to_vec(),
293 ))
294 }
295 BatchColumn::F64(a) => {
296 let view = f64_view(a)?;
297 let f32s: Vec<f32> = view.iter().map(|&x| x as f32).collect();
298 Ok(nan_to_zero_f32(
299 round_to_significant_figures(ArrayView1::from(f32s.as_slice()), 4).to_vec(),
300 ))
301 }
302 BatchColumn::Owned(ColumnVec::F32(v)) => Ok(nan_to_zero_f32(
303 round_to_significant_figures(ArrayView1::from(v.as_slice()), 4).to_vec(),
304 )),
305 BatchColumn::Owned(ColumnVec::F32Array(a)) => Ok(nan_to_zero_f32(
306 round_to_significant_figures(a.view(), 4).to_vec(),
307 )),
308 _ => Err("expected numeric column".into()),
309 }
310 }
311
312 fn col_as_f32(c: &ColumnVec) -> Result<Vec<f32>, String> {
313 match c {
314 ColumnVec::F32(v) => Ok(v.clone()),
315 ColumnVec::F32Array(a) => Ok(a.iter().copied().collect()),
316 ColumnVec::Utf8(_) => Err("expected numeric column".into()),
317 }
318 }
319
320 fn col_as_utf8(c: &ColumnVec) -> Result<Vec<String>, String> {
321 match c {
322 ColumnVec::Utf8(v) => Ok(v.clone()),
323 ColumnVec::F32(v) => Ok(v.iter().map(|x| x.to_string()).collect()),
324 ColumnVec::F32Array(a) => Ok(a.iter().map(|x| x.to_string()).collect()),
325 }
326 }
327
328 pub fn finish_map(&mut self, depth: &BinDepth) -> Result<(), String> {
329 for (&col, cp) in self.preprocess_map.iter_mut() {
330 let bin_depth = depth.depth_for(col);
331 cp.values.remove(&ValueKey::Str("nan".into()));
332
333 let items: Vec<(ValueKey, u64)> =
334 cp.values.iter().map(|(k, v)| (k.clone(), *v)).collect();
335 let selected = return_accounted_bins(&items, bin_depth);
336 let mut bins_map: HashMap<ValueKey, u64> = HashMap::new();
337 for (k, c) in &selected {
338 bins_map.insert(k.clone(), *c);
339 }
340
341 let mut bins_sorted: Vec<ValueKey> = bins_map.keys().cloned().collect();
342 bins_sorted.sort_by(key_cmp_for_sort);
343
344 match cp.bin_type {
345 BinType::Numerical => {
346 let mut nums: Vec<f32> = bins_sorted
347 .iter()
348 .filter_map(|k| k.as_num_bits().map(bits_to_f32))
349 .collect();
350 if nums.is_empty() {
351 cp.bins = bins_map;
352 cp.bin_labels = Vec::new();
353 cp.numerical_bins = Vec::new();
354 continue;
355 }
356 nums.sort_by(|x, y| x.partial_cmp(y).unwrap_or(std::cmp::Ordering::Equal));
357 let mut bins_sorted_f = nums;
358
359 if cp.min_v < bins_sorted_f[0] {
360 bins_sorted_f.insert(0, cp.min_v);
361 }
362 if cp.max_v > *bins_sorted_f.last().unwrap() {
363 let last = bins_sorted_f.len() - 1;
364 bins_sorted_f[last] = cp.max_v;
365 }
366 bins_sorted_f
367 .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
368
369 let mut nb: Vec<NumericalBin> = Vec::new();
370 for k in 0..bins_sorted_f.len().saturating_sub(1) {
371 let bottom = bins_sorted_f[k];
372 let top = bins_sorted_f[k + 1];
373 let label = format_bin_label_bottom_top(bottom, top, bottom);
374 nb.push(NumericalBin { bottom, top, label });
375 }
376 cp.bin_labels = nb.iter().map(|x| x.label.clone()).collect();
377 cp.numerical_bins = nb;
378 }
379 BinType::Categorical => {
380 cp.bin_labels = bins_sorted
381 .iter()
382 .filter_map(|k| match k {
383 ValueKey::Str(s) => Some(s.clone()),
384 _ => None,
385 })
386 .collect();
387 cp.numerical_bins = Vec::new();
388 }
389 }
390
391 cp.bins = bins_map;
392 }
393 self.finished = true;
394 Ok(())
395 }
396
397 pub fn use_map(&self, chunk: &ChunkTable) -> Result<HashMap<String, ColumnVec>, String> {
398 if !self.finished {
399 return Err("finish_map must be called before use_map".into());
400 }
401 chunk.validate()?;
402 let mut out = HashMap::new();
403 for col_idx in 0..chunk.names.len() {
404 let name = chunk.names[col_idx].clone();
405 let col = chunk
406 .cols
407 .get(col_idx)
408 .ok_or_else(|| format!("internal: missing column index {col_idx}"))?;
409 if let Some(cp) = self.preprocess_map.get(&col_idx) {
410 match cp.bin_type {
411 BinType::Numerical => {
412 let mut arr = Self::col_as_f32(col)?;
413 for x in &mut arr {
414 if x.is_nan() {
415 *x = 0.0;
416 }
417 }
418 let labels = vectorized_map_numerical_bins(cp, &arr);
419 out.insert(name, ColumnVec::Utf8(labels));
420 }
421 BinType::Categorical => {
422 let mut s = Self::col_as_utf8(col)?;
423 for t in &mut s {
424 if t == "nan" || t.eq_ignore_ascii_case("nan") {
425 *t = "empty".to_string();
426 }
427 }
428 let labels: Vec<String> = s
429 .into_iter()
430 .map(|val| map_categorical_bins(cp, &val))
431 .collect();
432 out.insert(name, ColumnVec::Utf8(labels));
433 }
434 }
435 } else {
436 out.insert(name, col.clone());
437 }
438 }
439 Ok(out)
440 }
441
442 pub fn use_map_batch(&self, chunk: &BatchChunk) -> Result<HashMap<String, ColumnVec>, String> {
443 if !self.finished {
444 return Err("finish_map must be called before use_map".into());
445 }
446 chunk.validate()?;
447 let mut out = HashMap::new();
448 for col_idx in 0..chunk.names.len() {
449 let name = chunk.names[col_idx].clone();
450 let col = chunk
451 .cols
452 .get(col_idx)
453 .ok_or_else(|| format!("internal: missing column index {col_idx}"))?;
454 if let Some(cp) = self.preprocess_map.get(&col_idx) {
455 match cp.bin_type {
456 BinType::Numerical => {
457 let arr = Self::batch_numeric_for_labels(col)?;
458 let labels = vectorized_map_numerical_bins(cp, &arr);
459 out.insert(name, ColumnVec::Utf8(labels));
460 }
461 BinType::Categorical => {
462 let labels = map_categorical_bins_arrow(cp, col)?;
463 out.insert(name, ColumnVec::Utf8(labels));
464 }
465 }
466 } else {
467 out.insert(name, batch_column_to_column_vec(col)?);
468 }
469 }
470 Ok(out)
471 }
472
473 fn batch_numeric_for_labels(col: &BatchColumn) -> Result<Vec<f32>, String> {
474 let mut arr = Self::batch_numeric_processed(col)?;
475 for x in &mut arr {
476 if x.is_nan() {
477 *x = 0.0;
478 }
479 }
480 Ok(arr)
481 }
482}
483
484fn key_cmp_for_sort(a: &ValueKey, b: &ValueKey) -> std::cmp::Ordering {
485 match (a, b) {
486 (ValueKey::Num(x), ValueKey::Num(y)) => bits_to_f32(*x)
487 .partial_cmp(&bits_to_f32(*y))
488 .unwrap_or(std::cmp::Ordering::Equal),
489 (ValueKey::Str(x), ValueKey::Str(y)) => x.cmp(y),
490 (ValueKey::Num(_), ValueKey::Str(_)) => std::cmp::Ordering::Less,
491 (ValueKey::Str(_), ValueKey::Num(_)) => std::cmp::Ordering::Greater,
492 }
493}
494
495fn merge_uniques_numerical(acc: &mut HashMap<ValueKey, u64>, arr: &[f32]) {
496 let mut local: HashMap<ValueKey, u64> = HashMap::new();
497 for &x in arr {
498 let k = ValueKey::from_f32(x);
499 *local.entry(k).or_insert(0) += 1;
500 }
501 for (k, c) in local {
502 *acc.entry(k).or_insert(0) += c;
503 }
504}
505
506fn merge_uniques_categorical(acc: &mut HashMap<ValueKey, u64>, arr: &[String]) {
507 for t in arr {
508 let k = ValueKey::Str(t.clone());
509 *acc.entry(k).or_insert(0) += 1;
510 }
511}
512
513fn normalize_nan_str(s: &mut String) {
514 if s == "nan" || s.eq_ignore_ascii_case("nan") {
515 *s = "empty".to_string();
516 }
517}
518
519fn merge_uniques_categorical_arrow(
520 acc: &mut HashMap<ValueKey, u64>,
521 col: &BatchColumn,
522) -> Result<(), String> {
523 match col {
524 BatchColumn::Utf8(a) => {
525 for i in 0..a.len() {
526 let mut s = utf8_value_at(a, i);
527 normalize_nan_str(&mut s);
528 *acc.entry(ValueKey::Str(s)).or_insert(0) += 1;
529 }
530 Ok(())
531 }
532 BatchColumn::Owned(ColumnVec::Utf8(v)) => {
533 for t in v {
534 let mut s = t.clone();
535 normalize_nan_str(&mut s);
536 *acc.entry(ValueKey::Str(s)).or_insert(0) += 1;
537 }
538 Ok(())
539 }
540 _ => {
541 let s = batch_column_to_strings(col)?;
542 merge_uniques_categorical(acc, &s);
543 Ok(())
544 }
545 }
546}
547
548fn batch_column_to_strings(col: &BatchColumn) -> Result<Vec<String>, String> {
549 match col {
550 BatchColumn::Utf8(a) => Ok((0..a.len()).map(|i| utf8_value_at(a, i)).collect()),
551 BatchColumn::Owned(c) => match c {
552 ColumnVec::Utf8(v) => Ok(v.clone()),
553 ColumnVec::F32(v) => Ok(v.iter().map(|x| x.to_string()).collect()),
554 ColumnVec::F32Array(a) => Ok(a.iter().map(|x| x.to_string()).collect()),
555 },
556 BatchColumn::F32(a) => Ok(crate::arrow_io::col_to_f32(a)?
557 .into_iter()
558 .map(|x| x.to_string())
559 .collect()),
560 BatchColumn::F64(a) => Ok(crate::arrow_io::col_to_f32(a)?
561 .into_iter()
562 .map(|x| x.to_string())
563 .collect()),
564 }
565}
566
567fn batch_column_to_column_vec(col: &BatchColumn) -> Result<ColumnVec, String> {
568 crate::arrow_io::batch_column_to_owned(col)
569}
570
571fn map_categorical_bins_arrow(
572 cp: &ColumnPreprocess,
573 col: &BatchColumn,
574) -> Result<Vec<String>, String> {
575 match col {
576 BatchColumn::Utf8(a) => {
577 let mut out = Vec::with_capacity(a.len());
578 for i in 0..a.len() {
579 let mut s = utf8_value_at(a, i);
580 normalize_nan_str(&mut s);
581 out.push(map_categorical_bins(cp, &s));
582 }
583 Ok(out)
584 }
585 _ => {
586 let mut s = batch_column_to_strings(col)?;
587 for t in &mut s {
588 normalize_nan_str(t);
589 }
590 Ok(s.into_iter()
591 .map(|val| map_categorical_bins(cp, &val))
592 .collect())
593 }
594 }
595}
596
597fn nan_to_zero_f32(mut v: Vec<f32>) -> Vec<f32> {
598 for x in &mut v {
599 if x.is_nan() {
600 *x = 0.0;
601 }
602 }
603 v
604}
605
606fn return_accounted_bins(items: &[(ValueKey, u64)], bin_depth: usize) -> Vec<(ValueKey, u64)> {
607 if items.is_empty() || bin_depth == 0 {
608 return Vec::new();
609 }
610 let mut sorted: Vec<(ValueKey, u64)> = items.to_vec();
611 sorted.sort_by_key(|b| std::cmp::Reverse(b.1));
612
613 let mut selected: Vec<(ValueKey, u64)> = sorted.iter().take(bin_depth).cloned().collect();
614
615 if bin_depth > sorted.len() {
616 return selected;
617 }
618
619 let last_count = sorted[bin_depth - 1].1;
620 let mut iterater = bin_depth;
621 while iterater < sorted.len() && sorted[iterater].1 == last_count {
622 selected = sorted.iter().take(iterater).cloned().collect();
623 iterater += 1;
624 }
625
626 selected
627}
628
629fn format_bin_label_bottom_top(_bottom: f32, top: f32, bin_val: f32) -> String {
630 let bin_1 = if bin_val.abs() < 10000.0 && bin_val.abs() > 0.1 {
631 format!("{:.4}", bin_val)
632 } else {
633 format!("{:.4e}", bin_val)
634 };
635 let bin_2 = if top.abs() < 10000.0 && top.abs() > 0.1 {
636 format!("{:.4}", top)
637 } else {
638 format!("{:.4e}", top)
639 };
640 format!("{bin_1} - {bin_2}")
641}
642
643fn vectorized_map_numerical_bins(cp: &ColumnPreprocess, arr: &[f32]) -> Vec<String> {
644 if cp.numerical_bins.is_empty() {
645 return vec!["other".to_string(); arr.len()];
646 }
647 let labels: Vec<&str> = cp.numerical_bins.iter().map(|b| b.label.as_str()).collect();
648
649 let mut out = Vec::with_capacity(arr.len());
650 for &x in arr {
651 let mut idx: Option<usize> = None;
652 for (i, nb) in cp.numerical_bins.iter().enumerate() {
653 if x >= nb.bottom && x < nb.top {
654 idx = Some(i);
655 break;
656 }
657 }
658 let label = idx
659 .map(|i| labels[i].to_string())
660 .unwrap_or_else(|| "other".to_string());
661 out.push(label);
662 }
663 out
664}
665
666fn map_categorical_bins(cp: &ColumnPreprocess, val: &str) -> String {
667 let k = ValueKey::Str(val.to_string());
668 if cp.values.contains_key(&k) {
669 val.to_string()
670 } else {
671 "other".to_string()
672 }
673}
674
675#[cfg(test)]
676mod tests {
677 use super::*;
678 use crate::arrow_io::{batch_chunk_to_table, split_batch_views};
679 use arrow::array::{Float32Array, StringArray};
680 use arrow::datatypes::{DataType, Field, Schema};
681 use arrow::record_batch::RecordBatch;
682 use std::sync::Arc;
683
684 #[test]
685 fn accounted_bins_ties_keep_extending() {
686 let items = vec![
687 (ValueKey::Str("a".into()), 10),
688 (ValueKey::Str("b".into()), 8),
689 (ValueKey::Str("c".into()), 8),
690 (ValueKey::Str("d".into()), 8),
691 (ValueKey::Str("e".into()), 3),
692 ];
693 let r = return_accounted_bins(&items, 2);
694 assert!(r.len() >= 2);
695 }
696
697 fn sample_batch() -> RecordBatch {
698 let id = Arc::new(StringArray::from(vec!["a", "b"]));
699 let feat = Arc::new(Float32Array::from(vec![1.0_f32, 20.0]));
700 let target = Arc::new(Float32Array::from(vec![0.5_f32, 1.5]));
701 let schema = Arc::new(Schema::new(vec![
702 Field::new("id", DataType::Utf8, false),
703 Field::new("feat", DataType::Float32, false),
704 Field::new("target", DataType::Float32, false),
705 ]));
706 RecordBatch::try_new(schema, vec![id, feat, target]).unwrap()
707 }
708
709 #[test]
710 fn preprocess_batch_matches_owned_path() {
711 let batch = sample_batch();
712 let (chunk, _, cg) = split_batch_views(&batch, "target", &["target".into()]).unwrap();
713 let table = batch_chunk_to_table(&chunk).unwrap();
714
715 let depth = BinDepth::new(4);
716 let mut via_batch = PreprocessStream::new(cg.clone());
717 via_batch.preprocess_batch(&chunk).unwrap();
718 via_batch.finish_map(&depth).unwrap();
719 let out_batch = via_batch.use_map_batch(&chunk).unwrap();
720
721 let mut via_table = PreprocessStream::new(cg);
722 via_table.preprocess(&table).unwrap();
723 via_table.finish_map(&depth).unwrap();
724 let out_table = via_table.use_map(&table).unwrap();
725
726 let mut batch_keys: Vec<_> = out_batch.keys().collect();
727 let mut table_keys: Vec<_> = out_table.keys().collect();
728 batch_keys.sort();
729 table_keys.sort();
730 assert_eq!(batch_keys, table_keys);
731 for name in out_batch.keys() {
732 match (out_batch.get(name).unwrap(), out_table.get(name).unwrap()) {
733 (ColumnVec::Utf8(a), ColumnVec::Utf8(b)) => assert_eq!(a, b),
734 (ColumnVec::F32(a), ColumnVec::F32(b)) => assert_eq!(a, b),
735 (ColumnVec::F32Array(a), ColumnVec::F32Array(b)) => assert_eq!(a, b),
736 (ColumnVec::F32Array(a), ColumnVec::F32(b)) => {
737 assert_eq!(a.iter().copied().collect::<Vec<_>>(), *b)
738 }
739 (ColumnVec::F32(a), ColumnVec::F32Array(b)) => {
740 assert_eq!(a, &b.iter().copied().collect::<Vec<_>>())
741 }
742 other => panic!("unexpected column type pairing for {name}: {other:?}"),
743 }
744 }
745 }
746}