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