Filtered valid category values

jpmml · Jan 7, 2024 · 7b37433 · 7b37433
1 parent 20cfa72
commit 7b37433
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sklearn2pmml/preprocessing/__init__.py b/sklearn2pmml/preprocessing/__init__.py
@@ -98,7 +98,8 @@ def __init__(self, dtype):
 
 	def fit(self, X, y = None):
 		if self.dtype == "category":
-			categories = numpy.unique(X)
+			nonmissing_mask = pandas.notnull(X)
+			categories = numpy.unique(X[nonmissing_mask])
 			self.dtype_ = CategoricalDtype(categories, ordered = False)
 		else:
 			self.dtype_ = self.dtype