aboutsummaryrefslogtreecommitdiffstats
path: root/predicate-search-core/src/main/java/com/yahoo/search/predicate/PredicateQueryParser.java
blob: 09487506ffe598c825eba81a7fa3aa52bda2bd4a (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.search.predicate;

import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonToken;

import java.io.IOException;
import java.util.Arrays;

/**
 * Parses predicate queries from JSON.
 *
 * Input JSON is assumed to have the following format:
 * {
 *      "features": [
 *          {"k": "key-name", "v":"value", "s":"0xDEADBEEFDEADBEEF"}
 *      ],
 *      "rangeFeatures": [
 *          {"k": "key-name", "v":42, "s":"0xDEADBEEFDEADBEEF"}
 *      ]
 *  }
 *
 * @author bjorncs
 */
public class PredicateQueryParser {

    private final JsonFactory factory = new JsonFactory();

    @FunctionalInterface
    public interface FeatureHandler<V> {
        void accept(String key, V value, long subqueryBitmap);
    }

    /**
     * Parses predicate query from JSON.
     * @param json JSON input.
     * @param featureHandler The handler is invoked when a feature is parsed.
     * @param rangeFeatureHandler The handler is invoked when a range feature is parsed.
     * @throws IllegalArgumentException If JSON is invalid.
     */
    public void parseJsonQuery(
            String json, FeatureHandler<String> featureHandler, FeatureHandler<Long> rangeFeatureHandler)
    throws IllegalArgumentException {

        try (JsonParser parser = factory.createParser(json)) {
            skipToken(parser, JsonToken.START_OBJECT);
            while (parser.nextToken() != JsonToken.END_OBJECT) {
                String fieldName = parser.getCurrentName();
                switch (fieldName) {
                    case "features":
                        parseFeatures(parser, JsonParser::getText, featureHandler);
                        break;
                    case "rangeFeatures":
                        parseFeatures(parser, JsonParser::getLongValue, rangeFeatureHandler);
                        break;
                    default:
                        throw new IllegalArgumentException("Invalid field name: " + fieldName);
                }
            }
        } catch (IllegalArgumentException e) {
            throw e;
        } catch (IOException e) {
            throw new AssertionError("This should never happen when parsing from a String", e);
        } catch (Exception e) {
            throw new IllegalArgumentException(String.format("Parsing query from JSON failed: '%s'", json), e);
        }
    }

    private static <V> void parseFeatures(
            JsonParser parser, ValueParser<V> valueParser, FeatureHandler<V> featureHandler) throws IOException {
        skipToken(parser, JsonToken.START_ARRAY);
        while (parser.nextToken() != JsonToken.END_ARRAY) {
            parseFeature(parser, valueParser, featureHandler);
        }
    }

    private static <V> void parseFeature(
            JsonParser parser, ValueParser<V> valueParser, FeatureHandler<V> featureHandler) throws IOException {
        String key = null;
        V value = null;
        long subqueryBitmap = SubqueryBitmap.DEFAULT_VALUE; // Specifying subquery bitmap is optional.

        while (parser.nextToken() != JsonToken.END_OBJECT) {
            String fieldName = parser.getCurrentName();
            skipToken(parser, JsonToken.VALUE_STRING, JsonToken.VALUE_NUMBER_INT);
            switch (fieldName) {
                case "k":
                    key = parser.getText();
                    break;
                case "v":
                    value = valueParser.parse(parser);
                    break;
                case "s":
                    subqueryBitmap = fromHexString(parser.getText());
                    break;
                default:
                    throw new IllegalArgumentException("Invalid field name: " + fieldName);
            }
        }
        if (key == null) {
            throw new IllegalArgumentException(
                    String.format("Feature key is missing! (%s)", parser.getCurrentLocation()));
        }
        if (value == null) {
            throw new IllegalArgumentException(
                    String.format("Feature value is missing! (%s)", parser.getCurrentLocation()));
        }
        featureHandler.accept(key, value, subqueryBitmap);
    }

    private static void skipToken(JsonParser parser, JsonToken... expected) throws IOException {
        JsonToken actual = parser.nextToken();
        if (Arrays.stream(expected).noneMatch(e -> e.equals(actual))) {
            throw new IllegalArgumentException(
                    String.format("Expected a token in %s, got %s (%s).",
                            Arrays.toString(expected), actual, parser.getTokenLocation()));
        }
    }

    private static long fromHexString(String subqueryBitmap) {
        if (!subqueryBitmap.startsWith("0x")) {
            throw new IllegalArgumentException("Not a valid subquery bitmap ('0x' prefix missing): " + subqueryBitmap);
        }
        return Long.parseUnsignedLong(subqueryBitmap.substring(2), 16);
    }

    @FunctionalInterface
    private interface ValueParser<V> {
        V parse(JsonParser parser) throws IOException;
    }

}