def test_type_error(self):
df = self.spark.createDataFrame([("a", 0),
("b", 0)]).toDF("features", "key")
keyedPCA = KeyedEstimator(sklearnEstimator=PCA())
self.assertRaises(TypeError, keyedPCA.fit, df)
df = self.spark.createDataFrame([(Vectors.dense([i]), [i], 0)
for i in range(10)])
df = df.toDF("features", "y", "key")
keyedLR = KeyedEstimator(sklearnEstimator=LinearRegression(), yCol="y")
self.assertRaises(TypeError, keyedLR.fit, df)
def test_correct_estimator_type(self):
self.checkEstimatorType(KeyedEstimator(sklearnEstimator=PCA()),
"transformer")
self.checkEstimatorType(
KeyedEstimator(sklearnEstimator=LinearRegression(), yCol="y"),
"predictor")
self.checkEstimatorType(KeyedEstimator(sklearnEstimator=DBSCAN()),
"clusterer")
self.checkEstimatorType(KeyedEstimator(sklearnEstimator=KMeans()),
"clusterer")
ke = KeyedEstimator(sklearnEstimator=KMeans(),
estimatorType="transformer")
self.checkEstimatorType(ke, "transformer")
custom = KeyedModelTests._CustomClusterer()
ke = KeyedEstimator(sklearnEstimator=custom)
self.checkEstimatorType(ke, "clusterer")
ke = KeyedEstimator(sklearnEstimator=custom,
estimatorType="transformer")
self.checkEstimatorType(ke, "transformer")
custom = KeyedModelTests._CustomTransformer()
self.checkEstimatorType(KeyedEstimator(sklearnEstimator=custom),
"transformer")
def test_invalid_argument(self):
self.assertRaises(ValueError, KeyedEstimator)
create = lambda: KeyedEstimator(sklearnEstimator=5)
self.assertRaises(ValueError, create)
class SomeUDC(object):
pass
create = lambda: KeyedEstimator(sklearnEstimator=SomeUDC())
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=PCA(), keyCols=[])
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=PCA(),
keyCols=["key", "estimator"])
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=PCA(),
xCol="estimator")
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=LinearRegression(),
yCol="estimator")
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=PCA(),
yCol="estimator")
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=LinearRegression())
self.assertRaises(ValueError, create)
def test_attr_error(self):
ke = KeyedEstimator(sklearnEstimator=PCA(), estimatorType="clusterer")
self.checkPredictionAttrError(ke)
ke = KeyedEstimator(sklearnEstimator=PCA(),
yCol="y",
estimatorType="predictor")
self.checkPredictionAttrError(ke)
ke = KeyedEstimator(sklearnEstimator=DBSCAN(),
estimatorType="transformer")
self.checkPredictionAttrError(ke)
ke = KeyedEstimator(sklearnEstimator=DBSCAN(),
yCol="y",
estimatorType="predictor")
self.checkPredictionAttrError(ke)
def test_surprise_key(self):
ke = KeyedEstimator(sklearnEstimator=PCA())
schema = StructType().add("features",
LongType()).add("key", LongType())
df = self.spark.createDataFrame([], schema)
km = ke.fit(df)
self.assertEqual(km.keyedModels.collect(), [])
self.assertEqual(km.keyedModels.dtypes,
[("key", LongType().simpleString()),
("estimator", "sklearn-estimator")])
df = self.spark.createDataFrame([(1, 2)], schema)
df = km.transform(df)
self.assertEqual(df.collect(), [(1, 2, None)])
self.assertEqual(df.dtypes, [("features", "bigint"), ("key", "bigint"),
("output", "vector")])
def test_surprise_key(self):
ke = KeyedEstimator(sklearnEstimator=PCA())
schema = StructType().add("features", LongType()).add("key", LongType())
df = self.spark.createDataFrame([], schema)
km = ke.fit(df)
self.assertEqual(km.keyedModels.collect(), [])
self.assertEqual(km.keyedModels.dtypes,
[("key", LongType().simpleString()),
("estimator", "sklearn-estimator")])
df = self.spark.createDataFrame([(1, 2)], schema)
df = km.transform(df)
self.assertEqual(df.collect(), [(1, 2, None)])
self.assertEqual(df.dtypes,
[("features", "bigint"),
("key", "bigint"),
("output", "vector")])
def checkKeyedModelEquivalent(self, minExamples, featureGen, labelGen, **kwargs):
NUSERS = 10
# featureGen() should generate a np rank-1 ndarray of equal length
# labelGen() should generate a scalar
assert (labelGen is not None) == ("yCol" in kwargs)
isPredictor = labelGen is not None
# sklearn's LinearRegression estimator is stable even if undetermined.
# User keys are just [0, NUSERS), repeated for each key if there are multiple columns.
# The i-th user has i examples.
keyCols = kwargs.get("keyCols", KeyedEstimator._paramSpecs["keyCols"]["default"])
outputCol = kwargs.get("outputCol", KeyedEstimator._paramSpecs["outputCol"]["default"])
xCol = kwargs.get("xCol", KeyedEstimator._paramSpecs["xCol"]["default"])
nExamplesPerUser = lambda i: max(minExamples, i + 1)
userKeys = [[i for _ in keyCols] for i in range(NUSERS)]
features = [[featureGen() for _ in range(nExamplesPerUser(i))] for i in range(NUSERS)]
useless = [["useless col" for _ in range(nExamplesPerUser(i))] for i in range(NUSERS)]
if isPredictor:
labels = [[labelGen() for _ in range(nExamplesPerUser(i))] for i in range(NUSERS)]
else:
labels = None
Xs = [np.vstack(x) for x in features]
ys = [np.array(y) for y in labels] if isPredictor else repeat(None)
localEstimators = [sklearn.base.clone(kwargs["sklearnEstimator"]).fit(X, y)
for X, y in zip(Xs, ys)]
expectedDF = pd.DataFrame(userKeys, columns=keyCols)
expectedDF["estimator"] = localEstimators
def flattenAndConvertNumpy(x):
return [Vectors.dense(i) if isinstance(i, np.ndarray) else i
for i in chain.from_iterable(x)]
inputDF = pd.DataFrame.from_dict(
{k: [i for i in range(NUSERS) for _ in range(nExamplesPerUser(i))] for k in keyCols})
inputDF[xCol] = flattenAndConvertNumpy(features)
inputDF["useless"] = flattenAndConvertNumpy(useless)
if labels:
inputDF[kwargs["yCol"]] = flattenAndConvertNumpy(labels)
inputDF = self.spark.createDataFrame(inputDF)
ke = KeyedEstimator(**kwargs)
km = ke.fit(inputDF)
actualDF = km.keyedModels.toPandas()
_assertPandasAlmostEqual(actualDF, expectedDF, keyCols)
# Test users with different amounts of points.
nTestPerUser = lambda i: NUSERS // 4 if i < NUSERS // 2 else NUSERS * 3 // 4
testFeatures = [[featureGen() for _ in range(nTestPerUser(i))] for i in range(NUSERS)]
# "useless" column has nothing to do with computation, but is essential for keeping order
# the same between the spark and non-spark versions
useless = [range(nTestPerUser(i)) for i in range(NUSERS)]
inputDF = pd.DataFrame.from_dict(
{k: [i for i in range(NUSERS) for _ in range(nTestPerUser(i))] for k in keyCols})
inputDF[xCol] = flattenAndConvertNumpy(testFeatures)
inputDF["useless"] = flattenAndConvertNumpy(useless)
estimatorType = km.sklearnEstimatorType # tested to be correct elsewhere
def makeOutput(estimator, X):
if estimatorType == "transformer":
return estimator.transform(X)
else:
assert estimatorType == "predictor" or estimatorType == "clusterer"
return estimator.predict(X).tolist()
Xs = [np.vstack(x) for x in testFeatures]
expectedOutput = map(makeOutput, localEstimators, Xs)
expectedDF = inputDF.copy(deep=True)
expectedDF[outputCol] = flattenAndConvertNumpy(expectedOutput)
inputDF = self.spark.createDataFrame(inputDF)
actualDF = km.transform(inputDF).toPandas()
_assertPandasAlmostEqual(actualDF, expectedDF, keyCols + ["useless"])
def checkKeyedModelEquivalent(self, minExamples, featureGen, labelGen,
**kwargs):
NUSERS = 10
# featureGen() should generate a np rank-1 ndarray of equal length
# labelGen() should generate a scalar
assert (labelGen is not None) == ("yCol" in kwargs)
isPredictor = labelGen is not None
# sklearn's LinearRegression estimator is stable even if undetermined.
# User keys are just [0, NUSERS), repeated for each key if there are multiple columns.
# The i-th user has i examples.
keyCols = kwargs.get("keyCols",
KeyedEstimator._paramSpecs["keyCols"]["default"])
outputCol = kwargs.get(
"outputCol", KeyedEstimator._paramSpecs["outputCol"]["default"])
xCol = kwargs.get("xCol",
KeyedEstimator._paramSpecs["xCol"]["default"])
nExamplesPerUser = lambda i: max(minExamples, i + 1)
userKeys = [[i for _ in keyCols] for i in range(NUSERS)]
features = [[featureGen() for _ in range(nExamplesPerUser(i))]
for i in range(NUSERS)]
useless = [["useless col" for _ in range(nExamplesPerUser(i))]
for i in range(NUSERS)]
if isPredictor:
labels = [[labelGen() for _ in range(nExamplesPerUser(i))]
for i in range(NUSERS)]
else:
labels = None
Xs = [np.vstack(x) for x in features]
ys = [np.array(y) for y in labels] if isPredictor else repeat(None)
localEstimators = [
sklearn.base.clone(kwargs["sklearnEstimator"]).fit(X, y)
for X, y in zip(Xs, ys)
]
expectedDF = pd.DataFrame(userKeys, columns=keyCols)
expectedDF["estimator"] = localEstimators
def flattenAndConvertNumpy(x):
return [
Vectors.dense(i) if isinstance(i, np.ndarray) else i
for i in chain.from_iterable(x)
]
inputDF = pd.DataFrame.from_dict({
k: [i for i in range(NUSERS) for _ in range(nExamplesPerUser(i))]
for k in keyCols
})
inputDF[xCol] = flattenAndConvertNumpy(features)
inputDF["useless"] = flattenAndConvertNumpy(useless)
if labels:
inputDF[kwargs["yCol"]] = flattenAndConvertNumpy(labels)
inputDF = self.spark.createDataFrame(inputDF)
ke = KeyedEstimator(**kwargs)
km = ke.fit(inputDF)
actualDF = km.keyedModels.toPandas()
_assertPandasAlmostEqual(actualDF, expectedDF, keyCols)
# Test users with different amounts of points.
nTestPerUser = lambda i: NUSERS // 4 if i < NUSERS // 2 else NUSERS * 3 // 4
testFeatures = [[featureGen() for _ in range(nTestPerUser(i))]
for i in range(NUSERS)]
# "useless" column has nothing to do with computation, but is essential for keeping order
# the same between the spark and non-spark versions
useless = [range(nTestPerUser(i)) for i in range(NUSERS)]
inputDF = pd.DataFrame.from_dict({
k: [i for i in range(NUSERS) for _ in range(nTestPerUser(i))]
for k in keyCols
})
inputDF[xCol] = flattenAndConvertNumpy(testFeatures)
inputDF["useless"] = flattenAndConvertNumpy(useless)
estimatorType = km.sklearnEstimatorType # tested to be correct elsewhere
def makeOutput(estimator, X):
if estimatorType == "transformer":
return estimator.transform(X)
else:
assert estimatorType == "predictor" or estimatorType == "clusterer"
return estimator.predict(X).tolist()
Xs = [np.vstack(x) for x in testFeatures]
expectedOutput = map(makeOutput, localEstimators, Xs)
expectedDF = inputDF.copy(deep=True)
expectedDF[outputCol] = flattenAndConvertNumpy(expectedOutput)
inputDF = self.spark.createDataFrame(inputDF)
actualDF = km.transform(inputDF).toPandas()
_assertPandasAlmostEqual(actualDF, expectedDF, keyCols + ["useless"])
def test_invalid_argument(self):
# Need to specify sklearnEstimator
self.assertRaises(ValueError, KeyedEstimator)
# sklearnEstimator must be a sklearn.base.Estimator
create = lambda: KeyedEstimator(sklearnEstimator=5)
self.assertRaises(ValueError, create)
class SomeUDC(object):
pass
create = lambda: KeyedEstimator(sklearnEstimator=SomeUDC())
self.assertRaises(ValueError, create)
# Must have fit()
create = lambda: KeyedEstimator(sklearnEstimator=KeyedModelTests.
_CustomMissingFit())
self.assertRaises(AttributeError, create)
# Must have key columns
create = lambda: KeyedEstimator(sklearnEstimator=PCA(), keyCols=[])
self.assertRaises(ValueError, create)
# Columns can't have "estimator" name in them
create = lambda: KeyedEstimator(sklearnEstimator=PCA(),
keyCols=["key", "estimator"])
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=PCA(),
xCol="estimator")
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=LinearRegression(),
yCol="estimator")
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=PCA(),
yCol="estimator")
self.assertRaises(ValueError, create)
# Presence of yCol requires predictor
create = lambda: KeyedEstimator(sklearnEstimator=LinearRegression(),
yCol="y",
estimatorType="transformer")
self.assertRaises(ValueError, create)
create = lambda: KeyedEstimator(sklearnEstimator=LinearRegression(),
yCol="y",
estimatorType="clusterer")
self.assertRaises(ValueError, create)
# estimatorType must be one of the three options
create = lambda: KeyedEstimator(sklearnEstimator=PCA(),
estimatorType="regressor")
self.assertRaises(ValueError, create)
def test_defaults(self):
ke = KeyedEstimator(sklearnEstimator=PCA())
for paramName, paramSpec in KeyedEstimator._paramSpecs.items():
if "default" in paramSpec:
self.assertEqual(paramSpec["default"],
ke.getOrDefault(paramName))
def test_create_no_errors(self):
KeyedEstimator(sklearnEstimator=PCA())
KeyedEstimator(sklearnEstimator=LinearRegression(), yCol="yCol")
KeyedEstimator(sklearnEstimator=KeyedModelTests._CustomTransformer())
def test_defaults(self):
ke = KeyedEstimator(sklearnEstimator=PCA())
for paramName, paramSpec in KeyedEstimator._paramSpecs.items():
if "default" in paramSpec:
self.assertEqual(paramSpec["default"], ke.getOrDefault(paramName))
