feat: flatten elasticsearch queries

Author: Codehardt

v2/evalostic.go

@@ -16,6 +16,7 @@ type Evalostic struct {
 	ahoCorasick  ahocorasick.AhoCorasick
 	strings      map[string]int
 	mapping      map[int][]int // which string can be found in which condition
+	orig         []node        // original conditions for export
 }
 
 // New builds a new Evalostic matcher that compiles all conditions to one big rule set that can be applied to strings.
@@ -37,6 +38,7 @@ func New(conditions []string) (*Evalostic, error) {
 		if err != nil {
 			return nil, fmt.Errorf("condition %d: %s", i, err)
 		}
+		e.orig = append(e.orig, root)
 		condStrings, _ := extractStrings(root)
 		for _, str := range condStrings {
 			strI, ok := e.strings[str]

v2/evalostic_test.go

@@ -271,7 +271,7 @@ func TestEvalosticAgainstStringContains(t *testing.T) {
 	}
 }
 
-func assert(t *testing.T, b bool) {
+func assertTrue(t *testing.T, b bool) {
 	_, f, l, _ := runtime.Caller(1)
 	if !b {
 		t.Fatalf("assertion failed in %s:%d", f, l)
@@ -298,23 +298,23 @@ func TestEvalostic(t *testing.T) {
 		`"1" AND NOT "2"`,
 		`"1" AND NOT "2"`,
 	})
-	assert(t, err == nil)
-	assert(t, sameIntegers(e.Match("foo"), []int{0}))
-	assert(t, sameIntegers(e.Match("bar"), []int{0}))
-	assert(t, sameIntegers(e.Match("foo bar"), []int{0}))
-	assert(t, sameIntegers(e.Match("baz"), []int{}))
-	assert(t, sameIntegers(e.Match("baz qux"), []int{1}))
-	assert(t, sameIntegers(e.Match("qux baz"), []int{1}))
-	assert(t, sameIntegers(e.Match("ab"), []int{}))
-	assert(t, sameIntegers(e.Match("ac"), []int{2}))
-	assert(t, sameIntegers(e.Match("ad"), []int{2}))
-	assert(t, sameIntegers(e.Match("bc"), []int{2}))
-	assert(t, sameIntegers(e.Match("bd"), []int{2}))
-	assert(t, sameIntegers(e.Match("cd"), []int{}))
-	assert(t, sameIntegers(e.Match("abcd"), []int{2}))
-	assert(t, sameIntegers(e.Match("1"), []int{3, 4}))
-	assert(t, sameIntegers(e.Match("2"), []int{}))
-	assert(t, sameIntegers(e.Match("12"), []int{}))
+	assertTrue(t, err == nil)
+	assertTrue(t, sameIntegers(e.Match("foo"), []int{0}))
+	assertTrue(t, sameIntegers(e.Match("bar"), []int{0}))
+	assertTrue(t, sameIntegers(e.Match("foo bar"), []int{0}))
+	assertTrue(t, sameIntegers(e.Match("baz"), []int{}))
+	assertTrue(t, sameIntegers(e.Match("baz qux"), []int{1}))
+	assertTrue(t, sameIntegers(e.Match("qux baz"), []int{1}))
+	assertTrue(t, sameIntegers(e.Match("ab"), []int{}))
+	assertTrue(t, sameIntegers(e.Match("ac"), []int{2}))
+	assertTrue(t, sameIntegers(e.Match("ad"), []int{2}))
+	assertTrue(t, sameIntegers(e.Match("bc"), []int{2}))
+	assertTrue(t, sameIntegers(e.Match("bd"), []int{2}))
+	assertTrue(t, sameIntegers(e.Match("cd"), []int{}))
+	assertTrue(t, sameIntegers(e.Match("abcd"), []int{2}))
+	assertTrue(t, sameIntegers(e.Match("1"), []int{3, 4}))
+	assertTrue(t, sameIntegers(e.Match("2"), []int{}))
+	assertTrue(t, sameIntegers(e.Match("12"), []int{}))
 }
 
 func ExampleMatch() {

v2/export.go

@@ -17,94 +17,114 @@ func (e *Evalostic) ExportElasticSearchQuery(wildcardField string, useMatchPhras
 // `"foo" OR "baz"` will be compiled to
 // {"bool":{"should":[{"wildcard":{"raw":{"case_insensitive":false,"value":"foo"}}},{"wildcard":{"raw":{"case_insensitive":false,"value":"bar"}}}]}}
 func (e *Evalostic) ExportElasticSearchQueryMap(wildcardField string, useMatchPhrase bool) map[string]interface{} {
-	indexToStrings := make(map[int]string)
-	for k, v := range e.strings {
-		indexToStrings[v] = k
+	var root node
+	for _, n := range e.orig {
+		if root == nil {
+			root = n
+		} else {
+			root = nodeOR{twoSubNodes{root, n}}
+		}
 	}
-	query := e.exportElasticSearchQuerySub(wildcardField, useMatchPhrase, indexToStrings, decisionTreeEntry{value: -1}, e.decisionTree, false)
-	if query == nil {
-		return make(map[string]interface{})
+	return nodeToElasticSearchQuery(root, useMatchPhrase)
+}
+
+func nodeToElasticSearchQuery(n node, useMatchPhrase bool) map[string]interface{} {
+	switch v := n.(type) {
+	case nodeVAL:
+		return leafToElasticSearchQuery(v, useMatchPhrase)
+	case nodeNOT:
+		return notToElasticSearchQuery(v, useMatchPhrase)
+	case nodeOR:
+		return orToElasticSearchQuery(v, useMatchPhrase)
+	case nodeAND:
+		return andToElasticSearchQuery(v, useMatchPhrase)
+	default:
+		return nil
 	}
-	return query
 }
 
 var wildcardReplacer = strings.NewReplacer("\\", "\\\\", "*", "\\*", "?", "\\?")
 
-func (e *Evalostic) exportElasticSearchQuerySub(wildcardField string, useMatchPhrase bool, indexToStrings map[int]string, entry decisionTreeEntry, node *decisionTreeNode, not bool) map[string]interface{} {
-	isLeaf := len(node.outputs) != 0
-	wildcard := map[string]interface{}{
-		"wildcard": map[string]interface{}{
-			wildcardField: map[string]interface{}{
-				"value":            "*" + wildcardReplacer.Replace(indexToStrings[entry.value]) + "*",
-				"case_insensitive": true,
+func notToElasticSearchQuery(n nodeNOT, useMatchPhrase bool) map[string]interface{} {
+	if not, ok := n.node.(nodeNOT); ok { // check for double negation
+		return nodeToElasticSearchQuery(not.node, useMatchPhrase)
+	}
+	return map[string]interface{}{
+		"bool": map[string]interface{}{
+			"must_not": []map[string]interface{}{
+				nodeToElasticSearchQuery(n.node, useMatchPhrase),
 			},
 		},
 	}
-	if useMatchPhrase {
-		wildcard = map[string]interface{}{
-			"match_phrase": map[string]interface{}{
-				wildcardField: indexToStrings[entry.value],
-			},
-		}
+}
+
+func flattenOr(n nodeOR) []node {
+	var nodes []node
+	if or, ok := n.node1.(nodeOR); ok {
+		nodes = append(nodes, flattenOr(or)...)
+	} else {
+		nodes = append(nodes, n.node1)
 	}
-	if not {
-		wildcard = map[string]interface{}{
-			"bool": map[string]interface{}{
-				"must_not": []interface{}{wildcard},
-			},
-		}
+	if or, ok := n.node2.(nodeOR); ok {
+		nodes = append(nodes, flattenOr(or)...)
+	} else {
+		nodes = append(nodes, n.node2)
 	}
-	if entry.value == -1 {
-		// special case: do not use root node as wildcard
-		wildcard = nil
+	return nodes
+}
+
+func flattenAnd(n nodeAND) []node {
+	var nodes []node
+	if and, ok := n.node1.(nodeAND); ok {
+		nodes = append(nodes, flattenAnd(and)...)
+	} else {
+		nodes = append(nodes, n.node1)
 	}
-	if isLeaf && wildcard != nil {
-		// special case: if it's a leaf, we don't need to process the sub tree
-		return wildcard
+	if and, ok := n.node2.(nodeAND); ok {
+		nodes = append(nodes, flattenAnd(and)...)
+	} else {
+		nodes = append(nodes, n.node2)
 	}
+	return nodes
+}
 
+func orToElasticSearchQuery(n nodeOR, useMatchPhrase bool) map[string]interface{} {
 	var should []map[string]interface{}
-
-	for subEntry, subNode := range node.children {
-		if subQuery := e.exportElasticSearchQuerySub(wildcardField, useMatchPhrase, indexToStrings, subEntry, subNode, false); subQuery != nil {
-			should = append(should, subQuery)
-		}
+	for _, node := range flattenOr(n) {
+		should = append(should, nodeToElasticSearchQuery(node, useMatchPhrase))
 	}
-	for subEntry, subNode := range node.notChildren {
-		if subQuery := e.exportElasticSearchQuerySub(wildcardField, useMatchPhrase, indexToStrings, subEntry, subNode, true); subQuery != nil {
-			should = append(should, subQuery)
-		}
+	return map[string]interface{}{
+		"bool": map[string]interface{}{
+			"should": should,
+		},
 	}
+}
 
-	toQuery := func(should []map[string]interface{}) map[string]interface{} {
-		if len(should) == 0 {
-			return nil
-		}
-		var res map[string]interface{}
-		if len(should) == 1 {
-			res = should[0]
-		} else {
-			res = map[string]interface{}{
-				"bool": map[string]interface{}{
-					"should": should,
-				},
-			}
-		}
-		return res
+func andToElasticSearchQuery(n nodeAND, useMatchPhrase bool) map[string]interface{} {
+	var must []map[string]interface{}
+	for _, node := range flattenAnd(n) {
+		must = append(must, nodeToElasticSearchQuery(node, useMatchPhrase))
 	}
-
-	childQuery := toQuery(should)
-	if childQuery == nil {
-		return nil
+	return map[string]interface{}{
+		"bool": map[string]interface{}{
+			"must": must,
+		},
 	}
-	if wildcard == nil {
-		return childQuery
+}
+
+func leafToElasticSearchQuery(n nodeVAL, useMatchPhrase bool) map[string]interface{} {
+	if useMatchPhrase {
+		return map[string]interface{}{
+			"match_phrase": map[string]interface{}{
+				"raw": n.nodeValue,
+			},
+		}
 	}
 	return map[string]interface{}{
-		"bool": map[string]interface{}{
-			"must": []interface{}{
-				wildcard,
-				childQuery,
+		"wildcard": map[string]interface{}{
+			"raw": map[string]interface{}{
+				"value":            "*" + wildcardReplacer.Replace(n.nodeValue) + "*",
+				"case_insensitive": true,
 			},
 		},
 	}