fix: Fix message to Json parse issue in Dataset class

PiperOrigin-RevId: 599675216

Author: vertex-sdk-bot

google/cloud/aiplatform/datasets/dataset.py

@@ -805,7 +805,7 @@ def export_data_for_custom_training(
                 saved_query_id,
                 annotation_schema_uri,
                 split,
-            )
+            )._pb
         )
 
     def update(

tests/system/aiplatform/test_dataset.py

@@ -405,9 +405,14 @@ def test_export_data_for_custom_training(self, staging_bucket):
         )
 
         # Ensure three output paths (training, validation and test) are provided
-        assert len(export_data_response["exported_files"]) == 3
+        assert len(export_data_response["exportedFiles"]) == 3
         # Ensure data stats are calculated and correct
-        assert export_data_response["data_stats"]["training_data_items_count"] == 40
+        assert int(export_data_response["dataStats"]["trainingDataItemsCount"]) == 40
+        assert int(export_data_response["dataStats"]["validationDataItemsCount"]) == 5
+        assert int(export_data_response["dataStats"]["testDataItemsCount"]) == 5
+        assert int(export_data_response["dataStats"]["trainingAnnotationsCount"]) == 40
+        assert int(export_data_response["dataStats"]["validationAnnotationsCount"]) == 5
+        assert int(export_data_response["dataStats"]["testAnnotationsCount"]) == 5
 
     def test_update_dataset(self):
         """Create a new dataset and use update() method to change its display_name, labels, and description.