#[cfg(feature = "sqlite")]
use crate::database::SQLiteQueryBuilder;
use nu_engine::command_prelude::*;
use nu_protocol::{
DeprecationEntry, DeprecationType, PipelineIterator, ReportMode, ast::PathMember,
casing::Casing, shell_error::generic::GenericError,
};
use std::collections::BTreeSet;
#[derive(Clone)]
pub struct Select;
impl Command for Select {
fn name(&self) -> &str {
"select"
}
fn signature(&self) -> Signature {
Signature::build("select")
.input_output_types(vec![
(Type::record(), Type::record()),
(Type::table(), Type::table()),
(Type::List(Box::new(Type::Any)), Type::Any),
#[cfg(feature = "sqlite")]
(
Type::Custom("SQLiteQueryBuilder".into()),
Type::Custom("SQLiteQueryBuilder".into()),
),
])
.switch(
"optional",
"Make all cell path members optional (returns `null` for missing values).",
Some('o'),
)
.switch(
"ignore-case",
"Make all cell path members case insensitive.",
None,
)
.switch(
"ignore-errors",
"Ignore missing data (make all cell path members optional) (deprecated).",
Some('i'),
)
.rest(
"rest",
SyntaxShape::CellPath,
"The columns to select from the table.",
)
.allow_variants_without_examples(true)
.category(Category::Filters)
}
fn description(&self) -> &str {
"Select only these columns or rows from the input. Opposite of `reject`."
}
fn extra_description(&self) -> &str {
"This differs from `get` in that, rather than accessing the given value in the data structure,
it removes all non-selected values from the structure. Hence, using `select` on a table will
produce a table, a list will produce a list, and a record will produce a record."
}
fn search_terms(&self) -> Vec<&str> {
vec!["pick", "choose", "get", "retain"]
}
fn run(
&self,
engine_state: &EngineState,
stack: &mut Stack,
call: &Call,
input: PipelineData,
) -> Result<PipelineData, ShellError> {
let columns: Vec<Value> = call.rest(engine_state, stack, 0)?;
let mut new_columns: Vec<CellPath> = vec![];
for col_val in columns {
let col_span = col_val.span();
match col_val {
Value::CellPath { val, .. } => {
new_columns.push(val);
}
Value::String { val, .. } => {
let cv = CellPath {
members: vec![PathMember::String {
val,
span: col_span,
optional: false,
casing: Casing::Sensitive,
}],
};
new_columns.push(cv);
}
Value::Int { val, .. } => {
if val < 0 {
return Err(ShellError::CantConvert {
to_type: "cell path".into(),
from_type: "negative number".into(),
span: col_span,
help: None,
});
}
let cv = CellPath {
members: vec![PathMember::Int {
val: val as usize,
span: col_span,
optional: false,
}],
};
new_columns.push(cv);
}
x => {
return Err(ShellError::CantConvert {
to_type: "cell path".into(),
from_type: x.get_type().to_string(),
span: x.span(),
help: None,
});
}
}
}
let optional = call.has_flag(engine_state, stack, "optional")?
|| call.has_flag(engine_state, stack, "ignore-errors")?;
let ignore_case = call.has_flag(engine_state, stack, "ignore-case")?;
let span = call.head;
if optional {
for cell_path in &mut new_columns {
cell_path.make_optional();
}
}
if ignore_case {
for cell_path in &mut new_columns {
cell_path.make_insensitive();
}
}
select(engine_state, span, new_columns, input)
}
fn deprecation_info(&self) -> Vec<DeprecationEntry> {
vec![DeprecationEntry {
ty: DeprecationType::Flag("ignore-errors".into()),
report_mode: ReportMode::FirstUse,
since: Some("0.106.0".into()),
expected_removal: None,
help: Some(
"This flag has been renamed to `--optional (-o)` to better reflect its behavior."
.into(),
),
}]
}
fn examples(&self) -> Vec<Example<'_>> {
vec![
Example {
description: "Select a column in a table.",
example: "[{a: a b: b}] | select a",
result: Some(Value::test_list(vec![Value::test_record(record! {
"a" => Value::test_string("a")
})])),
},
Example {
description: "Select a column even if some rows are missing that column.",
example: "[{a: a0 b: b0} {b: b1}] | select -o a",
result: Some(Value::test_list(vec![
Value::test_record(record! {
"a" => Value::test_string("a0")
}),
Value::test_record(record! {
"a" => Value::test_nothing()
}),
])),
},
Example {
description: "Select a field in a record.",
example: "{a: a b: b} | select a",
result: Some(Value::test_record(record! {
"a" => Value::test_string("a")
})),
},
Example {
description: "Select just the `name` column.",
example: "ls | select name",
result: None,
},
Example {
description: "Select the first four rows (this is the same as `first 4`).",
example: "ls | select 0 1 2 3",
result: None,
},
Example {
description: "Select multiple columns.",
example: "[[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb]] | select name type",
result: Some(Value::test_list(vec![
Value::test_record(record! {
"name" => Value::test_string("Cargo.toml"),
"type" => Value::test_string("toml"),
}),
Value::test_record(record! {
"name" => Value::test_string("Cargo.lock"),
"type" => Value::test_string("toml")
}),
])),
},
Example {
description: "Select multiple columns by spreading a list.",
example: "let cols = [name type]; [[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb]] | select ...$cols",
result: Some(Value::test_list(vec![
Value::test_record(record! {
"name" => Value::test_string("Cargo.toml"),
"type" => Value::test_string("toml")
}),
Value::test_record(record! {
"name" => Value::test_string("Cargo.lock"),
"type" => Value::test_string("toml")
}),
])),
},
]
}
}
fn select(
engine_state: &EngineState,
call_span: Span,
columns: Vec<CellPath>,
mut input: PipelineData,
) -> Result<PipelineData, ShellError> {
let mut unique_rows: BTreeSet<usize> = BTreeSet::new();
let mut new_columns = vec![];
for column in columns {
let CellPath { ref members } = column;
match members.first() {
Some(PathMember::Int { val, span, .. }) => {
if members.len() > 1 {
return Err(ShellError::Generic(GenericError::new(
"Select only allows row numbers for rows",
"extra after row number",
*span,
)));
}
unique_rows.insert(*val);
}
_ => {
if !new_columns.contains(&column) {
new_columns.push(column)
}
}
};
}
let columns = new_columns;
let input = if !unique_rows.is_empty() {
let metadata = input.take_metadata();
let pipeline_iter: PipelineIterator = input.into_iter();
NthIterator {
input: pipeline_iter,
rows: unique_rows.into_iter().peekable(),
current: 0,
}
.into_pipeline_data_with_metadata(
call_span,
engine_state.signals().clone(),
metadata,
)
} else {
input
};
#[cfg(feature = "sqlite")]
if let PipelineData::Value(Value::Custom { val, .. }, ..) = &input
&& let Some(table) = val.as_any().downcast_ref::<SQLiteQueryBuilder>()
{
let select_columns: Option<Vec<String>> = columns
.iter()
.map(|column| match column.members.as_slice() {
[PathMember::String { val, .. }] => Some(val.clone()),
_ => None,
})
.collect();
if let Some(select_columns) = select_columns.filter(|selected| !selected.is_empty())
&& let Some(new_table) = table.project_output_columns(&select_columns)
{
return Ok(Value::custom(Box::new(new_table), call_span).into_pipeline_data());
}
}
match input {
PipelineData::Value(v, metadata, ..) => {
let span = v.span();
match v {
Value::List {
vals: input_vals, ..
} => Ok(input_vals
.into_iter()
.map(move |input_val| {
if !columns.is_empty() {
let mut record = Record::new();
for path in &columns {
match input_val.follow_cell_path(&path.members) {
Ok(fetcher) => {
record.push(path.to_column_name(), fetcher.into_owned());
}
Err(e) => return Value::error(e, call_span),
}
}
Value::record(record, span)
} else {
input_val.clone()
}
})
.into_pipeline_data_with_metadata(
call_span,
engine_state.signals().clone(),
metadata,
)),
_ => {
if !columns.is_empty() {
let mut record = Record::new();
for cell_path in columns {
let result = v.follow_cell_path(&cell_path.members)?;
record.push(cell_path.to_column_name(), result.into_owned());
}
Ok(Value::record(record, call_span)
.into_pipeline_data_with_metadata(metadata))
} else {
Ok(v.into_pipeline_data_with_metadata(metadata))
}
}
}
}
PipelineData::ListStream(stream, metadata, ..) => Ok(stream
.map(move |x| {
if !columns.is_empty() {
let mut record = Record::new();
for path in &columns {
match x.follow_cell_path(&path.members) {
Ok(value) => {
record.push(path.to_column_name(), value.into_owned());
}
Err(e) => return Value::error(e, call_span),
}
}
Value::record(record, call_span)
} else {
x
}
})
.into_pipeline_data_with_metadata(call_span, engine_state.signals().clone(), metadata)),
_ => Ok(PipelineData::empty()),
}
}
struct NthIterator {
input: PipelineIterator,
rows: std::iter::Peekable<std::collections::btree_set::IntoIter<usize>>,
current: usize,
}
impl Iterator for NthIterator {
type Item = Value;
fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(row) = self.rows.peek() {
if self.current == *row {
self.rows.next();
self.current += 1;
return self.input.next();
} else {
self.current += 1;
let _ = self.input.next()?;
continue;
}
} else {
return None;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_examples() -> nu_test_support::Result {
nu_test_support::test().examples(Select)
}
}