2635c3894f
### What changes were proposed in this pull request? Now that we merged the Koalas main code into PySpark code base (#32036), we should enable doctests on the Spark's infrastructure. ### Why are the changes needed? Currently the pandas-on-Spark modules are not tested at all. We should enable doctests first, and we will port other unit tests separately later. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Enabled the whole doctests. Closes #32069 from ueshin/issues/SPARK-34972/pyspark-pandas_doctests. Authored-by: Takuya UESHIN <ueshin@databricks.com> Signed-off-by: HyukjinKwon <gurwls223@apache.org>
116 lines
4 KiB
Python
116 lines
4 KiB
Python
#
|
|
# 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.
|
|
#
|
|
|
|
from typing import List, Tuple, TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyspark
|
|
|
|
from pyspark.ml.feature import VectorAssembler
|
|
from pyspark.ml.stat import Correlation
|
|
|
|
from pyspark.pandas.utils import column_labels_level
|
|
|
|
if TYPE_CHECKING:
|
|
import pyspark.pandas as pp # noqa: F401 (SPARK-34943)
|
|
|
|
|
|
CORRELATION_OUTPUT_COLUMN = "__correlation_output__"
|
|
|
|
|
|
def corr(kdf: "pp.DataFrame", method: str = "pearson") -> pd.DataFrame:
|
|
"""
|
|
The correlation matrix of all the numerical columns of this dataframe.
|
|
|
|
Only accepts scalar numerical values for now.
|
|
|
|
:param kdf: the Koalas dataframe.
|
|
:param method: {'pearson', 'spearman'}
|
|
* pearson : standard correlation coefficient
|
|
* spearman : Spearman rank correlation
|
|
:return: :class:`pandas.DataFrame`
|
|
|
|
>>> pp.DataFrame({'A': [0, 1], 'B': [1, 0], 'C': ['x', 'y']}).corr()
|
|
A B
|
|
A 1.0 -1.0
|
|
B -1.0 1.0
|
|
"""
|
|
assert method in ("pearson", "spearman")
|
|
ndf, column_labels = to_numeric_df(kdf)
|
|
corr = Correlation.corr(ndf, CORRELATION_OUTPUT_COLUMN, method)
|
|
pcorr = corr.toPandas()
|
|
arr = pcorr.iloc[0, 0].toArray()
|
|
if column_labels_level(column_labels) > 1:
|
|
idx = pd.MultiIndex.from_tuples(column_labels)
|
|
else:
|
|
idx = pd.Index([label[0] for label in column_labels])
|
|
return pd.DataFrame(arr, columns=idx, index=idx)
|
|
|
|
|
|
def to_numeric_df(kdf: "pp.DataFrame") -> Tuple[pyspark.sql.DataFrame, List[Tuple]]:
|
|
"""
|
|
Takes a dataframe and turns it into a dataframe containing a single numerical
|
|
vector of doubles. This dataframe has a single field called '_1'.
|
|
|
|
TODO: index is not preserved currently
|
|
:param kdf: the Koalas dataframe.
|
|
:return: a pair of dataframe, list of strings (the name of the columns
|
|
that were converted to numerical types)
|
|
|
|
>>> to_numeric_df(pp.DataFrame({'A': [0, 1], 'B': [1, 0], 'C': ['x', 'y']}))
|
|
(DataFrame[__correlation_output__: vector], [('A',), ('B',)])
|
|
"""
|
|
# TODO, it should be more robust.
|
|
accepted_types = {
|
|
np.dtype(dt)
|
|
for dt in [np.int8, np.int16, np.int32, np.int64, np.float32, np.float64, np.bool_]
|
|
}
|
|
numeric_column_labels = [
|
|
label for label in kdf._internal.column_labels if kdf[label].dtype in accepted_types
|
|
]
|
|
numeric_df = kdf._internal.spark_frame.select(
|
|
*[kdf._internal.spark_column_for(idx) for idx in numeric_column_labels]
|
|
)
|
|
va = VectorAssembler(inputCols=numeric_df.columns, outputCol=CORRELATION_OUTPUT_COLUMN)
|
|
v = va.transform(numeric_df).select(CORRELATION_OUTPUT_COLUMN)
|
|
return v, numeric_column_labels
|
|
|
|
|
|
def _test():
|
|
import os
|
|
import doctest
|
|
import sys
|
|
from pyspark.sql import SparkSession
|
|
import pyspark.pandas.ml
|
|
|
|
os.chdir(os.environ["SPARK_HOME"])
|
|
|
|
globs = pyspark.pandas.ml.__dict__.copy()
|
|
globs["pp"] = pyspark.pandas
|
|
spark = SparkSession.builder.master("local[4]").appName("pyspark.pandas.ml tests").getOrCreate()
|
|
(failure_count, test_count) = doctest.testmod(
|
|
pyspark.pandas.ml, globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE
|
|
)
|
|
spark.stop()
|
|
if failure_count:
|
|
sys.exit(-1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
_test()
|