def _map_ktk_mps_to_groups(cube, datasets, label2gp): """ Map Kartothek metapartitions to groups. Parameters ---------- cube: Cube Cube specification. datasets: Dict[str, kartothek.core.dataset.DatasetMetadata] Datasets that are processed by the regrouper. label2gp: Dict[str, Dict[str, Tuple[int, int]]] Maps "dataset ID -> (label -> (group ID, partition ID))". Returns ------- groups: Dict[int, Dict[int, Dict[str, Tuple[kartothek.io_components.metapartition.MetaPartition, ...]]]] Maps "group ID -> (partition ID -> (dataset ID -> list of MetaPartitions))" """ groups = defaultdict(lambda: defaultdict(lambda: defaultdict(list))) for ktk_cube_dataset_id, ds in datasets.items(): label2gp_sub = label2gp[ktk_cube_dataset_id] for mp in dispatch_metapartitions_from_factory( dataset_factory=metadata_factory_from_dataset(ds), ): # FIXME: can this be simplified? if mp.label not in label2gp_sub: # filtered out by pre-condition continue for group_id, partition_id in label2gp_sub[mp.label]: groups[group_id][partition_id][ktk_cube_dataset_id].append(mp) return groups
def _multiplex_store_dataset_from_partitions_flat(mpss, cube, metadata, update, store, existing_datasets): dct = defaultdict(list) for sublist in mpss: for mp in sublist: for k, v in mp.items(): dct[k].append(v) result = {} for k, v in dct.items(): if update: ds_factory = metadata_factory_from_dataset(existing_datasets[k], with_schema=True, store=store) result[k] = update_dataset_from_partitions( v, dataset_uuid=cube.ktk_dataset_uuid(k), delete_scope=[], ds_factory=ds_factory, metadata=metadata[k], metadata_merger=None, store_factory=store, ) else: result[k] = store_dataset_from_partitions( v, dataset_metadata=metadata[k], dataset_uuid=cube.ktk_dataset_uuid(k), metadata_merger=None, metadata_storage_format=KTK_CUBE_METADATA_STORAGE_FORMAT, store=store, ) # list required for dask.bag return [result]
def test_metadata_factory_from_dataset_no_store(function_store, ds, load_schema): ds2 = DatasetMetadata.load_from_store( "uuid", function_store(), load_schema=load_schema ) factory = metadata_factory_from_dataset(ds2, with_schema=load_schema) assert factory.dataset_metadata is ds2 store = factory.store with pytest.raises(NotImplementedError): store.get("foo")
def test_metadata_factory_from_dataset_with_store(function_store, ds, load_schema): ds2 = DatasetMetadata.load_from_store( "uuid", function_store(), load_schema=load_schema ) factory = metadata_factory_from_dataset( ds2, with_schema=load_schema, store=function_store ) assert factory.dataset_metadata is ds2 store = factory.store store.put("foo", b"bar") assert store.get("foo") == b"bar"
def get_metapartitions_for_stats(datasets): """ Get all metapartitions that need to be scanned to gather cube stats. Parameters ---------- datasets: Dict[str, kartothek.core.dataset.DatasetMetadata] Datasets that are present. Returns ------- metapartitions: Tuple[Tuple[str, Tuple[kartothek.io_components.metapartition.MetaPartition, ...]], ...] Pre-aligned metapartitions (by primary index / physical partitions) and the ktk_cube dataset ID belonging to them. """ all_metapartitions = [] for ktk_cube_dataset_id, ds in datasets.items(): dataset_factory = metadata_factory_from_dataset(ds) for mp in dispatch_metapartitions_from_factory( dataset_factory=dataset_factory, dispatch_by=dataset_factory.partition_keys): all_metapartitions.append((ktk_cube_dataset_id, mp)) return all_metapartitions
def remove_partitions(cube, store, conditions=None, ktk_cube_dataset_ids=None, metadata=None): """ Remove given partition range from cube using a transaction. Parameters ---------- cube: kartothek.core.cube.cube.Cube Cube spec. store: Union[simplekv.KeyValueStore, Callable[[], simplekv.KeyValueStore]] Store. conditions: Union[None, Condition, Iterable[Condition], Conjunction] Conditions that should be applied, optional. Defaults to "entire cube". ktk_cube_dataset_ids: Optional[Union[Iterable[Union[Str, Bytes]], Union[Str, Bytes]]] Ktk_cube dataset IDs to apply the remove action to, optional. Default to "all". metadata: Optional[Dict[str, Dict[str, Any]]] Metadata for every the datasets, optional. Only given keys are updated/replaced. Deletion of metadata keys is not possible. Returns ------- datasets: Dict[str, kartothek.core.dataset.DatasetMetadata] Datasets, updated. """ if callable(store): store_instance = store() store_factory = store else: store_instance = store def store_factory(): return store existing_datasets = discover_datasets(cube, store) for ( ktk_cube_dataset_id, (ds, mp, delete_scope), ) in prepare_metapartitions_for_removal_action( cube=cube, store=store_instance, conditions=conditions, ktk_cube_dataset_ids=ktk_cube_dataset_ids, existing_datasets=existing_datasets, ).items(): mp = mp.store_dataframes( store=store_instance, dataset_uuid=ds.uuid, df_serializer=KTK_CUBE_DF_SERIALIZER, ) ds_factory = metadata_factory_from_dataset(ds, with_schema=True, store=store_factory) existing_datasets[ ktk_cube_dataset_id] = update_dataset_from_partitions( mp, store_factory=store_factory, dataset_uuid=ds.uuid, ds_factory=ds_factory, metadata=prepare_ktk_metadata(cube, ktk_cube_dataset_id, metadata), metadata_merger=None, delete_scope=delete_scope, ) return existing_datasets
def remove_partitions( cube: Cube, store: Union[simplekv.KeyValueStore, StoreFactory], conditions: Union[None, Condition, Sequence[Condition], Conjunction] = None, ktk_cube_dataset_ids: Optional[Union[Sequence[str], str]] = None, metadata: Optional[Dict[str, Dict[str, Any]]] = None, ): """ Remove given partition range from cube using a transaction. Remove the partitions selected by ``conditions``. If no ``conditions`` are given, remove all partitions. For each considered dataset, only the subset of ``conditions`` that refers to the partition columns of the respective dataset is used. In particular, a dataset that is not partitioned at all is always considered selected by ``conditions``. Parameters ---------- cube Cube spec. store Store. conditions Select the partitions to be removed. Must be a condition only on partition columns. ktk_cube_dataset_ids Ktk_cube dataset IDs to apply the remove action to, optional. Default to "all". metadata Metadata for every the datasets, optional. Only given keys are updated/replaced. Deletion of metadata keys is not possible. Returns ------- datasets: Dict[str, kartothek.core.dataset.DatasetMetadata] Datasets, updated. """ if callable(store): store_instance = store() store_factory = store else: store_instance = store def store_factory(): return store existing_datasets = discover_datasets(cube, store) for ( ktk_cube_dataset_id, (ds, mp, delete_scope), ) in prepare_metapartitions_for_removal_action( cube=cube, store=store_instance, conditions=conditions, ktk_cube_dataset_ids=ktk_cube_dataset_ids, existing_datasets=existing_datasets, ).items(): mp = mp.store_dataframes( store=store_instance, dataset_uuid=ds.uuid, df_serializer=KTK_CUBE_DF_SERIALIZER, ) ds_factory = metadata_factory_from_dataset(ds, with_schema=True, store=store_factory) existing_datasets[ ktk_cube_dataset_id] = update_dataset_from_partitions( mp, store_factory=store_factory, dataset_uuid=ds.uuid, ds_factory=ds_factory, metadata=prepare_ktk_metadata(cube, ktk_cube_dataset_id, metadata), metadata_merger=None, delete_scope=delete_scope, ) return existing_datasets