Skip to main content

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}