Repository: brimsec/zq (go)

Issue:
Z shaping: handle sets and arrays (for "deep json")
The shaper functions introduced in #1984 do not handle arrays and sets. This issue consists of adding that support. This is likely to be useful particular for "deep json" where arrays of objects are common.

Pull request description:
Add support for shaping arrays and sets
close #2038 

Patch:
diff --git a/expr/shaper.go b/expr/shaper.go
index 464d29ab0..60f30f8e0 100644
--- a/expr/shaper.go
+++ b/expr/shaper.go
@@ -27,7 +27,9 @@ const (
 	copyContainer
 	castPrimitive // cast field fromIndex from input record
 	null          // write null
-	record        // record into record below us
+	array         // build array
+	set           // build set
+	record        // build record
 )
 
 // A step is a recursive data structure encoding a series of
@@ -36,11 +38,13 @@ type step struct {
 	op        op
 	fromIndex int
 	castTypes struct{ from, to zng.Type } // for op == castPrimitive
-	record    []step
+	// if op == record, contains one op for each column.
+	// if op == array, contains one op for all array elements.
+	children []step
 }
 
 func (s *step) append(step step) {
-	s.record = append(s.record, step)
+	s.children = append(s.children, step)
 }
 
 // create the step needed to build a record of type out from a
@@ -50,7 +54,7 @@ func (s *step) append(step step) {
 // [a b] and the input type has fields [b a] that is ok). It is also
 // ok for leaf primitive types to be different; if they are a casting
 // step is inserted.
-func createStep(in, out *zng.TypeRecord) (step, error) {
+func createStepRecord(in, out *zng.TypeRecord) (step, error) {
 	s := step{op: record}
 	for _, outCol := range out.Columns {
 		ind, ok := in.ColumnOfField(outCol.Name)
@@ -58,32 +62,82 @@ func createStep(in, out *zng.TypeRecord) (step, error) {
 			s.append(step{op: null})
 			continue
 		}
-
 		inCol := in.Columns[ind]
-
-		switch {
-		case inCol.Type.ID() == outCol.Type.ID():
-			if zng.IsContainerType(inCol.Type) {
-				s.append(step{fromIndex: ind, op: copyContainer})
-			} else {
-				s.append(step{fromIndex: ind, op: copyPrimitive})
-			}
-		case zng.IsRecordType(inCol.Type) && zng.IsRecordType(outCol.Type):
-			child, err := createStep(inCol.Type.(*zng.TypeRecord), outCol.Type.(*zng.TypeRecord))
-			if err != nil {
-				return step{}, err
-			}
-			child.fromIndex = ind
-			s.append(child)
-		case zng.IsPrimitiveType(inCol.Type) && zng.IsPrimitiveType(outCol.Type):
-			s.append(step{fromIndex: ind, op: castPrimitive, castTypes: struct{ from, to zng.Type }{inCol.Type, outCol.Type}})
-		default:
-			return step{}, fmt.Errorf("createOp incompatible column types %s and %s\n", inCol.Type, outCol.Type)
+		child, err := createStep(inCol.Type, outCol.Type)
+		if err != nil {
+			return step{}, err
 		}
+		child.fromIndex = ind
+		s.append(child)
+	}
+	return s, nil
+}
+
+func createStepArray(in, out zng.Type) (step, error) {
+	s := step{op: array}
+	innerOp, err := createStep(in, out)
+	if err != nil {
+		return step{}, nil
+	}
+	s.append(innerOp)
+	return s, nil
+}
+
+func createStepSet(in, out zng.Type) (step, error) {
+	s := step{op: set}
+	innerOp, err := createStep(in, out)
+	if err != nil {
+		return step{}, nil
 	}
+	s.append(innerOp)
 	return s, nil
 }
 
+func isCollectionType(t zng.Type) bool {
+	if _, ok := t.(*zng.TypeArray); ok {
+		return true
+	}
+	if _, ok := t.(*zng.TypeSet); ok {
+		return true
+	}
+	return false
+}
+
+func innerType(t zng.Type) zng.Type {
+	switch t := t.(type) {
+	case *zng.TypeArray:
+		return t.Type
+	case *zng.TypeSet:
+		return t.Type
+	}
+	return nil
+}
+
+func createStep(in, out zng.Type) (step, error) {
+	switch {
+	case in.ID() == out.ID():
+		if zng.IsContainerType(in) {
+			return step{op: copyContainer}, nil
+		} else {
+			return step{op: copyPrimitive}, nil
+		}
+	case zng.IsRecordType(in) && zng.IsRecordType(out):
+		return createStepRecord(in.(*zng.TypeRecord), out.(*zng.TypeRecord))
+	case zng.IsPrimitiveType(in) && zng.IsPrimitiveType(out):
+		return step{op: castPrimitive, castTypes: struct{ from, to zng.Type }{in, out}}, nil
+	case isCollectionType(in):
+		if _, ok := out.(*zng.TypeArray); ok {
+			return createStepArray(innerType(in), innerType(out))
+		}
+		if _, ok := out.(*zng.TypeSet); ok {
+			return createStepSet(innerType(in), innerType(out))
+		}
+		fallthrough
+	default:
+		return step{}, fmt.Errorf("createStep incompatible column types %s and %s\n", in, out)
+	}
+}
+
 func (s *step) castPrimitive(in zcode.Bytes, b *zcode.Builder) {
 	if in == nil {
 		b.AppendNull()
@@ -98,12 +152,47 @@ func (s *step) castPrimitive(in zcode.Bytes, b *zcode.Builder) {
 	b.AppendPrimitive(v.Bytes)
 }
 
-func (s *step) buildRecord(in zcode.Bytes, b *zcode.Builder) {
-	if s.op != record {
-		panic("bad op")
+func (s *step) build(in zcode.Bytes, b *zcode.Builder) {
+	switch s.op {
+	case copyPrimitive:
+		b.AppendPrimitive(in)
+	case copyContainer:
+		b.AppendContainer(in)
+	case castPrimitive:
+		s.castPrimitive(in, b)
+	case record:
+		if in == nil {
+			b.AppendNull()
+			return
+		}
+		b.BeginContainer()
+		s.buildRecord(in, b)
+		b.EndContainer()
+	case array:
+		fallthrough
+	case set:
+		if in == nil {
+			b.AppendNull()
+			return
+		}
+		b.BeginContainer()
+		iter := in.Iter()
+		for !iter.Done() {
+			zv, _, err := iter.Next()
+			if err != nil {
+				panic(err)
+			}
+			s.children[0].build(zv, b)
+		}
+		if s.op == set {
+			b.TransformContainer(zng.NormalizeSet)
+		}
+		b.EndContainer()
 	}
-	for _, step := range s.record {
+}
 
+func (s *step) buildRecord(in zcode.Bytes, b *zcode.Builder) {
+	for _, step := range s.children {
 		switch step.op {
 		case null:
 			b.AppendNull()
@@ -119,23 +208,7 @@ func (s *step) buildRecord(in zcode.Bytes, b *zcode.Builder) {
 		if err != nil {
 			panic(err)
 		}
-
-		switch step.op {
-		case copyPrimitive:
-			b.AppendPrimitive(bytes)
-		case copyContainer:
-			b.AppendContainer(bytes)
-		case castPrimitive:
-			step.castPrimitive(bytes, b)
-		case record:
-			if bytes == nil {
-				b.AppendNull()
-				continue
-			}
-			b.BeginContainer()
-			step.buildRecord(bytes, b)
-			b.EndContainer()
-		}
+		step.build(bytes, b)
 	}
 }
 
@@ -254,7 +327,7 @@ func (s *Shaper) createShapeSpec(inType, spec *zng.TypeRecord) (shapeSpec, error
 			return shapeSpec{}, err
 		}
 	}
-	step, err := createStep(inType, typ)
+	step, err := createStepRecord(inType, typ)
 	return shapeSpec{typ, step}, err
 }
 
@@ -271,18 +344,35 @@ func (s *Shaper) cropRecordType(input, spec *zng.TypeRecord) (*zng.TypeRecord, e
 
 		specCol := spec.Columns[ind]
 		switch {
-		case !zng.IsRecordType(inCol.Type):
+		case zng.IsPrimitiveType(inCol.Type):
 			// 2. Field is non-record in input: keep (regardless of crop record-ness)
 			cols = append(cols, inCol)
-		case zng.IsRecordType(specCol.Type):
+		case zng.IsRecordType(inCol.Type) && zng.IsRecordType(specCol.Type):
 			// 3. Both records: recurse
 			out, err := s.cropRecordType(inCol.Type.(*zng.TypeRecord), specCol.Type.(*zng.TypeRecord))
 			if err != nil {
 				return nil, err
 			}
 			cols = append(cols, zng.Column{inCol.Name, out})
+		case isCollectionType(inCol.Type) && isCollectionType(specCol.Type) &&
+			zng.IsRecordType(innerType(inCol.Type)) && zng.IsRecordType(innerType(specCol.Type)):
+			// 4. array/set of records
+			inner, err := s.cropRecordType(innerType(inCol.Type).(*zng.TypeRecord), innerType(specCol.Type).(*zng.TypeRecord))
+			if err != nil {
+				return nil, err
+			}
+			var t zng.Type
+			if _, ok := inCol.Type.(*zng.TypeArray); ok {
+				t, err = s.zctx.LookupTypeArray(inner), nil
+			} else {
+				t, err = s.zctx.LookupTypeSet(inner), nil
+			}
+			if err != nil {
+				return nil, err
+			}
+			cols = append(cols, zng.Column{inCol.Name, t})
 		default:
-			// 4. record input but non-record in crop: keep crop
+			// 5. record input but non-record in crop: keep crop
 			cols = append(cols, specCol)
 
 		}
@@ -320,6 +410,24 @@ func (s *Shaper) orderRecordType(input, spec *zng.TypeRecord) (*zng.TypeRecord,
 				}
 				continue
 			}
+			if isCollectionType(inCol.Type) && isCollectionType(specCol.Type) && zng.IsRecordType(innerType(inCol.Type)) && zng.IsRecordType(innerType(specCol.Type)) {
+				if inner, err := s.orderRecordType(innerType(inCol.Type).(*zng.TypeRecord), innerType(specCol.Type).(*zng.TypeRecord)); err != nil {
+					return nil, err
+				} else {
+					var err error
+					var t zng.Type
+					if _, ok := inCol.Type.(*zng.TypeArray); ok {
+						t, err = s.zctx.LookupTypeArray(inner), nil
+					} else {
+						t, err = s.zctx.LookupTypeSet(inner), nil
+					}
+					if err != nil {
+						return nil, err
+					}
+					cols = append(cols, zng.Column{specCol.Name, t})
+				}
+				continue
+			}
 			cols = append(cols, inCol)
 		}
 	}
@@ -339,10 +447,11 @@ func (s *Shaper) fillRecordType(input, spec *zng.TypeRecord) (*zng.TypeRecord, e
 	copy(cols, input.Columns)
 	for _, specCol := range spec.Columns {
 		if i, ok := input.ColumnOfField(specCol.Name); ok {
+			inCol := input.Columns[i]
 			// Field is present both in input and spec: recurse if
 			// both records, or select appropriate type if not.
 			if specRecType, ok := specCol.Type.(*zng.TypeRecord); ok {
-				if inRecType, ok := input.Columns[i].Type.(*zng.TypeRecord); ok {
+				if inRecType, ok := inCol.Type.(*zng.TypeRecord); ok {
 					filled, err := s.fillRecordType(inRecType, specRecType)
 					if err != nil {
 						return nil, err
@@ -351,6 +460,24 @@ func (s *Shaper) fillRecordType(input, spec *zng.TypeRecord) (*zng.TypeRecord, e
 				} else {
 					cols[i] = specCol
 				}
+				continue
+			}
+			if isCollectionType(inCol.Type) && isCollectionType(specCol.Type) && zng.IsRecordType(innerType(inCol.Type)) && zng.IsRecordType(innerType(specCol.Type)) {
+				if inner, err := s.fillRecordType(innerType(inCol.Type).(*zng.TypeRecord), innerType(specCol.Type).(*zng.TypeRecord)); err != nil {
+					return nil, err
+				} else {
+					var err error
+					var t zng.Type
+					if _, ok := inCol.Type.(*zng.TypeArray); ok {
+						t, err = s.zctx.LookupTypeArray(inner), nil
+					} else {
+						t, err = s.zctx.LookupTypeSet(inner), nil
+					}
+					if err != nil {
+						return nil, err
+					}
+					cols[i] = zng.Column{specCol.Name, t}
+				}
 			}
 		} else {
 			cols = append(cols, specCol)
@@ -375,35 +502,51 @@ func (s *Shaper) castRecordType(input, spec *zng.TypeRecord) (*zng.TypeRecord, e
 		}
 		specCol := spec.Columns[ind]
 
-		inRec, inIsRec := inCol.Type.(*zng.TypeRecord)
-		castRec, castIsRec := specCol.Type.(*zng.TypeRecord)
+		if _, ok := specCol.Type.(*zng.TypeMap); ok {
+			return nil, fmt.Errorf("cannot yet use maps in shaping functions")
+		}
 
-		switch {
-		case inCol.Type.ID() == specCol.Type.ID():
-			// 2. Field has same type in cast: output type unmodified.
-			cols = append(cols, inCol)
-		case inIsRec && castIsRec:
-			// 3. Matching field is a record: recurse.
-			out, err := s.castRecordType(inRec, castRec)
-			if err != nil {
-				return nil, err
-			}
-			cols = append(cols, zng.Column{inCol.Name, out})
-		case zng.IsPrimitiveType(inCol.Type) && zng.IsPrimitiveType(specCol.Type):
-			// 4. Matching field is a primitive: output type is cast type.
-			if LookupPrimitiveCaster(specCol.Type) == nil {
-				return nil, fmt.Errorf("cast to %s not implemented", specCol.Type)
-			}
-			cols = append(cols, zng.Column{inCol.Name, specCol.Type})
-		default:
-			// 5. Non-castable type pair with at least one
-			// (non-record) container: output column is left
-			// unchanged.  Note that eventually, we should
-			// recognize and cast e.g. array[string] to array[ip].
-			// xxx before merge: file issue for this and mention
-			// it here.
+		if inCol.Type.ID() == specCol.Type.ID() {
+			// Field has same type in cast: output type unmodified.
 			cols = append(cols, inCol)
+			continue
+		}
+		castType, err := s.castType(inCol.Type, specCol.Type)
+		if err != nil {
+			return nil, err
 		}
+		cols = append(cols, zng.Column{inCol.Name, castType})
+
 	}
 	return s.zctx.LookupTypeRecord(cols)
 }
+
+func (c *Shaper) castType(inType, specType zng.Type) (zng.Type, error) {
+	switch {
+	case zng.IsRecordType(inType) && zng.IsRecordType(specType):
+		// Matching field is a record: recurse.
+		inRec := inType.(*zng.TypeRecord)
+		castRec := specType.(*zng.TypeRecord)
+		return c.castRecordType(inRec, castRec)
+	case zng.IsPrimitiveType(inType) && zng.IsPrimitiveType(specType):
+		// Matching field is a primitive: output type is cast type.
+		if LookupPrimitiveCaster(specType) == nil {
+			return nil, fmt.Errorf("cast to %s not implemented", specType)
+		}
+		return specType, nil
+	case isCollectionType(inType) && isCollectionType(specType):
+		out, err := c.castType(innerType(inType), innerType(specType))
+		if err != nil {
+			return nil, err
+		}
+		if _, ok := specType.(*zng.TypeArray); ok {
+			return c.zctx.LookupTypeArray(out), nil
+		}
+		return c.zctx.LookupTypeSet(out), nil
+	default:
+		// Non-castable type pair with at least one
+		// (non-record) container: output column is left
+		// unchanged.
+		return inType, nil
+	}
+}


Test patch:
diff --git a/expr/ztests/shape-cast-array-to-set.yaml b/expr/ztests/shape-cast-array-to-set.yaml
new file mode 100644
index 000000000..9889ed430
--- /dev/null
+++ b/expr/ztests/shape-cast-array-to-set.yaml
@@ -0,0 +1,17 @@
+zql: |
+  put . = cast(., {a:|[ip]|,b:|[{b:ip}]|})
+
+input: |
+  #0:record[a:array[string],b:array[record[b:string]]]
+  0:[-;-;]
+  0:[[][]]
+  0:[[1.1.1.1;2.2.2.2;2.2.2.2;][[1.1.1.1;][2.2.2.2;][2.2.2.2;]]]
+  0:[[-;1.1.1.1;][-;[-;][2.2.2.2;]]]
+
+
+output: |
+  #0:record[a:set[ip],b:set[record[b:ip]]]
+  0:[-;-;]
+  0:[[][]]
+  0:[[1.1.1.1;2.2.2.2;][[1.1.1.1;][2.2.2.2;]]]
+  0:[[-;1.1.1.1;][-;[-;][2.2.2.2;]]]
diff --git a/expr/ztests/shape-cast-arrays.yaml b/expr/ztests/shape-cast-arrays.yaml
new file mode 100644
index 000000000..2f75a7df2
--- /dev/null
+++ b/expr/ztests/shape-cast-arrays.yaml
@@ -0,0 +1,17 @@
+zql: |
+  put . = cast(., {a:[ip],b:[{b:ip}]})
+
+input: |
+  #0:record[a:array[string],b:array[record[b:string]]]
+  0:[-;-;]
+  0:[[][]]
+  0:[[1.1.1.1;2.2.2.2;][[1.1.1.1;][2.2.2.2;]]]
+  0:[[-;1.1.1.1;][-;[-;][2.2.2.2;]]]
+
+
+output: |
+  #0:record[a:array[ip],b:array[record[b:ip]]]
+  0:[-;-;]
+  0:[[][]]
+  0:[[1.1.1.1;2.2.2.2;][[1.1.1.1;][2.2.2.2;]]]
+  0:[[-;1.1.1.1;][-;[-;][2.2.2.2;]]]
diff --git a/expr/ztests/shape-cast-set-to-array.yaml b/expr/ztests/shape-cast-set-to-array.yaml
new file mode 100644
index 000000000..157d1e3ec
--- /dev/null
+++ b/expr/ztests/shape-cast-set-to-array.yaml
@@ -0,0 +1,17 @@
+zql: |
+  put . = cast(., {a:[ip],b:[{b:ip}]})
+
+input: |
+  #0:record[a:set[string],b:set[record[b:string]]]
+  0:[-;-;]
+  0:[[][]]
+  0:[[1.1.1.1;2.2.2.2;][[1.1.1.1;][2.2.2.2;]]]
+  0:[[-;1.1.1.1;][-;[-;][2.2.2.2;]]]
+
+
+output: |
+  #0:record[a:array[ip],b:array[record[b:ip]]]
+  0:[-;-;]
+  0:[[][]]
+  0:[[1.1.1.1;2.2.2.2;][[1.1.1.1;][2.2.2.2;]]]
+  0:[[-;1.1.1.1;][-;[-;][2.2.2.2;]]]
diff --git a/expr/ztests/shape-cast-sets.yaml b/expr/ztests/shape-cast-sets.yaml
new file mode 100644
index 000000000..7e464052f
--- /dev/null
+++ b/expr/ztests/shape-cast-sets.yaml
@@ -0,0 +1,17 @@
+zql: |
+  put . = cast(., {a:|[ip]|,b:|[{b:ip}]|})
+
+input: |
+  #0:record[a:set[string],b:set[record[b:string]]]
+  0:[-;-;]
+  0:[[][]]
+  0:[[1.1.1.1;2.2.2.2;][[1.1.1.1;][2.2.2.2;]]]
+  0:[[-;1.1.1.1;][-;[-;][2.2.2.2;]]]
+
+
+output: |
+  #0:record[a:set[ip],b:set[record[b:ip]]]
+  0:[-;-;]
+  0:[[][]]
+  0:[[1.1.1.1;2.2.2.2;][[1.1.1.1;][2.2.2.2;]]]
+  0:[[-;1.1.1.1;][-;[-;][2.2.2.2;]]]
diff --git a/expr/ztests/shape-crop-arrays-sets.yaml b/expr/ztests/shape-crop-arrays-sets.yaml
new file mode 100644
index 000000000..ad68979c6
--- /dev/null
+++ b/expr/ztests/shape-crop-arrays-sets.yaml
@@ -0,0 +1,10 @@
+zql: |
+  put . = crop(., {set:|[{b:string}]|,arr:[{a:string}]})
+
+input: |
+  #0:record[arr:array[record[a:string,b:string]],set:set[record[a:string,b:string]]]
+  0:[[[one;two;][three;four;]][[one;two;][three;four;]]]
+
+output: |
+  #0:record[arr:array[record[a:string]],set:set[record[b:string]]]
+  0:[[[one;][three;]][[two;][four;]]]
diff --git a/expr/ztests/shape-fill-arrays-sets.yaml b/expr/ztests/shape-fill-arrays-sets.yaml
new file mode 100644
index 000000000..47b361624
--- /dev/null
+++ b/expr/ztests/shape-fill-arrays-sets.yaml
@@ -0,0 +1,10 @@
+zql: |
+  put . = fill(., {arr:[{b:string,a:string}],set:|[{b:string,a:string}]|})
+
+input: |
+  #0:record[arr:array[record[a:string]],set:set[record[a:string]]]
+  0:[[[one;][three;]][[one;][three;]]]
+
+output: |
+  #0:record[arr:array[record[a:string,b:string]],set:set[record[a:string,b:string]]]
+  0:[[[one;-;][three;-;]][[one;-;][three;-;]]]
diff --git a/expr/ztests/shape-order-arrays-sets.yaml b/expr/ztests/shape-order-arrays-sets.yaml
new file mode 100644
index 000000000..ccce440ed
--- /dev/null
+++ b/expr/ztests/shape-order-arrays-sets.yaml
@@ -0,0 +1,10 @@
+zql: |
+  put . = order(., {set:|[{b:string,a:string}]|,arr:[{b:string,a:string}]})
+
+input: |
+  #0:record[arr:array[record[a:string,b:string]],set:set[record[a:string,b:string]]]
+  0:[[[one;two;][three;four;]][[one;two;][three;four;]]]
+
+output: |
+  #0:record[set:set[record[b:string,a:string]],arr:array[record[b:string,a:string]]]
+  0:[[[two;one;][four;three;]][[two;one;][four;three;]]]
diff --git a/ppl/zqd/ztests/pcap/postpcap-suricata-update.yaml b/ppl/zqd/ztests/pcap/postpcap-suricata-update.yaml
index ef8f35172..51e139266 100644
--- a/ppl/zqd/ztests/pcap/postpcap-suricata-update.yaml
+++ b/ppl/zqd/ztests/pcap/postpcap-suricata-update.yaml
@@ -12,12 +12,7 @@ script: |
   echo "suricataupdater --local rules --url http://notexist" >> ./suricataupdater
   chmod a+x ./suricataupdater
 
-  source services.sh
-  zqd listen -l=localhost:0 -portfile=zqd-suricata-port -data="$zqdroot" -suricataupdater=./suricataupdater -loglevel=info &> zqd.log &
-  zqdpid=$!
-  awaitfile zqd-suricata-port
-  trap "kill -9 $zqdpid" EXIT
-  export ZQD_HOST=localhost:$(cat zqd-suricata-port)
+  ZQD_EXTRA_FLAGS='-loglevel=info -suricataupdater=./suricataupdater' source services.sh
 
   # Wait for a zqd log line indicating that the suricata-update
   # background process has completed.
diff --git a/ppl/zqd/ztests/recruiter-root.sh b/ppl/zqd/ztests/recruiter-root.sh
index 56ca1e2e5..94d1fc8eb 100644
--- a/ppl/zqd/ztests/recruiter-root.sh
+++ b/ppl/zqd/ztests/recruiter-root.sh
@@ -16,7 +16,7 @@ awaitfile $portdir/zqd-w1
 
 zqd listen -personality=worker -l=localhost:0 -portfile=$portdir/zqd-w2 \
   -worker.recruiter=$recruiter -worker.node=$node_name &> zqd-w2.log &
-zqdw1pid=$!
+zqdw2pid=$!
 awaitfile $portdir/zqd-w2
 
 # (3) Start a zqd "root" process (default personality for now) that may recruit workers.