Added 'input_float' model transformation option. Fixes #128

jpmml · Apr 27, 2023 · 433ce1c · 433ce1c
1 parent 2487206
commit 433ce1c
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 0 deletions.
diff --git a/pmml-sparkml-xgboost/src/main/java/org/jpmml/sparkml/xgboost/BoosterUtil.java b/pmml-sparkml-xgboost/src/main/java/org/jpmml/sparkml/xgboost/BoosterUtil.java
@@ -23,12 +23,17 @@
 import java.io.InputStream;
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.function.Function;
 
 import ml.dmlc.xgboost4j.scala.Booster;
 import ml.dmlc.xgboost4j.scala.spark.params.GeneralParams;
 import org.apache.spark.ml.Model;
 import org.apache.spark.ml.param.shared.HasPredictionCol;
+import org.dmg.pmml.DataType;
+import org.dmg.pmml.Field;
 import org.dmg.pmml.mining.MiningModel;
+import org.jpmml.converter.ContinuousFeature;
+import org.jpmml.converter.Feature;
 import org.jpmml.converter.Schema;
 import org.jpmml.sparkml.ModelConverter;
 import org.jpmml.xgboost.HasXGBoostOptions;
@@ -60,6 +65,41 @@ public <M extends Model<M> & HasPredictionCol & GeneralParams, C extends ModelCo
 			throw new RuntimeException(ioe);
 		}
 
+		Boolean inputFloat = (Boolean)converter.getOption(HasSparkMLXGBoostOptions.OPTION_INPUT_FLOAT, null);
+		if((Boolean.TRUE).equals(inputFloat)){
+			Function<Feature, Feature> function = new Function<Feature, Feature>(){
+
+				@Override
+				public Feature apply(Feature feature){
+
+					if(feature instanceof ContinuousFeature){
+						ContinuousFeature continuousFeature = (ContinuousFeature)feature;
+
+						DataType dataType = continuousFeature.getDataType();
+						switch(dataType){
+							case INTEGER:
+							case FLOAT:
+								break;
+							case DOUBLE:
+								{
+									Field<?> field = continuousFeature.getField();
+
+									field.setDataType(DataType.FLOAT);
+
+									return new ContinuousFeature(continuousFeature.getEncoder(), field);
+								}
+							default:
+								break;
+						}
+					}
+
+					return feature;
+				}
+			};
+
+			schema = schema.toTransformedSchema(function);
+		}
+
 		Float missing = model.getMissing();
 		if(missing.isNaN()){
 			missing = null;

diff --git a/pmml-sparkml-xgboost/src/main/java/org/jpmml/sparkml/xgboost/HasSparkMLXGBoostOptions.java b/pmml-sparkml-xgboost/src/main/java/org/jpmml/sparkml/xgboost/HasSparkMLXGBoostOptions.java
@@ -22,4 +22,6 @@
 import org.jpmml.xgboost.HasXGBoostOptions;
 
 public interface HasSparkMLXGBoostOptions extends HasSparkMLOptions, HasXGBoostOptions {
+
+	String OPTION_INPUT_FLOAT = "input_float";
 }
diff --git a/pmml-sparkml-xgboost/src/test/java/org/jpmml/sparkml/xgboost/testing/XGBoostTest.java b/pmml-sparkml-xgboost/src/test/java/org/jpmml/sparkml/xgboost/testing/XGBoostTest.java
@@ -35,6 +35,7 @@
 import org.jpmml.model.visitors.AbstractVisitor;
 import org.jpmml.sparkml.testing.SparkMLEncoderBatch;
 import org.jpmml.sparkml.testing.SparkMLEncoderBatchTest;
+import org.jpmml.sparkml.xgboost.HasSparkMLXGBoostOptions;
 import org.jpmml.xgboost.HasXGBoostOptions;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -61,7 +62,10 @@ public XGBoostTest getArchiveBatchTest(){
 			public List<Map<String, Object>> getOptionsMatrix(){
 				Map<String, Object> options = new LinkedHashMap<>();
 
+				options.put(HasSparkMLXGBoostOptions.OPTION_INPUT_FLOAT, new Boolean[]{false, true});
+
 				options.put(HasXGBoostOptions.OPTION_COMPACT, new Boolean[]{false, true});
+				options.put(HasXGBoostOptions.OPTION_PRUNE, false);
 
 				return OptionsUtil.generateOptionsMatrix(options);
 			}

diff --git a/pmml-sparkml-xgboost/src/test/resources/pipeline/XGBoostHousing.zip b/pmml-sparkml-xgboost/src/test/resources/pipeline/XGBoostHousing.zip
diff --git a/pmml-sparkml-xgboost/src/test/resources/pipeline/XGBoostIris.zip b/pmml-sparkml-xgboost/src/test/resources/pipeline/XGBoostIris.zip