Add ability to slice sets

Closes #4938
brimdata · Dec 14, 2023 · dac2b5f · dac2b5f
1 parent 59b72d8
commit dac2b5f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/runtime/expr/slice.go b/runtime/expr/slice.go
@@ -39,14 +39,14 @@ func (s *Slice) Eval(ectx Context, this *zed.Value) *zed.Value {
 		length = len(elem.Bytes())
 	case *zed.TypeOfString:
 		length = utf8.RuneCount(elem.Bytes())
-	case *zed.TypeArray:
+	case *zed.TypeArray, *zed.TypeSet:
 		n, err := elem.ContainerLength()
 		if err != nil {
 			panic(err)
 		}
 		length = n
 	default:
-		return s.zctx.WrapError("sliced value is not array, bytes, or string", elem)
+		return s.zctx.WrapError("sliced value is not array, set, bytes, or string", elem)
 	}
 	if elem.IsNull() {
 		return elem
@@ -69,7 +69,7 @@ func (s *Slice) Eval(ectx Context, this *zed.Value) *zed.Value {
 	case *zed.TypeOfString:
 		bytes = bytes[utf8PrefixLen(bytes, from):]
 		bytes = bytes[:utf8PrefixLen(bytes, to-from)]
-	case *zed.TypeArray:
+	case *zed.TypeArray, *zed.TypeSet:
 		it := bytes.Iter()
 		for k := 0; k < to && !it.Done(); k++ {
 			if k == from {
@@ -81,6 +81,9 @@ func (s *Slice) Eval(ectx Context, this *zed.Value) *zed.Value {
 	default:
 		panic(elem.Type)
 	}
+	if _, ok := elem.Type.(*zed.TypeSet); ok {
+		bytes = zed.NormalizeSet(bytes)
+	}
 	return ectx.NewValue(elem.Type, bytes)
 }
 

diff --git a/runtime/expr/ztests/slice.yaml b/runtime/expr/ztests/slice.yaml
@@ -13,6 +13,7 @@ input: |
   {a:"0\u2071\u20723"}
   {a:"\u2070\u2071\u2072\u2073"}
   {a:[10(int32),11(int32),12(int32),13(int32)]}
+  {a:|[10(int32),11(int32),12(int32),13(int32)]|}
 
 output: |
   {a1:error({message:"sliced value is not array, bytes, or string",on:null}),a2:error({message:"sliced value is not array, bytes, or string",on:null}),a3:error({message:"sliced value is not array, bytes, or string",on:null}),a4:error({message:"sliced value is not array, bytes, or string",on:null}),a5:error({message:"sliced value is not array, bytes, or string",on:null}),a6:error({message:"sliced value is not array, bytes, or string",on:null}),a7:error({message:"sliced value is not array, bytes, or string",on:null}),a8:null}
@@ -27,3 +28,4 @@ output: |
   {a1:"ⁱ⁲",a2:"ⁱ⁲3",a3:"0",a4:"0ⁱ⁲",a5:"",a6:"3",a7:"⁲",a8:error("slice index is not a number")}
   {a1:"ⁱ⁲",a2:"ⁱ⁲⁳",a3:"⁰",a4:"⁰ⁱ⁲",a5:"",a6:"⁳",a7:"⁲",a8:error("slice index is not a number")}
   {a1:[11(int32),12(int32)],a2:[11(int32),12(int32),13(int32)],a3:[10(int32)],a4:[10(int32),11(int32),12(int32)],a5:[]([int32]),a6:[13(int32)],a7:[12(int32)],a8:[10(int32),11(int32)]}
+  {a1:|[11(int32),12(int32)]|,a2:|[11(int32),12(int32),13(int32)]|,a3:|[10(int32)]|,a4:|[10(int32),11(int32),12(int32)]|,a5:|[]|(|[int32]|),a6:|[13(int32)]|,a7:|[12(int32)]|,a8:|[10(int32),11(int32)]|}