Expanded commit 69fd0f7

jpmml · Sep 20, 2023 · d5357b2 · d5357b2
1 parent ea17da5
commit d5357b2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/sklearn2pmml/cross_reference/__init__.py b/sklearn2pmml/cross_reference/__init__.py
@@ -1,3 +1,4 @@
+from pandas import DataFrame
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.pipeline import FeatureUnion
 from sklearn2pmml.preprocessing import IdentityTransformer
@@ -36,7 +37,10 @@ def fit(self, X, y = None):
 
 	def transform(self, X):
 		for idx, name in enumerate(self.names):
-			x = X[:, idx]
+			if isinstance(X, DataFrame):
+				x = X.iloc[:, idx]
+			else:
+				x = X[:, idx]
 			self.memory[name] = x.copy()
 		return numpy.empty(shape = (X.shape[0], 0), dtype = int)
 

diff --git a/sklearn2pmml/cross_reference/tests/__init__.py b/sklearn2pmml/cross_reference/tests/__init__.py
@@ -6,7 +6,7 @@
 
 class MemorizerTest(TestCase):
 
-	def test_fit_transform(self):
+	def test_transform(self):
 		memory = dict()
 		self.assertEqual(0, len(memory))
 		memorizer = Memorizer(memory, ["int"])
@@ -31,7 +31,7 @@ def test_fit_transform(self):
 
 class RecallerTest(TestCase):
 
-	def test_fit_transform(self):
+	def test_transform(self):
 		memory = {
 			"int": [-1, 1]
 		}