Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
135 changes: 135 additions & 0 deletions openml-xgboost/pom.xml
Original file line number Diff line number Diff line change
@@ -0,0 +1,135 @@
<?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, but we only use the pure-Java API; scala-compiler is
not needed at all, so it is excluded to keep the classpath lean. scala-library is kept as a
harmless transitive runtime dependency.
-->
<dependency>
<groupId>ml.dmlc</groupId>
<artifactId>xgboost4j_2.13</artifactId>
<version>${xgboost.version}</version>
<exclusions>
<exclusion>
<groupId>org.scala-lang</groupId>
<artifactId>scala-compiler</artifactId>
</exclusion>
</exclusions>
</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>
Original file line number Diff line number Diff line change
@@ -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(
"XGBoost Binary Classifier",
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;
}
}
Original file line number Diff line number Diff line change
@@ -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.
*/
private 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;
}
}
Loading
Loading