rotoml/
data_operations.rs

1use polars::prelude::*;
2use std::collections::hash_map::DefaultHasher;
3use std::collections::{HashMap, HashSet};
4use std::error::Error;
5use std::hash::{Hash, Hasher};
6
7pub struct DataOperations;
8
9impl DataOperations {
10    /// Drop a single column from the DataFrame
11    ///
12    /// # Arguments
13    /// * `df` - The DataFrame to operate on
14    /// * `column_name` - The name of the column to drop
15    ///
16    /// # Returns
17    /// * `Result<DataFrame, Box<dyn Error>>` - A new DataFrame with the column removed
18    pub fn drop_column(mut df: DataFrame, column_name: &str) -> Result<DataFrame, Box<dyn Error>> {
19        // Check if the column exists
20        if !df.get_column_names().contains(&column_name) {
21            return Err(format!("Column '{}' does not exist in the DataFrame", column_name).into());
22        }
23        let _ = df.drop_in_place(column_name)?;
24        Ok(df)
25    }
26
27    /// Drop multiple columns from the DataFrame
28    ///
29    /// # Arguments
30    /// * `df` - The DataFrame to operate on
31    /// * `column_names` - A slice of column names to drop
32    ///
33    /// # Returns
34    /// * `Result<DataFrame, Box<dyn Error>>` - A new DataFrame with the columns removed
35    pub fn drop_columns(
36        mut df: DataFrame,
37        column_names: &[&str],
38    ) -> Result<DataFrame, Box<dyn Error>> {
39        // snapshot of existing columns
40        let existing: HashSet<&str> = df.get_column_names().into_iter().collect();
41
42        // Validate
43        let mut unique: Vec<&str> = Vec::new();
44        let mut seen: HashSet<&str> = HashSet::new();
45
46        for &col in column_names {
47            if !existing.contains(col) {
48                return Err(format!("Column '{}' does not exist in the DataFrame", col).into());
49            }
50
51            if seen.insert(col) {
52                unique.push(col);
53            }
54        }
55
56        for col in unique {
57            let _ = df.drop_in_place(col)?;
58        }
59
60        Ok(df)
61    }
62
63    /// Drop rows by their indexes from the DataFrame
64    ///
65    /// # Arguments
66    /// * `df` - The DataFrame to operate on
67    /// * ìndexes` - indexes to drop
68    ///
69    /// # Returns
70    /// * `Result<DataFrame, Box<dyn Error>>`
71    pub fn drop_rows(df: DataFrame, indexes: &[usize]) -> Result<DataFrame, Box<dyn Error>> {
72        let total_rows = df.height();
73
74        if indexes.is_empty() {
75            return Ok(df);
76        }
77
78        let mut drop_set: HashSet<usize> = HashSet::with_capacity(indexes.len());
79        for &idx in indexes {
80            if idx >= total_rows {
81                return Err(format!(
82                    "Index '{}' is out of bounds (DataFrame has {} rows)",
83                    idx, total_rows
84                )
85                .into());
86            }
87            drop_set.insert(idx);
88        }
89        let mask: Vec<bool> = (0..total_rows).map(|i| !drop_set.contains(&i)).collect();
90        let mask_series = BooleanChunked::new("mask", &mask);
91        let filtered_df = df.filter(&mask_series)?;
92
93        Ok(filtered_df)
94    }
95
96    /// Detect duplicate columns
97    ///
98    /// # Arguments
99    /// * `df` - The DataFrame to analyze
100    ///
101    /// # Returns
102    /// * `Result<Vec<(String, Vec<String>)>, Box<dyn Error>>` - A vector of tuples where
103    /// - the first element is the original column name
104    /// - the second element is a vector of duplicate column names that have the same data
105    pub fn detect_duplicate_columns(
106        df: &DataFrame,
107    ) -> Result<Vec<(String, Vec<String>)>, Box<dyn Error>> {
108        let mut duplicate_columns: Vec<(String, Vec<String>)> = Vec::new();
109        let mut seen_columns: HashMap<u64, String> = HashMap::new();
110        let mut duplicate_map: HashMap<String, Vec<String>> = HashMap::new();
111
112        for column_name in df.get_column_names() {
113            let column = df.column(column_name)?;
114            let mut hasher = DefaultHasher::new();
115
116            for i in 0..column.len() {
117                let value = column.get(i)?;
118                format!("{:?}", value).hash(&mut hasher);
119            }
120
121            let column_hash = hasher.finish();
122
123            if let Some(original_column) = seen_columns.get(&column_hash) {
124                if let Some(duplicates) = duplicate_map.get_mut(original_column) {
125                    duplicates.push(column_name.to_string());
126                } else {
127                    duplicate_map.insert(original_column.clone(), vec![column_name.to_string()]);
128                }
129            } else {
130                seen_columns.insert(column_hash, column_name.to_string());
131            }
132        }
133
134        for (original, duplicates) in duplicate_map {
135            duplicate_columns.push((original, duplicates));
136        }
137
138        Ok(duplicate_columns)
139    }
140
141    /// Count duplicate rows in the DataFrame
142    ///     
143    /// # Arguments
144    /// * `df` - The DataFrame to analyze
145    ///
146    /// # Returns
147    /// * `Result<(usize, Vec<usize>), Box<dyn Error>>` - A tuple where
148    ///   - the first element is the count of duplicate rows
149    ///   - the second element is a vector of indexes of duplicate rows
150    pub fn count_duplicate_rows(df: &DataFrame) -> Result<(usize, Vec<usize>), Box<dyn Error>> {
151        let total_rows = df.height();
152        let mut seen_rows: HashMap<u64, usize> = HashMap::new();
153        let mut duplicate_indexes: Vec<usize> = Vec::new();
154
155        for row_idx in 0..total_rows {
156            let mut hasher = DefaultHasher::new();
157
158            for column_name in df.get_column_names() {
159                let column = df.column(column_name)?;
160                let value = column.get(row_idx)?;
161                format!("{:?}", value).hash(&mut hasher);
162            }
163
164            let row_hash = hasher.finish();
165
166            if seen_rows.contains_key(&row_hash) {
167                duplicate_indexes.push(row_idx);
168            } else {
169                seen_rows.insert(row_hash, row_idx);
170            }
171        }
172        let duplicate_count = duplicate_indexes.len();
173        Ok((duplicate_count, duplicate_indexes))
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn test_drop_column() {
183        let df = df! {
184            "a" => &[1, 2, 3],
185            "b" => &[4, 5, 6],
186            "c" => &[7, 8, 9],
187        }
188        .unwrap();
189
190        let result = DataOperations::drop_column(df.clone(), "b").unwrap();
191
192        assert_eq!(result.width(), 2);
193        assert!(result.get_column_names().contains(&"a"));
194        assert!(result.get_column_names().contains(&"c"));
195        assert!(!result.get_column_names().contains(&"b"));
196    }
197
198    #[test]
199    fn test_drop_column_not_exists() {
200        let df = df! {
201            "a" => &[1, 2, 3],
202            "b" => &[4, 5, 6],
203        }
204        .unwrap();
205
206        let result = DataOperations::drop_column(df, "nonexistent");
207        assert!(result.is_err());
208    }
209
210    #[test]
211    fn test_drop_columns() {
212        let df = df! {
213            "a" => &[1, 2, 3],
214            "b" => &[4, 5, 6],
215            "c" => &[7, 8, 9],
216            "d" => &[10, 11, 12],
217        }
218        .unwrap();
219
220        let result = DataOperations::drop_columns(df.clone(), &["b", "d"]).unwrap();
221
222        assert_eq!(result.width(), 2);
223        assert!(result.get_column_names().contains(&"a"));
224        assert!(result.get_column_names().contains(&"c"));
225        assert!(!result.get_column_names().contains(&"b"));
226        assert!(!result.get_column_names().contains(&"d"));
227    }
228
229    #[test]
230    fn test_drop_columns_one_not_exists() {
231        let df = df! {
232            "a" => &[1, 2, 3],
233            "b" => &[4, 5, 6],
234        }
235        .unwrap();
236
237        let result = DataOperations::drop_columns(df, &["b", "nonexistent"]);
238        assert!(result.is_err());
239    }
240
241    #[test]
242    fn test_drop_rows() {
243        let df = df! {
244            "a" => &[1, 2, 3, 4, 5],
245            "b" => &[10, 20, 30, 40, 50],
246        }
247        .unwrap();
248
249        // Drop rows at indexes 1 and 3
250        let result = DataOperations::drop_rows(df.clone(), &[1, 3]).unwrap();
251
252        assert_eq!(result.height(), 3); // Should have 3 rows left
253        assert_eq!(result.width(), 2); // Should still have 2 columns
254
255        // Check remaining values
256        let col_a = result.column("a").unwrap();
257        let a_values: Vec<Option<i32>> = col_a.i32().unwrap().into_iter().collect();
258        assert_eq!(a_values, vec![Some(1), Some(3), Some(5)]);
259    }
260
261    #[test]
262    fn test_drop_rows_single() {
263        let df = df! {
264            "a" => &[1, 2, 3],
265            "b" => &[10, 20, 30],
266        }
267        .unwrap();
268
269        let result = DataOperations::drop_rows(df.clone(), &[0]).unwrap();
270
271        assert_eq!(result.height(), 2);
272        let col_a = result.column("a").unwrap();
273        let a_values: Vec<Option<i32>> = col_a.i32().unwrap().into_iter().collect();
274        assert_eq!(a_values, vec![Some(2), Some(3)]);
275    }
276
277    #[test]
278    fn test_drop_rows_out_of_bounds() {
279        let df = df! {
280            "a" => &[1, 2, 3],
281            "b" => &[10, 20, 30],
282        }
283        .unwrap();
284
285        let result = DataOperations::drop_rows(df, &[5]);
286        assert!(result.is_err());
287    }
288
289    #[test]
290    fn test_detect_duplicate_columns() {
291        let df = df! {
292            "a" => &[1, 2, 3],
293            "b" => &[1, 2, 3],  // duplicate of 'a'
294            "c" => &[4, 5, 6],
295            "d" => &[1, 2, 3],  // duplicate of 'a'
296        }
297        .unwrap();
298
299        let result = DataOperations::detect_duplicate_columns(&df).unwrap();
300
301        assert_eq!(result.len(), 1); // Should find one group of duplicates
302
303        let (original, duplicates) = &result[0];
304        assert_eq!(original, "a");
305        assert_eq!(duplicates.len(), 2);
306        assert!(duplicates.contains(&"b".to_string()));
307        assert!(duplicates.contains(&"d".to_string()));
308    }
309
310    #[test]
311    fn test_detect_duplicate_columns_no_duplicates() {
312        let df = df! {
313            "a" => &[1, 2, 3],
314            "b" => &[4, 5, 6],
315            "c" => &[7, 8, 9],
316        }
317        .unwrap();
318
319        let result = DataOperations::detect_duplicate_columns(&df).unwrap();
320
321        assert_eq!(result.len(), 0); // No duplicates found
322    }
323
324    #[test]
325    fn test_detect_duplicate_columns_multiple_groups() {
326        let df = df! {
327            "a" => &[1, 2, 3],
328            "b" => &[1, 2, 3],  // duplicate of 'a'
329            "c" => &[4, 5, 6],
330            "d" => &[4, 5, 6],  // duplicate of 'c'
331        }
332        .unwrap();
333
334        let result = DataOperations::detect_duplicate_columns(&df).unwrap();
335
336        assert_eq!(result.len(), 2); // Should find two groups of duplicates
337    }
338
339    #[test]
340    fn test_count_duplicate_rows() {
341        let df = df! {
342            "a" => &[1, 2, 1, 3, 2],
343            "b" => &[10, 20, 10, 30, 20],
344        }
345        .unwrap();
346
347        let (count, indexes) = DataOperations::count_duplicate_rows(&df).unwrap();
348
349        assert_eq!(count, 2); // Two duplicate rows
350        assert_eq!(indexes.len(), 2);
351        assert!(indexes.contains(&2)); // Row at index 2 is duplicate of row 0
352        assert!(indexes.contains(&4)); // Row at index 4 is duplicate of row 1
353    }
354
355    #[test]
356    fn test_count_duplicate_rows_no_duplicates() {
357        let df = df! {
358            "a" => &[1, 2, 3],
359            "b" => &[10, 20, 30],
360        }
361        .unwrap();
362
363        let (count, indexes) = DataOperations::count_duplicate_rows(&df).unwrap();
364
365        assert_eq!(count, 0);
366        assert_eq!(indexes.len(), 0);
367    }
368
369    #[test]
370    fn test_count_duplicate_rows_all_duplicates() {
371        let df = df! {
372            "a" => &[1, 1, 1],
373            "b" => &[10, 10, 10],
374        }
375        .unwrap();
376
377        let (count, indexes) = DataOperations::count_duplicate_rows(&df).unwrap();
378
379        assert_eq!(count, 2); // Two rows are duplicates (indexes 1 and 2)
380        assert_eq!(indexes, vec![1, 2]);
381    }
382}
rotoml/data_operations.rs

rotoml/
data_operations.rs
