def map(
self,
df: DataFrame,
map_func: Callable[[PartitionCursor, LocalDataFrame], LocalDataFrame],
output_schema: Any,
partition_spec: PartitionSpec,
metadata: Any = None,
on_init: Optional[Callable[[int, DataFrame], Any]] = None,
) -> DataFrame:
if partition_spec.num_partitions != "0":
self.log.warning(
"%s doesn't respect num_partitions %s",
self,
partition_spec.num_partitions,
)
cursor = partition_spec.get_cursor(df.schema, 0)
if on_init is not None:
on_init(0, df)
if len(partition_spec.partition_by) == 0: # no partition
df = to_local_df(df)
cursor.set(df.peek_array(), 0, 0)
output_df = map_func(cursor, df)
if (isinstance(output_df, PandasDataFrame)
and output_df.schema != output_schema):
output_df = PandasDataFrame(output_df.native, output_schema)
assert_or_throw(
output_df.schema == output_schema,
lambda: f"map output {output_df.schema} "
f"mismatches given {output_schema}",
)
output_df._metadata = ParamDict(metadata, deep=True)
output_df._metadata.set_readonly()
return self.to_df(output_df)
presort = partition_spec.presort
presort_keys = list(presort.keys())
presort_asc = list(presort.values())
output_schema = Schema(output_schema)
def _map(pdf: pd.DataFrame) -> pd.DataFrame:
if len(presort_keys) > 0:
pdf = pdf.sort_values(presort_keys, ascending=presort_asc)
input_df = PandasDataFrame(pdf.reset_index(drop=True),
df.schema,
pandas_df_wrapper=True)
cursor.set(input_df.peek_array(), cursor.partition_no + 1, 0)
output_df = map_func(cursor, input_df)
return output_df.as_pandas()
result = self.pl_utils.safe_groupby_apply(df.as_pandas(),
partition_spec.partition_by,
_map)
return PandasDataFrame(result, output_schema, metadata)
def to_input_data(self, df: DataFrame) -> List[Dict[str, Any]]:
return list(to_local_df(df).as_dict_iterable())
def to_input_data(self, df: DataFrame) -> LocalDataFrame:
return to_local_df(df)