[READ-ONLY] Mirror of https://github.com/probablykasper/csv-pipeline.
Add select method
+55
-1
+27
-1
src/pipeline.rs
+27
-1
src/pipeline.rs
···
1
1
use super::headers::Headers;
2
2
use crate::pipeline_iterators::{
3
-
AddCol, Flush, MapCol, MapRow, TransformInto, Validate, ValidateCol,
3
+
AddCol, Flush, MapCol, MapRow, Select, TransformInto, Validate, ValidateCol,
4
4
};
5
5
use crate::target::{StringTarget, Target};
6
6
use crate::transform::Transform;
···
132
132
name: col.to_string(),
133
133
index: self.headers.get_index(col),
134
134
});
135
+
self
136
+
}
137
+
138
+
/// Pick which columns to output, in the specified order. Panics if duplicate colums are specified.
139
+
///
140
+
/// ## Example
141
+
///
142
+
/// ```
143
+
/// use csv_pipeline::Pipeline;
144
+
///
145
+
/// let csv = Pipeline::from_path("test/AB.csv")
146
+
/// .unwrap()
147
+
/// .select(vec!["B"])
148
+
/// .collect_into_string()
149
+
/// .unwrap();
150
+
///
151
+
/// assert_eq!(csv, "B\n2\n");
152
+
/// ```
153
+
pub fn select(mut self, columns: Vec<&str>) -> Self {
154
+
let new_header_row = Row::from(columns.clone());
155
+
self.iterator = Box::new(Select {
156
+
iterator: self.iterator,
157
+
columns: columns.into_iter().map(String::from).collect(),
158
+
headers: self.headers.clone(),
159
+
});
160
+
self.headers = Headers::from_row(new_header_row).unwrap();
135
161
self
136
162
}
137
163
+28
src/pipeline_iterators.rs
+28
src/pipeline_iterators.rs
···
92
92
}
93
93
}
94
94
95
+
pub struct Select<I> {
96
+
pub iterator: I,
97
+
pub columns: Vec<String>,
98
+
pub headers: Headers,
99
+
}
100
+
impl<I> Iterator for Select<I>
101
+
where
102
+
I: Iterator<Item = RowResult>,
103
+
{
104
+
type Item = RowResult;
105
+
106
+
fn next(&mut self) -> Option<Self::Item> {
107
+
let row = match self.iterator.next()? {
108
+
Ok(row) => row,
109
+
Err(e) => return Some(Err(e)),
110
+
};
111
+
let mut selection = Vec::with_capacity(self.columns.len());
112
+
for col in &self.columns {
113
+
let field = match self.headers.get_field(&row, col) {
114
+
Some(field) => field,
115
+
None => return Some(Err(Error::MissingColumn(col.clone()))),
116
+
};
117
+
selection.push(field);
118
+
}
119
+
Some(Ok(selection.into()))
120
+
}
121
+
}
122
+
95
123
pub struct TransformInto<I, F>
96
124
where
97
125
F: FnMut() -> Vec<Box<dyn Transform>>,