python/pyspark/ml/tests/test_feature.py - master - githubmirror/spark

python/pyspark/ml/tests/test_feature.py
1 782 строки65 KB
Dongjoon Hyun
[SPARK-57977][PYTHON] Remove obsolete encoding declarations from Python files
07 июл 2026, 02:55
07 июл 2026, 02:556cb0e50
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import tempfile
from typing import List, Tuple, Any

import numpy as np

from pyspark.ml.feature import (
    DCT,
    Binarizer,
    Bucketizer,
    QuantileDiscretizer,
    CountVectorizer,
    CountVectorizerModel,
    OneHotEncoder,
    OneHotEncoderModel,
    FeatureHasher,
    ElementwiseProduct,
    HashingTF,
    IDF,
    IDFModel,
    Imputer,
    ImputerModel,
    NGram,
    Normalizer,
    Interaction,
    RFormula,
    RFormulaModel,
    Tokenizer,
    SQLTransformer,
    RegexTokenizer,
    StandardScaler,
    StandardScalerModel,
    MaxAbsScaler,
    MaxAbsScalerModel,
    MinMaxScaler,
    MinMaxScalerModel,
    RobustScaler,
    RobustScalerModel,
    ChiSqSelector,
    ChiSqSelectorModel,
    UnivariateFeatureSelector,
    UnivariateFeatureSelectorModel,
    VarianceThresholdSelector,
    VarianceThresholdSelectorModel,
    StopWordsRemover,
    StringIndexer,
    StringIndexerModel,
    VectorIndexer,
    VectorIndexerModel,
    TargetEncoder,
    TargetEncoderModel,
    VectorSizeHint,
    VectorSlicer,
    VectorAssembler,
    PCA,
    PCAModel,
    Word2Vec,
    Word2VecModel,
    BucketedRandomProjectionLSH,
    BucketedRandomProjectionLSHModel,
    MinHashLSH,
    MinHashLSHModel,
    IndexToString,
    PolynomialExpansion,
)
from pyspark.ml.linalg import DenseVector, SparseVector, Vectors
from pyspark.sql import Row
from pyspark.testing.sqlutils import ReusedSQLTestCase


class FeatureTestsMixin:
    def test_polynomial_expansion(self):
        df = self.spark.createDataFrame([(Vectors.dense([0.5, 2.0]),)], ["dense"])
        px = PolynomialExpansion(degree=2)
        px.setInputCol("dense")
        px.setOutputCol("expanded")
        self.assertTrue(
            np.allclose(
                px.transform(df).head().expanded.toArray(), [0.5, 0.25, 2.0, 1.0, 4.0], atol=1e-4
            )
        )

        def check(p: PolynomialExpansion) -> None:
            self.assertEqual(p.getInputCol(), "dense")
            self.assertEqual(p.getOutputCol(), "expanded")
            self.assertEqual(p.getDegree(), 2)

        check(px)

        # save & load
        with tempfile.TemporaryDirectory(prefix="px") as d:
            px.write().overwrite().save(d)
            px2 = PolynomialExpansion.load(d)
            self.assertEqual(str(px), str(px2))
            check(px2)

    def test_index_string(self):
        dataset = self.spark.createDataFrame(
            [
                (0, "a"),
                (1, "b"),
                (2, "c"),
                (3, "a"),
                (4, "a"),
                (5, "c"),
            ],
            ["id", "label"],
        )

        indexer = StringIndexer(inputCol="label", outputCol="labelIndex").fit(dataset)
        transformed = indexer.transform(dataset)
        idx2str = (
            IndexToString()
            .setInputCol("labelIndex")
            .setOutputCol("sameLabel")
            .setLabels(indexer.labels)
        )

        def check(t: IndexToString) -> None:
            self.assertEqual(t.getInputCol(), "labelIndex")
            self.assertEqual(t.getOutputCol(), "sameLabel")
            self.assertEqual(t.getLabels(), indexer.labels)

        check(idx2str)

        ret = idx2str.transform(transformed)
        self.assertEqual(
            sorted(ret.schema.names), sorted(["id", "label", "labelIndex", "sameLabel"])
        )

        rows = ret.select("label", "sameLabel").collect()
        for r in rows:
            self.assertEqual(r.label, r.sameLabel)

        # save & load
        with tempfile.TemporaryDirectory(prefix="index_string") as d:
            idx2str.write().overwrite().save(d)
            idx2str2 = IndexToString.load(d)
            self.assertEqual(str(idx2str), str(idx2str2))
            check(idx2str2)

    def test_dct(self):
        df = self.spark.createDataFrame([(Vectors.dense([5.0, 8.0, 6.0]),)], ["vec"])
        dct = DCT()
        dct.setInverse(False)
        dct.setInputCol("vec")
        dct.setOutputCol("resultVec")

        self.assertFalse(dct.getInverse())
        self.assertEqual(dct.getInputCol(), "vec")
        self.assertEqual(dct.getOutputCol(), "resultVec")

        output = dct.transform(df)
        self.assertEqual(output.columns, ["vec", "resultVec"])
        self.assertEqual(output.count(), 1)
        self.assertTrue(
            np.allclose(
                output.head().resultVec.toArray(),
                [10.96965511, -0.70710678, -2.04124145],
                atol=1e-4,
            )
        )

        # save & load
        with tempfile.TemporaryDirectory(prefix="dct") as d:
            dct.write().overwrite().save(d)
            dct2 = DCT.load(d)
            self.assertEqual(str(dct), str(dct2))

    def test_string_indexer(self):
        df = (
            self.spark.createDataFrame(
                [
                    (1, "a", "e"),
                    (2, "b", "f"),
                    (3, "c", "e"),
                    (4, "a", "f"),
                    (5, "a", "f"),
                    (6, "c", "f"),
                ],
                ["id", "label1", "label2"],
            )
            .coalesce(1)
            .sortWithinPartitions("id")
        )
        # single input
        si = StringIndexer(inputCol="label1", outputCol="index1")
        model = si.fit(df.select("label1"))
        self.assertEqual(si.uid, model.uid)
        self.assertEqual(model.labels, list(model.labelsArray[0]))

        # read/write
        with tempfile.TemporaryDirectory(prefix="string_indexer") as tmp_dir:
            si.write().overwrite().save(tmp_dir)
            si2 = StringIndexer.load(tmp_dir)
            self.assertEqual(str(si), str(si2))
            self.assertEqual(si.getInputCol(), "label1")
            self.assertEqual(si2.getInputCol(), "label1")

            model.write().overwrite().save(tmp_dir)
            model2 = StringIndexerModel.load(tmp_dir)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model.getInputCol(), "label1")
            self.assertEqual(model.getOutputCol(), "index1")
            self.assertEqual(model2.getInputCol(), "label1")

        indexed_df = model.transform(df.select("label1"))
        self.assertEqual(sorted(indexed_df.schema.names), sorted(["label1", "index1"]))

        def check_a_b(result: List[Tuple[Any, Any]]) -> None:
            self.assertTrue(result[0][0] == "a" and result[1][0] == "b" and result[2][0] == "c")
            sorted_value = sorted([v for _, v in result])
            self.assertEqual(sorted_value, [0.0, 1.0, 2.0])

        check_a_b(sorted(set([(i[0], i[1]) for i in indexed_df.collect()]), key=lambda x: x[0]))

        # multiple inputs
        input_cols = ["label1", "label2"]
        output_cols = ["index1", "index2"]
        si = StringIndexer(inputCols=input_cols, outputCols=output_cols)
        model = si.fit(df.select(*input_cols))
        self.assertEqual(model.getInputCols(), input_cols)
        self.assertEqual(model.getOutputCols(), output_cols)

        indexed_df = model.transform(df.select(*input_cols))
        self.assertEqual(
            sorted(indexed_df.schema.names), sorted(["label1", "index1", "label2", "index2"])
        )

        rows = indexed_df.collect()
        check_a_b(sorted(set([(i[0], i[2]) for i in rows]), key=lambda x: x[0]))

        # check e f
        result = sorted(set([(i[1], i[3]) for i in rows]), key=lambda x: x[0])
        self.assertTrue(result[0][0] == "e" and result[1][0] == "f")
        sorted_value = sorted([v for _, v in result])
        self.assertEqual(sorted_value, [0.0, 1.0])

    def test_vector_indexer(self):
        spark = self.spark
        df = spark.createDataFrame(
            [
                (Vectors.dense([-1.0, 0.0]),),
                (Vectors.dense([0.0, 1.0]),),
                (Vectors.dense([0.0, 2.0]),),
            ],
            ["a"],
        )

        indexer = VectorIndexer(maxCategories=2, inputCol="a")
        indexer.setOutputCol("indexed")
        self.assertEqual(indexer.getMaxCategories(), 2)
        self.assertEqual(indexer.getInputCol(), "a")
        self.assertEqual(indexer.getOutputCol(), "indexed")

        model = indexer.fit(df)
        self.assertEqual(indexer.uid, model.uid)
        self.assertEqual(model.numFeatures, 2)

        categoryMaps = model.categoryMaps
        self.assertEqual(categoryMaps, {0: {0.0: 0, -1.0: 1}}, categoryMaps)

        output = model.transform(df)
        self.assertEqual(output.columns, ["a", "indexed"])
        self.assertEqual(output.count(), 3)

        # save & load
        with tempfile.TemporaryDirectory(prefix="vector_indexer") as d:
            indexer.write().overwrite().save(d)
            indexer2 = VectorIndexer.load(d)
            self.assertEqual(str(indexer), str(indexer2))
            self.assertEqual(indexer2.getOutputCol(), "indexed")

            model.write().overwrite().save(d)
            model2 = VectorIndexerModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model2.getOutputCol(), "indexed")

    def test_elementwise_product(self):
        spark = self.spark
        df = spark.createDataFrame([(Vectors.dense([2.0, 1.0, 3.0]),)], ["values"])

        ep = ElementwiseProduct()
        ep.setScalingVec(Vectors.dense([1.0, 2.0, 3.0]))
        ep.setInputCol("values")
        ep.setOutputCol("eprod")

        self.assertEqual(ep.getScalingVec(), Vectors.dense([1.0, 2.0, 3.0]))
        self.assertEqual(ep.getInputCol(), "values")
        self.assertEqual(ep.getOutputCol(), "eprod")

        output = ep.transform(df)
        self.assertEqual(output.columns, ["values", "eprod"])
        self.assertEqual(output.count(), 1)

        # save & load
        with tempfile.TemporaryDirectory(prefix="elementwise_product") as d:
            ep.write().overwrite().save(d)
            ep2 = ElementwiseProduct.load(d)
            self.assertEqual(str(ep), str(ep2))
            self.assertEqual(ep2.getOutputCol(), "eprod")

    def test_pca(self):
        df = self.spark.createDataFrame(
            [
                (Vectors.sparse(5, [(1, 1.0), (3, 7.0)]),),
                (Vectors.dense([2.0, 0.0, 3.0, 4.0, 5.0]),),
                (Vectors.dense([4.0, 0.0, 0.0, 6.0, 7.0]),),
            ],
            ["features"],
        )
        pca = PCA(k=2, inputCol="features", outputCol="pca_features")

        model = pca.fit(df)
        self.assertTrue(np.allclose(model.pc.toArray()[0], [-0.44859172, -0.28423808], atol=1e-4))
        self.assertEqual(pca.uid, model.uid)
        self.assertEqual(model.getK(), 2)
        self.assertTrue(
            np.allclose(model.explainedVariance.toArray(), [0.79439, 0.20560], atol=1e-4)
        )
        model.setOutputCol("output")
        # Transform the data using the PCA model
        transformed_df = model.transform(df)
        self.assertTrue(
            np.allclose(
                transformed_df.collect()[0].output.toArray(), [1.64857, -4.013282], atol=1e-4
            )
        )

        # read/write
        with tempfile.TemporaryDirectory(prefix="pca") as tmp_dir:
            pca.write().overwrite().save(tmp_dir)
            pca2 = PCA.load(tmp_dir)
            self.assertEqual(str(pca), str(pca2))
            self.assertEqual(pca.getInputCol(), "features")
            self.assertEqual(pca2.getInputCol(), "features")

            model.write().overwrite().save(tmp_dir)
            model2 = PCAModel.load(tmp_dir)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model.getInputCol(), "features")
            self.assertEqual(model2.getInputCol(), "features")

    def test_vector_assembler(self):
        # Create a DataFrame
        df = (
            self.spark.createDataFrame(
                [
                    (1, 5.0, 6.0, 7.0),
                    (2, 1.0, 2.0, None),
                    (3, 3.0, float("nan"), 4.0),
                ],
                ["index", "a", "b", "c"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
        )

        # Initialize VectorAssembler
        vec_assembler = VectorAssembler(outputCol="features").setInputCols(["a", "b", "c"])
        output = vec_assembler.transform(df)
        self.assertEqual(output.columns, ["index", "a", "b", "c", "features"])
        self.assertEqual(output.head().features, Vectors.dense([5.0, 6.0, 7.0]))

        # Set custom parameters and transform the DataFrame
        params = {vec_assembler.inputCols: ["b", "a"], vec_assembler.outputCol: "vector"}
        self.assertEqual(
            vec_assembler.transform(df, params).head().vector, Vectors.dense([6.0, 5.0])
        )

        # read/write
        with tempfile.TemporaryDirectory(prefix="vector_assembler") as tmp_dir:
            vec_assembler.write().overwrite().save(tmp_dir)
            vec_assembler2 = VectorAssembler.load(tmp_dir)
            self.assertEqual(str(vec_assembler), str(vec_assembler2))

        # Initialize a new VectorAssembler with handleInvalid="keep"
        vec_assembler3 = VectorAssembler(
            inputCols=["a", "b", "c"], outputCol="features", handleInvalid="keep"
        )
        self.assertEqual(vec_assembler3.transform(df).count(), 3)

        # Update handleInvalid to "skip" and transform the DataFrame
        vec_assembler3.setParams(handleInvalid="skip")
        self.assertEqual(vec_assembler3.transform(df).count(), 1)

    def test_standard_scaler(self):
        df = (
            self.spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense([0.0])),
                    (2, 2.0, Vectors.dense([2.0])),
                    (3, 3.0, Vectors.sparse(1, [(0, 3.0)])),
                ],
                ["index", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("features")
        )
        scaler = StandardScaler(inputCol="features", outputCol="scaled")
        self.assertEqual(scaler.getInputCol(), "features")
        self.assertEqual(scaler.getOutputCol(), "scaled")

        model = scaler.fit(df)
        self.assertEqual(scaler.uid, model.uid)
        self.assertTrue(np.allclose(model.mean.toArray(), [1.66666667], atol=1e-4))
        self.assertTrue(np.allclose(model.std.toArray(), [1.52752523], atol=1e-4))

        output = model.transform(df)
        self.assertEqual(output.columns, ["features", "scaled"])
        self.assertEqual(output.count(), 3)

        # Model save & load
        with tempfile.TemporaryDirectory(prefix="standard_scaler") as d:
            scaler.write().overwrite().save(d)
            scaler2 = StandardScaler.load(d)
            self.assertEqual(str(scaler), str(scaler2))
            self.assertEqual(scaler2.getOutputCol(), "scaled")

            model.write().overwrite().save(d)
            model2 = StandardScalerModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model2.getOutputCol(), "scaled")

    def test_maxabs_scaler(self):
        df = (
            self.spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense([0.0])),
                    (2, 2.0, Vectors.dense([2.0])),
                    (3, 3.0, Vectors.sparse(1, [(0, 3.0)])),
                ],
                ["index", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("features")
        )

        scaler = MaxAbsScaler(inputCol="features", outputCol="scaled")
        self.assertEqual(scaler.getInputCol(), "features")
        self.assertEqual(scaler.getOutputCol(), "scaled")

        model = scaler.fit(df)
        self.assertEqual(scaler.uid, model.uid)
        self.assertTrue(np.allclose(model.maxAbs.toArray(), [3.0], atol=1e-4))

        output = model.transform(df)
        self.assertEqual(output.columns, ["features", "scaled"])
        self.assertEqual(output.count(), 3)

        # save & load
        with tempfile.TemporaryDirectory(prefix="standard_scaler") as d:
            scaler.write().overwrite().save(d)
            scaler2 = MaxAbsScaler.load(d)
            self.assertEqual(str(scaler), str(scaler2))
            self.assertEqual(scaler2.getOutputCol(), "scaled")

            model.write().overwrite().save(d)
            model2 = MaxAbsScalerModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model2.getOutputCol(), "scaled")

    def test_minmax_scaler(self):
        df = (
            self.spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense([0.0])),
                    (2, 2.0, Vectors.dense([2.0])),
                    (3, 3.0, Vectors.sparse(1, [(0, 3.0)])),
                ],
                ["index", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("features")
        )

        scaler = MinMaxScaler(inputCol="features", outputCol="scaled")
        self.assertEqual(scaler.getInputCol(), "features")
        self.assertEqual(scaler.getOutputCol(), "scaled")

        model = scaler.fit(df)
        self.assertEqual(scaler.uid, model.uid)
        self.assertTrue(np.allclose(model.originalMax.toArray(), [3.0], atol=1e-4))
        self.assertTrue(np.allclose(model.originalMin.toArray(), [0.0], atol=1e-4))

        output = model.transform(df)
        self.assertEqual(output.columns, ["features", "scaled"])
        self.assertEqual(output.count(), 3)

        # save & load
        with tempfile.TemporaryDirectory(prefix="min_max_scaler") as d:
            scaler.write().overwrite().save(d)
            scaler2 = MinMaxScaler.load(d)
            self.assertEqual(str(scaler), str(scaler2))
            self.assertEqual(scaler2.getOutputCol(), "scaled")

            model.write().overwrite().save(d)
            model2 = MinMaxScalerModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model2.getOutputCol(), "scaled")

    def test_robust_scaler(self):
        df = (
            self.spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense([0.0])),
                    (2, 2.0, Vectors.dense([2.0])),
                    (3, 3.0, Vectors.sparse(1, [(0, 3.0)])),
                ],
                ["index", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("features")
        )

        scaler = RobustScaler(inputCol="features", outputCol="scaled")
        self.assertEqual(scaler.getInputCol(), "features")
        self.assertEqual(scaler.getOutputCol(), "scaled")

        model = scaler.fit(df)
        self.assertEqual(scaler.uid, model.uid)
        self.assertTrue(np.allclose(model.range.toArray(), [3.0], atol=1e-4))
        self.assertTrue(np.allclose(model.median.toArray(), [2.0], atol=1e-4))

        output = model.transform(df)
        self.assertEqual(output.columns, ["features", "scaled"])
        self.assertEqual(output.count(), 3)

        # save & load
        with tempfile.TemporaryDirectory(prefix="robust_scaler") as d:
            scaler.write().overwrite().save(d)
            scaler2 = RobustScaler.load(d)
            self.assertEqual(str(scaler), str(scaler2))
            self.assertEqual(scaler2.getOutputCol(), "scaled")

            model.write().overwrite().save(d)
            model2 = RobustScalerModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model2.getOutputCol(), "scaled")

    def test_chi_sq_selector(self):
        df = self.spark.createDataFrame(
            [
                (Vectors.dense([0.0, 0.0, 18.0, 1.0]), 1.0),
                (Vectors.dense([0.0, 1.0, 12.0, 0.0]), 0.0),
                (Vectors.dense([1.0, 0.0, 15.0, 0.1]), 0.0),
            ],
            ["features", "label"],
        )

        selector = ChiSqSelector(numTopFeatures=1, outputCol="selectedFeatures")
        self.assertEqual(selector.getNumTopFeatures(), 1)
        self.assertEqual(selector.getOutputCol(), "selectedFeatures")

        model = selector.fit(df)
        self.assertEqual(selector.uid, model.uid)
        self.assertEqual(model.selectedFeatures, [2])

        output = model.transform(df)
        self.assertEqual(output.columns, ["features", "label", "selectedFeatures"])
        self.assertEqual(output.count(), 3)

        # save & load
        with tempfile.TemporaryDirectory(prefix="chi_sq_selector") as d:
            selector.write().overwrite().save(d)
            selector2 = ChiSqSelector.load(d)
            self.assertEqual(str(selector), str(selector2))

            model.write().overwrite().save(d)
            model2 = ChiSqSelectorModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_univariate_selector(self):
        df = self.spark.createDataFrame(
            [
                (Vectors.dense([0.0, 0.0, 18.0, 1.0]), 1.0),
                (Vectors.dense([0.0, 1.0, 12.0, 0.0]), 0.0),
                (Vectors.dense([1.0, 0.0, 15.0, 0.1]), 0.0),
            ],
            ["features", "label"],
        )

        selector = UnivariateFeatureSelector(outputCol="selectedFeatures")
        selector.setFeatureType("continuous").setLabelType("categorical").setSelectionThreshold(1)
        self.assertEqual(selector.getFeatureType(), "continuous")
        self.assertEqual(selector.getLabelType(), "categorical")
        self.assertEqual(selector.getOutputCol(), "selectedFeatures")
        self.assertEqual(selector.getSelectionThreshold(), 1)

        model = selector.fit(df)
        self.assertEqual(selector.uid, model.uid)
        self.assertEqual(model.selectedFeatures, [3])

        output = model.transform(df)
        self.assertEqual(output.columns, ["features", "label", "selectedFeatures"])
        self.assertEqual(output.count(), 3)

        # save & load
        with tempfile.TemporaryDirectory(prefix="univariate_selector") as d:
            selector.write().overwrite().save(d)
            selector2 = UnivariateFeatureSelector.load(d)
            self.assertEqual(str(selector), str(selector2))

            model.write().overwrite().save(d)
            model2 = UnivariateFeatureSelectorModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_variance_threshold_selector(self):
        df = self.spark.createDataFrame(
            [
                (Vectors.dense([0.0, 0.0, 18.0, 1.0]), 1.0),
                (Vectors.dense([0.0, 1.0, 12.0, 0.0]), 0.0),
                (Vectors.dense([1.0, 0.0, 15.0, 0.1]), 0.0),
            ],
            ["features", "label"],
        )

        selector = VarianceThresholdSelector(varianceThreshold=2, outputCol="selectedFeatures")
        self.assertEqual(selector.getVarianceThreshold(), 2)
        self.assertEqual(selector.getOutputCol(), "selectedFeatures")

        model = selector.fit(df)
        self.assertEqual(selector.uid, model.uid)
        self.assertEqual(model.selectedFeatures, [2])

        output = model.transform(df)
        self.assertEqual(output.columns, ["features", "label", "selectedFeatures"])
        self.assertEqual(output.count(), 3)

        # save & load
        with tempfile.TemporaryDirectory(prefix="variance_threshold_selector") as d:
            selector.write().overwrite().save(d)
            selector2 = VarianceThresholdSelector.load(d)
            self.assertEqual(str(selector), str(selector2))

            model.write().overwrite().save(d)
            model2 = VarianceThresholdSelectorModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_word2vec(self):
        sent = ("a b " * 100 + "a c " * 10).split(" ")
        df = self.spark.createDataFrame([(sent,), (sent,)], ["sentence"]).coalesce(1)

        w2v = Word2Vec(vectorSize=3, seed=42, inputCol="sentence", outputCol="model")
        w2v.setMaxIter(1)
        self.assertEqual(w2v.getInputCol(), "sentence")
        self.assertEqual(w2v.getOutputCol(), "model")
        self.assertEqual(w2v.getVectorSize(), 3)
        self.assertEqual(w2v.getSeed(), 42)
        self.assertEqual(w2v.getMaxIter(), 1)

        model = w2v.fit(df)
        self.assertEqual(w2v.uid, model.uid)
        self.assertEqual(model.getVectors().columns, ["word", "vector"])
        self.assertEqual(model.getVectors().count(), 3)

        synonyms = model.findSynonyms("a", 2)
        self.assertEqual(synonyms.columns, ["word", "similarity"])
        self.assertEqual(synonyms.count(), 2)

        synonyms = model.findSynonymsArray("a", 2)
        self.assertEqual(len(synonyms), 2)
        self.assertEqual(synonyms[0][0], "b")
        self.assertTrue(np.allclose(synonyms[0][1], -0.024012837558984756, atol=1e-4))
        self.assertEqual(synonyms[1][0], "c")
        self.assertTrue(np.allclose(synonyms[1][1], -0.19355154037475586, atol=1e-4))

        output = model.transform(df)
        self.assertEqual(output.columns, ["sentence", "model"])
        self.assertEqual(output.count(), 2)

        # save & load
        with tempfile.TemporaryDirectory(prefix="word2vec") as d:
            w2v.write().overwrite().save(d)
            w2v2 = Word2Vec.load(d)
            self.assertEqual(str(w2v), str(w2v2))

            model.write().overwrite().save(d)
            model2 = Word2VecModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_imputer(self):
        spark = self.spark
        df = spark.createDataFrame(
            [
                (1.0, float("nan")),
                (2.0, float("nan")),
                (float("nan"), 3.0),
                (4.0, 4.0),
                (5.0, 5.0),
            ],
            ["a", "b"],
        )

        imputer = Imputer(strategy="mean")
        imputer.setInputCols(["a", "b"])
        imputer.setOutputCols(["out_a", "out_b"])

        self.assertEqual(imputer.getStrategy(), "mean")
        self.assertEqual(imputer.getInputCols(), ["a", "b"])
        self.assertEqual(imputer.getOutputCols(), ["out_a", "out_b"])

        model = imputer.fit(df)
        self.assertEqual(imputer.uid, model.uid)
        self.assertEqual(model.surrogateDF.columns, ["a", "b"])
        self.assertEqual(model.surrogateDF.count(), 1)
        self.assertEqual(list(model.surrogateDF.head()), [3.0, 4.0])

        output = model.transform(df)
        self.assertEqual(output.columns, ["a", "b", "out_a", "out_b"])
        self.assertEqual(output.count(), 5)

        # save & load
        with tempfile.TemporaryDirectory(prefix="imputer") as d:
            imputer.write().overwrite().save(d)
            imputer2 = Imputer.load(d)
            self.assertEqual(str(imputer), str(imputer2))

            model.write().overwrite().save(d)
            model2 = ImputerModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_count_vectorizer(self):
        df = self.spark.createDataFrame(
            [(0, ["a", "b", "c"]), (1, ["a", "b", "b", "c", "a"])],
            ["label", "raw"],
        )

        cv = CountVectorizer()
        cv.setInputCol("raw")
        cv.setOutputCol("vectors")
        self.assertEqual(cv.getInputCol(), "raw")
        self.assertEqual(cv.getOutputCol(), "vectors")

        model = cv.fit(df)
        self.assertEqual(cv.uid, model.uid)
        self.assertEqual(sorted(model.vocabulary), ["a", "b", "c"])

        output = model.transform(df)
        self.assertEqual(output.columns, ["label", "raw", "vectors"])
        self.assertEqual(output.count(), 2)

        # save & load
        with tempfile.TemporaryDirectory(prefix="count_vectorizer") as d:
            cv.write().overwrite().save(d)
            cv2 = CountVectorizer.load(d)
            self.assertEqual(str(cv), str(cv2))

            model.write().overwrite().save(d)
            model2 = CountVectorizerModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_one_hot_encoder(self):
        df = self.spark.createDataFrame([(0.0,), (1.0,), (2.0,)], ["input"])

        encoder = OneHotEncoder()
        encoder.setInputCols(["input"])
        encoder.setOutputCols(["output"])
        self.assertEqual(encoder.getInputCols(), ["input"])
        self.assertEqual(encoder.getOutputCols(), ["output"])

        model = encoder.fit(df)
        self.assertEqual(encoder.uid, model.uid)
        self.assertEqual(model.categorySizes, [3])

        output = model.transform(df)
        self.assertEqual(output.columns, ["input", "output"])
        self.assertEqual(output.count(), 3)

        # save & load
        with tempfile.TemporaryDirectory(prefix="count_vectorizer") as d:
            encoder.write().overwrite().save(d)
            encoder2 = OneHotEncoder.load(d)
            self.assertEqual(str(encoder), str(encoder2))

            model.write().overwrite().save(d)
            model2 = OneHotEncoderModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_tokenizer(self):
        df = self.spark.createDataFrame([("a b c",)], ["text"])

        tokenizer = Tokenizer(outputCol="words")
        tokenizer.setInputCol("text")
        self.assertEqual(tokenizer.getInputCol(), "text")
        self.assertEqual(tokenizer.getOutputCol(), "words")

        output = tokenizer.transform(df)
        self.assertEqual(output.columns, ["text", "words"])
        self.assertEqual(output.count(), 1)
        self.assertEqual(output.head().words, ["a", "b", "c"])

        # save & load
        with tempfile.TemporaryDirectory(prefix="tokenizer") as d:
            tokenizer.write().overwrite().save(d)
            tokenizer2 = Tokenizer.load(d)
            self.assertEqual(str(tokenizer), str(tokenizer2))

    def test_regex_tokenizer(self):
        df = self.spark.createDataFrame([("A B  c",)], ["text"])

        tokenizer = RegexTokenizer(outputCol="words")
        tokenizer.setInputCol("text")
        self.assertEqual(tokenizer.getInputCol(), "text")
        self.assertEqual(tokenizer.getOutputCol(), "words")

        output = tokenizer.transform(df)
        self.assertEqual(output.columns, ["text", "words"])
        self.assertEqual(output.count(), 1)
        self.assertEqual(output.head().words, ["a", "b", "c"])

        # save & load
        with tempfile.TemporaryDirectory(prefix="regex_tokenizer") as d:
            tokenizer.write().overwrite().save(d)
            tokenizer2 = RegexTokenizer.load(d)
            self.assertEqual(str(tokenizer), str(tokenizer2))

    def test_sql_transformer(self):
        df = self.spark.createDataFrame([(0, 1.0, 3.0), (2, 2.0, 5.0)], ["id", "v1", "v2"])

        statement = "SELECT *, (v1 + v2) AS v3, (v1 * v2) AS v4 FROM __THIS__"
        sql = SQLTransformer(statement=statement)
        self.assertEqual(sql.getStatement(), statement)

        output = sql.transform(df)
        self.assertEqual(output.columns, ["id", "v1", "v2", "v3", "v4"])
        self.assertEqual(output.count(), 2)
        self.assertEqual(
            output.collect(),
            [
                Row(id=0, v1=1.0, v2=3.0, v3=4.0, v4=3.0),
                Row(id=2, v1=2.0, v2=5.0, v3=7.0, v4=10.0),
            ],
        )

        # save & load
        with tempfile.TemporaryDirectory(prefix="sql_transformer") as d:
            sql.write().overwrite().save(d)
            sql2 = SQLTransformer.load(d)
            self.assertEqual(str(sql), str(sql2))

    def test_stop_words_remover(self):
        df = self.spark.createDataFrame([(["a", "b", "c"],)], ["text"])

        remover = StopWordsRemover(stopWords=["b"])
        remover.setInputCol("text")
        remover.setOutputCol("words")

        self.assertEqual(remover.getStopWords(), ["b"])
        self.assertEqual(remover.getInputCol(), "text")
        self.assertEqual(remover.getOutputCol(), "words")

        output = remover.transform(df)
        self.assertEqual(output.columns, ["text", "words"])
        self.assertEqual(output.count(), 1)
        self.assertEqual(output.head().words, ["a", "c"])

        # save & load
        with tempfile.TemporaryDirectory(prefix="stop_words_remover") as d:
            remover.write().overwrite().save(d)
            remover2 = StopWordsRemover.load(d)
            self.assertEqual(str(remover), str(remover2))

    def test_stop_words_remover_with_given_words(self):
        spark = self.spark
        dataset = spark.createDataFrame([Row(input=["a", "panda"])])
        stopWordRemover = StopWordsRemover(inputCol="input", outputCol="output")
        # Default
        self.assertEqual(stopWordRemover.getInputCol(), "input")
        transformedDF = stopWordRemover.transform(dataset)
        self.assertEqual(transformedDF.head().output, ["panda"])
        self.assertEqual(type(stopWordRemover.getStopWords()), list)
        self.assertTrue(isinstance(stopWordRemover.getStopWords()[0], str))
        # Custom
        stopwords = ["panda"]
        stopWordRemover.setStopWords(stopwords)
        self.assertEqual(stopWordRemover.getInputCol(), "input")
        self.assertEqual(stopWordRemover.getStopWords(), stopwords)
        transformedDF = stopWordRemover.transform(dataset)
        self.assertEqual(transformedDF.head().output, ["a"])
        # with locale
        stopwords = ["BELKİ"]
        dataset = self.spark.createDataFrame([Row(input=["belki"])])
        stopWordRemover.setStopWords(stopwords).setLocale("tr")
        self.assertEqual(stopWordRemover.getStopWords(), stopwords)
        transformedDF = stopWordRemover.transform(dataset)
        self.assertEqual(transformedDF.head().output, [])

    def test_stop_words_remover_with_turkish(self):
        spark = self.spark
        dataset = spark.createDataFrame([Row(input=["acaba", "ama", "biri"])])
        stopWordRemover = StopWordsRemover(inputCol="input", outputCol="output")
        stopwords = StopWordsRemover.loadDefaultStopWords("turkish")
        stopWordRemover.setStopWords(stopwords)
        self.assertEqual(stopWordRemover.getStopWords(), stopwords)
        transformedDF = stopWordRemover.transform(dataset)
        self.assertEqual(transformedDF.head().output, [])

    def test_stop_words_remover_default(self):
        stopWordRemover = StopWordsRemover(inputCol="input", outputCol="output")

        # check the default value of local
        locale = stopWordRemover.getLocale()
        self.assertIsInstance(locale, str)
        self.assertTrue(len(locale) > 0)

        # check the default value of stop words
        stopwords = stopWordRemover.getStopWords()
        self.assertIsInstance(stopwords, list)
        self.assertTrue(len(stopwords) > 0)
        self.assertTrue(all(isinstance(word, str) for word in stopwords))

    def test_binarizer(self):
        b0 = Binarizer()
        self.assertListEqual(
            b0.params,
            [b0.inputCol, b0.inputCols, b0.outputCol, b0.outputCols, b0.threshold, b0.thresholds],
        )
        self.assertTrue(all([~b0.isSet(p) for p in b0.params]))
        self.assertTrue(b0.hasDefault(b0.threshold))
        self.assertEqual(b0.getThreshold(), 0.0)
        b0.setParams(inputCol="input", outputCol="output").setThreshold(1.0)
        self.assertTrue(not all([b0.isSet(p) for p in b0.params]))
        self.assertEqual(b0.getThreshold(), 1.0)
        self.assertEqual(b0.getInputCol(), "input")
        self.assertEqual(b0.getOutputCol(), "output")

        b0c = b0.copy({b0.threshold: 2.0})
        self.assertEqual(b0c.uid, b0.uid)
        self.assertListEqual(b0c.params, b0.params)
        self.assertEqual(b0c.getThreshold(), 2.0)

        b1 = Binarizer(threshold=2.0, inputCol="input", outputCol="output")
        self.assertNotEqual(b1.uid, b0.uid)
        self.assertEqual(b1.getThreshold(), 2.0)
        self.assertEqual(b1.getInputCol(), "input")
        self.assertEqual(b1.getOutputCol(), "output")

        df = self.spark.createDataFrame(
            [
                (0.1, 0.0),
                (0.4, 1.0),
                (1.2, 1.3),
                (1.5, float("nan")),
                (float("nan"), 1.0),
                (float("nan"), 0.0),
            ],
            ["v1", "v2"],
        )

        binarizer = Binarizer(threshold=1.0, inputCol="v1", outputCol="f1")
        output = binarizer.transform(df)
        self.assertEqual(output.columns, ["v1", "v2", "f1"])
        self.assertEqual(output.count(), 6)
        self.assertEqual(
            [r.f1 for r in output.select("f1").collect()],
            [0.0, 0.0, 1.0, 1.0, 0.0, 0.0],
        )

        binarizer = Binarizer(threshold=1.0, inputCols=["v1", "v2"], outputCols=["f1", "f2"])
        output = binarizer.transform(df)
        self.assertEqual(output.columns, ["v1", "v2", "f1", "f2"])
        self.assertEqual(output.count(), 6)
        self.assertEqual(
            [r.f1 for r in output.select("f1").collect()],
            [0.0, 0.0, 1.0, 1.0, 0.0, 0.0],
        )
        self.assertEqual(
            [r.f2 for r in output.select("f2").collect()],
            [0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        )

        # save & load
        with tempfile.TemporaryDirectory(prefix="binarizer") as d:
            binarizer.write().overwrite().save(d)
            binarizer2 = Binarizer.load(d)
            self.assertEqual(str(binarizer), str(binarizer2))

    def test_quantile_discretizer_single_column(self):
        spark = self.spark
        values = [(0.1,), (0.4,), (1.2,), (1.5,), (float("nan"),), (float("nan"),)]
        df = spark.createDataFrame(values, ["values"])

        qds = QuantileDiscretizer(inputCol="values", outputCol="buckets")
        qds.setNumBuckets(2)
        qds.setRelativeError(0.01)
        qds.setHandleInvalid("keep")

        self.assertEqual(qds.getInputCol(), "values")
        self.assertEqual(qds.getOutputCol(), "buckets")
        self.assertEqual(qds.getNumBuckets(), 2)
        self.assertEqual(qds.getRelativeError(), 0.01)
        self.assertEqual(qds.getHandleInvalid(), "keep")

        bucketizer = qds.fit(df)
        self.assertIsInstance(bucketizer, Bucketizer)
        # Bucketizer doesn't inherit uid from QuantileDiscretizer
        self.assertNotEqual(qds.uid, bucketizer.uid)
        self.assertTrue(qds.uid.startswith("QuantileDiscretizer"))
        self.assertTrue(bucketizer.uid.startswith("Bucketizer"))

        # check model coefficients
        self.assertEqual(bucketizer.getSplits(), [float("-inf"), 0.4, float("inf")])

        output = bucketizer.transform(df)
        self.assertEqual(output.columns, ["values", "buckets"])
        self.assertEqual(output.count(), 6)

        # save & load
        with tempfile.TemporaryDirectory(prefix="quantile_discretizer_single_column") as d:
            qds.write().overwrite().save(d)
            qds2 = QuantileDiscretizer.load(d)
            self.assertEqual(str(qds), str(qds2))

            bucketizer.write().overwrite().save(d)
            bucketizer2 = Bucketizer.load(d)
            self.assertEqual(str(bucketizer), str(bucketizer2))

    def test_quantile_discretizer_multiple_columns(self):
        spark = self.spark
        inputs = [
            (0.1, 0.0),
            (0.4, 1.0),
            (1.2, 1.3),
            (1.5, 1.5),
            (float("nan"), float("nan")),
            (float("nan"), float("nan")),
        ]
        df = spark.createDataFrame(inputs, ["input1", "input2"])

        qds = QuantileDiscretizer(
            relativeError=0.01,
            handleInvalid="keep",
            numBuckets=2,
            inputCols=["input1", "input2"],
            outputCols=["output1", "output2"],
        )

        self.assertEqual(qds.getInputCols(), ["input1", "input2"])
        self.assertEqual(qds.getOutputCols(), ["output1", "output2"])
        self.assertEqual(qds.getNumBuckets(), 2)
        self.assertEqual(qds.getRelativeError(), 0.01)
        self.assertEqual(qds.getHandleInvalid(), "keep")

        bucketizer = qds.fit(df)
        self.assertIsInstance(bucketizer, Bucketizer)
        # Bucketizer doesn't inherit uid from QuantileDiscretizer
        self.assertNotEqual(qds.uid, bucketizer.uid)
        self.assertTrue(qds.uid.startswith("QuantileDiscretizer"))
        self.assertTrue(bucketizer.uid.startswith("Bucketizer"))

        # check model coefficients
        self.assertEqual(
            bucketizer.getSplitsArray(),
            [
                [float("-inf"), 0.4, float("inf")],
                [float("-inf"), 1.0, float("inf")],
            ],
        )

        output = bucketizer.transform(df)
        self.assertEqual(output.columns, ["input1", "input2", "output1", "output2"])
        self.assertEqual(output.count(), 6)

        # save & load
        with tempfile.TemporaryDirectory(prefix="quantile_discretizer_multiple_columns") as d:
            qds.write().overwrite().save(d)
            qds2 = QuantileDiscretizer.load(d)
            self.assertEqual(str(qds), str(qds2))

            bucketizer.write().overwrite().save(d)
            bucketizer2 = Bucketizer.load(d)
            self.assertEqual(str(bucketizer), str(bucketizer2))

    def test_bucketizer(self):
        df = self.spark.createDataFrame(
            [
                (0.1, 0.0),
                (0.4, 1.0),
                (1.2, 1.3),
                (1.5, float("nan")),
                (float("nan"), 1.0),
                (float("nan"), 0.0),
            ],
            ["v1", "v2"],
        )

        splits = [-float("inf"), 0.5, 1.4, float("inf")]
        bucketizer = Bucketizer()
        bucketizer.setSplits(splits)
        bucketizer.setHandleInvalid("keep")
        bucketizer.setInputCol("v1")
        bucketizer.setOutputCol("b1")

        self.assertEqual(bucketizer.getSplits(), splits)
        self.assertEqual(bucketizer.getHandleInvalid(), "keep")
        self.assertEqual(bucketizer.getInputCol(), "v1")
        self.assertEqual(bucketizer.getOutputCol(), "b1")

        output = bucketizer.transform(df)
        self.assertEqual(output.columns, ["v1", "v2", "b1"])
        self.assertEqual(output.count(), 6)
        self.assertEqual(
            [r.b1 for r in output.select("b1").collect()],
            [0.0, 0.0, 1.0, 2.0, 3.0, 3.0],
        )

        splitsArray = [
            [-float("inf"), 0.5, 1.4, float("inf")],
            [-float("inf"), 0.5, float("inf")],
        ]
        bucketizer = Bucketizer(
            splitsArray=splitsArray,
            inputCols=["v1", "v2"],
            outputCols=["b1", "b2"],
        )
        bucketizer.setHandleInvalid("keep")
        self.assertEqual(bucketizer.getSplitsArray(), splitsArray)
        self.assertEqual(bucketizer.getHandleInvalid(), "keep")
        self.assertEqual(bucketizer.getInputCols(), ["v1", "v2"])
        self.assertEqual(bucketizer.getOutputCols(), ["b1", "b2"])

        output = bucketizer.transform(df)
        self.assertEqual(output.columns, ["v1", "v2", "b1", "b2"])
        self.assertEqual(output.count(), 6)
        self.assertEqual(
            [r.b1 for r in output.select("b1").collect()],
            [0.0, 0.0, 1.0, 2.0, 3.0, 3.0],
        )
        self.assertEqual(
            [r.b2 for r in output.select("b2").collect()],
            [0.0, 1.0, 1.0, 2.0, 1.0, 0.0],
        )

        # save & load
        with tempfile.TemporaryDirectory(prefix="bucketizer") as d:
            bucketizer.write().overwrite().save(d)
            bucketizer2 = Bucketizer.load(d)
            self.assertEqual(str(bucketizer), str(bucketizer2))

    def test_idf(self):
        df = self.spark.createDataFrame(
            [
                (DenseVector([1.0, 2.0]),),
                (DenseVector([0.0, 1.0]),),
                (DenseVector([3.0, 0.2]),),
            ],
            ["tf"],
        )
        idf = IDF(inputCol="tf")
        self.assertListEqual(idf.params, [idf.inputCol, idf.minDocFreq, idf.outputCol])

        model = idf.fit(df, {idf.outputCol: "idf"})
        self.assertEqual(idf.uid, model.uid)
        # self.assertEqual(
        #     model.uid, idf.uid, "Model should inherit the UID from its parent estimator."
        # )
        self.assertTrue(
            np.allclose(model.idf.toArray(), [0.28768207245178085, 0.0], atol=1e-4),
            model.idf,
        )
        self.assertEqual(model.docFreq, [2, 3])
        self.assertEqual(model.numDocs, 3)

        output = model.transform(df)
        self.assertEqual(output.columns, ["tf", "idf"])
        self.assertIsNotNone(output.head().idf)

        # save & load
        with tempfile.TemporaryDirectory(prefix="idf") as d:
            idf.write().overwrite().save(d)
            idf2 = IDF.load(d)
            self.assertEqual(str(idf), str(idf2))

            model.write().overwrite().save(d)
            model2 = IDFModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_ngram(self):
        spark = self.spark
        df = spark.createDataFrame([Row(input=["a", "b", "c", "d", "e"])])

        ngram = NGram(n=4, inputCol="input", outputCol="output")
        self.assertEqual(ngram.getN(), 4)
        self.assertEqual(ngram.getInputCol(), "input")
        self.assertEqual(ngram.getOutputCol(), "output")

        output = ngram.transform(df)
        self.assertEqual(output.head().output, ["a b c d", "b c d e"])

        # save & load
        with tempfile.TemporaryDirectory(prefix="ngram") as d:
            ngram.write().overwrite().save(d)
            ngram2 = NGram.load(d)
            self.assertEqual(str(ngram), str(ngram2))

    def test_normalizer(self):
        spark = self.spark
        df = spark.createDataFrame(
            [(Vectors.dense([3.0, -4.0]),), (Vectors.sparse(4, {1: 4.0, 3: 3.0}),)],
            ["input"],
        )

        normalizer = Normalizer(p=2.0, inputCol="input", outputCol="output")
        self.assertEqual(normalizer.getP(), 2.0)
        self.assertEqual(normalizer.getInputCol(), "input")
        self.assertEqual(normalizer.getOutputCol(), "output")

        output = normalizer.transform(df)
        self.assertEqual(output.columns, ["input", "output"])
        self.assertEqual(output.count(), 2)

        # save & load
        with tempfile.TemporaryDirectory(prefix="normalizer") as d:
            normalizer.write().overwrite().save(d)
            normalizer2 = Normalizer.load(d)
            self.assertEqual(str(normalizer), str(normalizer2))

    def test_interaction(self):
        spark = self.spark
        df = spark.createDataFrame([(0.0, 1.0), (2.0, 3.0)], ["a", "b"])

        interaction = Interaction()
        interaction.setInputCols(["a", "b"])
        interaction.setOutputCol("ab")
        self.assertEqual(interaction.getInputCols(), ["a", "b"])
        self.assertEqual(interaction.getOutputCol(), "ab")

        output = interaction.transform(df)
        self.assertEqual(output.columns, ["a", "b", "ab"])
        self.assertEqual(output.count(), 2)

        # save & load
        with tempfile.TemporaryDirectory(prefix="interaction") as d:
            interaction.write().overwrite().save(d)
            interaction2 = Interaction.load(d)
            self.assertEqual(str(interaction), str(interaction2))

    def test_count_vectorizer_with_binary(self):
        dataset = self.spark.createDataFrame(
            [
                (
                    0,
                    "a a a b b c".split(" "),
                    SparseVector(3, {0: 1.0, 1: 1.0, 2: 1.0}),
                ),
                (
                    1,
                    "a a".split(" "),
                    SparseVector(3, {0: 1.0}),
                ),
                (
                    2,
                    "a b".split(" "),
                    SparseVector(3, {0: 1.0, 1: 1.0}),
                ),
                (
                    3,
                    "c".split(" "),
                    SparseVector(3, {2: 1.0}),
                ),
            ],
            ["id", "words", "expected"],
        )
        cv = CountVectorizer(binary=True, inputCol="words", outputCol="features")
        model = cv.fit(dataset)
        self.assertEqual(cv.uid, model.uid)

        transformedList = model.transform(dataset).select("features", "expected").collect()

        for r in transformedList:
            feature, expected = r
            self.assertEqual(feature, expected)

    def test_count_vectorizer_with_maxDF(self):
        dataset = self.spark.createDataFrame(
            [
                (
                    0,
                    "a b c d".split(" "),
                    SparseVector(3, {0: 1.0, 1: 1.0, 2: 1.0}),
                ),
                (
                    1,
                    "a b c".split(" "),
                    SparseVector(3, {0: 1.0, 1: 1.0}),
                ),
                (
                    2,
                    "a b".split(" "),
                    SparseVector(3, {0: 1.0}),
                ),
                (
                    3,
                    "a".split(" "),
                    SparseVector(3, {}),
                ),
            ],
            ["id", "words", "expected"],
        )
        cv = CountVectorizer(inputCol="words", outputCol="features")
        model1 = cv.setMaxDF(3).fit(dataset)
        self.assertEqual(cv.uid, model1.uid)

        self.assertEqual(model1.vocabulary, ["b", "c", "d"])

        transformedList1 = model1.transform(dataset).select("features", "expected").collect()

        for r in transformedList1:
            feature, expected = r
            self.assertEqual(feature, expected)

        model2 = cv.setMaxDF(0.75).fit(dataset)
        self.assertEqual(model2.vocabulary, ["b", "c", "d"])

        transformedList2 = model2.transform(dataset).select("features", "expected").collect()

        for r in transformedList2:
            feature, expected = r
            self.assertEqual(feature, expected)

    def test_count_vectorizer_from_vocab(self):
        model = CountVectorizerModel.from_vocabulary(
            ["a", "b", "c"], inputCol="words", outputCol="features", minTF=2
        )
        self.assertEqual(model.vocabulary, ["a", "b", "c"])
        self.assertEqual(model.getMinTF(), 2)

        dataset = self.spark.createDataFrame(
            [
                (
                    0,
                    "a a a b b c".split(" "),
                    SparseVector(3, {0: 3.0, 1: 2.0}),
                ),
                (
                    1,
                    "a a".split(" "),
                    SparseVector(3, {0: 2.0}),
                ),
                (
                    2,
                    "a b".split(" "),
                    SparseVector(3, {}),
                ),
            ],
            ["id", "words", "expected"],
        )

        transformed_list = model.transform(dataset).select("features", "expected").collect()

        for r in transformed_list:
            feature, expected = r
            self.assertEqual(feature, expected)

        # Test an empty vocabulary
        with self.assertRaisesRegex(Exception, "Vocabulary list cannot be empty"):
            CountVectorizerModel.from_vocabulary([], inputCol="words")

        # Test model with default settings can transform
        model_default = CountVectorizerModel.from_vocabulary(["a", "b", "c"], inputCol="words")
        transformed_list = (
            model_default.transform(dataset)
            .select(model_default.getOrDefault(model_default.outputCol))
            .collect()
        )
        self.assertEqual(len(transformed_list), 3)

    def test_rformula_force_index_label(self):
        df = self.spark.createDataFrame(
            [(1.0, 1.0, "a"), (0.0, 2.0, "b"), (1.0, 0.0, "a")], ["y", "x", "s"]
        )
        # Does not index label by default since it's numeric type.
        rf = RFormula(formula="y ~ x + s")
        model = rf.fit(df)
        self.assertEqual(rf.uid, model.uid)

        transformedDF = model.transform(df)
        self.assertEqual(transformedDF.head().label, 1.0)
        # Force to index label.
        rf2 = RFormula(formula="y ~ x + s").setForceIndexLabel(True)
        model2 = rf2.fit(df)
        transformedDF2 = model2.transform(df)
        self.assertEqual(transformedDF2.head().label, 0.0)

    def test_rformula_string_indexer_order_type(self):
        df = self.spark.createDataFrame(
            [(1.0, 1.0, "a"), (0.0, 2.0, "b"), (1.0, 0.0, "a")], ["y", "x", "s"]
        )
        rf = RFormula(formula="y ~ x + s", stringIndexerOrderType="alphabetDesc")
        self.assertEqual(rf.getStringIndexerOrderType(), "alphabetDesc")
        model = rf.fit(df)
        self.assertEqual(rf.uid, model.uid)
        transformedDF = model.transform(df)
        observed = transformedDF.select("features").collect()
        expected = [[1.0, 0.0], [2.0, 1.0], [0.0, 0.0]]
        for i in range(0, len(expected)):
            self.assertTrue(all(observed[i]["features"].toArray() == expected[i]))

        # save & load
        with tempfile.TemporaryDirectory(prefix="rformula") as d:
            rf.write().overwrite().save(d)
            rf2 = RFormula.load(d)
            self.assertEqual(str(rf), str(rf2))

            model.write().overwrite().save(d)
            model2 = RFormulaModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model.getFormula(), model2.getFormula())

    def test_string_indexer_handle_invalid(self):
        df = self.spark.createDataFrame([(0, "a"), (1, "d"), (2, None)], ["id", "label"])

        si1 = StringIndexer(
            inputCol="label",
            outputCol="indexed",
            handleInvalid="keep",
            stringOrderType="alphabetAsc",
        )
        model1 = si1.fit(df)
        td1 = model1.transform(df)
        actual1 = td1.select("id", "indexed").collect()
        expected1 = [Row(id=0, indexed=0.0), Row(id=1, indexed=1.0), Row(id=2, indexed=2.0)]
        self.assertEqual(actual1, expected1)

        si2 = si1.setHandleInvalid("skip")
        model2 = si2.fit(df)
        td2 = model2.transform(df)
        actual2 = td2.select("id", "indexed").collect()
        expected2 = [Row(id=0, indexed=0.0), Row(id=1, indexed=1.0)]
        self.assertEqual(actual2, expected2)

    def test_string_indexer_from_labels(self):
        model = StringIndexerModel.from_labels(
            ["a", "b", "c"], inputCol="label", outputCol="indexed", handleInvalid="keep"
        )
        self.assertEqual(model.labels, ["a", "b", "c"])
        self.assertEqual(model.labelsArray, [["a", "b", "c"]])

        self.assertEqual(model.getInputCol(), "label")
        self.assertEqual(model.getOutputCol(), "indexed")
        self.assertEqual(model.getHandleInvalid(), "keep")

        df1 = self.spark.createDataFrame(
            [(0, "a"), (1, "c"), (2, None), (3, "b"), (4, "b")], ["id", "label"]
        )

        result1 = model.transform(df1)
        actual1 = result1.select("id", "indexed").collect()
        expected1 = [
            Row(id=0, indexed=0.0),
            Row(id=1, indexed=2.0),
            Row(id=2, indexed=3.0),
            Row(id=3, indexed=1.0),
            Row(id=4, indexed=1.0),
        ]
        self.assertEqual(actual1, expected1)

        model_empty_labels = StringIndexerModel.from_labels(
            [], inputCol="label", outputCol="indexed", handleInvalid="keep"
        )
        actual2 = model_empty_labels.transform(df1).select("id", "indexed").collect()
        expected2 = [
            Row(id=0, indexed=0.0),
            Row(id=1, indexed=0.0),
            Row(id=2, indexed=0.0),
            Row(id=3, indexed=0.0),
            Row(id=4, indexed=0.0),
        ]
        self.assertEqual(actual2, expected2)

        # Test model with default settings can transform
        model_default = StringIndexerModel.from_labels(["a", "b", "c"], inputCol="label")
        df2 = self.spark.createDataFrame(
            [(0, "a"), (1, "c"), (2, "b"), (3, "b"), (4, "b")], ["id", "label"]
        )
        transformed_list = (
            model_default.transform(df2)
            .select(model_default.getOrDefault(model_default.outputCol))
            .collect()
        )
        self.assertEqual(len(transformed_list), 5)

    def test_string_indexer_from_arrays_of_labels(self):
        model = StringIndexerModel.from_arrays_of_labels(
            [["a", "b", "c"], ["x", "y", "z"]],
            inputCols=["label1", "label2"],
            outputCols=["indexed1", "indexed2"],
            handleInvalid="keep",
        )

        self.assertEqual(model.labelsArray, [["a", "b", "c"], ["x", "y", "z"]])

        self.assertEqual(model.getInputCols(), ["label1", "label2"])
        self.assertEqual(model.getOutputCols(), ["indexed1", "indexed2"])
        self.assertEqual(model.getHandleInvalid(), "keep")

    def test_target_encoder_binary(self):
        df = self.spark.createDataFrame(
            [
                (0, 3, 5.0, 0.0),
                (1, 4, 5.0, 1.0),
                (2, 3, 5.0, 0.0),
                (0, 4, 6.0, 1.0),
                (1, 3, 6.0, 0.0),
                (2, 4, 6.0, 1.0),
                (0, 3, 7.0, 0.0),
                (1, 4, 8.0, 1.0),
                (2, 3, 9.0, 0.0),
            ],
            schema="input1 short, input2 int, input3 double, label double",
        )
        encoder = TargetEncoder(
            inputCols=["input1", "input2", "input3"],
            outputCols=["output", "output2", "output3"],
            labelCol="label",
            targetType="binary",
        )
        model = encoder.fit(df)
        output = model.transform(df)
        self.assertEqual(
            output.columns,
            ["input1", "input2", "input3", "label", "output", "output2", "output3"],
        )
        self.assertEqual(output.count(), 9)

        # save & load
        with tempfile.TemporaryDirectory(prefix="target_encoder") as d:
            encoder.write().overwrite().save(d)
            encoder2 = TargetEncoder.load(d)
            self.assertEqual(str(encoder), str(encoder2))

            model.write().overwrite().save(d)
            model2 = TargetEncoderModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_vector_size_hint(self):
        df = self.spark.createDataFrame(
            [
                (0, Vectors.dense([0.0, 10.0, 0.5])),
                (1, Vectors.dense([1.0, 11.0, 0.5, 0.6])),
                (2, Vectors.dense([2.0, 12.0])),
            ],
            ["id", "vector"],
        )

        sh = VectorSizeHint(inputCol="vector", handleInvalid="skip")
        sh.setSize(3)
        self.assertEqual(sh.getSize(), 3)

        output = sh.transform(df).head().vector
        expected = DenseVector([0.0, 10.0, 0.5])
        self.assertEqual(output, expected)

        # save & load
        with tempfile.TemporaryDirectory(prefix="vector_size_hint") as d:
            sh.write().overwrite().save(d)
            sh2 = VectorSizeHint.load(d)
            self.assertEqual(str(sh), str(sh2))

    def test_vector_slicer(self):
        spark = self.spark

        df = spark.createDataFrame(
            [
                (Vectors.dense([-2.0, 2.3, 0.0, 0.0, 1.0]),),
                (Vectors.dense([0.0, 0.0, 0.0, 0.0, 0.0]),),
                (Vectors.dense([0.6, -1.1, -3.0, 4.5, 3.3]),),
            ],
            ["features"],
        )

        vs = VectorSlicer(outputCol="sliced", indices=[1, 4])
        vs.setInputCol("features")
        self.assertEqual(vs.getIndices(), [1, 4])
        self.assertEqual(vs.getInputCol(), "features")
        self.assertEqual(vs.getOutputCol(), "sliced")

        output = vs.transform(df)
        self.assertEqual(output.columns, ["features", "sliced"])
        self.assertEqual(output.count(), 3)
        self.assertEqual(output.head().sliced, Vectors.dense([2.3, 1.0]))

        # save & load
        with tempfile.TemporaryDirectory(prefix="vector_slicer") as d:
            vs.write().overwrite().save(d)
            vs2 = VectorSlicer.load(d)
            self.assertEqual(str(vs), str(vs2))

    def test_feature_hasher(self):
        data = [(2.0, True, "1", "foo"), (3.0, False, "2", "bar")]
        cols = ["real", "bool", "stringNum", "string"]
        df = self.spark.createDataFrame(data, cols)

        hasher = FeatureHasher(numFeatures=2)
        hasher.setInputCols(cols)
        hasher.setOutputCol("features")

        self.assertEqual(hasher.getNumFeatures(), 2)
        self.assertEqual(hasher.getInputCols(), cols)
        self.assertEqual(hasher.getOutputCol(), "features")

        output = hasher.transform(df)
        self.assertEqual(output.columns, ["real", "bool", "stringNum", "string", "features"])
        self.assertEqual(output.count(), 2)

        features = output.head().features.toArray()
        self.assertTrue(
            np.allclose(features, [2.0, 3.0], atol=1e-4),
            features,
        )

        # save & load
        with tempfile.TemporaryDirectory(prefix="feature_hasher") as d:
            hasher.write().overwrite().save(d)
            hasher2 = FeatureHasher.load(d)
            self.assertEqual(str(hasher), str(hasher2))

    def test_hashing_tf(self):
        df = self.spark.createDataFrame([(0, ["a", "a", "b", "c", "c", "c"])], ["id", "words"])
        tf = HashingTF()
        tf.setInputCol("words").setOutputCol("features").setNumFeatures(10).setBinary(True)
        self.assertEqual(tf.getInputCol(), "words")
        self.assertEqual(tf.getOutputCol(), "features")
        self.assertEqual(tf.getNumFeatures(), 10)
        self.assertTrue(tf.getBinary())

        output = tf.transform(df)
        self.assertEqual(output.columns, ["id", "words", "features"])
        self.assertEqual(output.count(), 1)

        features = output.select("features").first().features.toArray()
        self.assertTrue(
            np.allclose(
                features,
                [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0],
                atol=1e-4,
            ),
            features,
        )

        # save & load
        with tempfile.TemporaryDirectory(prefix="hashing_tf") as d:
            tf.write().overwrite().save(d)
            tf2 = HashingTF.load(d)
            self.assertEqual(str(tf), str(tf2))

    def test_bucketed_random_projection_lsh(self):
        spark = self.spark

        data = [
            (0, Vectors.dense([-1.0, -1.0])),
            (1, Vectors.dense([-1.0, 1.0])),
            (2, Vectors.dense([1.0, -1.0])),
            (3, Vectors.dense([1.0, 1.0])),
        ]
        df = spark.createDataFrame(data, ["id", "features"])

        data2 = [
            (4, Vectors.dense([2.0, 2.0])),
            (5, Vectors.dense([2.0, 3.0])),
            (6, Vectors.dense([3.0, 2.0])),
            (7, Vectors.dense([3.0, 3.0])),
        ]
        df2 = spark.createDataFrame(data2, ["id", "features"])

        brp = BucketedRandomProjectionLSH()
        brp.setInputCol("features")
        brp.setOutputCol("hashes")
        brp.setSeed(12345)
        brp.setBucketLength(1.0)

        self.assertEqual(brp.getInputCol(), "features")
        self.assertEqual(brp.getOutputCol(), "hashes")
        self.assertEqual(brp.getBucketLength(), 1.0)
        self.assertEqual(brp.getSeed(), 12345)

        model = brp.fit(df)

        output = model.transform(df)
        self.assertEqual(output.columns, ["id", "features", "hashes"])
        self.assertEqual(output.count(), 4)

        output = model.approxNearestNeighbors(df2, Vectors.dense([1.0, 2.0]), 1)
        self.assertEqual(output.columns, ["id", "features", "hashes", "distCol"])
        self.assertEqual(output.count(), 1)

        output = model.approxSimilarityJoin(df, df2, 3)
        self.assertEqual(output.columns, ["datasetA", "datasetB", "distCol"])
        self.assertEqual(output.count(), 1)

        # save & load
        with tempfile.TemporaryDirectory(prefix="bucketed_random_projection_lsh") as d:
            brp.write().overwrite().save(d)
            brp2 = BucketedRandomProjectionLSH.load(d)
            self.assertEqual(str(brp), str(brp2))

            model.write().overwrite().save(d)
            model2 = BucketedRandomProjectionLSHModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_min_hash_lsh(self):
        spark = self.spark

        data = [
            (0, Vectors.dense([-1.0, -1.0])),
            (1, Vectors.dense([-1.0, 1.0])),
            (2, Vectors.dense([1.0, -1.0])),
            (3, Vectors.dense([1.0, 1.0])),
        ]
        df = spark.createDataFrame(data, ["id", "features"])

        data2 = [
            (4, Vectors.dense([2.0, 2.0])),
            (5, Vectors.dense([2.0, 3.0])),
            (6, Vectors.dense([3.0, 2.0])),
            (7, Vectors.dense([3.0, 3.0])),
        ]
        df2 = spark.createDataFrame(data2, ["id", "features"])

        mh = MinHashLSH()
        mh.setInputCol("features")
        mh.setOutputCol("hashes")
        mh.setSeed(12345)

        self.assertEqual(mh.getInputCol(), "features")
        self.assertEqual(mh.getOutputCol(), "hashes")
        self.assertEqual(mh.getSeed(), 12345)

        model = mh.fit(df)

        output = model.transform(df)
        self.assertEqual(output.columns, ["id", "features", "hashes"])
        self.assertEqual(output.count(), 4)

        output = model.approxNearestNeighbors(df2, Vectors.dense([1.0, 2.0]), 1)
        self.assertEqual(output.columns, ["id", "features", "hashes", "distCol"])
        self.assertEqual(output.count(), 1)

        output = model.approxSimilarityJoin(df, df2, 3)
        self.assertEqual(output.columns, ["datasetA", "datasetB", "distCol"])
        self.assertEqual(output.count(), 16)

        # save & load
        with tempfile.TemporaryDirectory(prefix="min_hash_lsh") as d:
            mh.write().overwrite().save(d)
            mh2 = MinHashLSH.load(d)
            self.assertEqual(str(mh), str(mh2))

            model.write().overwrite().save(d)
            model2 = MinHashLSHModel.load(d)
            self.assertEqual(str(model), str(model2))


class FeatureTests(FeatureTestsMixin, ReusedSQLTestCase):
    pass


if __name__ == "__main__":
    from pyspark.testing import main

    main()