1use nu_engine::command_prelude::*;
2use nu_protocol::{PipelineIterator, ast::PathMember, casing::Casing};
3use std::collections::BTreeSet;
4
5#[derive(Clone)]
6pub struct Select;
7
8impl Command for Select {
9 fn name(&self) -> &str {
10 "select"
11 }
12
13 fn signature(&self) -> Signature {
15 Signature::build("select")
16 .input_output_types(vec![
17 (Type::record(), Type::record()),
18 (Type::table(), Type::table()),
19 (Type::List(Box::new(Type::Any)), Type::Any),
20 ])
21 .switch(
22 "ignore-errors",
23 "ignore missing data (make all cell path members optional)",
24 Some('i'),
25 )
26 .rest(
27 "rest",
28 SyntaxShape::CellPath,
29 "The columns to select from the table.",
30 )
31 .allow_variants_without_examples(true)
32 .category(Category::Filters)
33 }
34
35 fn description(&self) -> &str {
36 "Select only these columns or rows from the input. Opposite of `reject`."
37 }
38
39 fn extra_description(&self) -> &str {
40 r#"This differs from `get` in that, rather than accessing the given value in the data structure,
41it removes all non-selected values from the structure. Hence, using `select` on a table will
42produce a table, a list will produce a list, and a record will produce a record."#
43 }
44
45 fn search_terms(&self) -> Vec<&str> {
46 vec!["pick", "choose", "get"]
47 }
48
49 fn run(
50 &self,
51 engine_state: &EngineState,
52 stack: &mut Stack,
53 call: &Call,
54 input: PipelineData,
55 ) -> Result<PipelineData, ShellError> {
56 let columns: Vec<Value> = call.rest(engine_state, stack, 0)?;
57 let mut new_columns: Vec<CellPath> = vec![];
58 for col_val in columns {
59 let col_span = col_val.span();
60 match col_val {
61 Value::CellPath { val, .. } => {
62 new_columns.push(val);
63 }
64 Value::String { val, .. } => {
65 let cv = CellPath {
66 members: vec![PathMember::String {
67 val,
68 span: col_span,
69 optional: false,
70 casing: Casing::Sensitive,
71 }],
72 };
73 new_columns.push(cv);
74 }
75 Value::Int { val, .. } => {
76 if val < 0 {
77 return Err(ShellError::CantConvert {
78 to_type: "cell path".into(),
79 from_type: "negative number".into(),
80 span: col_span,
81 help: None,
82 });
83 }
84 let cv = CellPath {
85 members: vec![PathMember::Int {
86 val: val as usize,
87 span: col_span,
88 optional: false,
89 }],
90 };
91 new_columns.push(cv);
92 }
93 x => {
94 return Err(ShellError::CantConvert {
95 to_type: "cell path".into(),
96 from_type: x.get_type().to_string(),
97 span: x.span(),
98 help: None,
99 });
100 }
101 }
102 }
103 let ignore_errors = call.has_flag(engine_state, stack, "ignore-errors")?;
104 let span = call.head;
105
106 if ignore_errors {
107 for cell_path in &mut new_columns {
108 cell_path.make_optional();
109 }
110 }
111
112 select(engine_state, span, new_columns, input)
113 }
114
115 fn examples(&self) -> Vec<Example> {
116 vec![
117 Example {
118 description: "Select a column in a table",
119 example: "[{a: a b: b}] | select a",
120 result: Some(Value::test_list(vec![Value::test_record(record! {
121 "a" => Value::test_string("a")
122 })])),
123 },
124 Example {
125 description: "Select a field in a record",
126 example: "{a: a b: b} | select a",
127 result: Some(Value::test_record(record! {
128 "a" => Value::test_string("a")
129 })),
130 },
131 Example {
132 description: "Select just the `name` column",
133 example: "ls | select name",
134 result: None,
135 },
136 Example {
137 description: "Select the first four rows (this is the same as `first 4`)",
138 example: "ls | select 0 1 2 3",
139 result: None,
140 },
141 Example {
142 description: "Select multiple columns",
143 example: "[[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb]] | select name type",
144 result: Some(Value::test_list(vec![
145 Value::test_record(record! {
146 "name" => Value::test_string("Cargo.toml"),
147 "type" => Value::test_string("toml"),
148 }),
149 Value::test_record(record! {
150 "name" => Value::test_string("Cargo.lock"),
151 "type" => Value::test_string("toml")
152 }),
153 ])),
154 },
155 Example {
156 description: "Select multiple columns by spreading a list",
157 example: r#"let cols = [name type]; [[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb]] | select ...$cols"#,
158 result: Some(Value::test_list(vec![
159 Value::test_record(record! {
160 "name" => Value::test_string("Cargo.toml"),
161 "type" => Value::test_string("toml")
162 }),
163 Value::test_record(record! {
164 "name" => Value::test_string("Cargo.lock"),
165 "type" => Value::test_string("toml")
166 }),
167 ])),
168 },
169 ]
170 }
171}
172
173fn select(
174 engine_state: &EngineState,
175 call_span: Span,
176 columns: Vec<CellPath>,
177 input: PipelineData,
178) -> Result<PipelineData, ShellError> {
179 let mut unique_rows: BTreeSet<usize> = BTreeSet::new();
180
181 let mut new_columns = vec![];
182
183 for column in columns {
184 let CellPath { ref members } = column;
185 match members.first() {
186 Some(PathMember::Int { val, span, .. }) => {
187 if members.len() > 1 {
188 return Err(ShellError::GenericError {
189 error: "Select only allows row numbers for rows".into(),
190 msg: "extra after row number".into(),
191 span: Some(*span),
192 help: None,
193 inner: vec![],
194 });
195 }
196 unique_rows.insert(*val);
197 }
198 _ => {
199 if !new_columns.contains(&column) {
200 new_columns.push(column)
201 }
202 }
203 };
204 }
205 let columns = new_columns;
206
207 let input = if !unique_rows.is_empty() {
208 let metadata = input.metadata();
209 let pipeline_iter: PipelineIterator = input.into_iter();
210
211 NthIterator {
212 input: pipeline_iter,
213 rows: unique_rows.into_iter().peekable(),
214 current: 0,
215 }
216 .into_pipeline_data_with_metadata(
217 call_span,
218 engine_state.signals().clone(),
219 metadata,
220 )
221 } else {
222 input
223 };
224
225 match input {
226 PipelineData::Value(v, metadata, ..) => {
227 let span = v.span();
228 match v {
229 Value::List {
230 vals: input_vals, ..
231 } => Ok(input_vals
232 .into_iter()
233 .map(move |input_val| {
234 if !columns.is_empty() {
235 let mut record = Record::new();
236 for path in &columns {
237 match input_val.follow_cell_path(&path.members) {
238 Ok(fetcher) => {
239 record.push(path.to_column_name(), fetcher.into_owned());
240 }
241 Err(e) => return Value::error(e, call_span),
242 }
243 }
244
245 Value::record(record, span)
246 } else {
247 input_val.clone()
248 }
249 })
250 .into_pipeline_data_with_metadata(
251 call_span,
252 engine_state.signals().clone(),
253 metadata,
254 )),
255 _ => {
256 if !columns.is_empty() {
257 let mut record = Record::new();
258
259 for cell_path in columns {
260 let result = v.follow_cell_path(&cell_path.members)?;
261 record.push(cell_path.to_column_name(), result.into_owned());
262 }
263
264 Ok(Value::record(record, call_span)
265 .into_pipeline_data_with_metadata(metadata))
266 } else {
267 Ok(v.into_pipeline_data_with_metadata(metadata))
268 }
269 }
270 }
271 }
272 PipelineData::ListStream(stream, metadata, ..) => Ok(stream
273 .map(move |x| {
274 if !columns.is_empty() {
275 let mut record = Record::new();
276 for path in &columns {
277 match x.follow_cell_path(&path.members) {
278 Ok(value) => {
279 record.push(path.to_column_name(), value.into_owned());
280 }
281 Err(e) => return Value::error(e, call_span),
282 }
283 }
284 Value::record(record, call_span)
285 } else {
286 x
287 }
288 })
289 .into_pipeline_data_with_metadata(call_span, engine_state.signals().clone(), metadata)),
290 _ => Ok(PipelineData::empty()),
291 }
292}
293
294struct NthIterator {
295 input: PipelineIterator,
296 rows: std::iter::Peekable<std::collections::btree_set::IntoIter<usize>>,
297 current: usize,
298}
299
300impl Iterator for NthIterator {
301 type Item = Value;
302
303 fn next(&mut self) -> Option<Self::Item> {
304 loop {
305 if let Some(row) = self.rows.peek() {
306 if self.current == *row {
307 self.rows.next();
308 self.current += 1;
309 return self.input.next();
310 } else {
311 self.current += 1;
312 let _ = self.input.next()?;
313 continue;
314 }
315 } else {
316 return None;
317 }
318 }
319 }
320}
321
322#[cfg(test)]
323mod tests {
324 use super::*;
325
326 #[test]
327 fn test_examples() {
328 use crate::test_examples;
329
330 test_examples(Select)
331 }
332}