Fds add num_partitions property to partitioners (#3095)

* Add num_partition property

* Trigger the partitioning in the num_partitions

---------

Co-authored-by: Daniel J. Beutel <[email protected]>

datasets/flwr_datasets/partitioner/dirichlet_partitioner.py

@@ -132,6 +132,13 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
         self._determine_node_id_to_indices_if_needed()
         return self.dataset.select(self._node_id_to_indices[node_id])
 
+    @property
+    def num_partitions(self) -> int:
+        """Total number of partitions."""
+        self._check_num_partitions_correctness_if_needed()
+        self._determine_node_id_to_indices_if_needed()
+        return self._num_partitions
+
     def _initialize_alpha(
         self, alpha: Union[int, float, List[float], NDArrayFloat]
     ) -> NDArrayFloat:

datasets/flwr_datasets/partitioner/iid_partitioner.py

@@ -50,3 +50,8 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
         return self.dataset.shard(
             num_shards=self._num_partitions, index=node_id, contiguous=True
         )
+
+    @property
+    def num_partitions(self) -> int:
+        """Total number of partitions."""
+        return self._num_partitions

datasets/flwr_datasets/partitioner/inner_dirichlet_partitioner.py

@@ -119,6 +119,17 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
         self._determine_node_id_to_indices_if_needed()
         return self.dataset.select(self._node_id_to_indices[node_id])
 
+    @property
+    def num_partitions(self) -> int:
+        """Total number of partitions."""
+        self._check_num_partitions_correctness_if_needed()
+        self._check_partition_sizes_correctness_if_needed()
+        self._check_the_sum_of_partition_sizes()
+        self._determine_num_unique_classes_if_needed()
+        self._alpha = self._initialize_alpha_if_needed(self._initial_alpha)
+        self._determine_node_id_to_indices_if_needed()
+        return self._num_partitions
+
     def _initialize_alpha_if_needed(
         self, alpha: Union[int, float, List[float], NDArrayFloat]
     ) -> NDArrayFloat:

datasets/flwr_datasets/partitioner/natural_id_partitioner.py

@@ -65,6 +65,13 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
             lambda row: row[self._partition_by] == self._node_id_to_natural_id[node_id]
         )
 
+    @property
+    def num_partitions(self) -> int:
+        """Total number of partitions."""
+        if len(self._node_id_to_natural_id) == 0:
+            self._create_int_node_id_to_natural_id()
+        return len(self._node_id_to_natural_id)
+
     @property
     def node_id_to_natural_id(self) -> Dict[int, str]:
         """Node id to corresponding natural id present.

datasets/flwr_datasets/partitioner/partitioner.py

@@ -79,3 +79,8 @@ def is_dataset_assigned(self) -> bool:
             True if a dataset is assigned, otherwise False.
         """
         return self._dataset is not None
+
+    @property
+    @abstractmethod
+    def num_partitions(self) -> int:
+        """Total number of partitions."""

datasets/flwr_datasets/partitioner/shard_partitioner.py

@@ -179,6 +179,15 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
         self._determine_node_id_to_indices_if_needed()
         return self.dataset.select(self._node_id_to_indices[node_id])
 
+    @property
+    def num_partitions(self) -> int:
+        """Total number of partitions."""
+        self._check_num_partitions_correctness_if_needed()
+        self._check_possibility_of_partitions_creation()
+        self._sort_dataset_if_needed()
+        self._determine_node_id_to_indices_if_needed()
+        return self._num_partitions
+
     def _determine_node_id_to_indices_if_needed(self) -> None:  # pylint: disable=R0914
         """Assign sample indices to each node id.
 

datasets/flwr_datasets/partitioner/size_partitioner.py

@@ -84,6 +84,12 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
         self._determine_node_id_to_indices_if_needed()
         return self.dataset.select(self._node_id_to_indices[node_id])
 
+    @property
+    def num_partitions(self) -> int:
+        """Total number of partitions."""
+        self._determine_node_id_to_indices_if_needed()
+        return self._num_partitions
+
     @property
     def node_id_to_size(self) -> Dict[int, int]:
         """Node id to the number of samples."""