ipfrs_tensorlogic/gradient/
arrow_ipc.rs

1//! Standalone Arrow IPC helpers for gradient serialization.
2//!
3//! These functions serialize/deserialize raw `f32` gradient slices to/from
4//! Apache Arrow IPC format so that gradients can be stored as content-addressed
5//! blocks in the IPFS block store.
6
7/// Serialize a gradient tensor as an Arrow IPC RecordBatch.
8///
9/// The output bytes contain a single record batch with two columns:
10/// - `"index"` (`Int32`): element indices 0..N
11/// - `"value"` (`Float32`): gradient values
12///
13/// This function is a standalone companion to
14/// `ComputationGraphStore::store_gradient_as_arrow`.  It does **not**
15/// require a node id or a graph — it operates purely on the raw gradient
16/// slice and is suitable for use in `DistributedGradientAccumulator`.
17pub fn store_gradient_as_arrow(gradient: &[f32]) -> anyhow::Result<Vec<u8>> {
18    use arrow::array::{Float32Array, Int32Array};
19    use arrow::datatypes::{DataType, Field, Schema};
20    use arrow::ipc::writer::FileWriter;
21    use arrow::record_batch::RecordBatch;
22    use std::io::Cursor;
23    use std::sync::Arc;
24
25    let n = gradient.len();
26    let indices: Int32Array = (0i32..(n as i32)).collect();
27    let values: Float32Array = gradient.iter().copied().collect();
28
29    let schema = Arc::new(Schema::new(vec![
30        Field::new("index", DataType::Int32, false),
31        Field::new("value", DataType::Float32, false),
32    ]));
33
34    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(indices), Arc::new(values)])?;
35
36    let mut buf = Vec::new();
37    {
38        let cursor = Cursor::new(&mut buf);
39        let mut writer = FileWriter::try_new(cursor, &schema)?;
40        writer.write(&batch)?;
41        writer.finish()?;
42    }
43
44    Ok(buf)
45}
46
47/// Deserialize gradient from Arrow IPC bytes produced by [`store_gradient_as_arrow`].
48///
49/// Reads the `"value"` column and returns it as `Vec<f32>`.
50pub fn load_gradient_from_arrow(bytes: &[u8]) -> anyhow::Result<Vec<f32>> {
51    use arrow::array::Float32Array;
52    use arrow::ipc::reader::FileReader;
53    use std::io::Cursor;
54
55    let cursor = Cursor::new(bytes);
56    let mut reader = FileReader::try_new(cursor, None)?;
57
58    let mut values: Vec<f32> = Vec::new();
59
60    for batch_result in &mut reader {
61        let batch = batch_result?;
62        // Locate the "value" column
63        let schema = batch.schema();
64        let col_idx = schema
65            .index_of("value")
66            .map_err(|_| anyhow::anyhow!("Arrow IPC block missing 'value' column"))?;
67
68        let col = batch
69            .column(col_idx)
70            .as_any()
71            .downcast_ref::<Float32Array>()
72            .ok_or_else(|| anyhow::anyhow!("'value' column is not Float32"))?;
73
74        values.extend_from_slice(col.values());
75    }
76
77    Ok(values)
78}
79
80#[cfg(test)]
81mod arrow_ipc_standalone_tests {
82    use super::*;
83
84    #[test]
85    fn test_arrow_ipc_roundtrip() {
86        // Generate 1000 f32 values in a deterministic pattern.
87        let original: Vec<f32> = (0u32..1000).map(|i| (i as f32) * 0.001).collect();
88
89        let bytes =
90            store_gradient_as_arrow(&original).expect("store_gradient_as_arrow should succeed");
91        assert!(!bytes.is_empty(), "IPC bytes must not be empty");
92
93        let loaded =
94            load_gradient_from_arrow(&bytes).expect("load_gradient_from_arrow should succeed");
95
96        assert_eq!(loaded.len(), original.len(), "element count must match");
97        for (i, (&orig, &val)) in original.iter().zip(loaded.iter()).enumerate() {
98            assert!(
99                (orig - val).abs() < 1e-6,
100                "value mismatch at index {i}: orig={orig}, loaded={val}"
101            );
102        }
103    }
104
105    #[test]
106    fn test_arrow_ipc_empty() {
107        let empty: Vec<f32> = vec![];
108
109        let bytes = store_gradient_as_arrow(&empty)
110            .expect("store_gradient_as_arrow on empty slice should succeed");
111
112        let loaded = load_gradient_from_arrow(&bytes)
113            .expect("load_gradient_from_arrow on empty IPC should succeed");
114
115        assert!(loaded.is_empty(), "loaded gradient must also be empty");
116    }
117}
ipfrs_tensorlogic/gradient/arrow_ipc.rs

ipfrs_tensorlogic/gradient/
arrow_ipc.rs
