Replace ExternalStream with new ByteStream type (#12774)

# Description
This PR introduces a `ByteStream` type which is a `Read`-able stream of
bytes. Internally, it has an enum over three different byte stream
sources:
```rust
pub enum ByteStreamSource {
    Read(Box<dyn Read + Send + 'static>),
    File(File),
    Child(ChildProcess),
}
```

This is in comparison to the current `RawStream` type, which is an
`Iterator<Item = Vec<u8>>` and has to allocate for each read chunk.

Currently, `PipelineData::ExternalStream` serves a weird dual role where
it is either external command output or a wrapper around `RawStream`.
`ByteStream` makes this distinction more clear (via `ByteStreamSource`)
and replaces `PipelineData::ExternalStream` in this PR:
```rust
pub enum PipelineData {
    Empty,
    Value(Value, Option<PipelineMetadata>),
    ListStream(ListStream, Option<PipelineMetadata>),
    ByteStream(ByteStream, Option<PipelineMetadata>),
}
```

The PR is relatively large, but a decent amount of it is just repetitive
changes.

This PR fixes #7017, fixes #10763, and fixes #12369.

This PR also improves performance when piping external commands. Nushell
should, in most cases, have competitive pipeline throughput compared to,
e.g., bash.
| Command | Before (MB/s) | After (MB/s) | Bash (MB/s) |
| -------------------------------------------------- | -------------:|
------------:| -----------:|
| `throughput \| rg 'x'` | 3059 | 3744 | 3739 |
| `throughput \| nu --testbin relay o> /dev/null` | 3508 | 8087 | 8136 |

# User-Facing Changes
- This is a breaking change for the plugin communication protocol,
because the `ExternalStreamInfo` was replaced with `ByteStreamInfo`.
Plugins now only have to deal with a single input stream, as opposed to
the previous three streams: stdout, stderr, and exit code.
- The output of `describe` has been changed for external/byte streams.
- Temporary breaking change: `bytes starts-with` no longer works with
byte streams. This is to keep the PR smaller, and `bytes ends-with`
already does not work on byte streams.
- If a process core dumped, then instead of having a `Value::Error` in
the `exit_code` column of the output returned from `complete`, it now is
a `Value::Int` with the negation of the signal number.

# After Submitting
- Update docs and book as necessary
- Release notes (e.g., plugin protocol changes)
- Adapt/convert commands to work with byte streams (high priority is
`str length`, `bytes starts-with`, and maybe `bytes ends-with`).
- Refactor the `tee` code, Devyn has already done some work on this.

---------

Co-authored-by: Devyn Cairns <devyn.cairns@gmail.com>
This commit is contained in:
Ian Manske
2024-05-16 14:11:18 +00:00
committed by GitHub
parent 1b8eb23785
commit 6fd854ed9f
210 changed files with 3955 additions and 4012 deletions

View File

@ -94,7 +94,7 @@ mod test {
.add_decl(Box::new(First))?
.add_decl(Box::new(Get))?
.eval("let df = ([[a b];[1 2] [3 4]] | polars into-df); polars store-ls | get key | first | polars store-rm $in")?;
let value = pipeline_data.into_value(Span::test_data());
let value = pipeline_data.into_value(Span::test_data())?;
let msg = value
.as_list()?
.first()

View File

@ -124,7 +124,7 @@ pub mod test {
assert!(tmp_file.exists());
let value = pipeline_data.into_value(Span::test_data());
let value = pipeline_data.into_value(Span::test_data())?;
let list = value.as_list()?;
assert_eq!(list.len(), 1);
let msg = list.first().expect("should have a value").as_str()?;

View File

@ -153,7 +153,7 @@ pub mod test {
assert!(tmp_file.exists());
let value = pipeline_data.into_value(Span::test_data());
let value = pipeline_data.into_value(Span::test_data())?;
let list = value.as_list()?;
assert_eq!(list.len(), 1);
let msg = list.first().expect("should have a value").as_str()?;

View File

@ -171,7 +171,7 @@ pub mod test {
assert!(tmp_file.exists());
let value = pipeline_data.into_value(Span::test_data());
let value = pipeline_data.into_value(Span::test_data())?;
let list = value.as_list()?;
assert_eq!(list.len(), 1);
let msg = list.first().expect("should have a value").as_str()?;

View File

@ -125,7 +125,7 @@ pub mod test {
assert!(tmp_file.exists());
let value = pipeline_data.into_value(Span::test_data());
let value = pipeline_data.into_value(Span::test_data())?;
let list = value.as_list()?;
assert_eq!(list.len(), 1);
let msg = list.first().expect("should have a value").as_str()?;

View File

@ -89,7 +89,7 @@ impl PluginCommand for ToNu {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) {
dataframe_command(plugin, call, value)
} else {

View File

@ -124,7 +124,7 @@ pub mod test {
assert!(tmp_file.exists());
let value = pipeline_data.into_value(Span::test_data());
let value = pipeline_data.into_value(Span::test_data())?;
let list = value.as_list()?;
assert_eq!(list.len(), 1);
let msg = list.first().expect("should have a value").as_str()?;

View File

@ -159,7 +159,7 @@ macro_rules! lazy_expr_command {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) {
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)
.map_err(LabeledError::from)?;
@ -239,7 +239,7 @@ macro_rules! lazy_expr_command {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) {
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)
.map_err(LabeledError::from)?;

View File

@ -114,8 +114,7 @@ impl PluginCommand for ExprIsIn {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df),
PolarsPluginObject::NuLazyFrame(lazy) => {

View File

@ -99,7 +99,7 @@ impl PluginCommand for ExprOtherwise {
let otherwise_predicate: Value = call.req(0)?;
let otherwise_predicate = NuExpression::try_from_value(plugin, &otherwise_predicate)?;
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let complete: NuExpression = match NuWhen::try_from_value(plugin, &value)?.when_type {
NuWhenType::Then(then) => then.otherwise(otherwise_predicate.into_polars()).into(),
NuWhenType::ChainedThen(chained_when) => chained_when

View File

@ -111,7 +111,7 @@ impl PluginCommand for ExprWhen {
let then_predicate: Value = call.req(1)?;
let then_predicate = NuExpression::try_from_value(plugin, &then_predicate)?;
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let when_then: NuWhen = match value {
Value::Nothing { .. } => when(when_predicate.into_polars())
.then(then_predicate.into_polars())

View File

@ -90,7 +90,7 @@ impl PluginCommand for CastDF {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuLazyFrame(lazy) => {
let (dtype, column_nm) = df_args(call)?;

View File

@ -61,7 +61,7 @@ impl PluginCommand for LazyCollect {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuLazyFrame(lazy) => {
let eager = lazy.collect(call.head)?;

View File

@ -50,7 +50,7 @@ impl PluginCommand for LazyExplode {
result: Some(
NuDataFrame::try_from_columns(vec![
Column::new(
"id".to_string(),
"id".to_string(),
vec![
Value::test_int(1),
Value::test_int(1),
@ -58,7 +58,7 @@ impl PluginCommand for LazyExplode {
Value::test_int(2),
]),
Column::new(
"name".to_string(),
"name".to_string(),
vec![
Value::test_string("Mercy"),
Value::test_string("Mercy"),
@ -66,7 +66,7 @@ impl PluginCommand for LazyExplode {
Value::test_string("Bob"),
]),
Column::new(
"hobbies".to_string(),
"hobbies".to_string(),
vec![
Value::test_string("Cycling"),
Value::test_string("Knitting"),
@ -84,7 +84,7 @@ impl PluginCommand for LazyExplode {
result: Some(
NuDataFrame::try_from_columns(vec![
Column::new(
"hobbies".to_string(),
"hobbies".to_string(),
vec![
Value::test_string("Cycling"),
Value::test_string("Knitting"),
@ -116,8 +116,7 @@ pub(crate) fn explode(
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, ShellError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => {
let lazy = df.lazy();

View File

@ -67,7 +67,7 @@ impl PluginCommand for LazyFetch {
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let rows: i64 = call.req(0)?;
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
let eager: NuDataFrame = lazy

View File

@ -92,7 +92,7 @@ impl PluginCommand for LazyFillNA {
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let fill: Value = call.req(0)?;
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => {

View File

@ -69,7 +69,7 @@ impl PluginCommand for LazyFillNull {
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let fill: Value = call.req(0)?;
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => cmd_lazy(plugin, engine, call, df.lazy(), fill),

View File

@ -72,7 +72,7 @@ impl PluginCommand for LazyFilter {
) -> Result<PipelineData, LabeledError> {
let expr_value: Value = call.req(0)?;
let filter_expr = NuExpression::try_from_value(plugin, &expr_value)?;
let pipeline_value = input.into_value(call.head);
let pipeline_value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &pipeline_value)?;
command(plugin, engine, call, lazy, filter_expr).map_err(LabeledError::from)
}

View File

@ -67,7 +67,7 @@ impl PluginCommand for FilterWith {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
command_lazy(plugin, engine, call, lazy).map_err(LabeledError::from)
}

View File

@ -97,7 +97,7 @@ impl PluginCommand for FirstDF {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
if NuLazyFrame::can_downcast(&value) || NuDataFrame::can_downcast(&value) {
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
command(plugin, engine, call, lazy).map_err(LabeledError::from)

View File

@ -138,7 +138,7 @@ impl PluginCommand for ToLazyGroupBy {
})?;
}
let pipeline_value = input.into_value(call.head);
let pipeline_value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &pipeline_value)?;
command(plugin, engine, call, lazy, expressions).map_err(LabeledError::from)
}

View File

@ -228,7 +228,7 @@ impl PluginCommand for LazyJoin {
let suffix: Option<String> = call.get_flag("suffix")?;
let suffix = suffix.unwrap_or_else(|| "_x".into());
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
let lazy = lazy.to_polars();

View File

@ -72,7 +72,7 @@ impl PluginCommand for LastDF {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) {
let df = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
command(plugin, engine, call, df).map_err(|e| e.into())

View File

@ -89,7 +89,7 @@ impl PluginCommand for LazyMedian {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => command(plugin, engine, call, df.lazy()),
PolarsPluginObject::NuLazyFrame(lazy) => command(plugin, engine, call, lazy),

View File

@ -97,7 +97,7 @@ impl PluginCommand for LazyQuantile {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let quantile: f64 = call.req(0)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => {

View File

@ -120,7 +120,7 @@ impl PluginCommand for RenameDF {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
command_lazy(plugin, engine, call, lazy).map_err(LabeledError::from)
}

View File

@ -65,7 +65,7 @@ impl PluginCommand for LazySelect {
let expr_value = Value::list(vals, call.head);
let expressions = NuExpression::extract_exprs(plugin, expr_value)?;
let pipeline_value = input.into_value(call.head);
let pipeline_value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &pipeline_value)?;
let lazy = NuLazyFrame::new(lazy.to_polars().select(&expressions));
lazy.to_pipeline_data(plugin, engine, call.head)

View File

@ -145,7 +145,7 @@ impl PluginCommand for LazySortBy {
maintain_order,
};
let pipeline_value = input.into_value(call.head);
let pipeline_value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &pipeline_value)?;
let lazy = NuLazyFrame::new(lazy.to_polars().sort_by_exprs(&expressions, sort_options));
lazy.to_pipeline_data(plugin, engine, call.head)

View File

@ -83,7 +83,7 @@ impl PluginCommand for WithColumn {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
command_lazy(plugin, engine, call, lazy).map_err(LabeledError::from)
}

View File

@ -78,8 +78,7 @@ impl PluginCommand for IsNotNull {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => command(plugin, engine, call, df),
PolarsPluginObject::NuLazyFrame(lazy) => {

View File

@ -80,8 +80,7 @@ impl PluginCommand for IsNull {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => command(plugin, engine, call, df),
PolarsPluginObject::NuLazyFrame(lazy) => {

View File

@ -70,8 +70,7 @@ impl PluginCommand for NUnique {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => command(plugin, engine, call, df),
PolarsPluginObject::NuLazyFrame(lazy) => {

View File

@ -92,7 +92,7 @@ impl PluginCommand for Shift {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
command_lazy(plugin, engine, call, lazy).map_err(LabeledError::from)
}

View File

@ -134,7 +134,7 @@ impl PluginCommand for Unique {
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
let value = input.into_value(call.head)?;
let df = NuLazyFrame::try_from_value_coerce(plugin, &value)?;
command_lazy(plugin, engine, call, df).map_err(LabeledError::from)
}

View File

@ -84,7 +84,7 @@ impl PolarsPluginObject {
input: PipelineData,
span: Span,
) -> Result<Self, ShellError> {
let value = input.into_value(span);
let value = input.into_value(span)?;
Self::try_from_value(plugin, &value)
}
@ -242,7 +242,7 @@ pub trait PolarsPluginCustomValue: CustomValue {
/// Handles the ability for a PolarsObjectType implementations to convert between
/// their respective CustValue type.
/// PolarsPluginObjectType's (NuDataFrame, NuLazyFrame) should
/// implement this trait.
/// implement this trait.
pub trait CustomValueSupport: Cacheable {
type CV: PolarsPluginCustomValue<PolarsPluginObjectType = Self> + CustomValue + 'static;
@ -301,7 +301,7 @@ pub trait CustomValueSupport: Cacheable {
input: PipelineData,
span: Span,
) -> Result<Self, ShellError> {
let value = input.into_value(span);
let value = input.into_value(span)?;
Self::try_from_value(plugin, &value)
}

View File

@ -519,7 +519,7 @@ impl NuDataFrame {
input: PipelineData,
span: Span,
) -> Result<Self, ShellError> {
let value = input.into_value(span);
let value = input.into_value(span)?;
Self::try_from_value_coerce(plugin, &value, span)
}
}

View File

@ -109,7 +109,7 @@ impl NuLazyFrame {
input: PipelineData,
span: Span,
) -> Result<Self, ShellError> {
let value = input.into_value(span);
let value = input.into_value(span)?;
Self::try_from_value_coerce(plugin, &value)
}
}