def store(self, store: StoreInput, dataset_uuid: str) -> str:
"""
Store the index as a parquet file
If compatible, the new keyname will be the name stored under the attribute `index_storage_key`.
If this attribute is None, a new key will be generated of the format
`{dataset_uuid}/indices/{column}/{timestamp}.by-dataset-index.parquet`
where the timestamp is in nanosecond accuracy and is created upon Index object initialization
Parameters
----------
store:
dataset_uuid:
"""
storage_key = None
store = ensure_store(store)
if (self.index_storage_key is not None and dataset_uuid
and dataset_uuid in self.index_storage_key):
storage_key = self.index_storage_key
if storage_key is None:
storage_key = "{dataset_uuid}/indices/{column}/{timestamp}{suffix}".format(
dataset_uuid=dataset_uuid,
suffix=naming.EXTERNAL_INDEX_SUFFIX,
column=quote(self.column),
timestamp=quote(self.creation_time.isoformat()),
)
# The arrow representation of index_dct requires a large amount of memory because strings are duplicated and
# flattened into the buffer. To avoid a high peak memory usage, split the index_dct into chunks and only convert
# one chunk a time to arrow.
parts_iter = partition_all(10_000, self.index_dct.items())
# Get first table explicit because its schema is required for ParquetWriter.
try:
table = _index_dct_to_table(dict(next(parts_iter)), self.column,
self.dtype)
except StopIteration:
# index_dct was empty, just pass it entirely
table = _index_dct_to_table(self.index_dct, self.column,
self.dtype)
buf = pa.BufferOutputStream()
with pq.ParquetWriter(buf, schema=table.schema) as writer:
writer.write_table(table)
del table
for part in parts_iter:
writer.write_table(
_index_dct_to_table(dict(part), self.column, self.dtype))
store.put(storage_key, buf.getvalue().to_pybytes())
return storage_key
def load(self, store: StoreInput):
"""
Load an external index into memory. Returns a new index object that
contains the index dictionary. Returns itself if the index is internal
or an already loaded index.
Parameters
----------
store
Object that implements the .get method for file/object loading.
Returns
-------
index: [kartothek.core.index.ExplicitSecondaryIndex]
"""
if self.loaded:
return self
store = ensure_store(store)
index_buffer = store.get(self.index_storage_key)
index_dct, column_type = _parquet_bytes_to_dict(
self.column, index_buffer)
return ExplicitSecondaryIndex(
column=self.column,
index_dct=index_dct,
dtype=column_type,
index_storage_key=self.index_storage_key,
normalize_dtype=False,
)
def _discover_dataset_meta_files(prefix: str, store: StoreInput) -> Set[str]:
"""
Get meta file names for all datasets.
Parameters
----------
prefix
the prefix.
store
KV store.
Returns
-------
names: Set[str]
The meta file names
"""
store = ensure_store(store)
names = {
name[: -len(METADATA_BASE_SUFFIX + suffix)]
for name in store.iter_prefixes(delimiter="/", prefix=prefix)
for suffix in [METADATA_FORMAT_JSON, METADATA_FORMAT_MSGPACK]
if name.endswith(METADATA_BASE_SUFFIX + suffix)
}
return names
def load_from_store(
uuid: str,
store: StoreInput,
load_schema: bool = True,
load_all_indices: bool = False,
) -> "DatasetMetadata":
"""
Load a dataset from a storage
Parameters
----------
uuid
UUID of the dataset.
store
Object that implements the .get method for file/object loading.
load_schema
Load table schema
load_all_indices
Load all registered indices into memory.
Returns
-------
dataset_metadata: :class:`~kartothek.core.dataset.DatasetMetadata`
Parsed metadata.
"""
key1 = naming.metadata_key_from_uuid(uuid)
store = ensure_store(store)
try:
value = store.get(key1)
metadata = load_json(value)
except KeyError:
key2 = naming.metadata_key_from_uuid(uuid, format="msgpack")
try:
value = store.get(key2)
metadata = unpackb(value)
except KeyError:
raise KeyError(
"Dataset does not exist. Tried {} and {}".format(
key1, key2))
ds = DatasetMetadata.load_from_dict(metadata,
store,
load_schema=load_schema)
if load_all_indices:
ds = ds.load_all_indices(store)
return ds
def storage_keys(uuid: str, store: StoreInput) -> List[str]:
"""
Retrieve all keys that belong to the given dataset.
Parameters
----------
uuid
UUID of the dataset.
store
Object that implements the .iter_keys method for key retrieval loading.
"""
store = ensure_store(store)
start_markers = ["{}.".format(uuid), "{}/".format(uuid)]
return list(
sorted(k for k in store.iter_keys(uuid) if any(
k.startswith(marker) for marker in start_markers)))
def get_parquet_metadata(self, store: StoreInput) -> pd.DataFrame:
"""
Retrieve the parquet metadata for the MetaPartition.
Especially relevant for calculating dataset statistics.
Parameters
----------
store
A factory function providing a KeyValueStore
table_name
Name of the kartothek table for which the statistics should be retrieved
Returns
-------
pd.DataFrame
A DataFrame with relevant parquet metadata
"""
store = ensure_store(store)
data = {}
with store.open(self.file) as fd: # type: ignore
pq_metadata = pa.parquet.ParquetFile(fd).metadata
data = {
"partition_label": self.label,
"serialized_size": pq_metadata.serialized_size,
"number_rows_total": pq_metadata.num_rows,
"number_row_groups": pq_metadata.num_row_groups,
"row_group_id": [],
"number_rows_per_row_group": [],
"row_group_compressed_size": [],
"row_group_uncompressed_size": [],
}
for rg_ix in range(pq_metadata.num_row_groups):
rg = pq_metadata.row_group(rg_ix)
data["row_group_id"].append(rg_ix)
data["number_rows_per_row_group"].append(rg.num_rows)
data["row_group_compressed_size"].append(rg.total_byte_size)
data["row_group_uncompressed_size"].append(
sum(
rg.column(col_ix).total_uncompressed_size
for col_ix in range(rg.num_columns)))
df = pd.DataFrame(data=data, columns=_METADATA_SCHEMA.keys())
df = df.astype(_METADATA_SCHEMA)
return df
def delete_from_store(self, dataset_uuid: Any,
store: StoreInput) -> "MetaPartition":
store = ensure_store(store)
# Delete data first
store.delete(self.file)
return self.copy(file=None, data=None)
def store_dataset_from_partitions(
partition_list,
store: StoreInput,
dataset_uuid,
dataset_metadata=None,
metadata_merger=None,
update_dataset=None,
remove_partitions=None,
metadata_storage_format=naming.DEFAULT_METADATA_STORAGE_FORMAT,
):
store = ensure_store(store)
if update_dataset:
dataset_builder = DatasetMetadataBuilder.from_dataset(update_dataset)
metadata_version = dataset_builder.metadata_version
else:
mp = next(iter(partition_list), None)
if mp is None:
raise ValueError(
"Cannot store empty datasets, partition_list must not be empty if in store mode."
)
metadata_version = mp.metadata_version
dataset_builder = DatasetMetadataBuilder(
uuid=dataset_uuid,
metadata_version=metadata_version,
partition_keys=mp.partition_keys,
)
dataset_builder.explicit_partitions = True
dataset_builder.table_meta = persist_common_metadata(
partition_list, update_dataset, store, dataset_uuid)
# We can only check for non unique partition labels here and if they occur we will
# fail hard. The resulting dataset may be corrupted or file may be left in the store
# without dataset metadata
partition_labels = partition_labels_from_mps(partition_list)
non_unique_labels = extract_duplicates(partition_labels)
if non_unique_labels:
raise ValueError(
"The labels {} are duplicated. Dataset metadata was not written.".
format(", ".join(non_unique_labels)))
if remove_partitions is None:
remove_partitions = []
if metadata_merger is None:
metadata_merger = combine_metadata
dataset_builder = update_metadata(dataset_builder, metadata_merger,
partition_list, dataset_metadata)
dataset_builder = update_partitions(dataset_builder, partition_list,
remove_partitions)
dataset_builder = update_indices(dataset_builder, store, partition_list,
remove_partitions)
if metadata_storage_format.lower() == "json":
store.put(*dataset_builder.to_json())
elif metadata_storage_format.lower() == "msgpack":
store.put(*dataset_builder.to_msgpack())
else:
raise ValueError(
"Unknown metadata storage format encountered: {}".format(
metadata_storage_format))
dataset = dataset_builder.to_dataset()
return dataset
def store_dataset_from_partitions(
partition_list,
store: StoreInput,
dataset_uuid,
dataset_metadata=None,
metadata_merger=None,
update_dataset=None,
remove_partitions=None,
metadata_storage_format=naming.DEFAULT_METADATA_STORAGE_FORMAT,
):
store = ensure_store(store)
schemas = set()
if update_dataset:
dataset_builder = DatasetMetadataBuilder.from_dataset(update_dataset)
metadata_version = dataset_builder.metadata_version
table_name = update_dataset.table_name
schemas.add(update_dataset.schema)
else:
mp = next(iter(partition_list), None)
if mp is None:
raise ValueError(
"Cannot store empty datasets, partition_list must not be empty if in store mode."
)
table_name = mp.table_name
metadata_version = mp.metadata_version
dataset_builder = DatasetMetadataBuilder(
uuid=dataset_uuid,
metadata_version=metadata_version,
partition_keys=mp.partition_keys,
)
for mp in partition_list:
if mp.schema:
schemas.add(mp.schema)
dataset_builder.schema = persist_common_metadata(
schemas=schemas,
update_dataset=update_dataset,
store=store,
dataset_uuid=dataset_uuid,
table_name=table_name,
)
# We can only check for non unique partition labels here and if they occur we will
# fail hard. The resulting dataset may be corrupted or file may be left in the store
# without dataset metadata
partition_labels = partition_labels_from_mps(partition_list)
# This could be safely removed since we do not allow to set this by the user
# anymore. It has implications on tests if mocks are used
non_unique_labels = extract_duplicates(partition_labels)
if non_unique_labels:
raise ValueError(
"The labels {} are duplicated. Dataset metadata was not written.".
format(", ".join(non_unique_labels)))
if remove_partitions is None:
remove_partitions = []
if metadata_merger is None:
metadata_merger = combine_metadata
dataset_builder = update_metadata(dataset_builder, metadata_merger,
dataset_metadata)
dataset_builder = update_partitions(dataset_builder, partition_list,
remove_partitions)
dataset_builder = update_indices(dataset_builder, store, partition_list,
remove_partitions)
if metadata_storage_format.lower() == "json":
store.put(*dataset_builder.to_json())
elif metadata_storage_format.lower() == "msgpack":
store.put(*dataset_builder.to_msgpack())
else:
raise ValueError(
"Unknown metadata storage format encountered: {}".format(
metadata_storage_format))
dataset = dataset_builder.to_dataset()
return dataset
