rs_arrow_array2freq2rbat/
lib.rs

1use std::collections::BTreeMap;
2use std::io;
3
4use arrow::array::PrimitiveArray;
5use arrow::record_batch::RecordBatch;
6
7use arrow::datatypes::ArrowPrimitiveType;
8
9use arrow::datatypes::SchemaRef;
10
11/// Computes the frequencies of each value in a [`PrimitiveArray`].
12///
13/// # Arguments
14///
15/// * `arr`: A [`PrimitiveArray`] of values.
16///
17/// # Returns
18///
19/// A [`Result`] containing a [`BTreeMap`] with the frequencies of each value,
20/// or an [`io::Error`] if an error occurs.
21///
22/// # Example
23///
24/// ```
25/// use arrow::array::Int32Array;
26/// use rs_arrow_array2freq2rbat::frequencies;
27///
28/// let arr = Int32Array::from(vec![Some(2), Some(3), Some(3), Some(5), Some(5), Some(5)]);
29/// let freqs = frequencies(arr).unwrap();
30///
31/// assert_eq!(freqs.get(&2), Some(&1));
32/// assert_eq!(freqs.get(&3), Some(&2));
33/// assert_eq!(freqs.get(&5), Some(&3));
34/// ```
35pub fn frequencies<T>(arr: PrimitiveArray<T>) -> Result<BTreeMap<T::Native, u64>, io::Error>
36where
37    T: ArrowPrimitiveType,
38    T::Native: Ord,
39{
40    let mut freqs = BTreeMap::new();
41    for item in arr.iter().flatten() {
42        *freqs.entry(item).or_insert(0) += 1;
43    }
44    Ok(freqs)
45}
46
47/// Converts a `BTreeMap` of frequencies into an Arrow `RecordBatch`.
48///
49/// # Arguments
50///
51/// * `m`: A `BTreeMap` where keys are the values and values are their frequencies.
52/// * `sch`: The `SchemaRef` for the output `RecordBatch`.
53///
54/// # Returns
55///
56/// A `Result` containing a `RecordBatch` with two columns, whose names are defined by the provided `SchemaRef`,
57/// or an `io::Error` if an error occurs.
58///
59/// # Example
60///
61/// ```
62/// use std::collections::BTreeMap;
63/// use std::sync::Arc;
64/// use arrow::datatypes::{DataType, Field, Schema};
65/// use rs_arrow_array2freq2rbat::map2batch;
66///
67/// let mut freqs = BTreeMap::new();
68/// freqs.insert(1, 1);
69/// freqs.insert(2, 2);
70/// freqs.insert(3, 3);
71///
72/// let schema = Arc::new(Schema::new(vec![
73///     Field::new("value", DataType::Int32, false),
74///     Field::new("frequency", DataType::UInt64, false),
75/// ]));
76///
77/// let batch = map2batch::<arrow::datatypes::Int32Type>(freqs, &schema).unwrap();
78///
79/// assert_eq!(batch.num_columns(), 2);
80/// assert_eq!(batch.num_rows(), 3);
81/// ```
82pub fn map2batch<T>(m: BTreeMap<T::Native, u64>, sch: &SchemaRef) -> Result<RecordBatch, io::Error>
83where
84    T: ArrowPrimitiveType,
85{
86    let keys: Vec<T::Native> = m.keys().cloned().collect();
87    let values: Vec<u64> = m.values().cloned().collect();
88
89    let key_array = PrimitiveArray::<T>::from_iter_values(keys);
90    let val_array = PrimitiveArray::<arrow::datatypes::UInt64Type>::from_iter(values);
91
92    let batch = RecordBatch::try_new(
93        sch.clone(),
94        vec![
95            std::sync::Arc::new(key_array),
96            std::sync::Arc::new(val_array),
97        ],
98    )
99    .map_err(|e| io::Error::other(e.to_string()))?;
100
101    Ok(batch)
102}
103
104/// Computes the frequencies of each value in a [`PrimitiveArray`] and converts the result into an Arrow [`RecordBatch`].
105///
106/// # Arguments
107///
108/// * `sch`: The `SchemaRef` for the output `RecordBatch`. This allows the caller to define custom column names and metadata for the resulting frequency and value columns. The schema is expected to have two fields, where the first field corresponds to the values from the `PrimitiveArray` (with `T::DATA_TYPE`) and the second field corresponds to the frequencies (with `DataType::UInt64`).
109///
110/// # Returns
111///
112/// A [`Result`] containing a [`RecordBatch`] with two columns, whose names are defined by the provided `SchemaRef`,
113/// or an [`io::Error`] if an error occurs.
114///
115/// # Example
116///
117/// ```
118/// use std::sync::Arc;
119/// use arrow::array::Int32Array;
120/// use arrow::datatypes::{DataType, Field, Schema};
121/// use rs_arrow_array2freq2rbat::array2frequency2batch;
122///
123/// let arr = Int32Array::from(vec![Some(2), Some(3), Some(3), Some(5), Some(5), Some(5)]);
124/// let schema = Arc::new(Schema::new(vec![
125///     Field::new("value", DataType::Int32, false),
126///     Field::new("frequency", DataType::UInt64, false),
127/// ]));
128/// let batch = array2frequency2batch(arr, schema).unwrap();
129///
130/// assert_eq!(batch.num_columns(), 2);
131/// assert_eq!(batch.num_rows(), 3);
132/// ```
133pub fn array2frequency2batch<T>(
134    arr: PrimitiveArray<T>,
135    sch: SchemaRef,
136) -> Result<RecordBatch, io::Error>
137where
138    T: ArrowPrimitiveType,
139    T::Native: Ord,
140{
141    let freqs = frequencies(arr)?;
142    map2batch::<T>(freqs, &sch)
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use arrow::array::Int32Array;
149
150    #[test]
151    fn test_frequencies() {
152        let arr = Int32Array::from(vec![1, 2, 2, 3, 3, 3]);
153        let freqs = frequencies(arr).unwrap();
154        assert_eq!(freqs.get(&1), Some(&1));
155        assert_eq!(freqs.get(&2), Some(&2));
156        assert_eq!(freqs.get(&3), Some(&3));
157    }
158
159    #[test]
160    fn test_map2batch() {
161        use arrow::datatypes::{DataType, Field, Schema};
162        use std::sync::Arc;
163
164        let mut freqs = BTreeMap::new();
165        freqs.insert(1, 1);
166        freqs.insert(2, 2);
167        freqs.insert(3, 3);
168
169        let schema = Arc::new(Schema::new(vec![
170            Field::new("value", DataType::Int32, false),
171            Field::new("frequency", DataType::UInt64, false),
172        ]));
173
174        let batch = map2batch::<arrow::datatypes::Int32Type>(freqs, &schema).unwrap();
175
176        assert_eq!(batch.num_columns(), 2);
177        assert_eq!(batch.num_rows(), 3);
178
179        let value_array = batch
180            .column(0)
181            .as_any()
182            .downcast_ref::<Int32Array>()
183            .unwrap();
184        let freq_array = batch
185            .column(1)
186            .as_any()
187            .downcast_ref::<arrow::array::UInt64Array>()
188            .unwrap();
189
190        assert_eq!(value_array.value(0), 1);
191        assert_eq!(value_array.value(1), 2);
192        assert_eq!(value_array.value(2), 3);
193
194        assert_eq!(freq_array.value(0), 1);
195        assert_eq!(freq_array.value(1), 2);
196        assert_eq!(freq_array.value(2), 3);
197    }
198
199    #[test]
200    fn test_array2frequency2batch() {
201        use arrow::datatypes::{DataType, Field, Schema};
202        use std::sync::Arc;
203
204        let arr = PrimitiveArray::<arrow::datatypes::Int32Type>::from(vec![1, 2, 2, 3, 3, 3]);
205        let schema = Arc::new(Schema::new(vec![
206            Field::new("value", DataType::Int32, false),
207            Field::new("frequency", DataType::UInt64, false),
208        ]));
209        let batch = array2frequency2batch(arr, schema).unwrap();
210
211        assert_eq!(batch.num_columns(), 2);
212        assert_eq!(batch.num_rows(), 3);
213
214        let value_array = batch
215            .column(0)
216            .as_any()
217            .downcast_ref::<Int32Array>()
218            .unwrap();
219        let freq_array = batch
220            .column(1)
221            .as_any()
222            .downcast_ref::<arrow::array::UInt64Array>()
223            .unwrap();
224
225        assert_eq!(value_array.value(0), 1);
226        assert_eq!(value_array.value(1), 2);
227        assert_eq!(value_array.value(2), 3);
228
229        assert_eq!(freq_array.value(0), 1);
230        assert_eq!(freq_array.value(1), 2);
231        assert_eq!(freq_array.value(2), 3);
232    }
233}
rs_arrow_array2freq2rbat/lib.rs

rs_arrow_array2freq2rbat/
lib.rs
