Add XGBoost to Vespa converter

2 files changed, 74 insertions, 0 deletions

diff --git a/config-model/src/test/integration/xgboost/models/xgboost.2.2.json b/config-model/src/test/integration/xgboost/models/xgboost.2.2.json
new file mode 100644
index 00000000000..f8949b47e52
--- /dev/null
+++ b/config-model/src/test/integration/xgboost/models/xgboost.2.2.json
@@ -0,0 +1,19 @@
+[
+  { "nodeid": 0, "depth": 0, "split": "f29", "split_condition": -0.1234567, "yes": 1, "no": 2, "missing": 1, "children": [
+    { "nodeid": 1, "depth": 1, "split": "f56", "split_condition": -0.242398, "yes": 3, "no": 4, "missing": 3, "children": [
+      { "nodeid": 3, "leaf": 1.71218 },
+      { "nodeid": 4, "leaf": -1.70044 }
+    ]},
+    { "nodeid": 2, "depth": 1, "split": "f109", "split_condition": 0.8723473, "yes": 5, "no": 6, "missing": 5, "children": [
+      { "nodeid": 5, "leaf": -1.94071 },
+      { "nodeid": 6, "leaf": 1.85965 }
+    ]}
+  ]},
+  { "nodeid": 0, "depth": 0, "split": "f60", "split_condition": -0.482947, "yes": 1, "no": 2, "missing": 1, "children": [
+    { "nodeid": 1, "depth": 1, "split": "f29", "split_condition": -4.2387498, "yes": 3, "no": 4, "missing": 3, "children": [
+      { "nodeid": 3, "leaf": 0.784718 },
+      { "nodeid": 4, "leaf": -0.96853 }
+    ]},
+    { "nodeid": 2, "leaf": -6.23624 }
+  ]}
+]
\ No newline at end of file
diff --git a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithXgboostTestCase.java b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithXgboostTestCase.java
new file mode 100644
index 00000000000..b65cb0b3d5f
--- /dev/null
+++ b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithXgboostTestCase.java
@@ -0,0 +1,55 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.searchdefinition.processing;
+
+import com.yahoo.path.Path;
+import com.yahoo.searchdefinition.parser.ParseException;
+import org.junit.Test;
+
+/**
+ * @author grace-lam
+ */
+public class RankingExpressionWithXgboostTestCase {
+
+    private final Path applicationDir = Path.fromString("src/test/integration/xgboost/");
+    private final static String vespaExpression = "if (f29 < -0.1234567, if (f56 < -0.242398, 1.71218, -1.70044), if (f109 < 0.8723473, -1.94071, 1.85965)) + " +
+            "if (f60 < -0.482947, if (f29 < -4.2387498, 0.784718, -0.96853), -6.23624)";
+
+    @Test
+    public void testXgboostReference() {
+        RankProfileSearchFixture search = fixtureWith("xgboost('xgboost.2.2.json')");
+        search.assertFirstPhaseExpression(vespaExpression, "my_profile");
+    }
+
+    @Test
+    public void testNestedXgboostReference() {
+        RankProfileSearchFixture search = fixtureWith("5 + sum(xgboost('xgboost.2.2.json'))");
+        search.assertFirstPhaseExpression("5 + reduce(" + vespaExpression + ", sum)", "my_profile");
+    }
+
+    private RankProfileSearchFixture fixtureWith(String firstPhaseExpression) {
+        return fixtureWith(firstPhaseExpression, null, null,
+                new RankingExpressionWithTensorFlowTestCase.StoringApplicationPackage(applicationDir));
+    }
+
+    private RankProfileSearchFixture fixtureWith(String firstPhaseExpression,
+                                                 String constant,
+                                                 String field,
+                                                 RankingExpressionWithTensorFlowTestCase.StoringApplicationPackage application) {
+        try {
+            return new RankProfileSearchFixture(
+                    application,
+                    application.getQueryProfiles(),
+                    "  rank-profile my_profile {\n" +
+                            "    first-phase {\n" +
+                            "      expression: " + firstPhaseExpression +
+                            "    }\n" +
+                            "  }",
+                    constant,
+                    field);
+        } catch (ParseException e) {
+            throw new IllegalArgumentException(e);
+        }
+    }
+
+}
+