refactor: add vec repeat support for logical ops

During #4622 lots of attention was given to Vector usage in binary ops since the literals that could be vectorized tended to show up there. Since boolean literals were not on the list, I somehow neglected to look at the logical operators (which depend on booleans to be used). This diff adds the necessary branching to cover the vec repeat case, looking forward to the time where boolean literals _can be vectorized_. It also adds skipped/commented out tests with reference to #4997 which should aim to enable these code paths by finally allowing bool literals to become vectorized.
influxdata · Aug 16, 2022 · 5531ee0 · 5531ee0
1 parent cf5caf0
commit 5531ee0
Show file tree

Hide file tree

Showing 6 changed files with 204 additions and 8 deletions.
diff --git a/array/binary.go b/array/binary.go
@@ -26,6 +26,38 @@ func And(l, r *Boolean, mem memory.Allocator) (*Boolean, error) {
 	return a, nil
 }
 
+func AndLConst(l bool, r *Boolean, mem memory.Allocator) (*Boolean, error) {
+	n := r.Len()
+	b := NewBooleanBuilder(mem)
+	b.Resize(n)
+	for i := 0; i < n; i++ {
+		if r.IsValid(i) {
+			b.Append(l && r.Value(i))
+		} else {
+			b.AppendNull()
+		}
+	}
+	a := b.NewBooleanArray()
+	b.Release()
+	return a, nil
+}
+
+func AndRConst(l *Boolean, r bool, mem memory.Allocator) (*Boolean, error) {
+	n := l.Len()
+	b := NewBooleanBuilder(mem)
+	b.Resize(n)
+	for i := 0; i < n; i++ {
+		if l.IsValid(i) {
+			b.Append(l.Value(i) && r)
+		} else {
+			b.AppendNull()
+		}
+	}
+	a := b.NewBooleanArray()
+	b.Release()
+	return a, nil
+}
+
 func Or(l, r *Boolean, mem memory.Allocator) (*Boolean, error) {
 	n := l.Len()
 	if n != r.Len() {
@@ -45,3 +77,35 @@ func Or(l, r *Boolean, mem memory.Allocator) (*Boolean, error) {
 	b.Release()
 	return a, nil
 }
+
+func OrLConst(l bool, r *Boolean, mem memory.Allocator) (*Boolean, error) {
+	n := r.Len()
+	b := NewBooleanBuilder(mem)
+	b.Resize(n)
+	for i := 0; i < n; i++ {
+		if r.IsValid(i) {
+			b.Append(l || r.Value(i))
+		} else {
+			b.AppendNull()
+		}
+	}
+	a := b.NewBooleanArray()
+	b.Release()
+	return a, nil
+}
+
+func OrRConst(l *Boolean, r bool, mem memory.Allocator) (*Boolean, error) {
+	n := l.Len()
+	b := NewBooleanBuilder(mem)
+	b.Resize(n)
+	for i := 0; i < n; i++ {
+		if l.IsValid(i) {
+			b.Append(l.Value(i) || r)
+		} else {
+			b.AppendNull()
+		}
+	}
+	a := b.NewBooleanArray()
+	b.Release()
+	return a, nil
+}
diff --git a/compiler/runtime.go b/compiler/runtime.go
@@ -350,7 +350,6 @@ func (e *logicalVectorEvaluator) Eval(ctx context.Context, scope Scope) (values.
 	} else if elemType := l.Vector().ElementType().Nature(); elemType != semantic.Bool {
 		return nil, errors.Newf(codes.Invalid, "cannot use operand of type %s with logical %s; expected boolean", elemType, e.operator)
 	}
-	lv := l.Vector().Arr().(*array.Boolean)
 
 	r, err := e.right.Eval(ctx, scope)
 	if err != nil {
@@ -362,26 +361,93 @@ func (e *logicalVectorEvaluator) Eval(ctx context.Context, scope Scope) (values.
 	} else if elemType := r.Vector().ElementType().Nature(); elemType != semantic.Bool {
 		return nil, errors.Newf(codes.Invalid, "cannot use operand of type %s with logical %s; expected boolean", elemType, e.operator)
 	}
-	rv := r.Vector().Arr().(*array.Boolean)
 
 	mem := memory.GetAllocator(ctx)
 
 	switch e.operator {
 	case ast.AndOperator:
-		res, err := array.And(lv, rv, mem)
+		return vectorAnd(l.Vector(), r.Vector(), mem)
+	case ast.OrOperator:
+		return vectorOr(l.Vector(), r.Vector(), mem)
+	default:
+		panic(errors.Newf(codes.Internal, "unknown logical operator %v", e.operator))
+	}
+
+}
+
+func vectorAnd(l, r values.Vector, mem memory.Allocator) (values.Vector, error) {
+	var lvr, rvr *values.Value
+
+	if vr, ok := l.(*values.VectorRepeatValue); ok {
+		v := vr.Value()
+		lvr = &v
+	}
+	if vr, ok := r.(*values.VectorRepeatValue); ok {
+		v := vr.Value()
+		rvr = &v
+	}
+
+	if lvr != nil && rvr != nil {
+		a := (*lvr).Bool()
+		b := (*rvr).Bool()
+		values.NewVectorRepeatValue(values.NewBool(a && b))
+	} else if lvr != nil {
+		res, err := array.AndLConst((*lvr).Bool(), r.Arr().(*array.Boolean), mem)
 		if err != nil {
 			return nil, err
 		}
 		return values.NewVectorValue(res, semantic.BasicBool), nil
-	case ast.OrOperator:
-		res, err := array.Or(lv, rv, mem)
+	} else if rvr != nil {
+		res, err := array.AndRConst(l.Arr().(*array.Boolean), (*rvr).Bool(), mem)
 		if err != nil {
 			return nil, err
 		}
 		return values.NewVectorValue(res, semantic.BasicBool), nil
-	default:
-		panic(errors.Newf(codes.Internal, "unknown logical operator %v", e.operator))
 	}
+
+	res, err := array.And(l.Arr().(*array.Boolean), r.Arr().(*array.Boolean), mem)
+	if err != nil {
+		return nil, err
+	}
+	return values.NewVectorValue(res, semantic.BasicBool), nil
+
+}
+
+func vectorOr(l, r values.Vector, mem memory.Allocator) (values.Vector, error) {
+	var lvr, rvr *values.Value
+
+	if vr, ok := l.(*values.VectorRepeatValue); ok {
+		v := vr.Value()
+		lvr = &v
+	}
+	if vr, ok := r.(*values.VectorRepeatValue); ok {
+		v := vr.Value()
+		rvr = &v
+	}
+
+	if lvr != nil && rvr != nil {
+		a := (*lvr).Bool()
+		b := (*rvr).Bool()
+		values.NewVectorRepeatValue(values.NewBool(a || b))
+	} else if lvr != nil {
+		res, err := array.OrLConst((*lvr).Bool(), r.Arr().(*array.Boolean), mem)
+		if err != nil {
+			return nil, err
+		}
+		return values.NewVectorValue(res, semantic.BasicBool), nil
+	} else if rvr != nil {
+		res, err := array.OrRConst(l.Arr().(*array.Boolean), (*rvr).Bool(), mem)
+		if err != nil {
+			return nil, err
+		}
+		return values.NewVectorValue(res, semantic.BasicBool), nil
+	}
+
+	res, err := array.Or(l.Arr().(*array.Boolean), r.Arr().(*array.Boolean), mem)
+	if err != nil {
+		return nil, err
+	}
+	return values.NewVectorValue(res, semantic.BasicBool), nil
 }
 
 type conditionalEvaluator struct {

diff --git a/compiler/vectorized_test.go b/compiler/vectorized_test.go
@@ -166,6 +166,14 @@ func TestVectorizedFns(t *testing.T) {
 			vectorizable: true,
 			skipComp:     true,
 		},
+		// FIXME: https://github.com/influxdata/flux/issues/4997
+		//  bool literals are not vectorized currently
+		//{
+		//	name:         "bool literals",
+		//	fn:           `(r) => ({r with c: true})`,
+		//	vectorizable: true,
+		//	skipComp:     true,
+		//},
 		{
 			name:         "no duration literals",
 			fn:           `(r) => ({r with c: 1h})`,

diff --git a/libflux/flux-core/src/semantic/tests/vectorize.rs b/libflux/flux-core/src/semantic/tests/vectorize.rs
@@ -330,8 +330,10 @@ fn vectorize_with_construction_using_literal_int() -> anyhow::Result<()> {
     Ok(())
 }
 
+// FIXME: https://github.com/influxdata/flux/issues/4997
+//  bool literals are not vectorized currently. This test gives "undefined identifier true"
 #[test]
-#[ignore] // FIXME: "undefined identifier true"
+#[ignore]
 fn vectorize_with_construction_using_literal_bool() -> anyhow::Result<()> {
     let pkg = vectorize(r#"(r) => ({ r with a: true })"#)?;
 

diff --git a/libflux/flux-core/src/semantic/vectorize.rs b/libflux/flux-core/src/semantic/vectorize.rs
@@ -158,6 +158,14 @@ impl Expression {
                 wrap_vec_repeat(expr.clone())
             }
             Expression::Call(expr) => Expression::Call(Box::new(expr.vectorize(env)?)),
+            // FIXME: https://github.com/influxdata/flux/issues/4997
+            //  bool literals are not vectorized currently.
+            //  This match arm doesn't seem to be hit when `true`/`false` appear in the function body.
+            // expr @ Expression::Boolean(_)
+            //     if env.config.features.contains(&Feature::VectorizedConst) =>
+            // {
+            //     wrap_vec_repeat(expr.clone())
+            // }
             _ => {
                 return Err(located(
                     self.loc().clone(),

diff --git a/stdlib/universe/map_vectorize_const_test.flux b/stdlib/universe/map_vectorize_const_test.flux
@@ -86,3 +86,51 @@ testcase vec_const_kitchen_sink_column_types {
 
     testing.diff(want: want, got: got)
 }
+
+testcase vec_const_bools {
+    option testing.tags = [
+        // FIXME: https://github.com/influxdata/flux/issues/4997
+        //  bool literals are not vectorized currently
+        "skip",
+    ]
+
+    input = array.from(rows: [{a: false, b: false}, {a: false, b: true}, {a: true, b: true}])
+    want =
+        array.from(
+            rows: [
+                {
+                    a: false,
+                    a_and_true: false,
+                    a_or_true: true,
+                    a_and_false: false,
+                    a_or_false: false,
+                },
+                {
+                    a: false,
+                    a_and_true: false,
+                    a_or_true: true,
+                    a_and_false: false,
+                    a_or_false: false,
+                },
+                {
+                    a: true,
+                    a_and_true: true,
+                    a_or_true: true,
+                    a_and_false: false,
+                    a_or_false: true,
+                },
+            ],
+        )
+    got =
+        input
+            |> map(
+                fn: (r) =>
+                    ({r with a_and_true: r.a and true,
+                        a_or_true: r.a or true,
+                        a_and_false: r.a and false,
+                        a_or_false: r.a or false,
+                    }),
+            )
+
+    testing.diff(want: want, got: got)
+}