def update_partitions(dataset_builder, add_partitions, remove_partitions): for mp in add_partitions: for sub_mp_dct in mp.metapartitions: # label is None in case of an empty partition if sub_mp_dct["label"] is not None: partition = Partition(label=sub_mp_dct["label"], files=sub_mp_dct["files"]) dataset_builder.add_partition(sub_mp_dct["label"], partition) for partition_name in remove_partitions: del dataset_builder.partitions[partition_name] return dataset_builder
def test_builder_full(metadata_version, frozen_time): expected = { "dataset_uuid": "uuid", "dataset_metadata_version": metadata_version, "partitions": { "run_id=1/L=1/P=1/part_1": { "files": { "core": "uuid/core/run_id=1/L=1/P=1/part_1.parquet", "helper": "uuid/helper/run_id=1/L=1/P=1/part_1.parquet", } } }, "metadata": { "key": "value", "creation_time": TIME_TO_FREEZE_ISO }, "indices": { "col1": { "a": ["run_id=1/L=1/P=1/part_1"], "b": ["run_id=2/L=1/P=1/part_1"], }, "col2": "uuid.col2.by-dataset-index.parquet", }, "partition_keys": ["L", "P"], } builder = DatasetMetadataBuilder("uuid", metadata_version=metadata_version, partition_keys=["L", "P"]) part_2 = Partition( label="run_id=1/L=1/P=1/part_1", files={ "core": "uuid/core/run_id=1/L=1/P=1/part_1.parquet", "helper": "uuid/helper/run_id=1/L=1/P=1/part_1.parquet", }, ) builder.add_partition("run_id=1/L=1/P=1/part_1", part_2) builder.add_metadata("key", "value") builder.add_external_index("col2") builder.add_embedded_index( "col1", ExplicitSecondaryIndex("col1", { "a": ["run_id=1/L=1/P=1/part_1"], "b": ["run_id=2/L=1/P=1/part_1"] }), ) key, result = builder.to_json() result = simplejson.loads(result) assert key == "uuid.by-dataset-metadata.json" assert result == expected
def test_eq(): assert not (Partition("label") == Partition("other_label")) assert not (Partition("label") == Partition("label", files={"some": "file"})) assert Partition(label="label", files={"some": "file"}) == Partition(label="label", files={"some": "file"})
def test_builder_to_dataset(metadata_version, frozen_time): expected = { "dataset_uuid": "uuid", "dataset_metadata_version": metadata_version, "partitions": { "part_2": { "files": { "core": "uuid/core/part_2.parquet" } } }, "metadata": { "key": "value", "creation_time": TIME_TO_FREEZE_ISO }, "indices": { "col1": { "a": ["part1"], "b": ["part2"] } }, } builder = DatasetMetadataBuilder("uuid", metadata_version=metadata_version) part_2 = Partition("part_2", {"core": "uuid/core/part_2.parquet"}) builder.add_partition("part_2", part_2) builder.add_metadata("key", "value") builder.add_embedded_index( "col1", ExplicitSecondaryIndex("col1", { "a": ["part1"], "b": ["part2"] })) result = builder.to_dataset() expected_from_dict = DatasetMetadata.from_dict(expected) assert result == expected_from_dict
def partition(self) -> Partition: return Partition(label=self.label, files={self.table_name: self.file})