def quick_concat(dfs, dimension_columns, partition_columns):
"""
Fast version of::
pd.concat(
dfs,
ignore_index=True,
sort=False,
).sort_values(dimension_columns + partition_columns).reset_index(drop=True)
if inputs are presorted.
Parameters
-----------
dfs: Iterable[pandas.DataFrame]
DataFrames to concat.
dimension_columns: Iterable[str]
Dimension columns in correct order.
partition_columns: Iterable[str]
Partition columns in correct order.
Returns
-------
df: pandas.DataFrame
Concatenated result.
"""
return sort_dataframe(
df=concat_dataframes(dfs),
columns=list(dimension_columns) + list(partition_columns),
)
def _load_partition_dfs(cube, group, partition_mps, store):
"""
Load partition Dataframes for seed, restrictive and other data.
The information about the merge strategy (seed, restricting, others) is taken from ``group``.
Parameters
----------
cube: Cube
Cube spec.
group: QueryGroup
Query group.
partition_mps: Dict[str, Iterable[MetaPartition]]
MetaPartitions for every dataset in this partition.
store: simplekv.KeyValueStore
Store to load data from.
Returns
-------
df_seed: pandas.DataFrame
Seed data.
dfs_restrict: List[pandas.DataFrame]
Restrictive data (for inner join).
dfs_other: List[pandas.DataFrame]
Other data (for left join).
"""
df_seed = None
dfs_restrict = []
dfs_other = []
for ktk_cube_dataset_id, empty in group.empty_df.items():
mps = partition_mps.get(ktk_cube_dataset_id, [])
df = _load_all_mps(
mps=mps,
store=store,
load_columns=list(group.load_columns[ktk_cube_dataset_id]),
predicates=group.predicates.get(ktk_cube_dataset_id, None),
empty=empty,
)
# de-duplicate and sort data
# PERF: keep order of dimensionality identical to group.dimension_columns
df_cols = set(df.columns)
dimensionality = [c for c in group.dimension_columns if c in df_cols]
df = sort_dataframe(df=df, columns=dimensionality)
df = drop_sorted_duplicates_keep_last(df, dimensionality)
if ktk_cube_dataset_id == cube.seed_dataset:
assert df_seed is None
df_seed = df
elif ktk_cube_dataset_id in group.restrictive_dataset_ids:
dfs_restrict.append(df)
else:
dfs_other.append(df)
assert df_seed is not None
return df_seed, dfs_restrict, dfs_other
def prepare_data_for_ktk(df,
ktk_cube_dataset_id,
cube,
existing_payload,
partition_on,
consume_df=False):
"""
Prepare data so it can be handed over to Kartothek.
Some checks will be applied to the data to ensure it is sane.
Parameters
----------
df: pandas.DataFrame
DataFrame to be passed to Kartothek.
ktk_cube_dataset_id: str
Ktk_cube dataset UUID (w/o cube prefix).
cube: kartothek.core.cube.cube.Cube
Cube specification.
existing_payload: Set[str]
Existing payload columns.
partition_on: Iterable[str]
Partition-on attribute for given dataset.
consume_df: bool
Whether the incoming DataFrame can be destroyed while processing it.
Returns
-------
mp: kartothek.io_components.metapartition.MetaPartition
Kartothek-ready MetaPartition, may be sentinel (aka empty and w/o label).
Raises
------
ValueError
In case anything is fishy.
"""
check_user_df(ktk_cube_dataset_id, df, cube, existing_payload,
partition_on)
if (df is None) or df.empty:
# fast-path for empty DF
return MetaPartition(
label=None,
metadata_version=KTK_CUBE_METADATA_VERSION,
partition_keys=list(partition_on),
)
# TODO: find a more elegant solution that works w/o copy
df_orig = df
df = df.copy()
if consume_df:
# the original df is still referenced in the parent scope, so drop it
df_orig.drop(columns=df_orig.columns,
index=df_orig.index,
inplace=True)
df_columns = list(df.columns)
df_columns_set = set(df_columns)
# normalize value order and reset index
sort_keys = [
col for col in itertools.chain(cube.partition_columns,
cube.dimension_columns)
if col in df_columns_set
]
df = sort_dataframe(df=df, columns=sort_keys)
# check duplicate cells
_check_duplicates(ktk_cube_dataset_id, df, sort_keys, cube)
# check+convert column names to unicode strings
df.rename(columns={c: converter_str(c) for c in df_columns}, inplace=True)
# create MetaPartition object for easier handling
mp = MetaPartition(
label=gen_uuid(),
data=df,
metadata_version=KTK_CUBE_METADATA_VERSION,
)
del df
# partition data
mp = mp.partition_on(list(partition_on))
# reset indices again (because partition_on breaks it)
for mp2 in mp:
mp2.data.reset_index(drop=True, inplace=True)
del mp2
# calculate indices
indices_to_build = set(cube.index_columns) & df_columns_set
if ktk_cube_dataset_id == cube.seed_dataset:
indices_to_build |= set(cube.dimension_columns) - set(
cube.suppress_index_on)
indices_to_build -= set(partition_on)
mp = mp.build_indices(indices_to_build)
return mp
def test_sort_dataframe(df, columns):
expected = df.sort_values(columns).reset_index(drop=True)
actual = sort_dataframe(df, columns)
pdt.assert_frame_equal(actual, expected)