alopex_dataframe/dataframe/
dataframe.rs

1use std::collections::HashSet;
2use std::sync::Arc;
3
4use arrow::datatypes::{Field, Schema, SchemaRef};
5use arrow::record_batch::RecordBatch;
6
7use crate::ops::{FillNull, JoinKeys, JoinType, SortOptions};
8use crate::{DataFrameError, Expr, Result, Series};
9
10/// An eager table backed by one or more Arrow `RecordBatch` values.
11#[derive(Debug, Clone)]
12pub struct DataFrame {
13    schema: SchemaRef,
14    batches: Vec<RecordBatch>,
15}
16
17impl DataFrame {
18    /// Construct a `DataFrame` from a list of `Series`.
19    ///
20    /// Chunk boundaries do not need to align across series as long as total lengths match.
21    pub fn new(columns: Vec<Series>) -> Result<Self> {
22        if columns.is_empty() {
23            return Ok(Self::empty());
24        }
25
26        let mut seen_names = HashSet::with_capacity(columns.len());
27        for c in &columns {
28            if !seen_names.insert(c.name().to_string()) {
29                return Err(DataFrameError::schema_mismatch(format!(
30                    "duplicate column name '{}'",
31                    c.name()
32                )));
33            }
34        }
35
36        let expected_len = columns[0].len();
37        for c in &columns[1..] {
38            if c.len() != expected_len {
39                return Err(DataFrameError::schema_mismatch(format!(
40                    "column length mismatch: '{}' has length {}, expected {}",
41                    c.name(),
42                    c.len(),
43                    expected_len
44                )));
45            }
46        }
47
48        let fields: Vec<Field> = columns
49            .iter()
50            .map(|c| Field::new(c.name(), c.dtype(), true))
51            .collect();
52        let schema: SchemaRef = Arc::new(Schema::new(fields));
53
54        let arrays = columns
55            .iter()
56            .map(|c| {
57                if c.chunks().is_empty() {
58                    Ok(arrow::array::new_empty_array(&c.dtype()))
59                } else if c.chunks().len() == 1 {
60                    Ok(c.chunks()[0].clone())
61                } else {
62                    let arrays = c
63                        .chunks()
64                        .iter()
65                        .map(|a| a.as_ref() as &dyn arrow::array::Array)
66                        .collect::<Vec<_>>();
67                    arrow::compute::concat(&arrays)
68                        .map_err(|source| DataFrameError::Arrow { source })
69                }
70            })
71            .collect::<Result<Vec<_>>>()?;
72
73        let batch = RecordBatch::try_new(schema.clone(), arrays).map_err(|e| {
74            DataFrameError::schema_mismatch(format!("failed to build RecordBatch: {e}"))
75        })?;
76
77        Ok(Self {
78            schema,
79            batches: vec![batch],
80        })
81    }
82
83    /// Construct a `DataFrame` from Arrow record batches (all batches must share the same schema).
84    pub fn from_batches(batches: Vec<RecordBatch>) -> Result<Self> {
85        if batches.is_empty() {
86            return Ok(Self::empty());
87        }
88
89        let schema = batches[0].schema();
90        for (i, b) in batches.iter().enumerate().skip(1) {
91            if b.schema().as_ref() != schema.as_ref() {
92                return Err(DataFrameError::schema_mismatch(format!(
93                    "schema mismatch between batches: batch 0 != batch {i}"
94                )));
95            }
96        }
97
98        Ok(Self { schema, batches })
99    }
100
101    /// Alias for `DataFrame::new`.
102    pub fn from_series(series: Vec<Series>) -> Result<Self> {
103        Self::new(series)
104    }
105
106    /// Return an empty `DataFrame` (no columns, no rows).
107    pub fn empty() -> Self {
108        Self {
109            schema: Arc::new(Schema::empty()),
110            batches: Vec::new(),
111        }
112    }
113
114    /// Return the number of rows.
115    pub fn height(&self) -> usize {
116        self.batches.iter().map(|b| b.num_rows()).sum()
117    }
118
119    /// Return the number of columns.
120    pub fn width(&self) -> usize {
121        self.schema.fields().len()
122    }
123
124    /// Return the Arrow schema.
125    pub fn schema(&self) -> SchemaRef {
126        self.schema.clone()
127    }
128
129    /// Get a column by name (case-sensitive).
130    pub fn column(&self, name: &str) -> Result<Series> {
131        let idx = self
132            .schema
133            .fields()
134            .iter()
135            .position(|f| f.name() == name)
136            .ok_or_else(|| DataFrameError::column_not_found(name.to_string()))?;
137
138        let chunks = self
139            .batches
140            .iter()
141            .map(|b| b.column(idx).clone())
142            .collect::<Vec<_>>();
143        Ok(Series::from_arrow_unchecked(name, chunks))
144    }
145
146    /// Return all columns in construction order.
147    pub fn columns(&self) -> Vec<Series> {
148        self.schema
149            .fields()
150            .iter()
151            .enumerate()
152            .map(|(idx, f)| {
153                let chunks = self
154                    .batches
155                    .iter()
156                    .map(|b| b.column(idx).clone())
157                    .collect::<Vec<_>>();
158                Series::from_arrow_unchecked(f.name(), chunks)
159            })
160            .collect()
161    }
162
163    /// Return the underlying Arrow batches.
164    pub fn to_arrow(&self) -> Vec<RecordBatch> {
165        self.batches.clone()
166    }
167
168    /// Convert this eager `DataFrame` to a `LazyFrame` for query planning/execution.
169    pub fn lazy(&self) -> crate::LazyFrame {
170        crate::LazyFrame::from_dataframe(self.clone())
171    }
172
173    /// Eager `select`, implemented by delegating to `LazyFrame`.
174    pub fn select(&self, exprs: Vec<Expr>) -> Result<Self> {
175        self.clone().lazy().select(exprs).collect()
176    }
177
178    /// Eager `filter`, implemented by delegating to `LazyFrame`.
179    pub fn filter(&self, predicate: Expr) -> Result<Self> {
180        self.clone().lazy().filter(predicate).collect()
181    }
182
183    /// Eager `with_columns`, implemented by delegating to `LazyFrame`.
184    pub fn with_columns(&self, exprs: Vec<Expr>) -> Result<Self> {
185        self.clone().lazy().with_columns(exprs).collect()
186    }
187
188    /// Start a group-by aggregation (eager API).
189    pub fn group_by(&self, by: Vec<Expr>) -> GroupBy {
190        GroupBy {
191            df: self.clone(),
192            by,
193        }
194    }
195
196    /// Join with another `DataFrame` using provided join keys.
197    pub fn join<K: Into<JoinKeys>>(
198        &self,
199        other: &DataFrame,
200        keys: K,
201        how: JoinType,
202    ) -> Result<Self> {
203        self.clone()
204            .lazy()
205            .join(other.clone().lazy(), keys, how)
206            .collect()
207    }
208
209    /// Sort by one or more columns.
210    pub fn sort(&self, by: Vec<String>, descending: Vec<bool>) -> Result<Self> {
211        let options = SortOptions {
212            by,
213            descending,
214            nulls_last: true,
215            stable: true,
216        };
217        self.clone().lazy().sort(options).collect()
218    }
219
220    /// Return the first `n` rows.
221    pub fn head(&self, n: usize) -> Result<Self> {
222        self.clone().lazy().head(n).collect()
223    }
224
225    /// Return the last `n` rows.
226    pub fn tail(&self, n: usize) -> Result<Self> {
227        self.clone().lazy().tail(n).collect()
228    }
229
230    /// Remove duplicate rows.
231    pub fn unique(&self, subset: Option<Vec<String>>) -> Result<Self> {
232        self.clone().lazy().unique(subset).collect()
233    }
234
235    /// Fill null values using a scalar or strategy.
236    pub fn fill_null<T: Into<FillNull>>(&self, fill: T) -> Result<Self> {
237        self.clone().lazy().fill_null(fill).collect()
238    }
239
240    /// Drop rows containing null values.
241    pub fn drop_nulls(&self, subset: Option<Vec<String>>) -> Result<Self> {
242        self.clone().lazy().drop_nulls(subset).collect()
243    }
244
245    /// Count null values per column.
246    pub fn null_count(&self) -> Result<Self> {
247        self.clone().lazy().null_count().collect()
248    }
249
250    /// Explode one `List<Utf8>` column into multiple rows.
251    pub fn explode(&self, column: impl Into<String>) -> Result<Self> {
252        self.clone().lazy().explode(column).collect()
253    }
254
255    /// Implode UTF-8 columns into one row of `List<Utf8>` columns.
256    pub fn implode(&self) -> Result<Self> {
257        self.clone().lazy().implode().collect()
258    }
259}
260
261/// Eager group-by handle that delegates execution to `LazyFrame`.
262#[derive(Debug, Clone)]
263pub struct GroupBy {
264    df: DataFrame,
265    by: Vec<Expr>,
266}
267
268impl GroupBy {
269    /// Perform aggregations for this group-by.
270    pub fn agg(self, aggs: Vec<Expr>) -> Result<DataFrame> {
271        self.df.lazy().group_by(self.by).agg(aggs).collect()
272    }
273
274    /// Return the underlying `DataFrame`.
275    pub fn into_df(self) -> DataFrame {
276        self.df
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use std::sync::Arc;
283
284    use arrow::array::{ArrayRef, Int32Array, StringArray};
285    use arrow::datatypes::{DataType, Field, Schema};
286    use arrow::record_batch::RecordBatch;
287
288    use super::DataFrame;
289    use crate::{DataFrameError, Series};
290
291    fn s_i32(name: &str, chunks: Vec<Vec<i32>>) -> Series {
292        let arrays: Vec<ArrayRef> = chunks
293            .into_iter()
294            .map(|v| Arc::new(Int32Array::from(v)) as ArrayRef)
295            .collect();
296        Series::from_arrow(name, arrays).unwrap()
297    }
298
299    #[test]
300    fn dataframe_new_accepts_misaligned_chunks_by_normalizing() {
301        let a = s_i32("a", vec![vec![1, 2], vec![3]]);
302        let b = s_i32("b", vec![vec![10], vec![20, 30]]);
303
304        let df = DataFrame::new(vec![a, b]).unwrap();
305        assert_eq!(df.height(), 3);
306        assert_eq!(df.width(), 2);
307        assert_eq!(df.schema().fields()[0].name(), "a");
308        assert_eq!(df.schema().fields()[1].name(), "b");
309
310        let batches = df.to_arrow();
311        assert_eq!(batches.len(), 1);
312        assert_eq!(batches[0].num_rows(), 3);
313    }
314
315    #[test]
316    fn dataframe_new_rejects_duplicate_column_names() {
317        let a1 = s_i32("a", vec![vec![1]]);
318        let a2 = s_i32("a", vec![vec![2]]);
319        let err = DataFrame::new(vec![a1, a2]).unwrap_err();
320        assert!(matches!(err, DataFrameError::SchemaMismatch { .. }));
321    }
322
323    #[test]
324    fn dataframe_new_rejects_length_mismatch() {
325        let a = s_i32("a", vec![vec![1, 2]]);
326        let b = s_i32("b", vec![vec![10]]);
327        let err = DataFrame::new(vec![a, b]).unwrap_err();
328        assert!(matches!(err, DataFrameError::SchemaMismatch { .. }));
329    }
330
331    #[test]
332    fn dataframe_new_accepts_different_chunk_counts() {
333        let a = s_i32("a", vec![vec![1], vec![2], vec![3]]);
334        let b = s_i32("b", vec![vec![10, 20, 30]]);
335        let df = DataFrame::new(vec![a, b]).unwrap();
336        assert_eq!(df.height(), 3);
337        assert_eq!(df.to_arrow().len(), 1);
338    }
339
340    #[test]
341    fn dataframe_column_is_case_sensitive() {
342        let a = s_i32("a", vec![vec![1]]);
343        let df = DataFrame::new(vec![a]).unwrap();
344        assert!(matches!(
345            df.column("A").unwrap_err(),
346            DataFrameError::ColumnNotFound { .. }
347        ));
348    }
349
350    #[test]
351    fn dataframe_from_batches_rejects_schema_mismatch() {
352        let a1: ArrayRef = Arc::new(Int32Array::from(vec![1]));
353        let a2: ArrayRef = Arc::new(StringArray::from(vec!["x"]));
354
355        let s1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
356        let s2 = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)]));
357
358        let b1 = RecordBatch::try_new(s1, vec![a1]).unwrap();
359        let b2 = RecordBatch::try_new(s2, vec![a2]).unwrap();
360
361        let err = DataFrame::from_batches(vec![b1, b2]).unwrap_err();
362        assert!(matches!(err, DataFrameError::SchemaMismatch { .. }));
363    }
364
365    #[test]
366    fn dataframe_columns_preserves_schema_order() {
367        let a = s_i32("a", vec![vec![1], vec![2]]);
368        let b = s_i32("b", vec![vec![10], vec![20]]);
369        let df = DataFrame::new(vec![b.clone(), a.clone()]).unwrap();
370
371        let cols = df.columns();
372        assert_eq!(cols[0].name(), "b");
373        assert_eq!(cols[1].name(), "a");
374        assert_eq!(cols[0].len(), 2);
375        assert_eq!(cols[1].len(), 2);
376    }
377}
alopex_dataframe/dataframe/dataframe.rs

alopex_dataframe/dataframe/
dataframe.rs
