napparent_tabular/
arrow_io.rs

1//! Load [`ChunkTable`] / [`ColGraph`] from Arrow [`RecordBatch`].
2
3use crate::ndarrow_bridge::{array1_f32_to_arrow, f32_view, f64_view};
4use crate::table::{BatchChunk, BatchColumn, ChunkTable, ColGraph, ColumnVec, TargetColumn};
5use arrow::array::{
6    Array, ArrayRef, AsArray, BooleanArray, Float32Array, Int16Array, Int32Array, Int64Array,
7    Int8Array, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
8};
9use arrow::datatypes::DataType;
10use arrow::record_batch::RecordBatch;
11use std::collections::HashSet;
12use std::sync::Arc;
13
14pub(crate) fn col_to_f32(col: &ArrayRef) -> Result<Vec<f32>, String> {
15    match col.data_type() {
16        DataType::Float32 => Ok(col
17            .as_primitive::<arrow::datatypes::Float32Type>()
18            .values()
19            .iter()
20            .copied()
21            .collect()),
22        DataType::Float64 => {
23            let a = col.as_primitive::<arrow::datatypes::Float64Type>();
24            Ok(a.values().iter().map(|x| *x as f32).collect())
25        }
26        DataType::Int8 => {
27            let a: &Int8Array = col.as_primitive();
28            Ok(a.values().iter().map(|x| *x as f32).collect())
29        }
30        DataType::Int16 => {
31            let a: &Int16Array = col.as_primitive();
32            Ok(a.values().iter().map(|x| *x as f32).collect())
33        }
34        DataType::Int32 => {
35            let a: &Int32Array = col.as_primitive();
36            Ok(a.values().iter().map(|x| *x as f32).collect())
37        }
38        DataType::Int64 => {
39            let a: &Int64Array = col.as_primitive();
40            Ok(a.values().iter().map(|x| *x as f32).collect())
41        }
42        DataType::UInt8 => {
43            let a: &UInt8Array = col.as_primitive();
44            Ok(a.values().iter().map(|x| *x as f32).collect())
45        }
46        DataType::UInt16 => {
47            let a: &UInt16Array = col.as_primitive();
48            Ok(a.values().iter().map(|x| *x as f32).collect())
49        }
50        DataType::UInt32 => {
51            let a: &UInt32Array = col.as_primitive();
52            Ok(a.values().iter().map(|x| *x as f32).collect())
53        }
54        DataType::UInt64 => {
55            let a: &UInt64Array = col.as_primitive();
56            Ok(a.values().iter().map(|x| *x as f32).collect())
57        }
58        DataType::Boolean => {
59            let a: &BooleanArray = col.as_boolean();
60            Ok((0..a.len())
61                .map(|i| {
62                    if a.is_null(i) {
63                        0.0
64                    } else if a.value(i) {
65                        1.0
66                    } else {
67                        0.0
68                    }
69                })
70                .collect())
71        }
72        other => Err(format!("unsupported numeric arrow type: {other}")),
73    }
74}
75
76fn col_to_utf8(col: &ArrayRef) -> Result<Vec<String>, String> {
77    match col.data_type() {
78        DataType::Utf8 => {
79            let a = col.as_string::<i32>();
80            Ok((0..a.len())
81                .map(|i| {
82                    if a.is_null(i) {
83                        "empty".to_string()
84                    } else {
85                        a.value(i).to_string()
86                    }
87                })
88                .collect())
89        }
90        DataType::LargeUtf8 => {
91            let a = col.as_string::<i64>();
92            Ok((0..a.len())
93                .map(|i| {
94                    if a.is_null(i) {
95                        "empty".to_string()
96                    } else {
97                        a.value(i).to_string()
98                    }
99                })
100                .collect())
101        }
102        DataType::Boolean => {
103            let a: &BooleanArray = col.as_boolean();
104            Ok((0..a.len())
105                .map(|i| {
106                    if a.is_null(i) {
107                        "empty".to_string()
108                    } else if a.value(i) {
109                        "true".to_string()
110                    } else {
111                        "false".to_string()
112                    }
113                })
114                .collect())
115        }
116        DataType::Float32
117        | DataType::Float64
118        | DataType::Int8
119        | DataType::Int16
120        | DataType::Int32
121        | DataType::Int64 => {
122            let v = col_to_f32(col)?;
123            Ok(v.into_iter().map(|x| x.to_string()).collect())
124        }
125        other => Err(format!("unsupported string-like arrow type: {other}")),
126    }
127}
128
129fn feature_column_from_arrow(col: ArrayRef, dt: &DataType) -> Result<BatchColumn, String> {
130    match dt {
131        DataType::Float32 => match f32_view(&col) {
132            Ok(_) => Ok(BatchColumn::F32(col)),
133            Err(_) => Ok(BatchColumn::Owned(ColumnVec::F32(col_to_f32(&col)?))),
134        },
135        DataType::Float64 => match f64_view(&col) {
136            Ok(_) => Ok(BatchColumn::F64(col)),
137            Err(_) => Ok(BatchColumn::Owned(ColumnVec::F32(col_to_f32(&col)?))),
138        },
139        DataType::Int8
140        | DataType::Int16
141        | DataType::Int32
142        | DataType::Int64
143        | DataType::UInt8
144        | DataType::UInt16
145        | DataType::UInt32
146        | DataType::UInt64
147        | DataType::Boolean => Ok(BatchColumn::Owned(ColumnVec::F32(col_to_f32(&col)?))),
148        DataType::Utf8 | DataType::LargeUtf8 => Ok(BatchColumn::Utf8(col)),
149        DataType::Dictionary(_, _) => {
150            let a = arrow::compute::cast(&col, &DataType::Utf8)
151                .map_err(|e| format!("dictionary decode: {e}"))?;
152            Ok(BatchColumn::Owned(ColumnVec::Utf8(col_to_utf8(&a)?)))
153        }
154        _ => Ok(BatchColumn::Owned(ColumnVec::Utf8(col_to_utf8(&col)?))),
155    }
156}
157
158fn target_from_arrow(col: ArrayRef) -> Result<TargetColumn, String> {
159    if col.data_type() == &DataType::Float32 && f32_view(&col).is_ok() {
160        return Ok(TargetColumn::F32(col));
161    }
162    let y = col_to_f32(&col).or_else(|_| {
163        col_to_utf8(&col).map(|v| {
164            v.into_iter()
165                .map(|s| s.parse::<f32>().unwrap_or(0.0))
166                .collect()
167        })
168    })?;
169    Ok(TargetColumn::Owned(y))
170}
171
172/// Zero-copy batch split: feature columns share Arrow buffers via `ArrayRef`.
173pub fn split_batch_views(
174    batch: &RecordBatch,
175    target: &str,
176    cols_to_drop: &[String],
177) -> Result<(BatchChunk, TargetColumn, ColGraph), String> {
178    let schema = batch.schema();
179    let fields: Vec<_> = schema.fields().iter().cloned().collect();
180    let mut target_idx: Option<usize> = None;
181    for (i, f) in fields.iter().enumerate() {
182        if f.name() == target {
183            target_idx = Some(i);
184            break;
185        }
186    }
187    let ti = target_idx.ok_or_else(|| format!("target column `{target}` not in batch schema"))?;
188
189    let drop_set: HashSet<&str> = cols_to_drop.iter().map(String::as_str).collect();
190    if !drop_set.contains(target) {
191        return Err(format!(
192            "`cols_to_drop` must include target `{target}` (see Python `load_numpy`)"
193        ));
194    }
195
196    let n = batch.num_rows();
197    let y_col = batch.column(ti);
198    let target_col = target_from_arrow(y_col.clone())?;
199
200    let mut names = Vec::new();
201    let mut cols = Vec::new();
202    let mut dropped = std::collections::HashSet::new();
203
204    for (i, field) in fields.iter().enumerate() {
205        if i == ti {
206            continue;
207        }
208        let name = field.name().clone();
209        let col = batch.column(i);
210        let logical = feature_column_from_arrow(col.clone(), field.data_type())?;
211        let idx = names.len();
212        if drop_set.contains(name.as_str()) {
213            dropped.insert(idx);
214        }
215        names.push(name);
216        cols.push(logical);
217    }
218
219    let chunk = BatchChunk { names, cols };
220    chunk.validate()?;
221    if target_col.len() != n {
222        return Err("Y length mismatch".into());
223    }
224    let col_graph = ColGraph {
225        names: chunk.names.clone(),
226        dropped,
227    };
228    Ok((chunk, target_col, col_graph))
229}
230
231pub(crate) fn batch_column_to_owned(col: &BatchColumn) -> Result<ColumnVec, String> {
232    match col {
233        BatchColumn::F32(a) => Ok(ColumnVec::F32(col_to_f32(a)?)),
234        BatchColumn::F64(a) => Ok(ColumnVec::F32(col_to_f32(a)?)),
235        BatchColumn::Utf8(a) => Ok(ColumnVec::Utf8(col_to_utf8(a)?)),
236        BatchColumn::Owned(c) => Ok(c.clone()),
237    }
238}
239
240pub fn batch_chunk_to_table(chunk: &BatchChunk) -> Result<ChunkTable, String> {
241    chunk.validate()?;
242    let cols = chunk
243        .cols
244        .iter()
245        .map(batch_column_to_owned)
246        .collect::<Result<Vec<_>, _>>()?;
247    Ok(ChunkTable {
248        names: chunk.names.clone(),
249        cols,
250    })
251}
252
253pub fn target_to_vec(target: &TargetColumn) -> Vec<f32> {
254    match target {
255        TargetColumn::F32(a) => f32_view(a)
256            .map(|v| v.iter().copied().collect())
257            .unwrap_or_else(|_| col_to_f32(a).unwrap_or_default()),
258        TargetColumn::Owned(v) => v.clone(),
259    }
260}
261
262pub fn target_as_outcomes(target: &TargetColumn) -> OutcomesRef<'_> {
263    match target {
264        TargetColumn::F32(a) => {
265            OutcomesRef::View(f32_view(a).expect("F32 target validated at split"))
266        }
267        TargetColumn::Owned(v) => OutcomesRef::Slice(v),
268    }
269}
270
271/// Borrowed outcomes for zero-copy aggregation when target is null-free Float32.
272pub enum OutcomesRef<'a> {
273    View(ndarray::ArrayView1<'a, f32>),
274    Slice(&'a [f32]),
275}
276
277impl OutcomesRef<'_> {
278    pub fn len(&self) -> usize {
279        match self {
280            OutcomesRef::View(v) => v.len(),
281            OutcomesRef::Slice(s) => s.len(),
282        }
283    }
284
285    pub fn is_empty(&self) -> bool {
286        self.len() == 0
287    }
288
289    pub fn get(&self, i: usize) -> f32 {
290        match self {
291            OutcomesRef::View(v) => v[i],
292            OutcomesRef::Slice(s) => s[i],
293        }
294    }
295
296    pub fn sum(&self) -> f32 {
297        match self {
298            OutcomesRef::View(v) => v.sum(),
299            OutcomesRef::Slice(s) => s.iter().sum(),
300        }
301    }
302
303    pub fn to_nan0_vec(&self) -> Vec<f32> {
304        match self {
305            OutcomesRef::View(v) => v
306                .iter()
307                .map(|&x| if x.is_nan() { 0.0 } else { x })
308                .collect(),
309            OutcomesRef::Slice(s) => s
310                .iter()
311                .map(|&x| if x.is_nan() { 0.0 } else { x })
312                .collect(),
313        }
314    }
315}
316
317/// Return the UTF-8 value at row `i`, matching [`col_to_utf8`] null/empty semantics.
318pub fn utf8_value_at(col: &ArrayRef, i: usize) -> String {
319    match col.data_type() {
320        DataType::Utf8 => {
321            let a = col.as_string::<i32>();
322            if a.is_null(i) {
323                "empty".to_string()
324            } else {
325                a.value(i).to_string()
326            }
327        }
328        DataType::LargeUtf8 => {
329            let a = col.as_string::<i64>();
330            if a.is_null(i) {
331                "empty".to_string()
332            } else {
333                a.value(i).to_string()
334            }
335        }
336        _ => "empty".to_string(),
337    }
338}
339
340/// Extract feature table `X`, target `Y`, and column metadata (materializes owned columns).
341///
342/// Prefer [`split_batch_views`] internally for zero-copy numeric ingest.
343pub fn split_batch_xy(
344    batch: &RecordBatch,
345    target: &str,
346    cols_to_drop: &[String],
347) -> Result<(ChunkTable, Vec<f32>, ColGraph), String> {
348    let (chunk, target_col, col_graph) = split_batch_views(batch, target, cols_to_drop)?;
349    let table = batch_chunk_to_table(&chunk)?;
350    let y = target_to_vec(&target_col);
351    Ok((table, y, col_graph))
352}
353
354/// Build a [`RecordBatch`] from string-name → column map (mixed float / utf8).
355pub fn batch_from_map(
356    schema: arrow::datatypes::SchemaRef,
357    mut columns_by_name: std::collections::HashMap<String, ColumnVec>,
358) -> Result<RecordBatch, String> {
359    use arrow::array::StringBuilder;
360
361    let n = schema
362        .fields()
363        .first()
364        .and_then(|f| columns_by_name.get(f.name()).map(|c| c.len()))
365        .unwrap_or(0);
366
367    let mut arrays: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());
368    for field in schema.fields() {
369        let name = field.name();
370        let col = columns_by_name
371            .remove(name)
372            .ok_or_else(|| format!("missing column `{name}` building RecordBatch"))?;
373        let arr: ArrayRef = match col {
374            ColumnVec::F32(v) => {
375                if v.len() != n {
376                    return Err(format!("column {name} length {}", v.len()));
377                }
378                Arc::new(Float32Array::from(v))
379            }
380            ColumnVec::F32Array(a) => {
381                if a.len() != n {
382                    return Err(format!("column {name} length {}", a.len()));
383                }
384                array1_f32_to_arrow(a)?
385            }
386            ColumnVec::Utf8(v) => {
387                if v.len() != n {
388                    return Err(format!("column {name} length {}", v.len()));
389                }
390                let mut b = StringBuilder::new();
391                for s in v {
392                    b.append_value(s);
393                }
394                Arc::new(b.finish())
395            }
396        };
397        arrays.push(arr);
398    }
399    RecordBatch::try_new(schema, arrays).map_err(|e| e.to_string())
400}
401
402pub fn concat_same_schema(batches: &[RecordBatch]) -> Result<RecordBatch, String> {
403    if batches.is_empty() {
404        return Err("empty batch list".into());
405    }
406    let schema = batches[0].schema();
407    arrow::compute::concat_batches(&schema, batches).map_err(|e| e.to_string())
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413    use arrow::array::{Float32Array, StringArray};
414    use arrow::datatypes::{DataType, Field, Schema};
415    use std::sync::Arc;
416
417    fn sample_batch() -> RecordBatch {
418        let id = Arc::new(StringArray::from(vec!["a", "b"]));
419        let feat = Arc::new(Float32Array::from(vec![1.0_f32, 20.0]));
420        let target = Arc::new(Float32Array::from(vec![0.5_f32, 1.5]));
421        let schema = Arc::new(Schema::new(vec![
422            Field::new("id", DataType::Utf8, false),
423            Field::new("feat", DataType::Float32, false),
424            Field::new("target", DataType::Float32, false),
425        ]));
426        RecordBatch::try_new(schema, vec![id, feat, target]).unwrap()
427    }
428
429    #[test]
430    fn split_batch_views_f32_zero_copy() {
431        let batch = sample_batch();
432        let (chunk, target, _cg) = split_batch_views(&batch, "target", &["target".into()]).unwrap();
433        assert!(matches!(target, TargetColumn::F32(_)));
434        assert!(matches!(chunk.cols[1], BatchColumn::F32(_)));
435    }
436
437    #[test]
438    fn split_batch_xy_matches_views_materialized() {
439        let batch = sample_batch();
440        let (table, y, cg) = split_batch_xy(&batch, "target", &["target".into()]).unwrap();
441        let (chunk, target, cg2) = split_batch_views(&batch, "target", &["target".into()]).unwrap();
442        assert_eq!(cg, cg2);
443        assert_eq!(y, target_to_vec(&target));
444        assert_eq!(table.names, chunk.names);
445    }
446}
napparent_tabular/arrow_io.rs

napparent_tabular/
arrow_io.rs
