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