def test_loc_timestamp_str(self):
df = pd.DataFrame({'A': np.random.randn(100), 'B': np.random.randn(100)},
index=pd.date_range('2011-01-01', freq='H', periods=100))
ddf = koalas.from_pandas(df)
# partial string slice
# TODO?: self.assert_eq(df.loc['2011-01-02'],
# TODO?: ddf.loc['2011-01-02'])
self.assert_eq(df.loc['2011-01-02':'2011-01-05'],
ddf.loc['2011-01-02':'2011-01-05'])
# series
# TODO?: self.assert_eq(df.A.loc['2011-01-02'],
# TODO?: ddf.A.loc['2011-01-02'])
self.assert_eq(df.A.loc['2011-01-02':'2011-01-05'],
ddf.A.loc['2011-01-02':'2011-01-05'])
df = pd.DataFrame({'A': np.random.randn(100), 'B': np.random.randn(100)},
index=pd.date_range('2011-01-01', freq='M', periods=100))
ddf = koalas.from_pandas(df)
# TODO?: self.assert_eq(df.loc['2011-01'], ddf.loc['2011-01'])
# TODO?: self.assert_eq(df.loc['2011'], ddf.loc['2011'])
self.assert_eq(df.loc['2011-01':'2012-05'], ddf.loc['2011-01':'2012-05'])
self.assert_eq(df.loc['2011':'2015'], ddf.loc['2011':'2015'])
# series
# TODO?: self.assert_eq(df.B.loc['2011-01'], ddf.B.loc['2011-01'])
# TODO?: self.assert_eq(df.B.loc['2011'], ddf.B.loc['2011'])
self.assert_eq(df.B.loc['2011-01':'2012-05'], ddf.B.loc['2011-01':'2012-05'])
self.assert_eq(df.B.loc['2011':'2015'], ddf.B.loc['2011':'2015'])
def test_csv(self):
pdf = self.pdf
kdf = self.kdf
self.assert_eq(kdf.to_csv(), pdf.to_csv())
pdf = pd.DataFrame({
'a': [1, np.nan, 3],
'b': ["one", "two", None],
}, index=[0, 1, 3])
kdf = koalas.from_pandas(pdf)
self.assert_eq(kdf.to_csv(na_rep='null'), pdf.to_csv(na_rep='null'))
pdf = pd.DataFrame({
'a': [1.0, 2.0, 3.0],
'b': [4.0, 5.0, 6.0],
}, index=[0, 1, 3])
kdf = koalas.from_pandas(pdf)
self.assert_eq(kdf.to_csv(float_format='%.1f'), pdf.to_csv(float_format='%.1f'))
self.assert_eq(kdf.to_csv(header=False), pdf.to_csv(header=False))
self.assert_eq(kdf.to_csv(index=False), pdf.to_csv(index=False))
def test_from_pandas_with_explicit_index(self):
pdf = self.pdf
df1 = koalas.from_pandas(pdf.set_index('month'))
self.assertPandasEqual(df1.toPandas(), pdf.set_index('month'))
df2 = koalas.from_pandas(pdf.set_index(['year', 'month']))
self.assertPandasEqual(df2.toPandas(), pdf.set_index(['year', 'month']))
def test_loc_non_informative_index(self):
df = pd.DataFrame({'x': [1, 2, 3, 4]}, index=[10, 20, 30, 40])
ddf = koalas.from_pandas(df)
self.assert_eq(ddf.loc[20:30], df.loc[20:30])
df = pd.DataFrame({'x': [1, 2, 3, 4]}, index=[10, 20, 20, 40])
ddf = koalas.from_pandas(df)
self.assert_eq(ddf.loc[20:20], df.loc[20:20])
def test_div(self):
pdf = self.pdf1
kdf = koalas.from_pandas(pdf)
for u in 'D', 's', 'ms':
duration = np.timedelta64(1, u)
self.assert_eq(
(kdf['end_date'] - kdf['start_date']) / duration,
(pdf['end_date'] - pdf['start_date']) / duration)
def test_to_excel(self):
with self.temp_dir() as dirpath:
pandas_location = dirpath + "/" + "output1.xlsx"
koalas_location = dirpath + "/" + "output2.xlsx"
pdf = self.pdf
kdf = self.kdf
kdf.to_excel(koalas_location)
pdf.to_excel(pandas_location)
dataframes = self.get_excel_dfs(koalas_location, pandas_location)
self.assert_eq(dataframes['got'], dataframes['expected'])
pdf = pd.DataFrame({
'a': [1, None, 3],
'b': ["one", "two", None],
}, index=[0, 1, 3])
kdf = koalas.from_pandas(pdf)
kdf.to_excel(koalas_location, na_rep='null')
pdf.to_excel(pandas_location, na_rep='null')
dataframes = self.get_excel_dfs(koalas_location, pandas_location)
self.assert_eq(dataframes['got'], dataframes['expected'])
pdf = pd.DataFrame({
'a': [1.0, 2.0, 3.0],
'b': [4.0, 5.0, 6.0],
}, index=[0, 1, 3])
kdf = koalas.from_pandas(pdf)
kdf.to_excel(koalas_location, float_format='%.1f')
pdf.to_excel(pandas_location, float_format='%.1f')
dataframes = self.get_excel_dfs(koalas_location, pandas_location)
self.assert_eq(dataframes['got'], dataframes['expected'])
kdf.to_excel(koalas_location, header=False)
pdf.to_excel(pandas_location, header=False)
dataframes = self.get_excel_dfs(koalas_location, pandas_location)
self.assert_eq(dataframes['got'], dataframes['expected'])
kdf.to_excel(koalas_location, index=False)
pdf.to_excel(pandas_location, index=False)
dataframes = self.get_excel_dfs(koalas_location, pandas_location)
self.assert_eq(dataframes['got'], dataframes['expected'])
def test_getitem_slice(self):
df = pd.DataFrame({'A': [1, 2, 3, 4, 5, 6, 7, 8, 9],
'B': [9, 8, 7, 6, 5, 4, 3, 2, 1],
'C': [True, False, True] * 3},
index=list('abcdefghi'))
ddf = koalas.from_pandas(df)
self.assert_eq(ddf['a':'e'], df['a':'e'])
self.assert_eq(ddf['a':'b'], df['a':'b'])
self.assert_eq(ddf['f':], df['f':])
def test_abs(self):
df = pd.DataFrame({'A': [1, -2, 3, -4, 5],
'B': [1., -2, 3, -4, 5],
'C': [-6., -7, -8, -9, 10],
'D': ['a', 'b', 'c', 'd', 'e']})
ddf = koalas.from_pandas(df)
self.assert_eq(ddf.A.abs(), df.A.abs())
self.assert_eq(ddf.B.abs(), df.B.abs())
self.assert_eq(ddf[['B', 'C']].abs(), df[['B', 'C']].abs())
def test_loc_datetime_no_freq(self):
datetime_index = pd.date_range('2016-01-01', '2016-01-31', freq='12h')
datetime_index.freq = None # FORGET FREQUENCY
df = pd.DataFrame({'num': range(len(datetime_index))}, index=datetime_index)
ddf = koalas.from_pandas(df)
slice_ = slice('2016-01-03', '2016-01-05')
result = ddf.loc[slice_, :]
expected = df.loc[slice_, :]
self.assert_eq(result, expected)
def test_getitem_period_str(self):
df = pd.DataFrame({'A': np.random.randn(100), 'B': np.random.randn(100)},
index=pd.period_range('2011-01-01', freq='H', periods=100))
ddf = koalas.from_pandas(df)
# partial string slice
# TODO?: self.assert_eq(df['2011-01-02'],
# TODO?: ddf['2011-01-02'])
self.assert_eq(df['2011-01-02':'2011-01-05'],
ddf['2011-01-02':'2011-01-05'])
df = pd.DataFrame({'A': np.random.randn(100), 'B': np.random.randn(100)},
index=pd.period_range('2011-01-01', freq='M', periods=100))
ddf = koalas.from_pandas(df)
# TODO?: self.assert_eq(df['2011-01'], ddf['2011-01'])
# TODO?: self.assert_eq(df['2011'], ddf['2011'])
self.assert_eq(df['2011-01':'2012-05'], ddf['2011-01':'2012-05'])
self.assert_eq(df['2011':'2015'], ddf['2011':'2015'])
def test_stats_on_boolean_series(self):
s = pd.Series([True, False, True])
ds = koalas.from_pandas(s)
self.assertEqual(ds.min(), s.min())
self.assertEqual(ds.max(), s.max())
self.assertEqual(ds.sum(), s.sum())
self.assertEqual(ds.mean(), s.mean())
self.assertAlmostEqual(ds.var(), s.var())
self.assertAlmostEqual(ds.std(), s.std())
def test_loc2d_with_known_divisions(self):
df = pd.DataFrame(np.random.randn(20, 5),
index=list('abcdefghijklmnopqrst'),
columns=list('ABCDE'))
ddf = koalas.from_pandas(df)
self.assert_eq(ddf.loc[['a'], 'A'], df.loc[['a'], 'A'])
self.assert_eq(ddf.loc[['a'], ['A']], df.loc[['a'], ['A']])
self.assert_eq(ddf.loc['a':'o', 'A'], df.loc['a':'o', 'A'])
self.assert_eq(ddf.loc['a':'o', ['A']], df.loc['a':'o', ['A']])
self.assert_eq(ddf.loc[['n'], ['A']], df.loc[['n'], ['A']])
self.assert_eq(ddf.loc[['a', 'c', 'n'], ['A']], df.loc[['a', 'c', 'n'], ['A']])
def test_stats_on_boolean_dataframe(self):
df = pd.DataFrame({'A': [True, False, True],
'B': [False, False, True]})
ddf = koalas.from_pandas(df)
pd.testing.assert_series_equal(ddf.min(), df.min())
pd.testing.assert_series_equal(ddf.max(), df.max())
pd.testing.assert_series_equal(ddf.sum(), df.sum())
pd.testing.assert_series_equal(ddf.mean(), df.mean())
pd.testing.assert_series_equal(ddf.var(), df.var())
pd.testing.assert_series_equal(ddf.std(), df.std())
def test_cov_corr_meta(self):
# Disable arrow execution since corr() is using UDT internally which is not supported.
with self.sql_conf({'spark.sql.execution.arrow.enabled': False}):
df = pd.DataFrame({'a': np.array([1, 2, 3], dtype='i1'),
'b': np.array([1, 2, 3], dtype='i2'),
'c': np.array([1, 2, 3], dtype='i4'),
'd': np.array([1, 2, 3]),
'e': np.array([1.0, 2.0, 3.0], dtype='f4'),
'f': np.array([1.0, 2.0, 3.0]),
'g': np.array([True, False, True]),
'h': np.array(list('abc'))},
index=pd.Index([1, 2, 3], name='myindex'))
ddf = koalas.from_pandas(df)
self.assert_eq(ddf.corr(), df.corr())
def test_to_json(self):
pdf = self.pdf
kdf = koalas.from_pandas(pdf)
self.assert_eq(kdf.to_json(), pdf.to_json())
self.assert_eq(kdf.to_json(orient='split'), pdf.to_json(orient='split'))
self.assert_eq(kdf.to_json(orient='records'), pdf.to_json(orient='records'))
self.assert_eq(kdf.to_json(orient='index'), pdf.to_json(orient='index'))
self.assert_eq(kdf.to_json(orient='values'), pdf.to_json(orient='values'))
self.assert_eq(kdf.to_json(orient='table'), pdf.to_json(orient='table'))
self.assert_eq(kdf.to_json(orient='records', lines=True),
pdf.to_json(orient='records', lines=True))
self.assert_eq(kdf.to_json(orient='split', index=False),
pdf.to_json(orient='split', index=False))
def test_getitem(self):
df = pd.DataFrame({'A': [1, 2, 3, 4, 5, 6, 7, 8, 9],
'B': [9, 8, 7, 6, 5, 4, 3, 2, 1],
'C': [True, False, True] * 3},
columns=list('ABC'))
ddf = koalas.from_pandas(df)
self.assert_eq(ddf['A'], df['A'])
self.assert_eq(ddf[['A', 'B']], df[['A', 'B']])
self.assert_eq(ddf[ddf.C], df[df.C])
self.assertRaises(KeyError, lambda: ddf['X'])
self.assertRaises(KeyError, lambda: ddf[['A', 'X']])
self.assertRaises(AttributeError, lambda: ddf.X)
def test_stat_functions(self):
df = pd.DataFrame({'A': [1, 2, 3, 4],
'B': [1.0, 2.1, 3, 4]})
ddf = koalas.from_pandas(df)
functions = ['max', 'min', 'mean', 'sum']
for funcname in functions:
self.assertEqual(getattr(ddf.A, funcname)(), getattr(df.A, funcname)())
self.assert_eq(getattr(ddf, funcname)(), getattr(df, funcname)())
functions = ['std', 'var']
for funcname in functions:
self.assertAlmostEqual(getattr(ddf.A, funcname)(), getattr(df.A, funcname)())
self.assertPandasAlmostEqual(getattr(ddf, funcname)(), getattr(df, funcname)())
# NOTE: To test skew and kurt, just make sure they run.
# The numbers are different in spark and pandas.
functions = ['skew', 'kurt']
for funcname in functions:
getattr(ddf.A, funcname)()
getattr(ddf, funcname)()
def test_loc2d_duplicated_columns(self):
df = pd.DataFrame(np.random.randn(20, 5),
index=list('abcdefghijklmnopqrst'),
columns=list('AABCD'))
ddf = koalas.from_pandas(df)
# TODO?: self.assert_eq(ddf.loc[['a'], 'A'], df.loc[['a'], 'A'])
# TODO?: self.assert_eq(ddf.loc[['a'], ['A']], df.loc[['a'], ['A']])
self.assert_eq(ddf.loc[['j'], 'B'], df.loc[['j'], 'B'])
self.assert_eq(ddf.loc[['j'], ['B']], df.loc[['j'], ['B']])
# TODO?: self.assert_eq(ddf.loc['a':'o', 'A'], df.loc['a':'o', 'A'])
# TODO?: self.assert_eq(ddf.loc['a':'o', ['A']], df.loc['a':'o', ['A']])
self.assert_eq(ddf.loc['j':'q', 'B'], df.loc['j':'q', 'B'])
self.assert_eq(ddf.loc['j':'q', ['B']], df.loc['j':'q', ['B']])
# TODO?: self.assert_eq(ddf.loc['a':'o', 'B':'D'], df.loc['a':'o', 'B':'D'])
# TODO?: self.assert_eq(ddf.loc['a':'o', 'B':'D'], df.loc['a':'o', 'B':'D'])
# TODO?: self.assert_eq(ddf.loc['j':'q', 'B':'A'], df.loc['j':'q', 'B':'A'])
# TODO?: self.assert_eq(ddf.loc['j':'q', 'B':'A'], df.loc['j':'q', 'B':'A'])
self.assert_eq(ddf.loc[ddf.B > 0, 'B'], df.loc[df.B > 0, 'B'])
def test_corr(self):
# Disable arrow execution since corr() is using UDT internally which is not supported.
with self.sql_conf({'spark.sql.execution.arrow.enabled': False}):
# DataFrame
# we do not handle NaNs for now
df = pd.util.testing.makeMissingDataframe(0.3, 42).fillna(0)
ddf = koalas.from_pandas(df)
res = ddf.corr()
sol = df.corr()
self.assertPandasAlmostEqual(res, sol)
# Series
a = df.A
b = df.B
da = ddf.A
db = ddf.B
res = da.corr(db)
sol = a.corr(b)
self.assertAlmostEqual(res, sol)
self.assertRaises(TypeError, lambda: da.corr(ddf))
def test_string_extractall(self):
kser = ks.from_pandas(self.pser)
with self.assertRaises(NotImplementedError):
kser.str.extractall("pat")
def kdf(self):
return koalas.from_pandas(self.pdf)
def test_subtraction(self):
pdf = self.pdf1
kdf = koalas.from_pandas(pdf)
kdf['diff_seconds'] = kdf['end_date'] - kdf['start_date'] - 1
self.assertEqual(list(kdf['diff_seconds'].toPandas()), [35545499, 33644699, 31571099])
def kdf1(self):
return ks.from_pandas(self.pdf1)
def test_assignment(self):
with self.assertRaisesRegex(ValueError,
"Cannot combine column argument"):
kdf = ks.from_pandas(self.pdf1)
kdf['c'] = self.kdf1.a
def test_cumprod(self):
pser = pd.Series([1.0, None, 1.0, 4.0, 9.0]).rename("a")
kser = koalas.from_pandas(pser)
self.assertEqual(repr(pser.cumprod()), repr(kser.cumprod()))
self.assertEqual(repr(pser.cumprod(skipna=False)),
repr(kser.cumprod(skipna=False)))
def transform_batch(self, func, *args, **kwargs):
"""
Transform chunks with a function that takes pandas DataFrame and outputs pandas DataFrame.
The pandas DataFrame given to the function is of a batch used internally. The length of
each input and output should be the same.
See also `Transform and apply a function
<https://koalas.readthedocs.io/en/latest/user_guide/transform_apply.html>`_.
.. note:: the `func` is unable to access to the whole input frame. Koalas internally
splits the input series into multiple batches and calls `func` with each batch multiple
times. Therefore, operations such as global aggregations are impossible. See the example
below.
>>> # This case does not return the length of whole frame but of the batch internally
... # used.
... def length(pdf) -> ks.DataFrame[int]:
... return pd.DataFrame([len(pdf)] * len(pdf))
...
>>> df = ks.DataFrame({'A': range(1000)})
>>> df.koalas.transform_batch(length) # doctest: +SKIP
c0
0 83
1 83
2 83
...
.. note:: this API executes the function once to infer the type which is
potentially expensive, for instance, when the dataset is created after
aggregations or sorting.
To avoid this, specify return type in ``func``, for instance, as below:
>>> def plus_one(x) -> ks.DataFrame[float, float]:
... return x + 1
If the return type is specified, the output column names become
`c0, c1, c2 ... cn`. These names are positionally mapped to the returned
DataFrame in ``func``.
To specify the column names, you can assign them in a pandas friendly style as below:
>>> def plus_one(x) -> ks.DataFrame['a': float, 'b': float]:
... return x + 1
>>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
>>> def plus_one(x) -> ks.DataFrame[zip(pdf.dtypes, pdf.columns)]:
... return x + 1
Parameters
----------
func : function
Function to transform each pandas frame.
*args
Positional arguments to pass to func.
**kwargs
Keyword arguments to pass to func.
Returns
-------
DataFrame
See Also
--------
DataFrame.koalas.apply_batch: For row/columnwise operations.
Series.koalas.transform_batch: transform the search as each pandas chunks.
Examples
--------
>>> df = ks.DataFrame([(1, 2), (3, 4), (5, 6)], columns=['A', 'B'])
>>> df
A B
0 1 2
1 3 4
2 5 6
>>> def plus_one_func(pdf) -> ks.DataFrame[int, int]:
... return pdf + 1
>>> df.koalas.transform_batch(plus_one_func)
c0 c1
0 2 3
1 4 5
2 6 7
>>> def plus_one_func(pdf) -> ks.DataFrame['A': int, 'B': int]:
... return pdf + 1
>>> df.koalas.transform_batch(plus_one_func)
A B
0 2 3
1 4 5
2 6 7
>>> def plus_one_func(pdf) -> ks.Series[int]:
... return pdf.B + 1
>>> df.koalas.transform_batch(plus_one_func)
0 3
1 5
2 7
dtype: int32
You can also omit the type hints so Koalas infers the return schema as below:
>>> df.koalas.transform_batch(lambda pdf: pdf + 1)
A B
0 2 3
1 4 5
2 6 7
>>> (df * -1).koalas.transform_batch(abs)
A B
0 1 2
1 3 4
2 5 6
Note that you should not transform the index. The index information will not change.
>>> df.koalas.transform_batch(lambda pdf: pdf.B + 1)
0 3
1 5
2 7
Name: B, dtype: int64
You can also specify extra arguments as below.
>>> df.koalas.transform_batch(lambda pdf, a, b, c: pdf.B + a + b + c, 1, 2, c=3)
0 8
1 10
2 12
Name: B, dtype: int64
"""
from databricks.koalas.groupby import GroupBy
from databricks.koalas.frame import DataFrame
from databricks.koalas.series import first_series
from databricks import koalas as ks
assert callable(func), "the first argument should be a callable function."
spec = inspect.getfullargspec(func)
return_sig = spec.annotations.get("return", None)
should_infer_schema = return_sig is None
original_func = func
func = lambda o: original_func(o, *args, **kwargs)
names = self._kdf._internal.to_internal_spark_frame.schema.names
should_by_pass = LooseVersion(pyspark.__version__) >= "3.0"
def pandas_concat(series):
# The input can only be a DataFrame for struct from Spark 3.0.
# This works around to make the input as a frame. See SPARK-27240
pdf = pd.concat(series, axis=1)
pdf = pdf.rename(columns=dict(zip(pdf.columns, names)))
return pdf
def pandas_extract(pdf, name):
# This is for output to work around a DataFrame for struct
# from Spark 3.0. See SPARK-23836
return pdf[name]
def pandas_series_func(f):
ff = f
return lambda *series: ff(pandas_concat(series))
def pandas_frame_func(f):
ff = f
return lambda *series: pandas_extract(ff(pandas_concat(series)), field.name)
if should_infer_schema:
# Here we execute with the first 1000 to get the return type.
# If the records were less than 1000, it uses pandas API directly for a shortcut.
limit = ks.get_option("compute.shortcut_limit")
pdf = self._kdf.head(limit + 1)._to_internal_pandas()
transformed = func(pdf)
if not isinstance(transformed, (pd.DataFrame, pd.Series)):
raise ValueError(
"The given function should return a frame; however, "
"the return type was %s." % type(transformed)
)
if len(transformed) != len(pdf):
raise ValueError("transform_batch cannot produce aggregated results")
kdf_or_kser = ks.from_pandas(transformed)
if isinstance(kdf_or_kser, ks.Series):
kser = kdf_or_kser
pudf = pandas_udf(
func if should_by_pass else pandas_series_func(func),
returnType=kser.spark.data_type,
functionType=PandasUDFType.SCALAR,
)
columns = self._kdf._internal.spark_columns
# TODO: Index will be lost in this case.
internal = self._kdf._internal.copy(
column_labels=kser._internal.column_labels,
data_spark_columns=[
(pudf(F.struct(*columns)) if should_by_pass else pudf(*columns)).alias(
kser._internal.data_spark_column_names[0]
)
],
column_label_names=kser._internal.column_label_names,
)
return first_series(DataFrame(internal))
else:
kdf = kdf_or_kser
if len(pdf) <= limit:
# only do the short cut when it returns a frame to avoid
# operations on different dataframes in case of series.
return kdf
return_schema = kdf._internal.to_internal_spark_frame.schema
# Force nullability.
return_schema = StructType(
[StructField(field.name, field.dataType) for field in return_schema.fields]
)
self_applied = DataFrame(self._kdf._internal.resolved_copy)
output_func = GroupBy._make_pandas_df_builder_func(
self_applied, func, return_schema, retain_index=True
)
columns = self_applied._internal.spark_columns
if should_by_pass:
pudf = pandas_udf(
output_func, returnType=return_schema, functionType=PandasUDFType.SCALAR
)
temp_struct_column = verify_temp_column_name(
self_applied._internal.spark_frame, "__temp_struct__"
)
applied = pudf(F.struct(*columns)).alias(temp_struct_column)
sdf = self_applied._internal.spark_frame.select(applied)
sdf = sdf.selectExpr("%s.*" % temp_struct_column)
else:
applied = []
for field in return_schema.fields:
applied.append(
pandas_udf(
pandas_frame_func(output_func),
returnType=field.dataType,
functionType=PandasUDFType.SCALAR,
)(*columns).alias(field.name)
)
sdf = self_applied._internal.spark_frame.select(*applied)
return DataFrame(kdf._internal.with_new_sdf(sdf))
else:
return_type = infer_return_type(original_func)
return_schema = return_type.tpe
is_return_series = isinstance(return_type, SeriesType)
is_return_dataframe = isinstance(return_type, DataFrameType)
if not is_return_dataframe and not is_return_series:
raise TypeError(
"The given function should specify a frame or series as its type "
"hints; however, the return type was %s." % return_sig
)
if is_return_series:
pudf = pandas_udf(
func if should_by_pass else pandas_series_func(func),
returnType=return_schema,
functionType=PandasUDFType.SCALAR,
)
columns = self._kdf._internal.spark_columns
internal = self._kdf._internal.copy(
column_labels=[None],
data_spark_columns=[
(pudf(F.struct(*columns)) if should_by_pass else pudf(*columns)).alias(
SPARK_DEFAULT_SERIES_NAME
)
],
column_label_names=None,
)
return first_series(DataFrame(internal))
else:
self_applied = DataFrame(self._kdf._internal.resolved_copy)
output_func = GroupBy._make_pandas_df_builder_func(
self_applied, func, return_schema, retain_index=False
)
columns = self_applied._internal.spark_columns
if should_by_pass:
pudf = pandas_udf(
output_func, returnType=return_schema, functionType=PandasUDFType.SCALAR
)
temp_struct_column = verify_temp_column_name(
self_applied._internal.spark_frame, "__temp_struct__"
)
applied = pudf(F.struct(*columns)).alias(temp_struct_column)
sdf = self_applied._internal.spark_frame.select(applied)
sdf = sdf.selectExpr("%s.*" % temp_struct_column)
else:
applied = []
for field in return_schema.fields:
applied.append(
pandas_udf(
pandas_frame_func(output_func),
returnType=field.dataType,
functionType=PandasUDFType.SCALAR,
)(*columns).alias(field.name)
)
sdf = self_applied._internal.spark_frame.select(*applied)
return DataFrame(sdf)
def test_to_numpy(self):
pser = pd.Series([1, 2, 3, 4, 5, 6, 7], name='x')
kser = ks.from_pandas(pser)
np.testing.assert_equal(kser.to_numpy(), pser.values)
def ks_start_date(self):
return ks.from_pandas(self.pd_start_date)
def test_aggregate(self):
pdf = pd.DataFrame({
'A': [1, 1, 2, 2],
'B': [1, 2, 3, 4],
'C': [0.362, 0.227, 1.267, -0.562]
})
kdf = koalas.from_pandas(pdf)
for as_index in [True, False]:
self.assert_eq(
kdf.groupby('A', as_index=as_index).agg({
'B': 'min',
'C': 'sum'
}),
pdf.groupby('A', as_index=as_index).agg({
'B': 'min',
'C': 'sum'
}))
self.assert_eq(
kdf.groupby('A', as_index=as_index).agg({
'B': ['min', 'max'],
'C': 'sum'
}),
pdf.groupby('A', as_index=as_index).agg({
'B': ['min', 'max'],
'C': 'sum'
}))
expected_error_message = (
r"aggs must be a dict mapping from column name \(string or "
r"tuple\) to aggregate functions \(string or list of strings\).")
with self.assertRaisesRegex(ValueError, expected_error_message):
kdf.groupby('A', as_index=as_index).agg(0)
# multi-index columns
columns = pd.MultiIndex.from_tuples([('X', 'A'), ('X', 'B'),
('Y', 'C')])
pdf.columns = columns
kdf.columns = columns
for as_index in [True, False]:
self.assert_eq(
kdf.groupby(('X', 'A'), as_index=as_index).agg({
('X', 'B'):
'min',
('Y', 'C'):
'sum'
}),
pdf.groupby(('X', 'A'), as_index=as_index).agg({
('X', 'B'):
'min',
('Y', 'C'):
'sum'
}))
self.assert_eq(
kdf.groupby(('X', 'A')).agg({
('X', 'B'): ['min', 'max'],
('Y', 'C'): 'sum'
}),
pdf.groupby(('X', 'A')).agg({
('X', 'B'): ['min', 'max'],
('Y', 'C'): 'sum'
}))
def test_loc(self):
kdf = self.kdf
pdf = self.pdf
self.assert_eq(kdf.loc[5:5], pdf.loc[5:5])
self.assert_eq(kdf.loc[3:8], pdf.loc[3:8])
self.assert_eq(kdf.loc[:8], pdf.loc[:8])
self.assert_eq(kdf.loc[3:], pdf.loc[3:])
self.assert_eq(kdf.loc[[5]], pdf.loc[[5]])
self.assert_eq(kdf.loc[:], pdf.loc[:])
# TODO?: self.assert_eq(kdf.loc[[3, 4, 1, 8]], pdf.loc[[3, 4, 1, 8]])
# TODO?: self.assert_eq(kdf.loc[[3, 4, 1, 9]], pdf.loc[[3, 4, 1, 9]])
# TODO?: self.assert_eq(kdf.loc[np.array([3, 4, 1, 9])], pdf.loc[np.array([3, 4, 1, 9])])
self.assert_eq(kdf.a.loc[5:5], pdf.a.loc[5:5])
self.assert_eq(kdf.a.loc[3:8], pdf.a.loc[3:8])
self.assert_eq(kdf.a.loc[:8], pdf.a.loc[:8])
self.assert_eq(kdf.a.loc[3:], pdf.a.loc[3:])
self.assert_eq(kdf.a.loc[[5]], pdf.a.loc[[5]])
# TODO?: self.assert_eq(kdf.a.loc[[3, 4, 1, 8]], pdf.a.loc[[3, 4, 1, 8]])
# TODO?: self.assert_eq(kdf.a.loc[[3, 4, 1, 9]], pdf.a.loc[[3, 4, 1, 9]])
# TODO?: self.assert_eq(kdf.a.loc[np.array([3, 4, 1, 9])],
# pdf.a.loc[np.array([3, 4, 1, 9])])
self.assert_eq(kdf.a.loc[[]], pdf.a.loc[[]])
self.assert_eq(kdf.a.loc[np.array([])], pdf.a.loc[np.array([])])
self.assert_eq(kdf.loc[1000:], pdf.loc[1000:])
self.assert_eq(kdf.loc[-2000:-1000], pdf.loc[-2000:-1000])
self.assert_eq(kdf.loc[5], pdf.loc[5])
self.assert_eq(kdf.loc[9], pdf.loc[9])
self.assert_eq(kdf.a.loc[5], pdf.a.loc[5])
self.assert_eq(kdf.a.loc[9], pdf.a.loc[9])
self.assertRaises(KeyError, lambda: kdf.loc[10])
self.assertRaises(KeyError, lambda: kdf.a.loc[10])
# monotonically increasing index test
pdf = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]},
index=[0, 1, 1, 2, 2, 2, 4, 5, 6])
kdf = ks.from_pandas(pdf)
self.assert_eq(kdf.loc[:2], pdf.loc[:2])
self.assert_eq(kdf.loc[:3], pdf.loc[:3])
self.assert_eq(kdf.loc[3:], pdf.loc[3:])
self.assert_eq(kdf.loc[4:], pdf.loc[4:])
self.assert_eq(kdf.loc[3:2], pdf.loc[3:2])
self.assert_eq(kdf.loc[-1:2], pdf.loc[-1:2])
self.assert_eq(kdf.loc[3:10], pdf.loc[3:10])
# monotonically decreasing index test
pdf = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]},
index=[6, 5, 5, 4, 4, 4, 2, 1, 0])
kdf = ks.from_pandas(pdf)
self.assert_eq(kdf.loc[:4], pdf.loc[:4])
self.assert_eq(kdf.loc[:3], pdf.loc[:3])
self.assert_eq(kdf.loc[3:], pdf.loc[3:])
self.assert_eq(kdf.loc[2:], pdf.loc[2:])
self.assert_eq(kdf.loc[2:3], pdf.loc[2:3])
self.assert_eq(kdf.loc[2:-1], pdf.loc[2:-1])
self.assert_eq(kdf.loc[10:3], pdf.loc[10:3])
# test when type of key is string and given value is not included in key
pdf = pd.DataFrame({"a": [1, 2, 3]}, index=["a", "b", "d"])
kdf = ks.from_pandas(pdf)
self.assert_eq(kdf.loc["a":"z"], pdf.loc["a":"z"])
# KeyError when index is not monotonic increasing or decreasing
# and specified values don't exist in index
kdf = ks.DataFrame([[1, 2], [4, 5], [7, 8]],
index=["cobra", "viper", "sidewinder"])
self.assertRaises(KeyError, lambda: kdf.loc["cobra":"koalas"])
self.assertRaises(KeyError, lambda: kdf.loc["koalas":"viper"])
kdf = ks.DataFrame([[1, 2], [4, 5], [7, 8]], index=[10, 30, 20])
self.assertRaises(KeyError, lambda: kdf.loc[0:30])
self.assertRaises(KeyError, lambda: kdf.loc[10:100])
def kidxs(self):
return [ks.from_pandas(pidx) for pidx in self.pidxs]
def test_stat_functions(self):
pdf = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [1.0, 2.1, 3, 4]})
kdf = ks.from_pandas(pdf)
self._test_stat_functions(pdf, kdf)
def ks(self):
return koalas.from_pandas(self.ps)
def test_assignment(self):
with self.assertRaisesRegex(ValueError,
"Cannot combine the series or dataframe"):
kdf = ks.from_pandas(self.pdf1)
kdf["c"] = self.kdf1.a
def kser(self):
return ks.from_pandas(self.pser)
def test_series_loc_setitem(self):
pdf = pd.DataFrame({
"x": [1, 2, 3],
"y": [4, 5, 6]
},
index=["cobra", "viper", "sidewinder"])
kdf = ks.from_pandas(pdf)
pser = pdf.x
psery = pdf.y
kser = kdf.x
ksery = kdf.y
pser_another = pd.Series([1, 2, 3],
index=["cobra", "viper", "sidewinder"])
kser_another = ks.from_pandas(pser_another)
kser.loc[kser % 2 == 1] = -kser_another
pser.loc[pser % 2 == 1] = -pser_another
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
pdf = pd.DataFrame({
"x": [1, 2, 3],
"y": [4, 5, 6]
},
index=["cobra", "viper", "sidewinder"])
kdf = ks.from_pandas(pdf)
pser = pdf.x
psery = pdf.y
kser = kdf.x
ksery = kdf.y
kser.loc[kser_another % 2 == 1] = -kser
pser.loc[pser_another % 2 == 1] = -pser
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
pdf = pd.DataFrame({
"x": [1, 2, 3],
"y": [4, 5, 6]
},
index=["cobra", "viper", "sidewinder"])
kdf = ks.from_pandas(pdf)
pser = pdf.x
psery = pdf.y
kser = kdf.x
ksery = kdf.y
kser.loc[kser_another % 2 == 1] = -kser
pser.loc[pser_another % 2 == 1] = -pser
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
pdf = pd.DataFrame({
"x": [1, 2, 3],
"y": [4, 5, 6]
},
index=["cobra", "viper", "sidewinder"])
kdf = ks.from_pandas(pdf)
pser = pdf.x
psery = pdf.y
kser = kdf.x
ksery = kdf.y
kser.loc[kser_another % 2 == 1] = -kser_another
pser.loc[pser_another % 2 == 1] = -pser_another
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
pdf = pd.DataFrame({
"x": [1, 2, 3],
"y": [4, 5, 6]
},
index=["cobra", "viper", "sidewinder"])
kdf = ks.from_pandas(pdf)
pser = pdf.x
psery = pdf.y
kser = kdf.x
ksery = kdf.y
kser.loc[["viper", "sidewinder"]] = -kser_another
pser.loc[["viper", "sidewinder"]] = -pser_another
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
pdf = pd.DataFrame({
"x": [1, 2, 3],
"y": [4, 5, 6]
},
index=["cobra", "viper", "sidewinder"])
kdf = ks.from_pandas(pdf)
pser = pdf.x
psery = pdf.y
kser = kdf.x
ksery = kdf.y
kser.loc[kser_another % 2 == 1] = 10
pser.loc[pser_another % 2 == 1] = 10
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
def ks(self):
return koalas.from_pandas(self.ps)
def test_add_suffix(self):
ps = pd.Series([1, 2, 3, 4], name='0')
ks = koalas.from_pandas(ps)
self.assert_eq(ps.add_suffix('_item'), ks.add_suffix('_item'))
def test_string_add_str_num(self):
pdf = pd.DataFrame(dict(col1=['a'], col2=[1]))
ds = koalas.from_pandas(pdf)
with self.assertRaises(TypeError):
ds['col1'] + ds['col2']
def test_loc_multiindex(self):
kdf = self.kdf
kdf = kdf.set_index("b", append=True)
pdf = self.pdf
pdf = pdf.set_index("b", append=True)
self.assert_eq(kdf.loc[:], pdf.loc[:])
self.assert_eq(kdf.loc[5:5], pdf.loc[5:5])
self.assert_eq(kdf.loc[5:9], pdf.loc[5:9])
self.assert_eq(kdf.loc[5], pdf.loc[5])
self.assert_eq(kdf.loc[9], pdf.loc[9])
# TODO: self.assert_eq(kdf.loc[(5, 3)], pdf.loc[(5, 3)])
# TODO: self.assert_eq(kdf.loc[(9, 0)], pdf.loc[(9, 0)])
self.assert_eq(kdf.a.loc[5], pdf.a.loc[5])
self.assert_eq(kdf.a.loc[9], pdf.a.loc[9])
self.assertTrue((kdf.a.loc[(5, 3)] == pdf.a.loc[(5, 3)]).all())
self.assert_eq(kdf.a.loc[(9, 0)], pdf.a.loc[(9, 0)])
# monotonically increasing index test
pdf = pd.DataFrame(
{"a": [1, 2, 3, 4, 5]},
index=pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"),
("y", "c"), ("y", "d"),
("z", "e")]),
)
kdf = ks.from_pandas(pdf)
for rows_sel in [
slice(None),
slice("y", None),
slice(None, "y"),
slice(("x", "b"), None),
slice(None, ("y", "c")),
slice(("x", "b"), ("y", "c")),
slice("x", ("y", "c")),
slice(("x", "b"), "y"),
]:
with self.subTest("monotonically increasing", rows_sel=rows_sel):
self.assert_eq(kdf.loc[rows_sel], pdf.loc[rows_sel])
self.assert_eq(kdf.a.loc[rows_sel], pdf.a.loc[rows_sel])
# monotonically increasing first index test
pdf = pd.DataFrame(
{"a": [1, 2, 3, 4, 5]},
index=pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"),
("y", "c"), ("y", "a"),
("z", "e")]),
)
kdf = ks.from_pandas(pdf)
for rows_sel in [
slice(None),
slice("y", None),
slice(None, "y"),
]:
with self.subTest("monotonically increasing first index",
rows_sel=rows_sel):
self.assert_eq(kdf.loc[rows_sel], pdf.loc[rows_sel])
self.assert_eq(kdf.a.loc[rows_sel], pdf.a.loc[rows_sel])
for rows_sel in [
slice(("x", "b"), None),
slice(None, ("y", "c")),
slice(("x", "b"), ("y", "c")),
slice("x", ("y", "c")),
slice(("x", "b"), "y"),
]:
with self.subTest("monotonically increasing first index",
rows_sel=rows_sel):
self.assertRaises(KeyError, lambda: kdf.loc[rows_sel])
self.assertRaises(KeyError, lambda: kdf.a.loc[rows_sel])
# not monotonically increasing index test
pdf = pd.DataFrame(
{"a": [1, 2, 3, 4, 5]},
index=pd.MultiIndex.from_tuples([("z", "e"), ("y", "d"),
("y", "c"), ("x", "b"),
("x", "a")]),
)
kdf = ks.from_pandas(pdf)
for rows_sel in [
slice("y", None),
slice(None, "y"),
slice(("x", "b"), None),
slice(None, ("y", "c")),
slice(("x", "b"), ("y", "c")),
slice("x", ("y", "c")),
slice(("x", "b"), "y"),
]:
with self.subTest("monotonically decreasing", rows_sel=rows_sel):
self.assertRaises(KeyError, lambda: kdf.loc[rows_sel])
self.assertRaises(KeyError, lambda: kdf.a.loc[rows_sel])
def test_string_add_str_num(self):
pdf = pd.DataFrame(dict(col1=["a"], col2=[1]))
kdf = ks.from_pandas(pdf)
with self.assertRaises(TypeError):
kdf["col1"] + kdf["col2"]
def kdf(self):
return koalas.from_pandas(self.pdf)
def df(self):
return koalas.from_pandas(self.full)
def test_string_add_str_lit(self):
pdf = pd.DataFrame(dict(col1=["a", "b", "c"]))
kdf = ks.from_pandas(pdf)
self.assert_eq(kdf["col1"] + "_lit", pdf["col1"] + "_lit")
self.assert_eq("_lit" + kdf["col1"], "_lit" + pdf["col1"])
def test_stat_functions(self):
pdf = pd.DataFrame({"A": [1, 2, 3, 4], "B": [1, 2, 3, 4]})
kdf = ks.from_pandas(pdf)
self._test_stat_functions(pdf, kdf)
def test_to_numpy(self):
s = pd.Series([1, 2, 3, 4, 5, 6, 7], name='x')
ddf = koalas.from_pandas(s)
np.testing.assert_equal(ddf.to_numpy(), s.values)
def test_index_nlevels(self):
pdf = pd.DataFrame({"a": [1, 2, 3]}, index=pd.Index(['a', 'b', 'c']))
kdf = ks.from_pandas(pdf)
self.assertEqual(kdf.index.nlevels, 1)
def test_to_datetime(self):
ps = pd.Series(['3/11/2000', '3/12/2000', '3/13/2000'] * 100)
ks = koalas.from_pandas(ps)
self.assert_eq(pd.to_datetime(ps, infer_datetime_format=True),
koalas.to_datetime(ks, infer_datetime_format=True))
def test_loc_on_pandas_datetimes(self):
df = pd.DataFrame({'x': [1, 2, 3]},
index=list(map(pd.Timestamp, ['2014', '2015', '2016'])))
a = koalas.from_pandas(df)
self.assert_eq(a.loc['2014':'2015'], df.loc['2014':'2015'])
def test_multiindex_nlevel(self):
pdf = pd.DataFrame({'a': [1, 2, 3]}, index=[list('abc'), list('def')])
kdf = ks.from_pandas(pdf)
self.assertEqual(kdf.index.nlevels, 2)
def test_string_encode(self):
kser = ks.from_pandas(self.pser)
with self.assertRaises(NotImplementedError):
kser.str.encode("utf-8")
def _test_groupby_expanding_func(self, f):
pser = pd.Series([1, 2, 3, 2], index=np.random.rand(4), name="a")
kser = ks.from_pandas(pser)
self.assert_eq(
getattr(kser.groupby(kser).expanding(2), f)().sort_index(),
getattr(pser.groupby(pser).expanding(2), f)().sort_index(),
almost=True,
)
self.assert_eq(
getattr(kser.groupby(kser).expanding(2), f)().sum(),
getattr(pser.groupby(pser).expanding(2), f)().sum(),
almost=True,
)
# Multiindex
pser = pd.Series(
[1, 2, 3, 2],
index=pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z"), ("c", "z")]),
name="a",
)
kser = ks.from_pandas(pser)
self.assert_eq(
getattr(kser.groupby(kser).expanding(2), f)().sort_index(),
getattr(pser.groupby(pser).expanding(2), f)().sort_index(),
almost=True,
)
pdf = pd.DataFrame({"a": [1.0, 2.0, 3.0, 2.0], "b": [4.0, 2.0, 3.0, 1.0]})
kdf = ks.from_pandas(pdf)
self.assert_eq(
getattr(kdf.groupby(kdf.a).expanding(2), f)().sort_index(),
getattr(pdf.groupby(pdf.a).expanding(2), f)().sort_index(),
almost=True,
)
self.assert_eq(
getattr(kdf.groupby(kdf.a).expanding(2), f)().sum(),
getattr(pdf.groupby(pdf.a).expanding(2), f)().sum(),
almost=True,
)
self.assert_eq(
getattr(kdf.groupby(kdf.a + 1).expanding(2), f)().sort_index(),
getattr(pdf.groupby(pdf.a + 1).expanding(2), f)().sort_index(),
almost=True,
)
self.assert_eq(
getattr(kdf.b.groupby(kdf.a).expanding(2), f)().sort_index(),
getattr(pdf.b.groupby(pdf.a).expanding(2), f)().sort_index(),
almost=True,
)
self.assert_eq(
getattr(kdf.groupby(kdf.a)["b"].expanding(2), f)().sort_index(),
getattr(pdf.groupby(pdf.a)["b"].expanding(2), f)().sort_index(),
almost=True,
)
self.assert_eq(
getattr(kdf.groupby(kdf.a)[["b"]].expanding(2), f)().sort_index(),
getattr(pdf.groupby(pdf.a)[["b"]].expanding(2), f)().sort_index(),
almost=True,
)
# Multiindex column
columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y")])
pdf.columns = columns
kdf.columns = columns
self.assert_eq(
getattr(kdf.groupby(("a", "x")).expanding(2), f)().sort_index(),
getattr(pdf.groupby(("a", "x")).expanding(2), f)().sort_index(),
almost=True,
)
self.assert_eq(
getattr(kdf.groupby([("a", "x"), ("a", "y")]).expanding(2), f)().sort_index(),
getattr(pdf.groupby([("a", "x"), ("a", "y")]).expanding(2), f)().sort_index(),
almost=True,
)
def test_series_iloc_setitem(self):
pdf = pd.DataFrame({
"x": [1, 2, 3],
"y": [4, 5, 6]
},
index=["cobra", "viper", "sidewinder"])
kdf = ks.from_pandas(pdf)
pser = pdf.x
psery = pdf.y
kser = kdf.x
ksery = kdf.y
pser1 = pser + 1
kser1 = kser + 1
pser_another = pd.Series([1, 2, 3],
index=["cobra", "viper", "sidewinder"])
kser_another = ks.from_pandas(pser_another)
kser.iloc[[1, 2]] = -kser_another
pser.iloc[[1, 2]] = -pser_another
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
kser.iloc[[0]] = 10 * kser_another
pser.iloc[[0]] = 10 * pser_another
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
kser1.iloc[[1, 2]] = -kser_another
pser1.iloc[[1, 2]] = -pser_another
self.assert_eq(kser1, pser1)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
pdf = pd.DataFrame({
"x": [1, 2, 3],
"y": [4, 5, 6]
},
index=["cobra", "viper", "sidewinder"])
kdf = ks.from_pandas(pdf)
pser = pdf.x
psery = pdf.y
kser = kdf.x
ksery = kdf.y
piloc = pser.iloc
kiloc = kser.iloc
kiloc[[1, 2]] = -kser_another
piloc[[1, 2]] = -pser_another
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
kiloc[[0]] = 10 * kser_another
piloc[[0]] = 10 * pser_another
self.assert_eq(kser, pser)
self.assert_eq(kdf, pdf)
self.assert_eq(ksery, psery)
def test_string_add_str_str(self):
pdf = pd.DataFrame(dict(col1=["a", "b", "c"], col2=["1", "2", "3"]))
kdf = ks.from_pandas(pdf)
self.assert_eq(kdf["col1"] + kdf["col2"], pdf["col1"] + pdf["col2"])
self.assert_eq(kdf["col2"] + kdf["col1"], pdf["col2"] + pdf["col1"])
