diff --git a/sklearn2pmml/cross_reference/__init__.py b/sklearn2pmml/cross_reference/__init__.py index d1138e4..afe6309 100644 --- 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 index b344b8a..e4839b7 100644 --- 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] }