[FSTORE-1975] Add row count endpoint for feature groups

Author: wolffbe

python/hopsworks/client/__init__.py

@@ -2,23 +2,17 @@
 # Do not edit it manually!
 # It is managed by hopsworks-apigen and is subject to periodical deletion.
 import hopsworks_common.client
-
-
 _is_external = hopsworks_common.client._is_external
 get_connection = hopsworks_common.client.get_connection
 get_instance = hopsworks_common.client.get_instance
 get_knative_domain = hopsworks_common.client.get_knative_domain
-get_serving_num_instances_limits = (
-    hopsworks_common.client.get_serving_num_instances_limits
-)
+get_serving_num_instances_limits = hopsworks_common.client.get_serving_num_instances_limits
 init = hopsworks_common.client.init
 is_kserve_installed = hopsworks_common.client.is_kserve_installed
 is_saas_connection = hopsworks_common.client.is_saas_connection
 is_scale_to_zero_required = hopsworks_common.client.is_scale_to_zero_required
 set_connection = hopsworks_common.client.set_connection
 set_knative_domain = hopsworks_common.client.set_knative_domain
 set_kserve_installed = hopsworks_common.client.set_kserve_installed
-set_serving_num_instances_limits = (
-    hopsworks_common.client.set_serving_num_instances_limits
-)
+set_serving_num_instances_limits = hopsworks_common.client.set_serving_num_instances_limits
 stop = hopsworks_common.client.stop

python/hopsworks/core/__init__.py

@@ -0,0 +1,8 @@
+# This file is generated by hopsworks-apigen.
+# Do not edit it manually!
+# It is managed by hopsworks-apigen and is subject to periodical deletion.
+import hopsworks_common.core.sink_job_configuration
+FeatureColumnMapping = hopsworks_common.core.sink_job_configuration.FeatureColumnMapping
+FullLoadConfig = hopsworks_common.core.sink_job_configuration.FullLoadConfig
+LoadingConfig = hopsworks_common.core.sink_job_configuration.LoadingConfig
+SinkJobConfiguration = hopsworks_common.core.sink_job_configuration.SinkJobConfiguration

python/hopsworks/mcp/tools/feature_group.py

@@ -54,6 +54,9 @@ def __init__(self, mcp: FastMCP):
         self.mcp.tool(tags=[TAGS.FEATURE_GROUP, TAGS.READ, TAGS.STATEFUL])(
             self.get_features
         )
+        self.mcp.tool(tags=[TAGS.FEATURE_GROUP, TAGS.READ, TAGS.STATEFUL])(
+            self.get_feature_group_row_count
+        )
 
     def _get_feature_group_versions(self, name: str | None = None):
         # Get the current project and its feature groups
@@ -180,6 +183,20 @@ async def preview_feature_group(  # docsig: disable
             f"Unable to convert preview to dictionary. Here's the raw preview:\n{preview}"
         )
 
+    async def get_feature_group_row_count(  # docsig: disable
+        self,
+        ctx: Context,
+        name: str,
+        version: int | None = None,
+    ) -> int:
+        """Get the row count of a feature group with the specified name and version (latest by default)."""
+        await ctx.info(
+            f"Retrieving row count of {name}{f' v{version}' if version else ''} feature group..."
+        )
+
+        fg = self._get_feature_group_version(name, version)
+        return fg.get_row_count()
+
     async def get_features(  # docsig: disable
         self,
         ctx: Context,

python/hsfs/core/feature_group_api.py

@@ -635,6 +635,33 @@ def get_generated_feature_groups(
             explicit_provenance.Links.Type.FEATURE_GROUP,
         )
 
+    def get_row_count(
+        self,
+        feature_group_instance: fg_mod.FeatureGroup
+        | fg_mod.ExternalFeatureGroup
+        | fg_mod.SpineGroup,
+    ) -> int:
+        """Get the row count of a feature group.
+
+        Parameters:
+            feature_group_instance:
+                metadata object of feature group.
+
+        Returns:
+            The number of rows in the feature group.
+        """
+        _client = client.get_instance()
+        path_params = [
+            "project",
+            _client._project_id,
+            "featurestores",
+            feature_group_instance.feature_store_id,
+            "featuregroups",
+            feature_group_instance.id,
+            "rowcount",
+        ]
+        return _client._send_request("GET", path_params)
+
     def _check_features(self, feature_group_instance) -> None:
         if not feature_group_instance._features:
             warnings.warn(

python/hsfs/core/feature_group_base_engine.py

@@ -41,6 +41,17 @@ def __init__(self, feature_store_id):
     def delete(self, feature_group):
         self._feature_group_api.delete(feature_group)
 
+    def get_row_count(self, feature_group: FeatureGroup) -> int:
+        """Get the row count of a feature group.
+
+        Parameters:
+            feature_group: The feature group to get the row count of.
+
+        Returns:
+            The row count of the feature group.
+        """
+        return self._feature_group_api.get_row_count(feature_group)
+
     def add_tag(self, feature_group: FeatureGroup, name: str, value: Any):
         """Attach a name/value tag to a feature group.
 

python/hsfs/feature_group.py

@@ -865,6 +865,28 @@ def get_generated_feature_groups(self) -> explicit_provenance.Links | None:
         """
         return self._feature_group_engine.get_generated_feature_groups(self)
 
+    def get_row_count(self) -> int:
+        """Get the number of rows in this feature group.
+
+        Example:
+            ```python
+            # connect to the Feature Store
+            fs = ...
+
+            # get the Feature Group instance
+            fg = fs.get_or_create_feature_group(...)
+
+            row_count = fg.get_row_count()
+            ```
+
+        Returns:
+            The number of rows in the feature group.
+
+        Raises:
+            hopsworks.client.exceptions.RestAPIError: If the backend encounters an error when handling the request.
+        """
+        return self._feature_group_engine.get_row_count(self)
+
     @public
     def get_feature(self, name: str) -> feature.Feature | None:
         """Retrieve a `Feature` object from the schema of the feature group.

python/tests/core/test_feature_group_api.py

@@ -64,6 +64,30 @@ def test_check_features(self, mocker, backend_fixtures):
         # Assert
         assert len(warning_record) == 0
 
+    def test_get_row_count(self, mocker, backend_fixtures):
+        # Arrange
+        fg_api = feature_group_api.FeatureGroupApi()
+
+        client_mock = Mock()
+        client_mock.configure_mock(
+            **{"_send_request.return_value": 42, "_project_id": 1}
+        )
+        mocker.patch(
+            "hopsworks_common.client.get_instance",
+            return_value=client_mock,
+        )
+        mocker.patch("hsfs.engine.get_instance")
+
+        json = backend_fixtures["feature_group"]["get"]["response"]
+        fg = fg_mod.FeatureGroup.from_response_json(json)
+
+        # Act
+        result = fg_api.get_row_count(fg)
+
+        # Assert
+        assert result == 42
+        assert client_mock._send_request.call_count == 1
+
     def test_check_features_no_features(self, mocker, backend_fixtures):
         # Arrange
         fg_api = feature_group_api.FeatureGroupApi()

python/tests/core/test_feature_group_base_engine.py

@@ -99,6 +99,22 @@ def test_get_tags(self, mocker):
         # Assert
         assert mock_tags_api.return_value.get.call_count == 1
 
+    def test_get_row_count(self, mocker):
+        # Arrange
+        feature_store_id = 99
+
+        mock_fg_api = mocker.patch("hsfs.core.feature_group_api.FeatureGroupApi")
+
+        fg_base_engine = feature_group_base_engine.FeatureGroupBaseEngine(
+            feature_store_id=feature_store_id
+        )
+
+        # Act
+        fg_base_engine.get_row_count(feature_group=None)
+
+        # Assert
+        assert mock_fg_api.return_value.get_row_count.call_count == 1
+
     def test_update_statistics_config(self, mocker):
         # Arrange
         feature_store_id = 99

python/tests/test_feature_group.py

@@ -385,6 +385,18 @@ def test_select_features(self):
         assert len(features) == 2
         assert {f.name for f in features} == {"f1", "f2"}
 
+    def test_get_row_count(self, mocker):
+        mock_engine = mocker.patch(
+            "hsfs.core.feature_group_base_engine.FeatureGroupBaseEngine.get_row_count",
+            return_value=123,
+        )
+
+        fg = get_test_feature_group()
+        result = fg.get_row_count()
+
+        assert result == 123
+        mock_engine.assert_called_once_with(fg)
+
     def test_materialization_job(self, mocker):
         mock_job = mocker.Mock()
         mock_job_api = mocker.patch(