feat: Implement new xgboost

openml-xgboost/pom.xml

@@ -0,0 +1,158 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Copyright 2026 Feedzai
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~     http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  ~
+  -->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <groupId>com.feedzai</groupId>
+        <artifactId>openml-java</artifactId>
+        <version>0.0.0-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>openml-xgboost</artifactId>
+    <name>OpenML XGBoost</name>
+    <description>Provider that imports, scores and trains XGBoost models using the native xgboost4j JVM package.</description>
+
+    <properties>
+        <!--
+          xgboost4j is the pure-JVM core of XGBoost. The published jar bundles the native library for
+          linux/x86_64, linux/aarch64, macos/x86_64, macos/aarch64 and windows/x86_64, so it works on
+          ARM (AWS Graviton / Apple Silicon) out of the box - unlike H2O-XGBoost (AMD64 only).
+          It is a thin JNI wrapper (no H2O-style JDK version gate), so it runs on Java 8-25.
+        -->
+        <xgboost.version>3.4.0</xgboost.version>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>com.feedzai</groupId>
+            <artifactId>openml-api</artifactId>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.feedzai</groupId>
+            <artifactId>openml-utils</artifactId>
+            <scope>provided</scope>
+        </dependency>
+
+        <!--
+          xgboost4j core (Java API: ml.dmlc.xgboost4j.java.{XGBoost,Booster,DMatrix}).
+          The artifact carries a Scala suffix; we use the _2.12 build to match Pulse's Scala 2.12
+          classpath, even though this provider only uses the pure-Java API. scala-compiler is not
+          needed at all, so it is excluded to keep the classpath lean.
+
+          kryo is excluded because xgboost4j (via its parent pom) pulls Kryo 5, whose
+          com.esotericsoftware.kryo.Kryo would clash on the class path with the Kryo 4 (kryo-shaded)
+          that Pulse and Spark rely on (Kryo 5 removed the nested Kryo.DefaultInstantiatorStrategy).
+          xgboost4j's Booster implements com.esotericsoftware.kryo.KryoSerializable, so kryo classes
+          must still be present at runtime - Pulse supplies them via kryo-shaded 4. To keep this module
+          self-contained for its own build/tests, kryo is re-declared below in 'provided' scope, which
+          is NOT propagated transitively to consumers.
+        -->
+        <dependency>
+            <groupId>ml.dmlc</groupId>
+            <artifactId>xgboost4j_2.12</artifactId>
+            <version>${xgboost.version}</version>
+            <exclusions>
+                <exclusion>
+                    <groupId>org.scala-lang</groupId>
+                    <artifactId>scala-compiler</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.esotericsoftware</groupId>
+                    <artifactId>kryo</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+        <!--
+          Needed only so xgboost4j's Booster (implements KryoSerializable) can load during this module's
+          own compilation and tests. 'provided' scope keeps it off consumers' transitive class paths so
+          it never clashes with the host's Kryo (e.g. Pulse's kryo-shaded 4).
+        -->
+        <dependency>
+            <groupId>com.esotericsoftware</groupId>
+            <artifactId>kryo</artifactId>
+            <version>5.6.2</version>
+            <scope>provided</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.guava</groupId>
+            <artifactId>guava</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-api</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.auto.service</groupId>
+            <artifactId>auto-service</artifactId>
+        </dependency>
+
+        <!--Testing-->
+        <dependency>
+            <groupId>com.feedzai</groupId>
+            <artifactId>openml-utils</artifactId>
+            <type>test-jar</type>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-csv</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>commons-io</groupId>
+            <artifactId>commons-io</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>ch.qos.logback</groupId>
+            <artifactId>logback-classic</artifactId>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <plugins>
+            <!-- allow to reuse the objects created in test directory from outside of this module -->
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-jar-plugin</artifactId>
+                <executions>
+                    <execution>
+                        <goals>
+                            <goal>test-jar</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+    </build>
+</project>

openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostAlgorithms.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2026 Feedzai
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.feedzai.openml.provider.xgboost;
+
+import com.feedzai.openml.provider.descriptor.MLAlgorithmDescriptor;
+import com.feedzai.openml.provider.descriptor.MachineLearningAlgorithmType;
+import com.feedzai.openml.util.algorithm.MLAlgorithmEnum;
+
+import static com.feedzai.openml.util.algorithm.MLAlgorithmEnum.createDescriptor;
+
+/**
+ * Specifies the XGBoost algorithms that can be imported and trained through this provider.
+ *
+ * @since 1.0.0
+ */
+public enum XgboostAlgorithms implements MLAlgorithmEnum {
+
+    /**
+     * XGBoost binary classifier.
+     */
+    XGBOOST_BINARY_CLASSIFIER(createDescriptor(
+            "DMLC - XGBoost",
+            XgboostDescriptorUtil.PARAMS,
+            MachineLearningAlgorithmType.SUPERVISED_BINARY_CLASSIFICATION,
+            "https://xgboost.readthedocs.io/"
+    ));
+
+    /**
+     * {@link MLAlgorithmDescriptor} for this algorithm.
+     */
+    private final MLAlgorithmDescriptor descriptor;
+
+    /**
+     * Constructor.
+     *
+     * @param descriptor {@link MLAlgorithmDescriptor} for this algorithm.
+     */
+    XgboostAlgorithms(final MLAlgorithmDescriptor descriptor) {
+        this.descriptor = descriptor;
+    }
+
+    @Override
+    public MLAlgorithmDescriptor getAlgorithmDescriptor() {
+        return this.descriptor;
+    }
+}

openml-xgboost/src/main/java/com/feedzai/openml/provider/xgboost/XgboostClassificationModel.java

@@ -0,0 +1,159 @@
+/*
+ * Copyright 2026 Feedzai
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.feedzai.openml.provider.xgboost;
+
+import com.feedzai.openml.data.Instance;
+import com.feedzai.openml.data.schema.DatasetSchema;
+import com.feedzai.openml.model.ClassificationMLModel;
+import com.feedzai.openml.provider.exception.ModelLoadingException;
+import ml.dmlc.xgboost4j.java.Booster;
+import ml.dmlc.xgboost4j.java.XGBoostError;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.nio.file.Path;
+
+/**
+ * A classification model backed by a native XGBoost {@link Booster}, used for real-time single-instance
+ * scoring.
+ *
+ * <p>Scoring uses {@link Booster#inplace_predict(float[], int, int, float)} on a single-row feature
+ * vector, which avoids allocating a {@code DMatrix} per prediction. The native booster handle is not
+ * thread-safe, so predictions are serialized on a private lock (mirrors the H2O provider's approach).
+ *
+ * @since 1.0.0
+ */
+public class XgboostClassificationModel implements ClassificationMLModel {
+
+    /**
+     * Logger for this class.
+     */
+    private static final Logger logger = LoggerFactory.getLogger(XgboostClassificationModel.class);
+
+    /**
+     * Value used to signal a missing feature to XGBoost.
+     */
+    private static final float MISSING_VALUE = Float.NaN;
+
+    /**
+     * The native XGBoost booster.
+     */
+    private final Booster booster;
+
+    /**
+     * The schema the model uses.
+     */
+    private final DatasetSchema schema;
+
+    /**
+     * The number of predictive features expected by the model.
+     */
+    private final int numFeatures;
+
+    /**
+     * Lock serializing access to the non-thread-safe native booster during prediction.
+     */
+    private final Object predictLock = new Object();
+
+    /**
+     * Constructor.
+     *
+     * @param booster The trained/loaded native XGBoost booster.
+     * @param schema  The {@link DatasetSchema} the model uses.
+     */
+    XgboostClassificationModel(final Booster booster, final DatasetSchema schema) {
+        this.booster = booster;
+        this.schema = schema;
+        this.numFeatures = XgboostSchemaUtils.numFeatures(schema);
+    }
+
+    @Override
+    public double[] getClassDistribution(final Instance instance) {
+        final float[] row = XgboostSchemaUtils.featureRow(instance, this.schema);
+
+        final float[][] predictions;
+        try {
+            // The native booster handle is not thread-safe; serialize predictions.
+            synchronized (this.predictLock) {
+                predictions = this.booster.inplace_predict(row, 1, this.numFeatures, MISSING_VALUE);
+            }
+        } catch (final XGBoostError e) {
+            throw new RuntimeException("XGBoost failed to score the instance.", e);
+        }
+
+        return toClassDistribution(predictions[0]);
+    }
+
+    @Override
+    public int classify(final Instance instance) {
+        final double[] distribution = getClassDistribution(instance);
+
+        int argMax = 0;
+        for (int i = 1; i < distribution.length; i++) {
+            if (distribution[i] > distribution[argMax]) {
+                argMax = i;
+            }
+        }
+        return argMax;
+    }
+
+    @Override
+    public boolean save(final Path dir, final String name) {
+        try {
+            this.booster.saveModel(dir.resolve(XgboostModelCreator.MODEL_BINARY_RESOURCE_FILE_NAME).toString());
+            return true;
+        } catch (final XGBoostError e) {
+            logger.error("Failed to save XGBoost model {} to {}.", name, dir, e);
+            return false;
+        }
+    }
+
+    @Override
+    public DatasetSchema getSchema() {
+        return this.schema;
+    }
+
+    @Override
+    public void close() {
+        this.booster.dispose();
+    }
+
+    /**
+     * Converts a raw XGBoost prediction row into a class distribution aligned with the schema's target
+     * classes.
+     *
+     * <p>For binary objectives XGBoost outputs a single value - the probability of the positive class -
+     * which is expanded to {@code [1 - p, p]}. For multi-class objectives ({@code multi:softprob}) the
+     * per-class probability vector is returned as-is.
+     *
+     * @param prediction The raw prediction row for a single instance.
+     * @return The class distribution.
+     */
+    static double[] toClassDistribution(final float[] prediction) {
+        if (prediction.length == 1) {
+            final double positiveProbability = prediction[0];
+            return new double[]{1.0 - positiveProbability, positiveProbability};
+        }
+
+        final double[] distribution = new double[prediction.length];
+        for (int i = 0; i < prediction.length; i++) {
+            distribution[i] = prediction[i];
+        }
+        return distribution;
+    }
+}