1 files changed, 241 insertions, 0 deletions
diff --git a/container-search/src/main/java/com/yahoo/prelude/querytransform/QueryRewrite.java b/container-search/src/main/java/com/yahoo/prelude/querytransform/QueryRewrite.java
new file mode 100644
index 00000000000..fe680bd5ad0
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/prelude/querytransform/QueryRewrite.java
@@ -0,0 +1,241 @@
+// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.prelude.querytransform;
+
+import com.yahoo.prelude.query.AndItem;
+import com.yahoo.prelude.query.CompositeItem;
+import com.yahoo.prelude.query.EquivItem;
+import com.yahoo.prelude.query.Item;
+import com.yahoo.prelude.query.NearItem;
+import com.yahoo.prelude.query.NotItem;
+import com.yahoo.prelude.query.NullItem;
+import com.yahoo.prelude.query.OrItem;
+import com.yahoo.prelude.query.RankItem;
+import com.yahoo.prelude.query.SimpleIndexedItem;
+import com.yahoo.prelude.query.SubstringItem;
+import com.yahoo.search.Query;
+import com.yahoo.search.query.Model;
+import com.yahoo.search.result.Hit;
+
+/**
+ * @author balder
+ */
+public class QueryRewrite {
+
+    private static enum Recall {
+        RECALLS_EVERYTHING,
+        RECALLS_NOTHING,
+        UNKNOWN_RECALL
+    }
+
+    /**
+     * Optimize multiple NotItems under and or by collapsing them in to one and leaving
+     * the positive ones behind in its place and moving itself with the original and as its positive item
+     * and the union of all the negative items of all the original NotItems as its negative items.
+     *
+     * @param query to optimize
+     */
+    public static void optimizeAndNot(Query query) {
+        Item root = query.getModel().getQueryTree().getRoot();
+        Item possibleNewRoot = optimizeAndNot(root);
+        if (root != possibleNewRoot) {
+            query.getModel().getQueryTree().setRoot(possibleNewRoot);
+        }
+    }
+    private static Item optimizeAndNot(Item node) {
+        if (node instanceof CompositeItem) {
+            return extractAndNotRecursively((CompositeItem) node);
+        }
+        return node;
+    }
+    private static CompositeItem extractAndNotRecursively(CompositeItem parent) {
+        for (int i = 0; i < parent.getItemCount(); i++) {
+            Item child = parent.getItem(i);
+            Item possibleNewChild = optimizeAndNot(child);
+            if (child != possibleNewChild) {
+                parent.setItem(i, possibleNewChild);
+            }
+        }
+        if (parent instanceof AndItem) {
+            return extractAndNot((AndItem) parent);
+        }
+        return parent;
+    }
+    private static CompositeItem extractAndNot(AndItem parent) {
+        NotItem theOnlyNot = null;
+        for (int i = 0; i < parent.getItemCount(); i++) {
+            Item child = parent.getItem(i);
+            if (child instanceof NotItem) {
+                NotItem thisNot = (NotItem) child;
+                parent.setItem(i, thisNot.getPositiveItem());
+                if (theOnlyNot == null) {
+                    theOnlyNot = thisNot;
+                    theOnlyNot.setPositiveItem(parent);
+                } else {
+                    for (int j=1; j < thisNot.getItemCount(); j++) {
+                        theOnlyNot.addNegativeItem(thisNot.getItem(j));
+                    }
+                }
+            }
+        }
+        return (theOnlyNot != null) ? theOnlyNot : parent;
+    }
+    /**
+     * Optimizes the given query tree based on its {@link Model#getRestrict()} parameter, if any.
+     *
+     * @param query to optimize.
+     */
+    public static void optimizeByRestrict(Query query) {
+        if (query.getModel().getRestrict().size() != 1) {
+            return;
+        }
+        Item root = query.getModel().getQueryTree().getRoot();
+        if (optimizeByRestrict(root, query.getModel().getRestrict().iterator().next()) == Recall.RECALLS_NOTHING) {
+            query.getModel().getQueryTree().setRoot(new NullItem());
+        }
+    }
+
+    private static Recall optimizeByRestrict(Item item, String restrictParam) {
+        if (item instanceof SimpleIndexedItem) {
+            return optimizeIndexedItemByRestrict((SimpleIndexedItem)item, restrictParam);
+        } else if (item instanceof NotItem) {
+            return optimizeNotItemByRestrict((NotItem)item, restrictParam);
+        } else if (item instanceof CompositeItem) {
+            return optimizeCompositeItemByRestrict((CompositeItem)item, restrictParam);
+        } else {
+            return Recall.UNKNOWN_RECALL;
+        }
+    }
+
+    private static Recall optimizeIndexedItemByRestrict(SimpleIndexedItem item, String restrictParam) {
+        if (!Hit.SDDOCNAME_FIELD.equals(item.getIndexName())) {
+            return Recall.UNKNOWN_RECALL;
+        }
+        // a query term searching for sddocname will either recall everything or nothing, depending on whether
+        // the term matches the restrict parameter or not
+        return restrictParam.equals(item.getIndexedString())
+               ? Recall.RECALLS_EVERYTHING
+               : Recall.RECALLS_NOTHING;
+    }
+
+    private static Recall optimizeNotItemByRestrict(NotItem item, String restrictParam) {
+        // first item is the positive one
+        if (optimizeByRestrict(item.getItem(0), restrictParam) == Recall.RECALLS_NOTHING) {
+            return Recall.RECALLS_NOTHING;
+        }
+        // all the remaining items are negative ones
+        for (int i = item.getItemCount(); --i >= 1; ) {
+            Item child = item.getItem(i);
+            switch (optimizeByRestrict(child, restrictParam)) {
+            case RECALLS_EVERYTHING:
+                return Recall.RECALLS_NOTHING;
+            case RECALLS_NOTHING:
+                item.removeItem(i);
+                break;
+            }
+        }
+        return Recall.UNKNOWN_RECALL;
+    }
+
+    private static Recall optimizeCompositeItemByRestrict(CompositeItem item, String restrictParam) {
+        for (int i = item.getItemCount(); --i >= 0; ) {
+            switch (optimizeByRestrict(item.getItem(i), restrictParam)) {
+            case RECALLS_EVERYTHING:
+                if ((item instanceof OrItem) || (item instanceof EquivItem)) {
+                    retainChild(item, i);
+                    return Recall.RECALLS_EVERYTHING;
+                } else if ((item instanceof AndItem) || (item instanceof NearItem)) {
+                    item.removeItem(i);
+                } else if (item instanceof RankItem) {
+                    // empty
+                } else {
+                    throw new UnsupportedOperationException(item.getClass().getName());
+                }
+                break;
+            case RECALLS_NOTHING:
+                if ((item instanceof OrItem) || (item instanceof EquivItem)) {
+                    item.removeItem(i);
+                } else if ((item instanceof AndItem) || (item instanceof NearItem)) {
+                    return Recall.RECALLS_NOTHING;
+                } else if (item instanceof RankItem) {
+                    item.removeItem(i);
+                } else {
+                    throw new UnsupportedOperationException(item.getClass().getName());
+                }
+                break;
+            }
+        }
+        return Recall.UNKNOWN_RECALL;
+    }
+
+    private static void retainChild(CompositeItem item, int childIdx) {
+        Item child = item.removeItem(childIdx);
+        for (int i = item.getItemCount(); --i >= 0; ) {
+            item.removeItem(i);
+        }
+        item.addItem(child);
+    }
+
+    /**
+     * Collapses all single-child {@link CompositeItem}s into their parent item.
+     *
+     * @param query The query whose composites to collapse.
+     */
+    public static void collapseSingleComposites(Query query) {
+        Item oldRoot = query.getModel().getQueryTree().getRoot();
+        Item newRoot = collapseSingleComposites(oldRoot);
+        if (oldRoot != newRoot) {
+            query.getModel().getQueryTree().setRoot(newRoot);
+        }
+    }
+
+    private static Item collapseSingleComposites(Item item) {
+        if (!(item instanceof CompositeItem)) {
+            return item;
+        }
+        CompositeItem parent = (CompositeItem)item;
+        int numChildren = parent.getItemCount();
+        for (int i = 0; i < numChildren; ++i) {
+            Item oldChild = parent.getItem(i);
+            Item newChild = collapseSingleComposites(oldChild);
+            if (oldChild != newChild) {
+                parent.setItem(i, newChild);
+            }
+        }
+        return numChildren == 1 ? parent.getItem(0) : item;
+    }
+
+    /**
+     * Replaces and {@link SimpleIndexedItem} searching in the {@link Hit#SDDOCNAME_FIELD} with an item
+     * appropriate for the search node.
+     *
+     * @param query The query to rewrite.
+     */
+    public static void rewriteSddocname(Query query) {
+        Item oldRoot = query.getModel().getQueryTree().getRoot();
+        Item newRoot = rewriteSddocname(oldRoot);
+        if (oldRoot != newRoot) {
+            query.getModel().getQueryTree().setRoot(newRoot);
+        }
+    }
+
+    private static Item rewriteSddocname(Item item) {
+        if (item instanceof CompositeItem) {
+            CompositeItem parent = (CompositeItem)item;
+            for (int i = 0, len = parent.getItemCount(); i < len; ++i) {
+                Item oldChild = parent.getItem(i);
+                Item newChild = rewriteSddocname(oldChild);
+                if (oldChild != newChild) {
+                    parent.setItem(i, newChild);
+                }
+            }
+        } else if (item instanceof SimpleIndexedItem) {
+            SimpleIndexedItem oldItem = (SimpleIndexedItem)item;
+            if (Hit.SDDOCNAME_FIELD.equals(oldItem.getIndexName())) {
+                SubstringItem newItem = new SubstringItem(oldItem.getIndexedString());
+                newItem.setIndexName("[documentmetastore]");
+                return newItem;
+            }
+        }
+        return item;
+    }
+}