Repository: brimsec/zq
Pull request #2421 (2021-03-24 22:59:35)

update "in" to work with maps
This commit updates "in" to operate on maps and brings the code
into conformance with "missing" and in-line errors.


Patch:
diff --git a/expr/eval.go b/expr/eval.go
index cb36fc5a1..9793c72c2 100644
--- a/expr/eval.go
+++ b/expr/eval.go
@@ -129,27 +129,39 @@ func (i *In) Eval(rec *zng.Record) (zng.Value, error) {
 	if err != nil {
 		return container, err
 	}
-	if typ := zng.AliasOf(container.Type); typ == zng.TypeNet {
-		n, err := zng.DecodeNet(container.Bytes)
-		if err != nil {
-			return zng.Value{}, err
-		}
-		if typ := zng.AliasOf(elem.Type); typ != zng.TypeIP {
-			return zng.Value{}, ErrIncompatibleTypes
-		}
-		a, err := zng.DecodeIP(elem.Bytes)
-		if err != nil {
-			return zng.Value{}, err
-		}
-		if n.IP.Equal(a.Mask(n.Mask)) {
-			return zng.True, nil
-		}
-		return zng.False, nil
+	switch typ := zng.AliasOf(container.Type).(type) {
+	case *zng.TypeOfNet:
+		return inNet(elem, container)
+	case *zng.TypeArray:
+		return i.inContainer(zng.AliasOf(typ.Type), elem, container)
+	case *zng.TypeSet:
+		return i.inContainer(zng.AliasOf(typ.Type), elem, container)
+	case *zng.TypeMap:
+		return i.inMap(typ, elem, container)
+	default:
+		return zng.NewErrorf("'in' operator applied to non-container type"), nil
+	}
+}
+
+func inNet(elem, net zng.Value) (zng.Value, error) {
+	n, err := zng.DecodeNet(net.Bytes)
+	if err != nil {
+		return zng.Value{}, err
+	}
+	if typ := zng.AliasOf(elem.Type); typ != zng.TypeIP {
+		return zng.NewErrorf("'in' operator applied to non-container type"), nil
+	}
+	a, err := zng.DecodeIP(elem.Bytes)
+	if err != nil {
+		return zng.Value{}, err
 	}
-	typ := zng.InnerType(container.Type)
-	if typ == nil {
-		return zng.Value{}, ErrNotContainer
+	if n.IP.Equal(a.Mask(n.Mask)) {
+		return zng.True, nil
 	}
+	return zng.False, nil
+}
+
+func (i *In) inContainer(typ zng.Type, elem, container zng.Value) (zng.Value, error) {
 	iter := container.Bytes.Iter()
 	for {
 		if iter.Done() {
@@ -160,13 +172,35 @@ func (i *In) Eval(rec *zng.Record) (zng.Value, error) {
 			return zng.Value{}, err
 		}
 		_, err = i.vals.Coerce(elem, zng.Value{typ, zv})
+		if err == nil && i.vals.Equal() {
+			return zng.True, nil
+		}
+	}
+}
+
+func (i *In) inMap(typ *zng.TypeMap, elem, container zng.Value) (zng.Value, error) {
+	keyType := zng.AliasOf(typ.KeyType)
+	valType := zng.AliasOf(typ.ValType)
+	iter := container.Bytes.Iter()
+	for !iter.Done() {
+		zv, _, err := iter.Next()
+		if err != nil {
+			return zng.Value{}, err
+		}
+		_, err = i.vals.Coerce(elem, zng.Value{keyType, zv})
+		if err == nil && i.vals.Equal() {
+			return zng.True, nil
+		}
+		zv, _, err = iter.Next()
 		if err != nil {
 			return zng.Value{}, err
 		}
-		if i.vals.Equal() {
+		_, err = i.vals.Coerce(elem, zng.Value{valType, zv})
+		if err == nil && i.vals.Equal() {
 			return zng.True, nil
 		}
 	}
+	return zng.False, nil
 }
 
 type Equal struct {
@@ -555,7 +589,7 @@ func getNthFromContainer(container zcode.Bytes, idx uint) (zcode.Bytes, error) {
 			return zv, nil
 		}
 	}
-	return nil, ErrIndexOutOfBounds
+	return nil, zng.ErrMissing
 }
 
 func lookupKey(mapBytes, target zcode.Bytes) (zcode.Bytes, bool) {
@@ -618,7 +652,7 @@ func indexArray(typ *zng.TypeArray, array zcode.Bytes, index zng.Value) (zng.Val
 	if zng.IsSigned(id) {
 		v, _ := zng.DecodeInt(index.Bytes)
 		if idx < 0 {
-			return zng.NewErrorf("array index out of bounds"), nil
+			return zng.Value{}, zng.ErrMissing
 		}
 		idx = uint(v)
 	} else {
@@ -653,7 +687,7 @@ func indexMap(typ *zng.TypeMap, mapBytes zcode.Bytes, key zng.Value) (zng.Value,
 	if valBytes, ok := lookupKey(mapBytes, key.Bytes); ok {
 		return zng.Value{typ.ValType, valBytes}, nil
 	}
-	return zng.NewErrorf("key not found in map: %s", key), nil
+	return zng.Value{}, zng.ErrMissing
 }
 
 type Conditional struct {
diff --git a/zng/map.go b/zng/map.go
index 12a03acab..a3f6439ac 100644
--- a/zng/map.go
+++ b/zng/map.go
@@ -129,6 +129,7 @@ func (t *TypeMap) ZSONOf(zv zcode.Bytes) string {
 		b.WriteString(sep)
 		b.WriteByte('{')
 		b.WriteString(t.KeyType.ZSONOf(val))
+		b.WriteByte(',')
 		val, _, err = it.Next()
 		if err != nil {
 			return badZng(err, t, zv)


Test patch:
diff --git a/expr/expr_test.go b/expr/expr_test.go
index e66890e0d..f13259cb6 100644
--- a/expr/expr_test.go
+++ b/expr/expr_test.go
@@ -61,6 +61,10 @@ func zbool(b bool) zng.Value {
 	return zng.Value{zng.TypeBool, zng.EncodeBool(b)}
 }
 
+func zerr(s string) zng.Value {
+	return zng.Value{zng.TypeError, zng.EncodeString(s)}
+}
+
 func zint32(v int32) zng.Value {
 	return zng.Value{zng.TypeInt32, zng.EncodeInt(int64(v))}
 }
@@ -418,9 +422,9 @@ func TestIn(t *testing.T) {
 	testSuccessful(t, "1 in s", record, zbool(false))
 	testSuccessful(t, "4 in s", record, zbool(true))
 
-	testError(t, `"boo" in a`, record, expr.ErrIncompatibleTypes, "in operator with mismatched type")
-	testError(t, `"boo" in s`, record, expr.ErrIncompatibleTypes, "in operator with mismatched type")
-	testError(t, "1 in 2", record, expr.ErrNotContainer, "in operator with non-container")
+	testSuccessful(t, `"boo" in a`, record, zbool(false))
+	testSuccessful(t, `"boo" in s`, record, zbool(false))
+	testSuccessful(t, "1 in 2", record, zerr("'in' operator applied to non-container type"))
 }
 
 func TestArithmetic(t *testing.T) {
@@ -571,9 +575,6 @@ func TestArrayIndex(t *testing.T) {
 	testSuccessful(t, "x[i]", record, zint64(2))
 	testSuccessful(t, "i+1", record, zint64(2))
 	testSuccessful(t, "x[i+1]", record, zint64(3))
-
-	testError(t, "x[-1]", record, expr.ErrIndexOutOfBounds, "negative array index")
-	testError(t, "x[3]", record, expr.ErrIndexOutOfBounds, "array index too large")
 }
 
 func TestFieldReference(t *testing.T) {
diff --git a/expr/ztests/array-bounds-2.yaml b/expr/ztests/array-bounds-2.yaml
new file mode 100644
index 000000000..ef482520f
--- /dev/null
+++ b/expr/ztests/array-bounds-2.yaml
@@ -0,0 +1,7 @@
+zql: cut m1=missing(a[-1]),m2=missing(a[3])
+
+input: |
+  {x:[1,2,3],i:1 (uint16)} (=0)
+
+output: |
+  {m1:true,m2:true}
diff --git a/expr/ztests/array-bounds.yaml b/expr/ztests/array-bounds.yaml
new file mode 100644
index 000000000..443bbb316
--- /dev/null
+++ b/expr/ztests/array-bounds.yaml
@@ -0,0 +1,9 @@
+zql: put b=missing(a[2])
+
+input: |
+  {a:[1,2,3]}
+  {a:[1,2]}
+
+output: |
+  {a:[1,2,3],b:false}
+  {a:[1,2],b:true}
diff --git a/expr/ztests/in-map.yaml b/expr/ztests/in-map.yaml
new file mode 100644
index 000000000..6f53a93d3
--- /dev/null
+++ b/expr/ztests/in-map.yaml
@@ -0,0 +1,16 @@
+zql: |
+  const m = |{"a": 1, "b": 2}|
+  cut b1=match(x in m),b2=match("a" in a)
+
+input: |
+  {x:"a"}
+  {x:1}
+  {x:10}
+  {y:1}
+  {a:|{{"a","hello"}}|}
+
+output: |
+  {b1:true}
+  {b1:true}
+  {b1:false}
+  {b2:true}
diff --git a/expr/ztests/type-map.yaml b/expr/ztests/type-map.yaml
new file mode 100644
index 000000000..c64674c4b
--- /dev/null
+++ b/expr/ztests/type-map.yaml
@@ -0,0 +1,25 @@
+zql: |
+  type conn = {_path:string, id:{src:ip,dst:ip}, etc:string}
+  type dns = {_path:string, id:{src:ip,dst:ip}, query:string, query_size:uint64}
+  const schemas = |{
+    "conn": conn,
+    "dns": dns
+  }|
+  split (
+    => cut schema=schemas[_path]
+    => missing(schemas[_path]) | put _UNCLASSIFIED=true
+  ) | sort .
+
+input: |
+  {_path:"conn", id:{src:"192.168.1.1",dst:"192.168.1.2"}, etc:"foo"}
+  {_path:"dns", id:{src:"192.168.1.1",dst:"192.168.1.2"}, query:"www.google.com", query_size:23}
+  {_path:"conn", id:{src:"192.168.1.1",dst:"192.168.1.3"}, etc:"foo"}
+  {_path:"unknown", x:"has unknown _path"}
+  {y:"has no _path"}
+
+output: |
+  {_path:"unknown",x:"has unknown _path",_UNCLASSIFIED:true}
+  {schema:(conn=({_path:string,id:{src:ip,dst:ip},etc:string}))}
+  {schema:(conn=({_path:string,id:{src:ip,dst:ip},etc:string}))}
+  {schema:(dns=({_path:string,id:{src:ip,dst:ip},query:string,query_size:uint64}))}
+  {y:"has no _path",_UNCLASSIFIED:true}