Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix DataFrame.apply to support additional dtypes. #2125

Merged
merged 2 commits into from
Mar 30, 2021
Merged

Fix DataFrame.apply to support additional dtypes. #2125

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5446ec7
Fix DataFrame.apply to support additional dtypes.
ueshin Mar 29, 2021
3d78a82
Merge branch 'master' into frame_apply
ueshin Mar 30, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion databricks/koalas/frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2588,6 +2588,7 @@ def apply_func(pdf):
self_applied.columns, [return_schema] * len(self_applied.columns)
)
return_schema = StructType([StructField(c, t) for c, t in fields_types])
data_dtypes = [cast(SeriesType, return_type).dtype] * len(self_applied.columns)
elif require_column_axis:
if axis != 1:
raise TypeError(
Expand All @@ -2596,11 +2597,13 @@ def apply_func(pdf):
"was %s" % return_sig
)
return_schema = cast(DataFrameType, return_type).spark_type
data_dtypes = cast(DataFrameType, return_type).dtypes
else:
# any axis is fine.
should_return_series = True
return_schema = cast(ScalarType, return_type).spark_type
return_schema = StructType([StructField(SPARK_DEFAULT_SERIES_NAME, return_schema)])
data_dtypes = [cast(ScalarType, return_type).dtype]
column_labels = [None]

if should_use_map_in_pandas:
Expand All @@ -2621,7 +2624,10 @@ def apply_func(pdf):

# Otherwise, it loses index.
internal = InternalFrame(
spark_frame=sdf, index_spark_columns=None, column_labels=column_labels
spark_frame=sdf,
index_spark_columns=None,
column_labels=column_labels,
data_dtypes=data_dtypes,
)

result = DataFrame(internal) # type: "DataFrame"
Expand Down
29 changes: 28 additions & 1 deletion databricks/koalas/tests/test_categorical.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def kdf(self):

@property
def df_pair(self):
return (self.pdf, self.kdf)
return self.pdf, self.kdf

def test_categorical_frame(self):
pdf, kdf = self.df_pair
Expand Down Expand Up @@ -106,6 +106,33 @@ def test_factorize(self):
self.assert_eq(kcodes.tolist(), pcodes.tolist())
self.assert_eq(kuniques, puniques)

def test_frame_apply(self):
pdf, kdf = self.df_pair

self.assert_eq(kdf.apply(lambda x: x).sort_index(), pdf.apply(lambda x: x).sort_index())
self.assert_eq(
kdf.apply(lambda x: x, axis=1).sort_index(), pdf.apply(lambda x: x, axis=1).sort_index()
)

def test_frame_apply_without_shortcut(self):
with ks.option_context("compute.shortcut_limit", 0):
self.test_frame_apply()

pdf = pd.DataFrame(
{"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]}
)
kdf = ks.from_pandas(pdf)

dtype = CategoricalDtype(categories=["a", "b", "c"])

def categorize(ser) -> ks.Series[dtype]:
return ser.astype(dtype)

self.assert_eq(
kdf.apply(categorize).sort_values(["a", "b"]).reset_index(drop=True),
pdf.apply(categorize).sort_values(["a", "b"]).reset_index(drop=True),
)

def test_groupby_apply(self):
pdf, kdf = self.df_pair

Expand Down