diff --git a/datasets/flwr_datasets/partitioner/__init__.py b/datasets/flwr_datasets/partitioner/__init__.py index 59f647f44b16..583c48efee93 100644 --- a/datasets/flwr_datasets/partitioner/__init__.py +++ b/datasets/flwr_datasets/partitioner/__init__.py @@ -30,6 +30,7 @@ from .size_partitioner import SizePartitioner from .square_partitioner import SquarePartitioner from .vertical_even_partitioner import VerticalEvenPartitioner +from .vertical_size_partitioner import VerticalSizePartitioner __all__ = [ "DirichletPartitioner", @@ -47,4 +48,5 @@ "SizePartitioner", "SquarePartitioner", "VerticalEvenPartitioner", + "VerticalSizePartitioner", ] diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py new file mode 100644 index 000000000000..de6161a51c67 --- /dev/null +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -0,0 +1,297 @@ +# Copyright 2024 Flower Labs GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""VerticalSizePartitioner class.""" +# flake8: noqa: E501 +# pylint: disable=C0301, R0902, R0913 +from math import floor +from typing import Literal, Optional, Union, cast + +import numpy as np + +import datasets +from flwr_datasets.partitioner.partitioner import Partitioner +from flwr_datasets.partitioner.vertical_partitioner_utils import ( + _add_active_party_columns, +) + + +class VerticalSizePartitioner(Partitioner): + """Creates vertical partitions by spliting features (columns) based on sizes. + + The sizes refer to the number of columns after the `drop_columns` are + dropped. `shared_columns` and `active_party_column` are excluded and + added only after the size-based division. + + Enables selection of "active party" column(s) and palcement into + a specific partition or creation of a new partition just for it. + Also enables droping columns and sharing specified columns across + all partitions. + + Parameters + ---------- + partition_sizes : Union[list[int], list[float]] + A list where each value represents the size of a partition. + list[int] -> each value represent an absolute number of columns. Size zero is + allowed and will result in an empty partition if no shared columns are present. + list of floats -> each value represent a fraction total number of columns. + Note that applies to collums without `active_party_columns` or `shared_columns`. + They are additionally included in to the partition(s). + active_party_column : Optional[Union[str, list[str]]] + Column(s) (typically representing labels) associated with the + "active party" (which can be the server). + active_party_columns_mode : Union[Literal[["add_to_first", "add_to_last", "create_as_first", "create_as_last", "add_to_all"], int] + Determines how to assign the active party columns: + - "add_to_first": Append active party columns to the first partition. + - "add_to_last": Append active party columns to the last partition. + - int: Append active party columns to the specified partition index. + - "create_as_first": Create a new partition at the start containing only + these columns. + - "create_as_last": Create a new partition at the end containing only + these columns. + - "add_to_all": Append active party columns to all partitions. + drop_columns : Optional[list[str]] + Columns to remove entirely from the dataset before partitioning. + shared_columns : Optional[list[str]] + Columns to duplicate into every partition after initial partitioning. + shuffle : bool + Whether to shuffle the order of columns before partitioning. + seed : Optional[int] + Random seed for shuffling columns. Has no effect if `shuffle=False`. + + Examples + -------- + >>> partitioner = VerticalEvenPartitioner( + ... partition_sizes=[8, 4, 2], + ... active_party_columns=["income"], + ... active_party_columns_mode="create_as_last" + ... ) + >>> fds = FederatedDataset( + ... dataset="scikit-learn/adult-census-income", + ... partitioners={"train": partitioner} + ... ) + >>> partitions = [fds.load_partition(i) for i in range(partitioner.num_partitions)] + >>> print([partition.column_names for partition in partitions]) + """ + + def __init__( + self, + partition_sizes: Union[list[int], list[float]], + active_party_column: Optional[Union[str, list[str]]] = None, + active_party_columns_mode: Union[ + Literal[ + "add_to_first", + "add_to_last", + "create_as_first", + "create_as_last", + "add_to_all", + ], + int, + ] = "add_to_last", + drop_columns: Optional[list[str]] = None, + shared_columns: Optional[list[str]] = None, + shuffle: bool = True, + seed: Optional[int] = 42, + ) -> None: + super().__init__() + + self._partition_sizes = partition_sizes + self._active_party_columns = self._init_active_party_column(active_party_column) + self._active_party_columns_mode = active_party_columns_mode + self._drop_columns = drop_columns or [] + self._shared_columns = shared_columns or [] + self._shuffle = shuffle + self._seed = seed + self._rng = np.random.default_rng(seed=self._seed) + + self._partition_columns: Optional[list[list[str]]] = None + self._partitions_determined = False + + self._validate_parameters_in_init() + + def _determine_partitions_if_needed(self) -> None: + if self._partitions_determined: + return + + if self.dataset is None: + raise ValueError("No dataset is set for this partitioner.") + + all_columns = list(self.dataset.column_names) + self._validate_parameters_while_partitioning( + all_columns, self._shared_columns, self._active_party_columns + ) + columns = [column for column in all_columns if column not in self._drop_columns] + columns = [column for column in columns if column not in self._shared_columns] + columns = [ + column for column in columns if column not in self._active_party_columns + ] + + if self._shuffle: + self._rng.shuffle(columns) + if all(isinstance(fraction, float) for fraction in self._partition_sizes): + partition_columns = _fraction_split( + columns, cast(list[float], self._partition_sizes) + ) + else: + partition_columns = _count_split( + columns, cast(list[int], self._partition_sizes) + ) + + partition_columns = _add_active_party_columns( + self._active_party_columns, + self._active_party_columns_mode, + partition_columns, + ) + + # Add shared columns to all partitions + for partition in partition_columns: + for column in self._shared_columns: + partition.append(column) + + self._partition_columns = partition_columns + self._partitions_determined = True + + def load_partition(self, partition_id: int) -> datasets.Dataset: + """Load a partition based on the partition index. + + Parameters + ---------- + partition_id : int + The index that corresponds to the requested partition. + + Returns + ------- + dataset_partition : Dataset + Single partition of a dataset. + """ + self._determine_partitions_if_needed() + assert self._partition_columns is not None + if partition_id < 0 or partition_id >= len(self._partition_columns): + raise ValueError(f"Invalid partition_id {partition_id}.") + columns = self._partition_columns[partition_id] + return self.dataset.select_columns(columns) + + @property + def num_partitions(self) -> int: + """Number of partitions.""" + self._determine_partitions_if_needed() + assert self._partition_columns is not None + return len(self._partition_columns) + + def _validate_parameters_in_init(self) -> None: + if not isinstance(self._partition_sizes, list): + raise ValueError("partition_sizes must be a list.") + if all(isinstance(fraction, float) for fraction in self._partition_sizes): + fraction_sum = sum(self._partition_sizes) + if fraction_sum != 1.0: + raise ValueError("Float ratios in column_distribution must sum to 1.0.") + if any( + fraction < 0.0 or fraction > 1.0 for fraction in self._partition_sizes + ): + raise ValueError( + "All floats in column_distribution must be >= 0.0 and <= 1.0." + ) + elif all( + isinstance(coulumn_count, int) for coulumn_count in self._partition_sizes + ): + if any(coulumn_count < 0 for coulumn_count in self._partition_sizes): + raise ValueError("All integers in column_distribution must be >= 0.") + else: + raise ValueError("partition_sizes list must be all floats or all ints.") + + # Validate columns lists + for parameter_name, parameter_list in [ + ("drop_columns", self._drop_columns), + ("shared_columns", self._shared_columns), + ("active_party_columns", self._active_party_columns), + ]: + if not all(isinstance(column, str) for column in parameter_list): + raise ValueError(f"All entries in {parameter_name} must be strings.") + + valid_modes = { + "add_to_first", + "add_to_last", + "create_as_first", + "create_as_last", + "add_to_all", + } + if not ( + isinstance(self._active_party_columns_mode, int) + or self._active_party_columns_mode in valid_modes + ): + raise ValueError( + "active_party_columns_mode must be an int or one of " + "'add_to_first', 'add_to_last', 'create_as_first', 'create_as_last', " + "'add_to_all'." + ) + + def _validate_parameters_while_partitioning( + self, + all_columns: list[str], + shared_columns: list[str], + active_party_columns: list[str], + ) -> None: + # Shared columns existance check + for column in shared_columns: + if column not in all_columns: + raise ValueError(f"Shared column '{column}' not found in the dataset.") + # Active party columns existence check + for column in active_party_columns: + if column not in all_columns: + raise ValueError( + f"Active party column '{column}' not found in the dataset." + ) + num_columns = len(all_columns) + if all(isinstance(size, int) for size in self._partition_sizes): + if sum(self._partition_sizes) > num_columns: + raise ValueError( + "Sum of partition sizes cannot exceed the total number of columns." + ) + else: + pass + + def _init_active_party_column( + self, active_party_column: Optional[Union[str, list[str]]] + ) -> list[str]: + if active_party_column is None: + return [] + if isinstance(active_party_column, str): + return [active_party_column] + if isinstance(active_party_column, list): + return active_party_column + raise ValueError("active_party_column must be a string or a list of strings.") + + +def _count_split(columns: list[str], counts: list[int]) -> list[list[str]]: + partition_columns = [] + start = 0 + for count in counts: + end = start + count + partition_columns.append(columns[start:end]) + start = end + return partition_columns + + +def _fraction_split(columns: list[str], fractions: list[float]) -> list[list[str]]: + num_columns = len(columns) + partitions = [] + cumulative = 0 + for index, fraction in enumerate(fractions): + count = int(floor(fraction * num_columns)) + if index == len(fractions) - 1: + # Last partition takes the remainder + count = num_columns - cumulative + partitions.append(columns[cumulative : cumulative + count]) + cumulative += count + return partitions diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner_test.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner_test.py new file mode 100644 index 000000000000..bc6b8324ac52 --- /dev/null +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner_test.py @@ -0,0 +1,186 @@ +# Copyright 2024 Flower Labs GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""VerticalSizePartitioner class tests.""" +# mypy: disable-error-code=arg-type +# pylint: disable=R0902, R0913 +import unittest + +import numpy as np + +from datasets import Dataset +from flwr_datasets.partitioner.vertical_size_partitioner import VerticalSizePartitioner + + +def _create_dummy_dataset(column_names: list[str], num_rows: int = 100) -> Dataset: + """Create a dataset with random integer data.""" + rng = np.random.default_rng(seed=42) + data = {col: rng.integers(0, 100, size=num_rows).tolist() for col in column_names} + return Dataset.from_dict(data) + + +class TestVerticalSizePartitioner(unittest.TestCase): + """Tests for VerticalSizePartitioner.""" + + def test_init_invalid_partition_sizes_type(self) -> None: + """Check ValueError if partition_sizes is not a list.""" + with self.assertRaises(ValueError): + VerticalSizePartitioner(partition_sizes="not_a_list") + + def test_init_mixed_partition_sizes_types(self) -> None: + """Check ValueError if partition_sizes mix int and float.""" + with self.assertRaises(ValueError): + VerticalSizePartitioner(partition_sizes=[0.5, 1]) + + def test_init_float_partitions_sum_not_one(self) -> None: + """Check ValueError if float partitions do not sum to 1.""" + with self.assertRaises(ValueError): + VerticalSizePartitioner(partition_sizes=[0.3, 0.3]) + + def test_init_float_partitions_out_of_range(self) -> None: + """Check ValueError if any float partition <0 or >1.""" + with self.assertRaises(ValueError): + VerticalSizePartitioner(partition_sizes=[-0.5, 1.5]) + + def test_init_int_partitions_negative(self) -> None: + """Check ValueError if any int partition size is negative.""" + with self.assertRaises(ValueError): + VerticalSizePartitioner(partition_sizes=[5, -1]) + + def test_init_invalid_mode(self) -> None: + """Check ValueError if active_party_columns_mode is invalid.""" + with self.assertRaises(ValueError): + VerticalSizePartitioner( + partition_sizes=[2, 2], active_party_columns_mode="invalid" + ) + + def test_init_active_party_column_invalid_type(self) -> None: + """Check ValueError if active_party_column is not str/list.""" + with self.assertRaises(ValueError): + VerticalSizePartitioner(partition_sizes=[2, 2], active_party_column=123) + + def test_partitioning_with_int_sizes(self) -> None: + """Check correct partitioning with integer sizes.""" + columns = ["f1", "f2", "f3", "f4", "f5"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner(partition_sizes=[2, 3], shuffle=False) + partitioner.dataset = dataset + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + self.assertEqual(len(p0.column_names), 2) + self.assertEqual(len(p1.column_names), 3) + + def test_partitioning_with_fraction_sizes(self) -> None: + """Check correct partitioning with fraction sizes.""" + columns = ["f1", "f2", "f3", "f4"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner(partition_sizes=[0.5, 0.5], shuffle=False) + partitioner.dataset = dataset + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + self.assertEqual(len(p0.column_names), 2) + self.assertEqual(len(p1.column_names), 2) + + def test_partitioning_with_drop_columns(self) -> None: + """Check dropping specified columns before partitioning.""" + columns = ["f1", "drop_me", "f2", "f3"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner( + partition_sizes=[2, 1], drop_columns=["drop_me"], shuffle=False + ) + partitioner.dataset = dataset + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + all_cols = p0.column_names + p1.column_names + self.assertNotIn("drop_me", all_cols) + + def test_partitioning_with_shared_columns(self) -> None: + """Check shared columns added to every partition.""" + columns = ["f1", "f2", "shared"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner( + partition_sizes=[1, 1], shared_columns=["shared"], shuffle=False + ) + partitioner.dataset = dataset + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + self.assertIn("shared", p0.column_names) + self.assertIn("shared", p1.column_names) + + def test_partitioning_with_active_party_add_to_last(self) -> None: + """Check active party columns added to the last partition.""" + columns = ["f1", "f2", "label"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner( + partition_sizes=[2], + active_party_column="label", + active_party_columns_mode="add_to_last", + shuffle=False, + ) + partitioner.dataset = dataset + p0 = partitioner.load_partition(0) + self.assertIn("label", p0.column_names) + + def test_partitioning_with_active_party_create_as_first(self) -> None: + """Check creating a new first partition for active party cols.""" + columns = ["f1", "f2", "label"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner( + partition_sizes=[2], + active_party_column="label", + active_party_columns_mode="create_as_first", + shuffle=False, + ) + partitioner.dataset = dataset + self.assertEqual(partitioner.num_partitions, 2) + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + self.assertEqual(p0.column_names, ["label"]) + self.assertIn("f1", p1.column_names) + self.assertIn("f2", p1.column_names) + + def test_partitioning_with_nonexistent_shared_column(self) -> None: + """Check ValueError if shared column does not exist.""" + columns = ["f1", "f2"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner( + partition_sizes=[1], shared_columns=["nonexistent"], shuffle=False + ) + partitioner.dataset = dataset + with self.assertRaises(ValueError): + partitioner.load_partition(0) + + def test_partitioning_with_nonexistent_active_party_column(self) -> None: + """Check ValueError if active party column does not exist.""" + columns = ["f1", "f2"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner( + partition_sizes=[1], active_party_column="missing_label", shuffle=False + ) + partitioner.dataset = dataset + with self.assertRaises(ValueError): + partitioner.load_partition(0) + + def test_sum_of_int_partition_sizes_exceeds_num_columns(self) -> None: + """Check ValueError if sum of int sizes > total columns.""" + columns = ["f1", "f2"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner(partition_sizes=[3], shuffle=False) + partitioner.dataset = dataset + with self.assertRaises(ValueError): + partitioner.load_partition(0) + + +if __name__ == "__main__": + unittest.main()