From 6ac5a7b7532881dfedbd932a8e7cbdfc55cb8ee1 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Thu, 12 Dec 2024 17:46:54 +0100 Subject: [PATCH 01/15] Add even vertical partitioner --- .../flwr_datasets/partitioner/__init__.py | 2 + .../partitioner/vertical_even_partitioner.py | 226 ++++++++++++++++++ .../vertical_even_partitioner_test.py | 201 ++++++++++++++++ .../partitioner/vertical_partitioner_utils.py | 102 ++++++++ .../vertical_partitioner_utils_test.py | 144 +++++++++++ 5 files changed, 675 insertions(+) create mode 100644 datasets/flwr_datasets/partitioner/vertical_even_partitioner.py create mode 100644 datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py create mode 100644 datasets/flwr_datasets/partitioner/vertical_partitioner_utils.py create mode 100644 datasets/flwr_datasets/partitioner/vertical_partitioner_utils_test.py diff --git a/datasets/flwr_datasets/partitioner/__init__.py b/datasets/flwr_datasets/partitioner/__init__.py index a14efa1cc905..59f647f44b16 100644 --- a/datasets/flwr_datasets/partitioner/__init__.py +++ b/datasets/flwr_datasets/partitioner/__init__.py @@ -29,6 +29,7 @@ from .shard_partitioner import ShardPartitioner from .size_partitioner import SizePartitioner from .square_partitioner import SquarePartitioner +from .vertical_even_partitioner import VerticalEvenPartitioner __all__ = [ "DirichletPartitioner", @@ -45,4 +46,5 @@ "ShardPartitioner", "SizePartitioner", "SquarePartitioner", + "VerticalEvenPartitioner", ] diff --git a/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py new file mode 100644 index 000000000000..6a6df3df35a0 --- /dev/null +++ b/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py @@ -0,0 +1,226 @@ +# 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. +# ============================================================================== +"""VerticalEvenPartitioner class.""" +# flake8: noqa: E501 +from typing import Literal, Optional, Union + +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, + _list_split, +) + + +class VerticalEvenPartitioner(Partitioner): + """Partitioner that splits features (columns) evenly into vertical partitions. + + 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. + + The number and nature of partitions can be defined in various ways: + - By specifying a simple integer for even splitting. + - By providing ratios or absolute counts for each partition. + - By explicitly listing the columns for each partition. + (see `column_distribution` and `mode` parameters for more details) + + Parameters + ---------- + num_partitions : int + Number of partitions to create. + active_party_columns : Optional[list[str]] + Columns 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( + ... num_partitions=3, + ... active_party_columns=["income"], + ... active_party_columns_mode="add_to_last", + ... shuffle=True, + ... seed=42 + ... ) + >>> 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, + num_partitions: int, + active_party_columns: Optional[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._num_partitions = num_partitions + self._active_party_columns = active_party_columns or [] + 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) + partition_columns = _list_split(columns, self._num_partitions) + 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 self._num_partitions < 1: + raise ValueError("column_distribution as int must be >= 1.") + + # 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." + ) diff --git a/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py b/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py new file mode 100644 index 000000000000..3b35208706c2 --- /dev/null +++ b/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py @@ -0,0 +1,201 @@ +# 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. +# ============================================================================== +"""VerticalEvenPartitioner class tests.""" +# mypy: disable-error-code=list-item,arg-type +import unittest + +import numpy as np + +from datasets import Dataset +from flwr_datasets.partitioner.vertical_even_partitioner import VerticalEvenPartitioner + + +def _create_dummy_dataset(column_names: list[str], num_rows: int = 100) -> Dataset: + """Create a dummy dataset with random data for testing.""" + data = {} + rng = np.random.default_rng(seed=42) + for col in column_names: + # Just numeric data; could also be strings, categoricals, etc. + data[col] = rng.integers(0, 100, size=num_rows).tolist() + return Dataset.from_dict(data) + + +class TestVerticalEvenPartitioner(unittest.TestCase): + """Unit tests for VerticalEvenPartitioner.""" + + def test_init_with_invalid_num_partitions(self) -> None: + """Test that initializing with an invalid number of partitions.""" + with self.assertRaises(ValueError): + VerticalEvenPartitioner(num_partitions=0) + + def test_init_with_invalid_active_party_mode(self) -> None: + """Test initialization with invalid active_party_columns_mode.""" + with self.assertRaises(ValueError): + VerticalEvenPartitioner( + num_partitions=2, active_party_columns_mode="invalid_mode" + ) + + def test_init_with_non_string_drop_columns(self) -> None: + """Test initialization with non-string elements in drop_columns.""" + with self.assertRaises(ValueError): + VerticalEvenPartitioner(num_partitions=2, drop_columns=[1, "a", 3]) + + def test_init_with_non_string_shared_columns(self) -> None: + """Test initialization with non-string elements in shared_columns.""" + with self.assertRaises(ValueError): + VerticalEvenPartitioner(num_partitions=2, shared_columns=["col1", 123]) + + def test_init_with_non_string_active_party_columns(self) -> None: + """Test initialization with non-string elements in active_party_columns.""" + with self.assertRaises(ValueError): + VerticalEvenPartitioner( + num_partitions=2, active_party_columns=["col1", None] + ) + + def test_partitioning_basic(self) -> None: + """Test basic partitioning with no special columns or dropping.""" + columns = ["feature1", "feature2", "feature3", "feature4"] + dataset = _create_dummy_dataset(columns, num_rows=50) + partitioner = VerticalEvenPartitioner(num_partitions=2, shuffle=False) + partitioner.dataset = dataset + + self.assertEqual(partitioner.num_partitions, 2) + + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + + self.assertEqual(len(p0.column_names), 2) + self.assertEqual(len(p1.column_names), 2) + self.assertIn("feature1", p0.column_names) + self.assertIn("feature2", p0.column_names) + self.assertIn("feature3", p1.column_names) + self.assertIn("feature4", p1.column_names) + + def test_partitioning_with_drop_columns(self) -> None: + """Test partitioning while dropping some columns.""" + columns = ["feature1", "feature2", "drop_me", "feature3", "feature4"] + dataset = _create_dummy_dataset(columns, num_rows=50) + partitioner = VerticalEvenPartitioner( + num_partitions=2, drop_columns=["drop_me"], shuffle=False, seed=42 + ) + partitioner.dataset = dataset + + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + all_partition_columns = p0.column_names + p1.column_names + + # The drop_me should not be in any partition + self.assertNotIn("drop_me", all_partition_columns) + # The rest of columns should be distributed + self.assertIn("feature1", all_partition_columns) + self.assertIn("feature2", all_partition_columns) + self.assertIn("feature3", all_partition_columns) + self.assertIn("feature4", all_partition_columns) + + def test_partitioning_with_shared_columns(self) -> None: + """Test that shared columns are present in all partitions.""" + columns = ["f1", "f2", "f3", "f4", "shared_col"] + dataset = _create_dummy_dataset(columns, num_rows=50) + partitioner = VerticalEvenPartitioner( + num_partitions=2, shared_columns=["shared_col"], shuffle=False, seed=42 + ) + partitioner.dataset = dataset + + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + + self.assertIn("shared_col", p0.column_names) + self.assertIn("shared_col", p1.column_names) + + def test_partitioning_with_active_party_columns_add_to_last(self) -> None: + """Test active party columns are appended to the last partition.""" + columns = ["f1", "f2", "f3", "f4", "income"] + dataset = _create_dummy_dataset(columns, num_rows=50) + partitioner = VerticalEvenPartitioner( + num_partitions=2, + active_party_columns=["income"], + active_party_columns_mode="add_to_last", + shuffle=False, + seed=42, + ) + partitioner.dataset = dataset + + p0 = partitioner.load_partition(0) + p1 = partitioner.load_partition(1) + + # The income should be only in the last partition + self.assertNotIn("income", p0.column_names) + self.assertIn("income", p1.column_names) + + def test_partitioning_with_active_party_columns_create_as_first(self) -> None: + """Test creating a new partition solely for active party columns.""" + columns = ["f1", "f2", "f3", "f4", "income"] + dataset = _create_dummy_dataset(columns, num_rows=50) + partitioner = VerticalEvenPartitioner( + num_partitions=2, + active_party_columns=["income"], + active_party_columns_mode="create_as_first", + shuffle=False, + ) + partitioner.dataset = dataset + + # The first partition should be just the active party columns + # and then two more partitions from original splitting. + self.assertEqual(partitioner.num_partitions, 3) + + p0 = partitioner.load_partition(0) # active party partition + p1 = partitioner.load_partition(1) + p2 = partitioner.load_partition(2) + + self.assertEqual(p0.column_names, ["income"]) + self.assertIn("f1", p1.column_names) + self.assertIn("f2", p1.column_names) + self.assertIn("f3", p2.column_names) + self.assertIn("f4", p2.column_names) + + def test_partitioning_with_nonexistent_active_party_columns(self) -> None: + """Test that a ValueError is raised if active party column does not exist.""" + columns = ["f1", "f2", "f3", "f4"] + dataset = _create_dummy_dataset(columns, num_rows=50) + partitioner = VerticalEvenPartitioner( + num_partitions=2, + active_party_columns=["income"], # Not present in dataset + active_party_columns_mode="add_to_last", + shuffle=False, + ) + partitioner.dataset = dataset + + with self.assertRaises(ValueError) as context: + partitioner.load_partition(0) + self.assertIn("Active party column 'income' not found", str(context.exception)) + + def test_partitioning_with_nonexistent_shared_columns(self) -> None: + """Test that a ValueError is raised if shared column does not exist.""" + columns = ["f1", "f2", "f3"] + dataset = _create_dummy_dataset(columns, num_rows=50) + partitioner = VerticalEvenPartitioner( + num_partitions=2, shared_columns=["nonexistent_col"], shuffle=False + ) + partitioner.dataset = dataset + + with self.assertRaises(ValueError) as context: + partitioner.load_partition(0) + self.assertIn( + "Shared column 'nonexistent_col' not found", str(context.exception) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/datasets/flwr_datasets/partitioner/vertical_partitioner_utils.py b/datasets/flwr_datasets/partitioner/vertical_partitioner_utils.py new file mode 100644 index 000000000000..8859bec6c675 --- /dev/null +++ b/datasets/flwr_datasets/partitioner/vertical_partitioner_utils.py @@ -0,0 +1,102 @@ +# 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. +# ============================================================================== +"""VerticalPartitioner utils.py.""" +# flake8: noqa: E501 +from typing import Any, Literal, Union + + +def _list_split(lst: list[Any], num_sublists: int) -> list[list[Any]]: + """Split a list into n nearly equal-sized sublists. + + Parameters + ---------- + lst : list[Any] + The list to split. + num_sublists : int + Number of sublists to create. + + Returns + ------- + subslist: list[list[Any]] + A list containing num_sublists sublists. + """ + if num_sublists <= 0: + raise ValueError("Number of splits must be greater than 0") + chunk_size, remainder = divmod(len(lst), num_sublists) + sublists = [] + start_index = 0 + for i in range(num_sublists): + end_index = start_index + chunk_size + if i < remainder: + end_index += 1 + sublists.append(lst[start_index:end_index]) + start_index = end_index + return sublists + + +def _add_active_party_columns( + active_party_columns: list[str], + active_party_columns_mode: Union[ + Literal[ + "add_to_first", + "add_to_last", + "create_as_first", + "create_as_last", + "add_to_all", + ], + int, + ], + partition_columns: list[list[str]], +) -> list[list[str]]: + """Add active party columns to the partition columns based on the mode. + + Parameters + ---------- + active_party_columns : list[str] + List of active party columns. + active_party_columns_mode : Union[Literal["add_to_first", "add_to_last", "create_as_first", "create_as_last", "add_to_all"], int] + Mode to add active party columns to partition columns. + + Returns + ------- + partition_columns: list[list[str]] + List of partition columns after the modyfication. + """ + if isinstance(active_party_columns_mode, int): + partition_id = active_party_columns_mode + if partition_id < 0 or partition_id >= len(partition_columns): + raise ValueError( + f"Invalid partition index {partition_id} for active_party_columns_mode." + f"Must be in the range [0, {len(partition_columns) - 1}]" + f"but given {partition_id}" + ) + for column in active_party_columns: + partition_columns[partition_id].append(column) + else: + if active_party_columns_mode == "add_to_first": + for column in active_party_columns: + partition_columns[0].append(column) + elif active_party_columns_mode == "add_to_last": + for column in active_party_columns: + partition_columns[-1].append(column) + elif active_party_columns_mode == "create_as_first": + partition_columns.insert(0, active_party_columns) + elif active_party_columns_mode == "create_as_last": + partition_columns.append(active_party_columns) + elif active_party_columns_mode == "add_to_all": + for column in active_party_columns: + for partition in partition_columns: + partition.append(column) + return partition_columns diff --git a/datasets/flwr_datasets/partitioner/vertical_partitioner_utils_test.py b/datasets/flwr_datasets/partitioner/vertical_partitioner_utils_test.py new file mode 100644 index 000000000000..f85d027fe444 --- /dev/null +++ b/datasets/flwr_datasets/partitioner/vertical_partitioner_utils_test.py @@ -0,0 +1,144 @@ +# 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. +# ============================================================================== +"""Tests for vertical partitioner utilities.""" +import unittest +from typing import Any, Literal + +from flwr_datasets.partitioner.vertical_partitioner_utils import ( + _add_active_party_columns, + _list_split, +) + + +class TestVerticalPartitionerUtils(unittest.TestCase): + """Tests for _list_split and _add_active_party_columns utilities.""" + + def test_list_split_basic_splitting(self) -> None: + """Check equal splitting with divisible lengths.""" + lst = [1, 2, 3, 4, 5, 6] + result = _list_split(lst, 3) + expected = [[1, 2], [3, 4], [5, 6]] + self.assertEqual(result, expected) + + def test_list_split_uneven_splitting(self) -> None: + """Check uneven splitting with non-divisible lengths.""" + lst = [10, 20, 30, 40, 50] + result = _list_split(lst, 2) + expected = [[10, 20, 30], [40, 50]] + self.assertEqual(result, expected) + + def test_list_split_single_sublist(self) -> None: + """Check that single sublist returns the full list.""" + lst = [1, 2, 3] + result = _list_split(lst, 1) + expected = [[1, 2, 3]] + self.assertEqual(result, expected) + + def test_list_split_more_sublists_than_elements(self) -> None: + """Check extra sublists are empty when count exceeds length.""" + lst = [42] + result = _list_split(lst, 3) + expected = [[42], [], []] + self.assertEqual(result, expected) + + def test_list_split_empty_list(self) -> None: + """Check splitting empty list produces empty sublists.""" + lst: list[Any] = [] + result = _list_split(lst, 3) + expected: list[list[Any]] = [[], [], []] + self.assertEqual(result, expected) + + def test_list_split_invalid_num_sublists(self) -> None: + """Check ValueError when sublist count is zero or negative.""" + lst = [1, 2, 3] + with self.assertRaises(ValueError): + _list_split(lst, 0) + + def test_add_to_first(self) -> None: + """Check adding active cols to the first partition.""" + partition_columns = [["col1", "col2"], ["col3"], ["col4"]] + active_party_columns = ["active1", "active2"] + mode: Literal["add_to_first"] = "add_to_first" + result = _add_active_party_columns( + active_party_columns, mode, partition_columns + ) + self.assertEqual( + result, [["col1", "col2", "active1", "active2"], ["col3"], ["col4"]] + ) + + def test_add_to_last(self) -> None: + """Check adding active cols to the last partition.""" + partition_columns = [["col1", "col2"], ["col3"], ["col4"]] + active_party_columns = ["active"] + mode: Literal["add_to_last"] = "add_to_last" + result = _add_active_party_columns( + active_party_columns, mode, partition_columns + ) + self.assertEqual(result, [["col1", "col2"], ["col3"], ["col4", "active"]]) + + def test_create_as_first(self) -> None: + """Check creating a new first partition for active cols.""" + partition_columns = [["col1"], ["col2"]] + active_party_columns = ["active1", "active2"] + mode: Literal["create_as_first"] = "create_as_first" + result = _add_active_party_columns( + active_party_columns, mode, partition_columns + ) + self.assertEqual(result, [["active1", "active2"], ["col1"], ["col2"]]) + + def test_create_as_last(self) -> None: + """Check creating a new last partition for active cols.""" + partition_columns = [["col1"], ["col2"]] + active_party_columns = ["active1", "active2"] + mode: Literal["create_as_last"] = "create_as_last" + result = _add_active_party_columns( + active_party_columns, mode, partition_columns + ) + self.assertEqual(result, [["col1"], ["col2"], ["active1", "active2"]]) + + def test_add_to_all(self) -> None: + """Check adding active cols to all partitions.""" + partition_columns = [["col1"], ["col2", "col3"], ["col4"]] + active_party_columns = ["active"] + mode: Literal["add_to_all"] = "add_to_all" + result = _add_active_party_columns( + active_party_columns, mode, partition_columns + ) + self.assertEqual( + result, [["col1", "active"], ["col2", "col3", "active"], ["col4", "active"]] + ) + + def test_add_to_specific_partition_valid_index(self) -> None: + """Check adding active cols to a specific valid partition.""" + partition_columns = [["col1"], ["col2"], ["col3"]] + active_party_columns = ["active1", "active2"] + mode: int = 1 + result = _add_active_party_columns( + active_party_columns, mode, partition_columns + ) + self.assertEqual(result, [["col1"], ["col2", "active1", "active2"], ["col3"]]) + + def test_add_to_specific_partition_invalid_index(self) -> None: + """Check ValueError when partition index is invalid.""" + partition_columns = [["col1"], ["col2"]] + active_party_columns = ["active"] + mode: int = 5 + with self.assertRaises(ValueError) as context: + _add_active_party_columns(active_party_columns, mode, partition_columns) + self.assertIn("Invalid partition index", str(context.exception)) + + +if __name__ == "__main__": + unittest.main() From 7b56c1fe191d05282617ecd3b174658330390582 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Fri, 13 Dec 2024 09:34:13 +0100 Subject: [PATCH 02/15] Fix formatting errors --- .../flwr_datasets/partitioner/vertical_even_partitioner.py | 1 + .../partitioner/vertical_even_partitioner_test.py | 5 +++-- .../flwr_datasets/partitioner/vertical_partitioner_utils.py | 1 + 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py index 6a6df3df35a0..180c4bd07347 100644 --- a/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py @@ -14,6 +14,7 @@ # ============================================================================== """VerticalEvenPartitioner class.""" # flake8: noqa: E501 +# pylint: disable=C0301, R0902, R0913 from typing import Literal, Optional, Union import numpy as np diff --git a/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py b/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py index 3b35208706c2..8e766617d609 100644 --- a/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py +++ b/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== """VerticalEvenPartitioner class tests.""" -# mypy: disable-error-code=list-item,arg-type +# mypy: disable-error-code=list-item import unittest import numpy as np @@ -44,7 +44,8 @@ def test_init_with_invalid_active_party_mode(self) -> None: """Test initialization with invalid active_party_columns_mode.""" with self.assertRaises(ValueError): VerticalEvenPartitioner( - num_partitions=2, active_party_columns_mode="invalid_mode" + num_partitions=2, + active_party_columns_mode="invalid_mode", # type: ignore[arg-type] ) def test_init_with_non_string_drop_columns(self) -> None: diff --git a/datasets/flwr_datasets/partitioner/vertical_partitioner_utils.py b/datasets/flwr_datasets/partitioner/vertical_partitioner_utils.py index 8859bec6c675..e9e7e3855ef4 100644 --- a/datasets/flwr_datasets/partitioner/vertical_partitioner_utils.py +++ b/datasets/flwr_datasets/partitioner/vertical_partitioner_utils.py @@ -14,6 +14,7 @@ # ============================================================================== """VerticalPartitioner utils.py.""" # flake8: noqa: E501 +# pylint: disable=C0301 from typing import Any, Literal, Union From b0b1d50c921faab16d541cfdb82c99eccb0a3b4b Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Fri, 13 Dec 2024 12:16:24 +0100 Subject: [PATCH 03/15] Add VerticalSizePartitioner --- .../flwr_datasets/partitioner/__init__.py | 2 + .../partitioner/vertical_size_partitioner.py | 297 ++++++++++++++++++ .../vertical_size_partitioner_test.py | 186 +++++++++++ 3 files changed, 485 insertions(+) create mode 100644 datasets/flwr_datasets/partitioner/vertical_size_partitioner.py create mode 100644 datasets/flwr_datasets/partitioner/vertical_size_partitioner_test.py 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() From 99780e093fcfd8589854d6e7f5a0cf78716f6e75 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Fri, 13 Dec 2024 12:27:23 +0100 Subject: [PATCH 04/15] Update example --- .../partitioner/vertical_size_partitioner.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index de6161a51c67..ba847963d994 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -72,16 +72,19 @@ class VerticalSizePartitioner(Partitioner): Examples -------- - >>> partitioner = VerticalEvenPartitioner( + >>> from flwr_datasets import FederatedDataset + >>> from flwr_datasets.partitioner import VerticalSizePartitioner + >>> + >>> partitioner = VerticalSizePartitioner( ... partition_sizes=[8, 4, 2], - ... active_party_columns=["income"], + ... active_party_column="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)] + >>> partitions = [fds.load_partition(i) for i in range(fds.partitioners["train"].num_partitions)] >>> print([partition.column_names for partition in partitions]) """ From 08cb0653843b41d9e15e5dc4032119f313a3484a Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Fri, 13 Dec 2024 12:42:19 +0100 Subject: [PATCH 05/15] Remove vertical even partitioner --- .../partitioner/vertical_even_partitioner.py | 227 ------------------ .../vertical_even_partitioner_test.py | 202 ---------------- 2 files changed, 429 deletions(-) delete mode 100644 datasets/flwr_datasets/partitioner/vertical_even_partitioner.py delete mode 100644 datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py diff --git a/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py deleted file mode 100644 index 180c4bd07347..000000000000 --- a/datasets/flwr_datasets/partitioner/vertical_even_partitioner.py +++ /dev/null @@ -1,227 +0,0 @@ -# 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. -# ============================================================================== -"""VerticalEvenPartitioner class.""" -# flake8: noqa: E501 -# pylint: disable=C0301, R0902, R0913 -from typing import Literal, Optional, Union - -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, - _list_split, -) - - -class VerticalEvenPartitioner(Partitioner): - """Partitioner that splits features (columns) evenly into vertical partitions. - - 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. - - The number and nature of partitions can be defined in various ways: - - By specifying a simple integer for even splitting. - - By providing ratios or absolute counts for each partition. - - By explicitly listing the columns for each partition. - (see `column_distribution` and `mode` parameters for more details) - - Parameters - ---------- - num_partitions : int - Number of partitions to create. - active_party_columns : Optional[list[str]] - Columns 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( - ... num_partitions=3, - ... active_party_columns=["income"], - ... active_party_columns_mode="add_to_last", - ... shuffle=True, - ... seed=42 - ... ) - >>> 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, - num_partitions: int, - active_party_columns: Optional[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._num_partitions = num_partitions - self._active_party_columns = active_party_columns or [] - 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) - partition_columns = _list_split(columns, self._num_partitions) - 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 self._num_partitions < 1: - raise ValueError("column_distribution as int must be >= 1.") - - # 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." - ) diff --git a/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py b/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py deleted file mode 100644 index 8e766617d609..000000000000 --- a/datasets/flwr_datasets/partitioner/vertical_even_partitioner_test.py +++ /dev/null @@ -1,202 +0,0 @@ -# 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. -# ============================================================================== -"""VerticalEvenPartitioner class tests.""" -# mypy: disable-error-code=list-item -import unittest - -import numpy as np - -from datasets import Dataset -from flwr_datasets.partitioner.vertical_even_partitioner import VerticalEvenPartitioner - - -def _create_dummy_dataset(column_names: list[str], num_rows: int = 100) -> Dataset: - """Create a dummy dataset with random data for testing.""" - data = {} - rng = np.random.default_rng(seed=42) - for col in column_names: - # Just numeric data; could also be strings, categoricals, etc. - data[col] = rng.integers(0, 100, size=num_rows).tolist() - return Dataset.from_dict(data) - - -class TestVerticalEvenPartitioner(unittest.TestCase): - """Unit tests for VerticalEvenPartitioner.""" - - def test_init_with_invalid_num_partitions(self) -> None: - """Test that initializing with an invalid number of partitions.""" - with self.assertRaises(ValueError): - VerticalEvenPartitioner(num_partitions=0) - - def test_init_with_invalid_active_party_mode(self) -> None: - """Test initialization with invalid active_party_columns_mode.""" - with self.assertRaises(ValueError): - VerticalEvenPartitioner( - num_partitions=2, - active_party_columns_mode="invalid_mode", # type: ignore[arg-type] - ) - - def test_init_with_non_string_drop_columns(self) -> None: - """Test initialization with non-string elements in drop_columns.""" - with self.assertRaises(ValueError): - VerticalEvenPartitioner(num_partitions=2, drop_columns=[1, "a", 3]) - - def test_init_with_non_string_shared_columns(self) -> None: - """Test initialization with non-string elements in shared_columns.""" - with self.assertRaises(ValueError): - VerticalEvenPartitioner(num_partitions=2, shared_columns=["col1", 123]) - - def test_init_with_non_string_active_party_columns(self) -> None: - """Test initialization with non-string elements in active_party_columns.""" - with self.assertRaises(ValueError): - VerticalEvenPartitioner( - num_partitions=2, active_party_columns=["col1", None] - ) - - def test_partitioning_basic(self) -> None: - """Test basic partitioning with no special columns or dropping.""" - columns = ["feature1", "feature2", "feature3", "feature4"] - dataset = _create_dummy_dataset(columns, num_rows=50) - partitioner = VerticalEvenPartitioner(num_partitions=2, shuffle=False) - partitioner.dataset = dataset - - self.assertEqual(partitioner.num_partitions, 2) - - p0 = partitioner.load_partition(0) - p1 = partitioner.load_partition(1) - - self.assertEqual(len(p0.column_names), 2) - self.assertEqual(len(p1.column_names), 2) - self.assertIn("feature1", p0.column_names) - self.assertIn("feature2", p0.column_names) - self.assertIn("feature3", p1.column_names) - self.assertIn("feature4", p1.column_names) - - def test_partitioning_with_drop_columns(self) -> None: - """Test partitioning while dropping some columns.""" - columns = ["feature1", "feature2", "drop_me", "feature3", "feature4"] - dataset = _create_dummy_dataset(columns, num_rows=50) - partitioner = VerticalEvenPartitioner( - num_partitions=2, drop_columns=["drop_me"], shuffle=False, seed=42 - ) - partitioner.dataset = dataset - - p0 = partitioner.load_partition(0) - p1 = partitioner.load_partition(1) - all_partition_columns = p0.column_names + p1.column_names - - # The drop_me should not be in any partition - self.assertNotIn("drop_me", all_partition_columns) - # The rest of columns should be distributed - self.assertIn("feature1", all_partition_columns) - self.assertIn("feature2", all_partition_columns) - self.assertIn("feature3", all_partition_columns) - self.assertIn("feature4", all_partition_columns) - - def test_partitioning_with_shared_columns(self) -> None: - """Test that shared columns are present in all partitions.""" - columns = ["f1", "f2", "f3", "f4", "shared_col"] - dataset = _create_dummy_dataset(columns, num_rows=50) - partitioner = VerticalEvenPartitioner( - num_partitions=2, shared_columns=["shared_col"], shuffle=False, seed=42 - ) - partitioner.dataset = dataset - - p0 = partitioner.load_partition(0) - p1 = partitioner.load_partition(1) - - self.assertIn("shared_col", p0.column_names) - self.assertIn("shared_col", p1.column_names) - - def test_partitioning_with_active_party_columns_add_to_last(self) -> None: - """Test active party columns are appended to the last partition.""" - columns = ["f1", "f2", "f3", "f4", "income"] - dataset = _create_dummy_dataset(columns, num_rows=50) - partitioner = VerticalEvenPartitioner( - num_partitions=2, - active_party_columns=["income"], - active_party_columns_mode="add_to_last", - shuffle=False, - seed=42, - ) - partitioner.dataset = dataset - - p0 = partitioner.load_partition(0) - p1 = partitioner.load_partition(1) - - # The income should be only in the last partition - self.assertNotIn("income", p0.column_names) - self.assertIn("income", p1.column_names) - - def test_partitioning_with_active_party_columns_create_as_first(self) -> None: - """Test creating a new partition solely for active party columns.""" - columns = ["f1", "f2", "f3", "f4", "income"] - dataset = _create_dummy_dataset(columns, num_rows=50) - partitioner = VerticalEvenPartitioner( - num_partitions=2, - active_party_columns=["income"], - active_party_columns_mode="create_as_first", - shuffle=False, - ) - partitioner.dataset = dataset - - # The first partition should be just the active party columns - # and then two more partitions from original splitting. - self.assertEqual(partitioner.num_partitions, 3) - - p0 = partitioner.load_partition(0) # active party partition - p1 = partitioner.load_partition(1) - p2 = partitioner.load_partition(2) - - self.assertEqual(p0.column_names, ["income"]) - self.assertIn("f1", p1.column_names) - self.assertIn("f2", p1.column_names) - self.assertIn("f3", p2.column_names) - self.assertIn("f4", p2.column_names) - - def test_partitioning_with_nonexistent_active_party_columns(self) -> None: - """Test that a ValueError is raised if active party column does not exist.""" - columns = ["f1", "f2", "f3", "f4"] - dataset = _create_dummy_dataset(columns, num_rows=50) - partitioner = VerticalEvenPartitioner( - num_partitions=2, - active_party_columns=["income"], # Not present in dataset - active_party_columns_mode="add_to_last", - shuffle=False, - ) - partitioner.dataset = dataset - - with self.assertRaises(ValueError) as context: - partitioner.load_partition(0) - self.assertIn("Active party column 'income' not found", str(context.exception)) - - def test_partitioning_with_nonexistent_shared_columns(self) -> None: - """Test that a ValueError is raised if shared column does not exist.""" - columns = ["f1", "f2", "f3"] - dataset = _create_dummy_dataset(columns, num_rows=50) - partitioner = VerticalEvenPartitioner( - num_partitions=2, shared_columns=["nonexistent_col"], shuffle=False - ) - partitioner.dataset = dataset - - with self.assertRaises(ValueError) as context: - partitioner.load_partition(0) - self.assertIn( - "Shared column 'nonexistent_col' not found", str(context.exception) - ) - - -if __name__ == "__main__": - unittest.main() From c227f94410fb04693e9b2f8b3c58dda60f2d36bc Mon Sep 17 00:00:00 2001 From: Adam Narozniak <51029327+adam-narozniak@users.noreply.github.com> Date: Fri, 13 Dec 2024 13:56:32 +0100 Subject: [PATCH 06/15] Apply suggestions from code review Co-authored-by: Javier --- .../partitioner/vertical_size_partitioner.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index ba847963d994..f321be673ed0 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -198,20 +198,20 @@ def _validate_parameters_in_init(self) -> None: 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.") + raise ValueError("Float ratios in `partition_sizes` 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." + "All floats in `partition_sizes` 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.") + raise ValueError("All integers in `partition_sizes` must be >= 0.") else: - raise ValueError("partition_sizes list must be all floats or all ints.") + raise ValueError("`partition_sizes` list must be all floats or all ints.") # Validate columns lists for parameter_name, parameter_list in [ From a2e81b716759519adab2600fe3c45a24baef7fec Mon Sep 17 00:00:00 2001 From: Adam Narozniak <51029327+adam-narozniak@users.noreply.github.com> Date: Fri, 13 Dec 2024 13:56:50 +0100 Subject: [PATCH 07/15] Apply suggestions from code review Co-authored-by: Javier --- datasets/flwr_datasets/partitioner/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/datasets/flwr_datasets/partitioner/__init__.py b/datasets/flwr_datasets/partitioner/__init__.py index 583c48efee93..8770d5b8b76e 100644 --- a/datasets/flwr_datasets/partitioner/__init__.py +++ b/datasets/flwr_datasets/partitioner/__init__.py @@ -29,7 +29,6 @@ from .shard_partitioner import ShardPartitioner from .size_partitioner import SizePartitioner from .square_partitioner import SquarePartitioner -from .vertical_even_partitioner import VerticalEvenPartitioner from .vertical_size_partitioner import VerticalSizePartitioner __all__ = [ @@ -47,6 +46,5 @@ "ShardPartitioner", "SizePartitioner", "SquarePartitioner", - "VerticalEvenPartitioner", "VerticalSizePartitioner", ] From 8634269ab2ac0ad92a6f59e83cd936a4a6543b59 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Fri, 13 Dec 2024 14:16:15 +0100 Subject: [PATCH 08/15] Update index error message --- .../flwr_datasets/partitioner/vertical_size_partitioner.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index f321be673ed0..d70279e542ff 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -181,7 +181,9 @@ def load_partition(self, partition_id: int) -> datasets.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}.") + raise IndexError( + f"partition_id: {partition_id} out of range <0, {len(self.num_partitions) - 1}>." + ) columns = self._partition_columns[partition_id] return self.dataset.select_columns(columns) From ac3bc68bce61968e24499b8d0bf93611aa1d5e55 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Fri, 13 Dec 2024 14:17:52 +0100 Subject: [PATCH 09/15] Change the check to allow using only all columns --- .../flwr_datasets/partitioner/vertical_size_partitioner.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index d70279e542ff..e1c494eb433a 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -259,9 +259,9 @@ def _validate_parameters_while_partitioning( ) num_columns = len(all_columns) if all(isinstance(size, int) for size in self._partition_sizes): - if sum(self._partition_sizes) > num_columns: + if sum(self._partition_sizes) != num_columns: raise ValueError( - "Sum of partition sizes cannot exceed the total number of columns." + "Sum of partition sizes cannot differ from the total number of columns." ) else: pass From eaf9ac3e86f71ad2f3d8fe94d155130aea422e7e Mon Sep 17 00:00:00 2001 From: jafermarq Date: Mon, 16 Dec 2024 17:29:08 +0000 Subject: [PATCH 10/15] fix --- datasets/flwr_datasets/partitioner/vertical_size_partitioner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index e1c494eb433a..c2af7def660a 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -182,7 +182,7 @@ def load_partition(self, partition_id: int) -> datasets.Dataset: assert self._partition_columns is not None if partition_id < 0 or partition_id >= len(self._partition_columns): raise IndexError( - f"partition_id: {partition_id} out of range <0, {len(self.num_partitions) - 1}>." + f"partition_id: {partition_id} out of range <0, {self.num_partitions - 1}>." ) columns = self._partition_columns[partition_id] return self.dataset.select_columns(columns) From 4e4d8c176d8d1cee2bc2db586e97b6d72cc69fe6 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Tue, 17 Dec 2024 13:34:01 +0100 Subject: [PATCH 11/15] Update size check validation --- .../partitioner/vertical_size_partitioner.py | 25 +++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index c2af7def660a..ffd84feb4302 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -45,9 +45,12 @@ class VerticalSizePartitioner(Partitioner): 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). + A list of floats -> each value represent a fraction total number of columns. + Note that these values apply to collums without `active_party_columns`, `shared_columns`. + They are additionally included in to the partition(s). `drop_columns` are also not counted + toward the partition sizes. + In case fo list[int]: sum(partition_sizes) == len(columns) - len(drop_columns) - + len(shared_columns) - len(active_party_columns) active_party_column : Optional[Union[str, list[str]]] Column(s) (typically representing labels) associated with the "active party" (which can be the server). @@ -258,13 +261,21 @@ def _validate_parameters_while_partitioning( f"Active party column '{column}' not found in the dataset." ) num_columns = len(all_columns) + num_cols_unused_in_core_div = 0 + if self._active_party_columns is not None: + num_cols_unused_in_core_div += len(self._active_party_columns) + if self._shared_columns is not None: + num_cols_unused_in_core_div += len(self._shared_columns) + if self._drop_columns is not None: + num_cols_unused_in_core_div += len(self._drop_columns) + num_core_div_columns = num_columns - num_cols_unused_in_core_div if all(isinstance(size, int) for size in self._partition_sizes): - if sum(self._partition_sizes) != num_columns: + if sum(self._partition_sizes) != num_core_div_columns: raise ValueError( - "Sum of partition sizes cannot differ from the total number of columns." + "Sum of partition sizes cannot differ from the total number of columns" + "used in the division. Note that shared_columns, drop_columns and" + "active_party_columns are not included in the division." ) - else: - pass def _init_active_party_column( self, active_party_column: Optional[Union[str, list[str]]] From fc2aa812750741dd931a151dcdb0be05c14ed698 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Tue, 17 Dec 2024 13:34:45 +0100 Subject: [PATCH 12/15] Extend tests --- .../vertical_size_partitioner_test.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner_test.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner_test.py index bc6b8324ac52..d2c483c2be88 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner_test.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner_test.py @@ -181,6 +181,26 @@ def test_sum_of_int_partition_sizes_exceeds_num_columns(self) -> None: with self.assertRaises(ValueError): partitioner.load_partition(0) + def test_sum_of_int_partition_sizes_indirectly_exceeds_num_columns(self) -> None: + """Check ValueError if sum of int sizes > total columns.""" + columns = ["f1", "f2", "f3"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner( + partition_sizes=[1, 1], drop_columns=["f3", "f2"], shuffle=False + ) + partitioner.dataset = dataset + with self.assertRaises(ValueError): + partitioner.load_partition(0) + + def test_sum_of_int_partition_sizes_is_smaller_than_num_columns(self) -> None: + """Check ValueError if sum of int sizes < total columns.""" + columns = ["f1", "f2", "f3"] + dataset = _create_dummy_dataset(columns) + partitioner = VerticalSizePartitioner(partition_sizes=[2], shuffle=False) + partitioner.dataset = dataset + with self.assertRaises(ValueError): + partitioner.load_partition(0) + if __name__ == "__main__": unittest.main() From eb4043c0a33fc1225f0c3cca5dfe94f5ddfa4934 Mon Sep 17 00:00:00 2001 From: jafermarq Date: Wed, 18 Dec 2024 16:08:39 +0000 Subject: [PATCH 13/15] fix docstrings render --- .../flwr_datasets/partitioner/vertical_size_partitioner.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index ffd84feb4302..7e00d85da7b1 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -56,13 +56,12 @@ class VerticalSizePartitioner(Partitioner): "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. + - "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. From 79a60a59f3d6d65f5664271ac3e4e2eb6e252314 Mon Sep 17 00:00:00 2001 From: jafermarq Date: Wed, 18 Dec 2024 16:11:08 +0000 Subject: [PATCH 14/15] better docstrings render --- .../partitioner/vertical_size_partitioner.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index 7e00d85da7b1..74e6f81b580c 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -57,12 +57,12 @@ class VerticalSizePartitioner(Partitioner): 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. + - `"add_to_first"`: Append active party columns to the first partition. + - `"add_to_last"`: Append active party columns to the last partition. + - `"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. - 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]] From 8d5e6690f5dc52f7ab137e91fd04112fdc176173 Mon Sep 17 00:00:00 2001 From: jafermarq Date: Wed, 18 Dec 2024 16:24:03 +0000 Subject: [PATCH 15/15] fix error message + format --- datasets/flwr_datasets/partitioner/vertical_size_partitioner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py index 74e6f81b580c..462a76a2e3f5 100644 --- a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py +++ b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py @@ -271,7 +271,7 @@ def _validate_parameters_while_partitioning( if all(isinstance(size, int) for size in self._partition_sizes): if sum(self._partition_sizes) != num_core_div_columns: raise ValueError( - "Sum of partition sizes cannot differ from the total number of columns" + "Sum of partition sizes cannot differ from the total number of columns " "used in the division. Note that shared_columns, drop_columns and" "active_party_columns are not included in the division." )