def _get_meta_partitions_with_dataframe(metadata_version):
df = pd.DataFrame(
OrderedDict(
[
("P", [1]),
("L", [1]),
("TARGET", [1]),
("DATE", [datetime.date(2010, 1, 1)]),
]
)
)
df_2 = pd.DataFrame(OrderedDict([("P", [1]), ("info", ["a"])]))
mp = MetaPartition(
label="cluster_1",
data={SINGLE_TABLE: df, "helper": df_2},
metadata_version=metadata_version,
)
df_3 = pd.DataFrame(
OrderedDict(
[
("P", [2]),
("L", [2]),
("TARGET", [2]),
("DATE", [datetime.date(2009, 12, 31)]),
]
)
)
df_4 = pd.DataFrame(OrderedDict([("P", [2]), ("info", ["b"])]))
mp2 = MetaPartition(
label="cluster_2",
data={SINGLE_TABLE: df_3, "helper": df_4},
metadata_version=metadata_version,
)
return [mp, mp2]
def test_nested_copy():
mp = MetaPartition(
label="label_1",
file="file",
data=pd.DataFrame({"test": [1, 2, 3]}),
indices={"test": {
1: "label_1",
2: "label_2",
3: "label_3"
}},
)
mp_2 = MetaPartition(
label="label_2",
data=pd.DataFrame({"test": [4, 5, 6]}),
indices={"test": [4, 5, 6]},
)
mp = mp.add_metapartition(mp_2)
assert len(mp.metapartitions) == 2
new_mp = mp.copy()
# Check if the copy is identical
assert len(new_mp.metapartitions) == len(mp.metapartitions)
assert new_mp == mp
# ... but not the same object
assert id(new_mp) != id(mp)
def test_partition_on_one_level_ts():
original_df = pd.DataFrame({
"test": [
pd.Timestamp("2001-01-01"),
pd.Timestamp("2001-01-02"),
pd.Timestamp("2001-01-03"),
],
"some_values": [1, 2, 3],
})
mp = MetaPartition(label="label_1",
file="file",
data=original_df,
metadata_version=4)
new_mp = mp.partition_on(["test"])
assert len(new_mp.metapartitions) == 3
labels = set()
for mp in new_mp:
labels.add(mp.label)
assert len(mp.data) == 1
assert mp.data is not None
df = mp.data
assert df._is_view
# try to be agnostic about the order
assert len(df) == 1
assert "test" not in df
expected_labels = set([
"test=2001-01-01%2000%3A00%3A00/label_1",
"test=2001-01-02%2000%3A00%3A00/label_1",
"test=2001-01-03%2000%3A00%3A00/label_1",
])
assert labels == expected_labels
def test_partition_on_one_level():
original_df = pd.DataFrame({"test": [1, 2, 3], "some_values": [1, 2, 3]})
mp = MetaPartition(
label="label_1",
files={"core": "file"},
data={"core": original_df},
dataset_metadata={"dataset": "metadata"},
metadata_version=4,
)
new_mp = mp.partition_on(["test"])
assert len(new_mp.metapartitions) == 3
labels = set()
for mp in new_mp:
labels.add(mp.label)
assert len(mp.data) == 1
assert "core" in mp.data
df = mp.data["core"]
assert df._is_view
# try to be agnostic about the order
assert len(df) == 1
assert "test" not in df
expected_labels = set(
["test=1/label_1", "test=2/label_1", "test=3/label_1"])
assert labels == expected_labels
def test_to_dict(metadata_version):
df = pd.DataFrame({"A": [1]})
schema = make_meta(df, origin="test")
mp = MetaPartition(
label="label_1",
file="file",
data=df,
indices={"test": [1, 2, 3]},
metadata_version=metadata_version,
schema=schema,
)
mp_dct = mp.to_dict()
assert mp_dct == {
"label": "label_1",
"data": df,
"file": "file",
"indices": {
"test": [1, 2, 3]
},
"metadata_version": metadata_version,
"schema": schema,
"partition_keys": [],
"logical_conjunction": None,
"table_name": SINGLE_TABLE,
}
def test_store_single_dataframe_as_partition_no_metadata(
store, metadata_version):
df = pd.DataFrame({
"P": np.arange(0, 10),
"L": np.arange(0, 10),
"TARGET": np.arange(10, 20)
})
mp = MetaPartition(label="test_label",
data={"core": df},
metadata_version=metadata_version)
partition = mp.store_dataframes(
store=store,
df_serializer=ParquetSerializer(),
dataset_uuid="dataset_uuid",
store_metadata=False,
)
assert len(partition.data) == 0
expected_file = "dataset_uuid/core/test_label.parquet"
assert partition.files == {"core": expected_file}
assert partition.label == "test_label"
# One meta one actual file
files_in_store = list(store.keys())
assert len(files_in_store) == 1
stored_df = DataFrameSerializer.restore_dataframe(store=store,
key=expected_file)
pdt.assert_frame_equal(df, stored_df)
def test_to_dict(metadata_version):
mp = MetaPartition(
label="label_1",
files={"core": "file"},
data={"core": "placeholder"},
indices={"test": [1, 2, 3]},
metadata_version=metadata_version,
table_meta={"core": {
"test": "int8"
}},
)
mp_dct = mp.to_dict()
assert mp_dct == {
"label": "label_1",
"data": {
"core": "placeholder"
},
"files": {
"core": "file"
},
"indices": {
"test": [1, 2, 3]
},
"dataset_metadata": {},
"metadata_version": metadata_version,
"table_meta": {
"core": {
"test": "int8"
}
},
"partition_keys": [],
}
def test_store_single_dataframe_as_partition(store, metadata_storage_format,
metadata_version, expected_key):
df = pd.DataFrame({
"P": np.arange(0, 10),
"L": np.arange(0, 10),
"TARGET": np.arange(10, 20)
})
mp = MetaPartition(label="test_label",
data={"core": df},
metadata_version=metadata_version)
meta_partition = mp.store_dataframes(
store=store,
df_serializer=ParquetSerializer(),
dataset_uuid="dataset_uuid",
store_metadata=True,
metadata_storage_format=metadata_storage_format,
)
assert len(meta_partition.data) == 0
assert meta_partition.files == {"core": expected_key}
assert meta_partition.label == "test_label"
files_in_store = list(store.keys())
expected_num_files = 1
assert len(files_in_store) == expected_num_files
stored_df = DataFrameSerializer.restore_dataframe(store=store,
key=expected_key)
pdt.assert_frame_equal(df, stored_df)
files_in_store.remove(expected_key)
assert len(files_in_store) == expected_num_files - 1
def test_load_dataframes(meta_partitions_files_only, store_session,
predicate_pushdown_to_io):
expected_df = pd.DataFrame(
OrderedDict([
("P", [1]),
("L", [1]),
("TARGET", [1]),
("DATE", pd.to_datetime([date(2010, 1, 1)])),
]))
expected_df_2 = pd.DataFrame(OrderedDict([("P", [1]), ("info", ["a"])]))
mp = meta_partitions_files_only[0]
assert len(mp.files) > 0
assert len(mp.data) == 0
mp = meta_partitions_files_only[0].load_dataframes(
store=store_session, predicate_pushdown_to_io=predicate_pushdown_to_io)
assert len(mp.data) == 2
data = mp.data
pdt.assert_frame_equal(data["core"], expected_df, check_dtype=False)
pdt.assert_frame_equal(data["helper"], expected_df_2, check_dtype=False)
empty_mp = MetaPartition("empty_mp", metadata_version=mp.metadata_version)
empty_mp.load_dataframes(store_session,
predicate_pushdown_to_io=predicate_pushdown_to_io)
assert empty_mp.data == {}
def test_reconstruct_index_categories(store):
ser = ParquetSerializer()
df = pd.DataFrame({
"index_col": [1, 1],
"second_index_col": [2, 2],
"column": list("ab")
})
label = "dontcare"
key_prefix = "uuid/table/index_col=2/second_index_col=2/{}".format(label)
key = ser.store(store, key_prefix, df)
schema = make_meta(df, origin="1", partition_keys="index_col")
store_schema_metadata(schema, "uuid", store, "table")
mp = MetaPartition(
label="index_col=2/dontcare",
files={"table": key},
metadata_version=4,
table_meta={"table": schema},
partition_keys=["index_col", "second_index_col"],
)
categories = ["second_index_col", "index_col"]
mp = mp.load_dataframes(store, categoricals={"table": categories})
df_actual = mp.data["table"]
df_expected = pd.DataFrame(
OrderedDict([
("index_col", [2, 2]),
("second_index_col", [2, 2]),
("column", list("ab")),
]))
df_expected = df_expected.astype({col: "category" for col in categories})
pdt.assert_frame_equal(df_actual, df_expected)
def test_table_meta(store):
mp = MetaPartition(
label="label_1",
data=pd.DataFrame(
{
"i32": np.array([1, 2, 3, 1, 2, 3], dtype="int32"),
"float": np.array([1, 1, 1, 2, 2, 2], dtype="float64"),
}
),
metadata_version=4,
)
assert mp.schema is not None
expected_meta = make_meta(
pd.DataFrame(
{"i32": np.array([], dtype="int32"), "float": np.array([], dtype="float64")}
),
origin="1",
)
actual_meta = mp.schema
assert actual_meta == expected_meta
mp = mp.store_dataframes(store, "dataset_uuid")
actual_meta = mp.schema
assert actual_meta == expected_meta
class BuildIndex(AsvBenchmarkConfig):
params = ([-1, 1], [10**3, 10**4], [10, 100])
param_names = ["cardinality", "num_values", "partitions_to_merge"]
def setup(self, cardinality, num_values, partitions_to_merge):
self.column = "column"
self.table = "table"
self.merge_indices = []
for n in range(partitions_to_merge):
if cardinality < 0:
array = [
"{:010f}".format(x) for x in np.random.randn(num_values)
]
else:
unique_vals = ["{:010d}".format(n) for n in range(cardinality)]
array = [
unique_vals[x % len(unique_vals)]
for x in range(num_values)
]
self.df = pd.DataFrame({self.column: array})
self.mp = MetaPartition(label=self.table,
data={"core": self.df},
metadata_version=4)
self.mp_indices = self.mp.build_indices([self.column])
self.merge_indices.append(self.mp_indices)
def time_metapartition_build_index(self, cardinality, num_values,
partitions_to_merge):
self.mp.build_indices([self.column])
def time_merge_indices(self, cardinality, num_values, partitions_to_merge):
MetaPartition.merge_indices(self.merge_indices)
def test_partition_urlencode():
original_df = pd.DataFrame({"ÖŒå": [1, 2, 3], "some_values": [1, 2, 3]})
mp = MetaPartition(label="label_1", data=original_df, metadata_version=4)
new_mp = mp.partition_on(["ÖŒå"])
assert len(new_mp.metapartitions) == 3
labels = set()
for mp in new_mp:
labels.add(mp.label)
assert len(mp.data) == 1
assert mp.data is not None
df = mp.data
assert df._is_view
# try to be agnostic about the order
assert len(df) == 1
assert "ÖŒå" not in df
expected_labels = set(
[
"%C3%96%C5%92%C3%A5=1/label_1",
"%C3%96%C5%92%C3%A5=2/label_1",
"%C3%96%C5%92%C3%A5=3/label_1",
]
)
assert labels == expected_labels
def test_reconstruct_index_empty_df(store, categoricals):
ser = ParquetSerializer()
df = pd.DataFrame({"index_col": [1, 1], "column": list("ab")})
df = df[0:0]
label = "dontcare"
key_prefix = "uuid/table/index_col=2/{}".format(label)
key = ser.store(store, key_prefix, df)
schema = make_meta(df, origin="1", partition_keys="index_col")
store_schema_metadata(schema, "uuid", store)
mp = MetaPartition(
label="index_col=2/dontcare",
file=key,
metadata_version=4,
schema=schema,
partition_keys=["index_col"],
)
categoricals = None
if categoricals:
categoricals = ["index_col"]
mp = mp.load_dataframes(store, categoricals=categoricals)
df_actual = mp.data
df_expected = pd.DataFrame(
OrderedDict([("index_col", [2, 2]), ("column", list("ab"))])
)
df_expected = df_expected[0:0]
if categoricals:
df_expected = df_expected.astype({"index_col": "category"})
pdt.assert_frame_equal(df_actual, df_expected)
def meta_partitions_evaluation_dataframe(metadata_version):
"""
Create a list of MetaPartitions for testing. The partitions
include in-memory pd.DataFrames without external references, i.e. files
are empty
"""
df = pd.DataFrame(
OrderedDict([("P", [1]), ("L", [1]), ("HORIZON", [1]), ("PRED", [10])])
)
mp = MetaPartition(
label="cluster_1_1", data={"PRED": df}, metadata_version=metadata_version
)
df_2 = pd.DataFrame(
OrderedDict([("P", [1]), ("L", [1]), ("HORIZON", [2]), ("PRED", [20])])
)
mp2 = MetaPartition(
label="cluster_1_2", data={"PRED": df_2}, metadata_version=metadata_version
)
df_3 = pd.DataFrame(
OrderedDict([("P", [2]), ("L", [2]), ("HORIZON", [1]), ("PRED", [10])])
)
mp3 = MetaPartition(
label="cluster_2_1", data={"PRED": df_3}, metadata_version=metadata_version
)
df_4 = pd.DataFrame(
OrderedDict([("P", [2]), ("L", [2]), ("HORIZON", [2]), ("PRED", [20])])
)
mp4 = MetaPartition(
label="cluster_2_2", data={"PRED": df_4}, metadata_version=metadata_version
)
return [mp, mp2, mp3, mp4]
def test_load_dataframe_logical_conjunction(store, metadata_version):
df = pd.DataFrame(
{"P": np.arange(0, 10), "L": np.arange(0, 10), "TARGET": np.arange(10, 20)}
)
mp = MetaPartition(
label="cluster_1",
data=df,
metadata_version=metadata_version,
logical_conjunction=[("P", ">", 4)],
)
meta_partition = mp.store_dataframes(
store=store, df_serializer=None, dataset_uuid="dataset_uuid",
)
predicates = None
loaded_mp = meta_partition.load_dataframes(store=store, predicates=predicates)
data = pd.DataFrame(
{"P": [5, 6, 7, 8, 9], "L": [5, 6, 7, 8, 9], "TARGET": [15, 16, 17, 18, 19]}
).set_index(np.arange(5, 10))
pdt.assert_frame_equal(loaded_mp.data, data)
predicates = [[("L", ">", 6), ("TARGET", "<", 18)]]
loaded_mp = meta_partition.load_dataframes(store=store, predicates=predicates)
data = pd.DataFrame({"P": [7], "L": [7], "TARGET": [17]}).set_index(np.array([7]))
pdt.assert_frame_equal(loaded_mp.data, data)
predicates = [[("L", ">", 2), ("TARGET", "<", 17)], [("TARGET", "==", 19)]]
loaded_mp = meta_partition.load_dataframes(store=store, predicates=predicates)
data = pd.DataFrame(
{"P": [5, 6, 9], "L": [5, 6, 9], "TARGET": [15, 16, 19]}
).set_index(np.array([5, 6, 9]))
pdt.assert_frame_equal(loaded_mp.data, data)
def test_raises_other_index_missing(self, cube, function_store):
store_data(
cube=cube,
function_store=function_store,
df=MetaPartition(
label=gen_uuid(),
data={
SINGLE_TABLE: pd.DataFrame({"x": [0], "y": [0], "p": [0], "q": [0]})
},
metadata_version=KTK_CUBE_METADATA_VERSION,
).build_indices(["x", "y"]),
name=cube.seed_dataset,
)
store_data(
cube=cube,
function_store=function_store,
df=MetaPartition(
label=gen_uuid(),
data={
SINGLE_TABLE: pd.DataFrame(
{"x": [0], "y": [0], "p": [0], "q": [0], "i1": [1337]}
)
},
metadata_version=KTK_CUBE_METADATA_VERSION,
),
name="enrich",
)
with pytest.raises(ValueError) as exc:
discover_datasets(cube, function_store)
assert (
str(exc.value)
== 'ExplicitSecondaryIndex or PartitionIndex "i1" is missing in dataset "enrich".'
)
def test_store_single_dataframe_as_partition(store, metadata_version):
df = pd.DataFrame(
{"P": np.arange(0, 10), "L": np.arange(0, 10), "TARGET": np.arange(10, 20)}
)
mp = MetaPartition(label="test_label", data=df, metadata_version=metadata_version)
meta_partition = mp.store_dataframes(
store=store, df_serializer=ParquetSerializer(), dataset_uuid="dataset_uuid",
)
assert meta_partition.data is None
expected_key = "dataset_uuid/table/test_label.parquet"
assert meta_partition.file == expected_key
assert meta_partition.label == "test_label"
files_in_store = list(store.keys())
expected_num_files = 1
assert len(files_in_store) == expected_num_files
stored_df = DataFrameSerializer.restore_dataframe(store=store, key=expected_key)
pdt.assert_frame_equal(df, stored_df)
files_in_store.remove(expected_key)
assert len(files_in_store) == expected_num_files - 1
def test_merge_indices():
indices = [
MetaPartition(
label="label1",
indices={"location": {"Loc1": ["label1"], "Loc2": ["label1"]}},
),
MetaPartition(
label="label2",
indices={
"location": {"Loc3": ["label2"], "Loc2": ["label2"]},
"product": {"Product1": ["label2"], "Product2": ["label2"]},
},
),
]
result = MetaPartition.merge_indices(indices)
expected = {
"location": ExplicitSecondaryIndex(
"location",
{"Loc1": ["label1"], "Loc2": ["label1", "label2"], "Loc3": ["label2"]},
),
"product": ExplicitSecondaryIndex(
"product", {"Product1": ["label2"], "Product2": ["label2"]}
),
}
assert result == expected
def test_get_parquet_metadata_row_group_size(store):
df = pd.DataFrame({"P": np.arange(0, 10), "L": np.arange(0, 10)})
mp = MetaPartition(label="test_label", data=df)
ps = ParquetSerializer(chunk_size=5)
meta_partition = mp.store_dataframes(
store=store, dataset_uuid="dataset_uuid", df_serializer=ps
)
actual = meta_partition.get_parquet_metadata(store=store)
actual.drop(
columns=[
"serialized_size",
"row_group_compressed_size",
"row_group_uncompressed_size",
],
axis=1,
inplace=True,
)
expected = pd.DataFrame(
{
"partition_label": ["test_label", "test_label"],
"row_group_id": [0, 1],
"number_rows_total": [10, 10],
"number_row_groups": [2, 2],
"number_rows_per_row_group": [5, 5],
}
)
pd.testing.assert_frame_equal(actual, expected)
def test_get_parquet_metadata_empty_df(store):
df = pd.DataFrame()
mp = MetaPartition(label="test_label", data=df)
meta_partition = mp.store_dataframes(store=store, dataset_uuid="dataset_uuid")
actual = meta_partition.get_parquet_metadata(store=store)
actual.drop(
columns=[
"serialized_size",
"row_group_compressed_size",
"row_group_uncompressed_size",
],
axis=1,
inplace=True,
)
expected = pd.DataFrame(
{
"partition_label": ["test_label"],
"row_group_id": 0,
"number_rows_total": 0,
"number_row_groups": 1,
"number_rows_per_row_group": 0,
}
)
pd.testing.assert_frame_equal(actual, expected)
def test_table_meta(store):
mp = MetaPartition(
label="label_1",
data={
"core":
pd.DataFrame({
"i32":
np.array([1, 2, 3, 1, 2, 3], dtype="int32"),
"float":
np.array([1, 1, 1, 2, 2, 2], dtype="float64"),
})
},
metadata_version=4,
)
assert len(mp.table_meta) == 1
assert "core" in mp.table_meta
expected_meta = make_meta(
pd.DataFrame({
"i32": np.array([], dtype="int32"),
"float": np.array([], dtype="float64")
}),
origin="1",
)
actual_meta = mp.table_meta["core"]
assert actual_meta == expected_meta
mp = mp.store_dataframes(store, "dataset_uuid")
actual_meta = mp.table_meta["core"]
assert actual_meta == expected_meta
def test_partition_on_scalar_intermediate(df_not_nested, col):
"""
Test against a bug where grouping leaves a scalar value
"""
assert len(df_not_nested) == 1
mp = MetaPartition(label="somelabel", data=df_not_nested, metadata_version=4)
new_mp = mp.partition_on(col)
assert len(new_mp) == 1
def test_concat_metapartition_wrong_types(df_all_types):
mp1 = MetaPartition(label="first", data=df_all_types, metadata_version=4)
df_corrupt = df_all_types.copy()
df_corrupt["int8"] = "NoInteger"
mp2 = MetaPartition(label="second", data=df_corrupt, metadata_version=4)
with pytest.raises(ValueError, match="Schema violation"):
MetaPartition.concat_metapartitions([mp1, mp2])
def store_data(
cube,
function_store,
df,
name,
partition_on="default",
metadata_version=KTK_CUBE_METADATA_VERSION,
metadata_storage_format=KTK_CUBE_METADATA_STORAGE_FORMAT,
metadata=None,
overwrite=False,
new_ktk_cube_metadata=True,
write_suppress_index_on=True,
):
if partition_on == "default":
partition_on = cube.partition_columns
if isinstance(df, pd.DataFrame):
mp = MetaPartition(
label=gen_uuid(), data={SINGLE_TABLE: df}, metadata_version=metadata_version
)
indices_to_build = set(cube.index_columns) & set(df.columns)
if name == cube.seed_dataset:
indices_to_build |= set(cube.dimension_columns) - set(
cube.suppress_index_on
)
mp = mp.build_indices(indices_to_build)
dfs = mp
else:
assert isinstance(df, MetaPartition)
assert df.metadata_version == metadata_version
dfs = df
if metadata is None:
metadata = {
KTK_CUBE_METADATA_DIMENSION_COLUMNS: cube.dimension_columns,
KTK_CUBE_METADATA_KEY_IS_SEED: (name == cube.seed_dataset),
}
if new_ktk_cube_metadata:
metadata.update(
{KTK_CUBE_METADATA_PARTITION_COLUMNS: cube.partition_columns}
)
if write_suppress_index_on:
metadata.update(
{KTK_CUBE_METADATA_SUPPRESS_INDEX_ON: list(cube.suppress_index_on)}
)
return store_dataframes_as_dataset(
store=function_store,
dataset_uuid=cube.ktk_dataset_uuid(name),
dfs=dfs,
partition_on=list(partition_on) if partition_on else None,
metadata_storage_format=metadata_storage_format,
metadata_version=metadata_version,
df_serializer=KTK_CUBE_DF_SERIALIZER,
metadata=metadata,
overwrite=overwrite,
)
def test_partition_on_explicit_index():
original_df = pd.DataFrame({
"level1": [1, 2, 1, 2, 1, 2],
"level2": [1, 1, 1, 2, 2, 2],
"explicit_index_col": np.arange(0, 6),
})
mp = MetaPartition(
label="label_1",
file="file",
data=original_df,
indices={
"explicit_index_col":
{value: ["label_1"]
for value in np.arange(0, 6)}
},
metadata_version=4,
)
new_mp = mp.partition_on(["level1", "level2"])
assert len(new_mp) == 4
expected_indices = {
"explicit_index_col":
ExplicitSecondaryIndex(
"explicit_index_col",
{
0: ["level1=1/level2=1/label_1"],
2: ["level1=1/level2=1/label_1"]
},
)
}
assert expected_indices == new_mp["level1=1/level2=1/label_1"].indices
expected_indices = {
"explicit_index_col":
ExplicitSecondaryIndex("explicit_index_col",
{4: ["level1=1/level2=2/label_1"]})
}
assert expected_indices == new_mp["level1=1/level2=2/label_1"].indices
expected_indices = {
"explicit_index_col":
ExplicitSecondaryIndex("explicit_index_col",
{1: ["level1=2/level2=1/label_1"]})
}
assert expected_indices == new_mp["level1=2/level2=1/label_1"].indices
expected_indices = {
"explicit_index_col":
ExplicitSecondaryIndex(
"explicit_index_col",
{
3: ["level1=2/level2=2/label_1"],
5: ["level1=2/level2=2/label_1"]
},
)
}
assert expected_indices == new_mp["level1=2/level2=2/label_1"].indices
def test_partition_on_raises_no_cols_left(empty):
original_df = pd.DataFrame({"test": [1, 2, 3]})
if empty:
original_df = original_df.loc[[]]
mp = MetaPartition(
label="label_1", file="file", data=original_df, metadata_version=4
)
with pytest.raises(ValueError) as e:
mp.partition_on(["test"])
assert str(e.value) == "No data left to save outside partition columns"
def test_partition_on_raises_pocols_missing(empty):
original_df = pd.DataFrame({"test": [1, 2, 3]})
if empty:
original_df = original_df.loc[[]]
mp = MetaPartition(
label="label_1", file="file", data=original_df, metadata_version=4
)
with pytest.raises(ValueError) as e:
mp.partition_on(["test", "foo", "bar"])
assert str(e.value) == "Partition column(s) missing: bar, foo"
def test_partition_on_valid_schemas():
"""
Ensure that partitioning is possible even if the output schemas of the
sub partitions may be different
"""
df = pd.DataFrame({"partition_col": [0, 1], "values": [None, "str"]})
mp = MetaPartition(label="base_label", data=df, metadata_version=4)
mp = mp.partition_on(["partition_col"])
assert len(mp) == 2
expected_meta = make_meta(df, origin="1", partition_keys="partition_col")
assert mp.schema == expected_meta
def test_concat_metapartition(df_all_types):
mp1 = MetaPartition(label="first", data=df_all_types, metadata_version=4)
mp2 = MetaPartition(label="second", data=df_all_types, metadata_version=4)
new_mp = MetaPartition.concat_metapartitions([mp1, mp2])
# what the label actually is, doesn't matter so much
assert new_mp.label is not None
df_expected = pd.concat([df_all_types, df_all_types])
df_actual = new_mp.data
pdt.assert_frame_equal(df_actual, df_expected)
