steru 0.8.0

steru
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

use arrow::array::RecordBatch as ArrowRecordBatch;
use arrow::datatypes::Schema as ArrowSchema;
use arrow::ipc::reader::StreamReader;
use polars::error::PolarsError;
use polars::frame::DataFrame;
use polars::prelude::*;
use polars_io::ipc::{IpcCompression, IpcStreamWriter};
use std::io::Cursor;
use std::sync::Arc;

pub fn write_to_arrow(
    df: &mut DataFrame,
    compression: Option<IpcCompression>,
) -> PolarsResult<(Arc<ArrowSchema>, Vec<ArrowRecordBatch>)> {
    let mut buffer = Vec::new();
    IpcStreamWriter::new(&mut buffer)
        .with_compression(compression)
        .finish(df)?;

    let cursor = Cursor::new(buffer);
    let stream_reader = StreamReader::try_new(cursor, None)
        .map_err(|e| PolarsError::ComputeError(format!("{e}").into()))?;

    let schema = stream_reader.schema();
    let batches: Result<Vec<_>, _> = stream_reader.collect();
    let batches = batches.map_err(|e| PolarsError::ComputeError(format!("{e}").into()))?;

    Ok((schema, batches))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_df() -> DataFrame {
        df!(
            "ticker" => ["AAPL", "NVDA", "MSFT", "GOOG", "AMZN"],
            "price" => [229.9, 138.93, 420.56, 166.41, 188.4],
            "high" => [231.31, 139.6, 424.04, 167.62, 189.83],
            "low" => [228.6, 136.3, 417.52, 164.78, 188.44],
        )
        .unwrap()
    }

    #[test]
    fn compressed_write() {
        let mut df = create_df();
        let arrow_data = write_to_arrow(&mut df, Some(IpcCompression::ZSTD));
        let (arrow_schema, record_batches) = arrow_data.unwrap();

        let df_schema = df.schema();
        let arrow_schema_fields = arrow_schema.fields();
        let df_schema_fields: Vec<Field> = df_schema.iter_fields().collect();
        let column_names = df.get_column_names();
        let column_names_owned = df.get_column_names_owned();

        assert_eq!(arrow_schema_fields.len(), df_schema_fields.len());
        assert_eq!(arrow_schema_fields.len(), df.width());

        for (i, field) in arrow_schema_fields.iter().enumerate() {
            assert_eq!(field.name().to_owned(), column_names_owned[i].to_string());
        }

        assert!(!record_batches.is_empty());
        assert_eq!(record_batches.len(), 1);

        let batch = &record_batches[0];
        assert_eq!(batch.num_rows(), df.height());
        assert_eq!(batch.num_columns(), df.width());

        assert_eq!(column_names.len(), batch.num_columns());
        assert_eq!(column_names_owned.len(), batch.num_columns());
    }

    #[test]
    fn uncompressed_write() {
        let mut df = create_df();
        let arrow_data = write_to_arrow(&mut df, None);
        let (arrow_schema, record_batches) = arrow_data.unwrap();

        let df_schema = df.schema();
        let arrow_schema_fields = arrow_schema.fields();
        let df_schema_fields: Vec<Field> = df_schema.iter_fields().collect();
        let column_names = df.get_column_names();
        let column_names_owned = df.get_column_names_owned();

        assert_eq!(arrow_schema_fields.len(), df_schema_fields.len());
        assert_eq!(arrow_schema_fields.len(), df.width());

        for (i, field) in arrow_schema_fields.iter().enumerate() {
            assert_eq!(field.name().to_owned(), column_names_owned[i].to_string());
        }

        assert!(!record_batches.is_empty());
        assert_eq!(record_batches.len(), 1);

        let batch = &record_batches[0];
        assert_eq!(batch.num_rows(), df.height());
        assert_eq!(batch.num_columns(), df.width());

        assert_eq!(column_names.len(), batch.num_columns());
        assert_eq!(column_names_owned.len(), batch.num_columns());
    }
}