Skip to content

Commit

Permalink
enhancement(stdlib): Allow zip function to take two array parameters (
Browse files Browse the repository at this point in the history
#1159)

* enhancement(stdlib): Allow `zip` function to take two array parameters

The 2-parameter micro benchmark shows that this two-parameter mode runs 50%
faster than a single array of two arrays.

* Fix clippy lint

* Adjust function parameter names
  • Loading branch information
bruceg authored Dec 2, 2024
1 parent 6c55b02 commit 62f0b4f
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 24 deletions.
12 changes: 10 additions & 2 deletions benches/stdlib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2938,8 +2938,16 @@ bench_function! {
bench_function! {
zip => vrl::stdlib::Zip;

default {
args: func_args![array: value!([["one", "two", "three", "four"], ["one", 2, null, true]])],
one_parameter {
args: func_args![array_0: value!([["one", "two", "three", "four"], ["one", 2, null, true]])],
want: Ok(value!([["one","one"], ["two",2], ["three",null], ["four",true]])),
}

two_parameters {
args: func_args![
array_0: value!(["one", "two", "three", "four"]),
array_1: value!(["one", 2, null, true]),
],
want: Ok(value!([["one","one"], ["two",2], ["three",null], ["four",true]])),
}
}
79 changes: 57 additions & 22 deletions src/stdlib/zip.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,15 @@
use crate::compiler::prelude::*;

fn zip(value: Value) -> Resolved {
fn zip2(value0: Value, value1: Value) -> Resolved {
Ok(value0
.try_array()?
.into_iter()
.zip(value1.try_array()?)
.map(|(v0, v1)| Value::Array(vec![v0, v1]))
.collect())
}

fn zip_all(value: Value) -> Resolved {
Ok(MultiZip(
value
.try_array()?
Expand Down Expand Up @@ -30,19 +39,33 @@ impl Function for Zip {
}

fn parameters(&self) -> &'static [Parameter] {
&[Parameter {
keyword: "array",
kind: kind::ARRAY,
required: true,
}]
&[
Parameter {
keyword: "array_0",
kind: kind::ARRAY,
required: true,
},
Parameter {
keyword: "array_1",
kind: kind::ARRAY,
required: false,
},
]
}

fn examples(&self) -> &'static [Example] {
&[Example {
title: "merge three arrays into an array of 3-tuples",
source: r#"zip([["a", "b", "c"], [1, null, true], [4, 5, 6]])"#,
result: Ok(r#"[["a", 1, 4], ["b", null, 5], ["c", true, 6]]"#),
}]
&[
Example {
title: "merge an array of three arrays into an array of 3-tuples",
source: r#"zip([["a", "b", "c"], [1, null, true], [4, 5, 6]])"#,
result: Ok(r#"[["a", 1, 4], ["b", null, 5], ["c", true, 6]]"#),
},
Example {
title: "merge two array parameters",
source: "zip([1, 2, 3, 4], [5, 6, 7])",
result: Ok("[[1, 5], [2, 6], [3, 7]]"),
},
]
}

fn compile(
Expand All @@ -51,19 +74,25 @@ impl Function for Zip {
_ctx: &mut FunctionCompileContext,
arguments: ArgumentList,
) -> Compiled {
let array = arguments.required("array");
Ok(ZipFn { array }.as_expr())
let array_0 = arguments.required("array_0");
let array_1 = arguments.optional("array_1");
Ok(ZipFn { array_0, array_1 }.as_expr())
}
}

#[derive(Debug, Clone)]
struct ZipFn {
array: Box<dyn Expression>,
array_0: Box<dyn Expression>,
array_1: Option<Box<dyn Expression>>,
}

impl FunctionExpression for ZipFn {
fn resolve(&self, ctx: &mut Context) -> Resolved {
zip(self.array.resolve(ctx)?)
let array_0 = self.array_0.resolve(ctx)?;
match &self.array_1 {
None => zip_all(array_0),
Some(array_1) => zip2(array_0, array_1.resolve(ctx)?),
}
}

fn type_def(&self, _state: &TypeState) -> TypeDef {
Expand All @@ -81,43 +110,49 @@ mod tests {
zip => Zip;

zips_two_arrays {
args: func_args![array: value!([[1, 2, 3], [4, 5, 6]])],
args: func_args![array_0: value!([[1, 2, 3], [4, 5, 6]])],
want: Ok(value!([[1, 4], [2, 5], [3, 6]])),
tdef: TypeDef::array(Collection::any()),
}

zips_three_arrays {
args: func_args![array: value!([[1, 2, 3], [4, 5, 6], [7, 8, 9]])],
args: func_args![array_0: value!([[1, 2, 3], [4, 5, 6], [7, 8, 9]])],
want: Ok(value!([[1, 4, 7], [2, 5, 8], [3, 6, 9]])),
tdef: TypeDef::array(Collection::any()),
}

zips_two_parameters {
args: func_args![array_0: value!([1, 2, 3]), array_1: value!([4, 5, 6])],
want: Ok(value!([[1, 4], [2, 5], [3, 6]])),
tdef: TypeDef::array(Collection::any()),
}

uses_shortest_length1 {
args: func_args![array: value!([[1, 2, 3], [4, 5]])],
args: func_args![array_0: value!([[1, 2, 3], [4, 5]])],
want: Ok(value!([[1, 4], [2, 5]])),
tdef: TypeDef::array(Collection::any()),
}

uses_shortest_length2 {
args: func_args![array: value!([[1, 2], [4, 5, 6]])],
args: func_args![array_0: value!([[1, 2], [4, 5, 6]])],
want: Ok(value!([[1, 4], [2, 5]])),
tdef: TypeDef::array(Collection::any()),
}

requires_outer_array {
args: func_args![array: 1],
args: func_args![array_0: 1],
want: Err("expected array, got integer"),
tdef: TypeDef::array(Collection::any()),
}

requires_inner_arrays1 {
args: func_args![array: value!([true, []])],
args: func_args![array_0: value!([true, []])],
want: Err("expected array, got boolean"),
tdef: TypeDef::array(Collection::any()),
}

requires_inner_arrays2 {
args: func_args![array: value!([[], null])],
args: func_args![array_0: value!([[], null])],
want: Err("expected array, got null"),
tdef: TypeDef::array(Collection::any()),
}
Expand Down

0 comments on commit 62f0b4f

Please sign in to comment.