entrenar/config/train/arrow.rs
1//! Arrow array conversion utilities
2
3use crate::error::{Error, Result};
4
5/// Convert Arrow array to f32 vector
6#[cfg(all(not(target_arch = "wasm32"), feature = "parquet"))]
7pub fn arrow_array_to_f32(array: &arrow::array::ArrayRef) -> Result<Vec<f32>> {
8 use arrow::array::{Float32Array, Float64Array, Int32Array, Int64Array};
9 use arrow::datatypes::DataType;
10
11 match array.data_type() {
12 DataType::Float32 => {
13 let arr = array
14 .as_any()
15 .downcast_ref::<Float32Array>()
16 .ok_or_else(|| Error::ConfigError("Failed to downcast Float32Array".to_string()))?;
17 Ok(arr.values().to_vec())
18 }
19 DataType::Float64 => {
20 let arr = array
21 .as_any()
22 .downcast_ref::<Float64Array>()
23 .ok_or_else(|| Error::ConfigError("Failed to downcast Float64Array".to_string()))?;
24 Ok(arr.values().iter().map(|&x| x as f32).collect())
25 }
26 DataType::Int32 => {
27 let arr = array
28 .as_any()
29 .downcast_ref::<Int32Array>()
30 .ok_or_else(|| Error::ConfigError("Failed to downcast Int32Array".to_string()))?;
31 Ok(arr.values().iter().map(|&x| x as f32).collect())
32 }
33 DataType::Int64 => {
34 let arr = array
35 .as_any()
36 .downcast_ref::<Int64Array>()
37 .ok_or_else(|| Error::ConfigError("Failed to downcast Int64Array".to_string()))?;
38 Ok(arr.values().iter().map(|&x| x as f32).collect())
39 }
40 other => Err(Error::ConfigError(format!(
41 "Unsupported Arrow data type: {other:?}. Use Float32, Float64, Int32, or Int64."
42 ))),
43 }
44}