Implement vector IP comparison

brimdata · Dec 9, 2024 · 0444a85 · 0444a85
1 parent 3aed16e
commit 0444a85
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 3 deletions.
diff --git a/runtime/sam/expr/eval.go b/runtime/sam/expr/eval.go
@@ -319,6 +319,8 @@ func (c *Compare) Eval(ectx Context, this super.Value) super.Value {
 		return c.result(bytes.Compare(super.DecodeBytes(lhs.Bytes()), super.DecodeBytes(rhs.Bytes())))
 	case lid == super.IDString:
 		return c.result(cmp.Compare(super.DecodeString(lhs.Bytes()), super.DecodeString(rhs.Bytes())))
+	case lid == super.IDIP:
+		return c.result(super.DecodeIP(lhs.Bytes()).Compare(super.DecodeIP(rhs.Bytes())))
 	default:
 		if bytes.Equal(lhs.Bytes(), rhs.Bytes()) {
 			return c.result(0)

diff --git a/runtime/vam/expr/compare.go b/runtime/vam/expr/compare.go
@@ -45,7 +45,10 @@ func (c *Compare) eval(vecs ...vector.Any) vector.Any {
 	if kind != vector.KindOf(rhs) {
 		panic("vector kind mismatch after coerce")
 	}
-	if kind == vector.KindType {
+	switch kind {
+	case vector.KindIP:
+		return c.compareIPs(lhs, rhs, nulls)
+	case vector.KindType:
 		return c.compareTypeVals(lhs, rhs)
 	}
 	lform, ok := vector.FormOf(lhs)
@@ -64,6 +67,42 @@ func (c *Compare) eval(vecs ...vector.Any) vector.Any {
 	return vector.CopyAndSetNulls(out, nulls)
 }
 
+func (c *Compare) compareIPs(lhs, rhs vector.Any, nulls *vector.Bool) vector.Any {
+	out := vector.NewBoolEmpty(lhs.Len(), nulls)
+	for i := range lhs.Len() {
+		l, null := vector.IPValue(lhs, i)
+		if null {
+			continue
+		}
+		r, null := vector.IPValue(rhs, i)
+		if null {
+			continue
+		}
+		if isCompareOpSatisfied(c.opCode, l.Compare(r)) {
+			out.Set(i)
+		}
+	}
+	return out
+}
+
+func isCompareOpSatisfied(opCode, i int) bool {
+	switch opCode {
+	case vector.CompLT:
+		return i < 0
+	case vector.CompLE:
+		return i <= 0
+	case vector.CompGT:
+		return i > 0
+	case vector.CompGE:
+		return i >= 0
+	case vector.CompEQ:
+		return i == 0
+	case vector.CompNE:
+		return i != 0
+	}
+	panic(opCode)
+}
+
 func (c *Compare) compareTypeVals(lhs, rhs vector.Any) vector.Any {
 	if c.opCode == vector.CompLT || c.opCode == vector.CompGT {
 		return vector.NewConst(super.False, lhs.Len(), nil)

diff --git a/vector/kind.go b/vector/kind.go
@@ -16,8 +16,9 @@ const (
 	KindFloat   = 3
 	KindString  = 4
 	KindBytes   = 5
-	KindType    = 6
-	KindError   = 7
+	KindIP      = 6
+	KindType    = 7
+	KindError   = 8
 )
 
 const (
@@ -43,6 +44,8 @@ func KindOf(v Any) Kind {
 		return KindString
 	case *Error:
 		return KindError
+	case *IP:
+		return KindIP
 	case *TypeValue:
 		return KindType
 	case *Dict:
@@ -87,6 +90,8 @@ func KindOfType(typ super.Type) Kind {
 		return KindString
 	case *super.TypeOfBytes:
 		return KindBytes
+	case *super.TypeOfIP:
+		return KindIP
 	case *super.TypeOfType:
 		return KindType
 	}