diff --git a/sdks/python/apache_beam/dataframe/partitionings.py b/sdks/python/apache_beam/dataframe/partitionings.py index 1fe760fe8589..f1346708b1d9 100644 --- a/sdks/python/apache_beam/dataframe/partitionings.py +++ b/sdks/python/apache_beam/dataframe/partitionings.py @@ -114,9 +114,15 @@ def _hash_index(self, df): for level in levels) def partition_fn(self, df, num_partitions): - hashes = self._hash_index(df) + if not len(df): + for key in range(num_partitions): + yield key, df + return + pids = self._hash_index(df) % num_partitions + empty_df = df.iloc[0:0] + groups = dict(tuple(df.groupby(pids, sort=False))) for key in range(num_partitions): - yield key, df[hashes % num_partitions == key] + yield key, groups.get(key, empty_df) def check(self, dfs): # Drop empty DataFrames diff --git a/sdks/python/apache_beam/typehints/pandas_type_compatibility.py b/sdks/python/apache_beam/typehints/pandas_type_compatibility.py index 45ae27baffe7..724aff336ab4 100644 --- a/sdks/python/apache_beam/typehints/pandas_type_compatibility.py +++ b/sdks/python/apache_beam/typehints/pandas_type_compatibility.py @@ -189,17 +189,18 @@ def explode_batch(self, batch: pd.DataFrame): # nullable types def make_null_checking_generator(series): nulls = pd.isnull(series) + if not nulls.any(): + return iter(series) return (None if isnull else value for isnull, value in zip(nulls, series)) all_series = self._get_series(batch) iterators = [make_null_checking_generator(series) for series in all_series] - - for values in zip(*iterators): - yield self._element_type.user_type( - **{ - column: value - for column, value in zip(self._columns, values) - }) + user_type = self._element_type.user_type + if issubclass(user_type, tuple) and hasattr(user_type, '_make'): + yield from map(user_type._make, zip(*iterators)) + else: + for values in zip(*iterators): + yield user_type(*values) def combine_batches(self, batches: List[pd.DataFrame]): return pd.concat(batches) @@ -267,8 +268,12 @@ def __init__( if is_nullable(element_type): def unbatch(series): - for isnull, value in zip(pd.isnull(series), series): - yield None if isnull else value + nulls = pd.isnull(series) + if not nulls.any(): + yield from series + else: + for isnull, value in zip(nulls, series): + yield None if isnull else value else: def unbatch(series):