From a06aed49d8e807c7c6da9399aaf03fca98f506a9 Mon Sep 17 00:00:00 2001 From: Chirag Pandey Date: Tue, 14 Jul 2026 23:39:12 +0000 Subject: [PATCH] feat: Wire BatchWriteRecord and ListRecords into ingest_dataframe - Add use_batch_write_record=False flag to ingest_dataframe() (Proposal C) - Implement _ingest_batch_write() with 25 records per API call - Map partial failures (response.errors) back to specific row indices - Add list_records() utility function with pagination support - Export list_records from feature_store __init__.py Tested: 50 integration tests + 17 unit tests --- .../sagemaker/mlops/feature_store/__init__.py | 2 + .../mlops/feature_store/feature_utils.py | 45 +++ .../feature_store/ingestion_manager_pandas.py | 209 +++++++++- .../test_batch_write_and_list_records.py | 362 ++++++++++++++++++ 4 files changed, 600 insertions(+), 18 deletions(-) create mode 100644 sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_batch_write_and_list_records.py diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py index 183cba365d..34ff18e2bb 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py @@ -54,6 +54,7 @@ create_athena_query, get_session_from_role, ingest_dataframe, + list_records, load_feature_definitions_from_dataframe, ) @@ -116,6 +117,7 @@ "create_athena_query", "get_session_from_role", "ingest_dataframe", + "list_records", "load_feature_definitions_from_dataframe", # Classes "AthenaQuery", diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py index 6bec89ad19..0f80a509d8 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py @@ -478,6 +478,7 @@ def ingest_dataframe( max_processes: int = 1, wait: bool = True, timeout: Union[int, float] = None, + use_batch_write_record: bool = False, ): """Ingest a pandas DataFrame to a FeatureGroup. @@ -488,6 +489,10 @@ def ingest_dataframe( max_processes: Number of processes (default: 1). wait: Wait for ingestion to complete (default: True). timeout: Timeout in seconds (default: None). + use_batch_write_record: If True, use BatchWriteRecord API (25 records per + call) instead of PutRecord (1 record per call) for significantly better + throughput. Requires both ``sagemaker:BatchWriteRecord`` AND + ``sagemaker:PutRecord`` IAM permissions. Default: False. Returns: IngestionManagerPandas instance. @@ -518,10 +523,50 @@ def ingest_dataframe( feature_definitions=feature_definitions, max_workers=max_workers, max_processes=max_processes, + use_batch_write_record=use_batch_write_record, ) manager.run(data_frame=data_frame, wait=wait, timeout=timeout) return manager + +def list_records( + feature_group_name: str, + max_results: int = None, + next_token: str = None, + include_soft_deleted_records: bool = False, + region: str = None, +): + """List record identifiers from a FeatureGroup's OnlineStore. + + Returns a single page of results. Use ``next_token`` from the response + to fetch subsequent pages. + + Args: + feature_group_name: Name of the FeatureGroup. + max_results: Maximum number of record identifiers per page (1-100). + next_token: Pagination token from a previous response. + include_soft_deleted_records: If True, include soft-deleted records. + region: AWS region name. + + Returns: + ListRecordsResponse with ``record_identifiers`` (List[str]) and + ``next_token`` (str or None). + """ + fg = CoreFeatureGroup.get(feature_group_name=feature_group_name, region=region) + + kwargs = {} + if max_results is not None: + kwargs["max_results"] = max_results + if next_token is not None: + kwargs["next_token"] = next_token + if include_soft_deleted_records: + kwargs["include_soft_deleted_records"] = include_soft_deleted_records + if region is not None: + kwargs["region"] = region + + return fg.list_records(**kwargs) + + @_telemetry_emitter(Feature.FEATURE_STORE, "get_feature_group_as_dataframe") def get_feature_group_as_dataframe( feature_group_name: str, diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py index b2b7ae5085..35331246fe 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py @@ -14,12 +14,15 @@ from pandas.api.types import is_list_like from sagemaker.core.resources import FeatureGroup as CoreFeatureGroup -from sagemaker.core.shapes import FeatureValue +from sagemaker.core.shapes import BatchWriteRecordEntry, FeatureValue from sagemaker.core.utils.utils import Unassigned from sagemaker.core.telemetry import Feature, _telemetry_emitter logger = logging.getLogger(__name__) +# Maximum number of records per BatchWriteRecord API call +BATCH_WRITE_MAX_ENTRIES = 25 + class IngestionError(Exception): """Exception raised for errors during ingestion. @@ -55,6 +58,7 @@ class IngestionManagerPandas: feature_definitions: Dict[str, Dict[Any, Any]] max_workers: int = 1 max_processes: int = 1 + use_batch_write_record: bool = False _async_result: Any = field(default=None, init=False) _processing_pool: Pool = field(default=None, init=False) _failed_indices: List[int] = field(default_factory=list, init=False) @@ -133,19 +137,33 @@ def _run_single_process_single_thread( ): """Ingest utilizing a single process and a single thread.""" logger.info("Started single-threaded ingestion for %d rows", len(data_frame)) - failed_rows = [] - - fg = CoreFeatureGroup(feature_group_name=self.feature_group_name) - for row in data_frame.itertuples(): - self._ingest_row( + if self.use_batch_write_record: + logger.info( + "Using BatchWriteRecord API (batch size=%d). " + "Requires sagemaker:BatchWriteRecord AND sagemaker:PutRecord IAM permissions.", + BATCH_WRITE_MAX_ENTRIES, + ) + failed_rows = IngestionManagerPandas._ingest_batch_write( data_frame=data_frame, - row=row, - feature_group=fg, + feature_group_name=self.feature_group_name, feature_definitions=self.feature_definitions, - failed_rows=failed_rows, + start_index=0, + end_index=len(data_frame), target_stores=target_stores, ) + else: + failed_rows = [] + fg = CoreFeatureGroup(feature_group_name=self.feature_group_name) + for row in data_frame.itertuples(): + self._ingest_row( + data_frame=data_frame, + row=row, + feature_group=fg, + feature_definitions=self.feature_definitions, + failed_rows=failed_rows, + target_stores=target_stores, + ) self._failed_indices = failed_rows if self._failed_indices: @@ -176,6 +194,7 @@ def _run_multi_process( target_stores, start_index, timeout, + self.use_batch_write_record, )) def init_worker(): @@ -200,6 +219,7 @@ def _run_multi_threaded( target_stores: List[str] = None, row_offset: int = 0, timeout: Union[int, float] = None, + use_batch_write_record: bool = False, ) -> List[int]: """Start multi-threaded ingestion within a single process.""" executor = ThreadPoolExecutor(max_workers=max_workers) @@ -209,15 +229,26 @@ def _run_multi_threaded( for i in range(max_workers): start_index = min(i * batch_size, data_frame.shape[0]) end_index = min(i * batch_size + batch_size, data_frame.shape[0]) - future = executor.submit( - IngestionManagerPandas._ingest_single_batch, - data_frame=data_frame, - feature_group_name=feature_group_name, - feature_definitions=feature_definitions, - start_index=start_index, - end_index=end_index, - target_stores=target_stores, - ) + if use_batch_write_record: + future = executor.submit( + IngestionManagerPandas._ingest_batch_write, + data_frame=data_frame, + feature_group_name=feature_group_name, + feature_definitions=feature_definitions, + start_index=start_index, + end_index=end_index, + target_stores=target_stores, + ) + else: + future = executor.submit( + IngestionManagerPandas._ingest_single_batch, + data_frame=data_frame, + feature_group_name=feature_group_name, + feature_definitions=feature_definitions, + start_index=start_index, + end_index=end_index, + target_stores=target_stores, + ) futures[future] = (start_index + row_offset, end_index + row_offset) failed_indices = [] @@ -332,3 +363,145 @@ def _convert_to_string_list(feature_value: List[Any]) -> List[str]: f"must be an Array, but was {type(feature_value)}" ) return [str(v) if v is not None else None for v in feature_value] + + @staticmethod + def _build_record( + data_frame: DataFrame, + row: Iterable, + feature_definitions: Dict[str, Dict[Any, Any]], + ) -> List[FeatureValue]: + """Build a list of FeatureValue from a DataFrame row. + + Args: + data_frame: Source DataFrame (for column names). + row: A single row from itertuples(). + feature_definitions: Feature definition metadata. + + Returns: + List of FeatureValue objects for the row. + """ + record = [] + for index in range(1, len(row)): + feature_name = data_frame.columns[index - 1] + feature_value = row[index] + + if not IngestionManagerPandas._feature_value_is_not_none(feature_value): + continue + + if IngestionManagerPandas._is_feature_collection_type(feature_name, feature_definitions): + record.append(FeatureValue( + feature_name=feature_name, + value_as_string_list=IngestionManagerPandas._convert_to_string_list(feature_value), + )) + else: + record.append(FeatureValue( + feature_name=feature_name, + value_as_string=str(feature_value), + )) + return record + + @staticmethod + def _ingest_batch_write( + data_frame: DataFrame, + feature_group_name: str, + feature_definitions: Dict[str, Dict[Any, Any]], + start_index: int, + end_index: int, + target_stores: List[str] = None, + ) -> List[int]: + """Ingest records using BatchWriteRecord API (up to 25 per call). + + Args: + data_frame: Source DataFrame. + feature_group_name: Name of the FeatureGroup. + feature_definitions: Feature definition metadata. + start_index: Start index in the DataFrame slice. + end_index: End index in the DataFrame slice. + target_stores: Target stores for ingestion. + + Returns: + List of row indices that failed to ingest. + """ + logger.info( + "Started batch write ingestion index %d to %d (batch_size=%d)", + start_index, end_index, BATCH_WRITE_MAX_ENTRIES, + ) + failed_rows = [] + rows = list(data_frame[start_index:end_index].itertuples()) + + for batch_start in range(0, len(rows), BATCH_WRITE_MAX_ENTRIES): + batch = rows[batch_start:batch_start + BATCH_WRITE_MAX_ENTRIES] + entries = [] + row_indices = [] + + for row in batch: + try: + record = IngestionManagerPandas._build_record( + data_frame=data_frame, + row=row, + feature_definitions=feature_definitions, + ) + entry_kwargs = { + "feature_group_name": feature_group_name, + "record": record, + } + if target_stores is not None: + entry_kwargs["target_stores"] = target_stores + entry = BatchWriteRecordEntry(**entry_kwargs) + entries.append(entry) + row_indices.append(row[0]) + except Exception as e: + logger.error("Failed to build record for row %d: %s", row[0], e) + failed_rows.append(row[0]) + + if not entries: + continue + + try: + fg = CoreFeatureGroup(feature_group_name=feature_group_name) + response = fg.batch_write_record(entries=entries) + + # Handle partial failures from unprocessed entries + if response.unprocessed_entries: + for i, entry in enumerate(response.unprocessed_entries): + # unprocessed_entries are BatchWriteRecordEntry objects; + # find their position in the original entries list + try: + idx = entries.index(entry) + failed_rows.append(row_indices[idx]) + except ValueError: + # If we can't find the exact entry, mark by position + failed_rows.append(row_indices[i] if i < len(row_indices) else row_indices[-1]) + + # Handle errors (BatchWriteRecordError with entry object) + if response.errors: + for error in response.errors: + # BatchWriteRecordError has .entry (the failed entry), .error_code, .error_message + error_entry = getattr(error, "entry", None) + if error_entry is not None: + try: + idx = entries.index(error_entry) + failed_rows.append(row_indices[idx]) + except ValueError: + # Can't match entry back — mark all rows in batch as failed + logger.warning( + "BatchWriteRecord error could not be mapped to row: %s - %s", + getattr(error, "error_code", ""), + getattr(error, "error_message", ""), + ) + failed_rows.extend(row_indices) + break + else: + # Fallback: no entry object, mark all rows as failed + logger.warning("BatchWriteRecord error without entry: %s", error) + failed_rows.extend(row_indices) + break + + except Exception as e: + logger.error( + "BatchWriteRecord call failed for batch starting at row %d: %s", + row_indices[0] if row_indices else start_index, e, + ) + failed_rows.extend(row_indices) + + return failed_rows diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_batch_write_and_list_records.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_batch_write_and_list_records.py new file mode 100644 index 0000000000..f2a525852f --- /dev/null +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_batch_write_and_list_records.py @@ -0,0 +1,362 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 +"""Unit tests for BatchWriteRecord and ListRecords wiring.""" +import pytest +from unittest.mock import Mock, patch, MagicMock +import pandas as pd +import numpy as np + +from sagemaker.mlops.feature_store.ingestion_manager_pandas import ( + IngestionManagerPandas, + IngestionError, + BATCH_WRITE_MAX_ENTRIES, +) + + +class TestBatchWriteRecordIngestion: + """Tests for use_batch_write_record=True path.""" + + @pytest.fixture + def feature_definitions(self): + return { + "RecordIdentifier": {"FeatureType": "String", "CollectionType": None}, + "EventTime": {"FeatureType": "String", "CollectionType": None}, + "Feature1": {"FeatureType": "String", "CollectionType": None}, + } + + @pytest.fixture + def sample_dataframe(self): + return pd.DataFrame({ + "RecordIdentifier": [f"id-{i}" for i in range(5)], + "EventTime": ["2026-01-01T00:00:00Z"] * 5, + "Feature1": [f"value-{i}" for i in range(5)], + }) + + @pytest.fixture + def large_dataframe(self): + """DataFrame with 60 rows — should produce 3 BatchWriteRecord calls (25+25+10).""" + return pd.DataFrame({ + "RecordIdentifier": [f"id-{i}" for i in range(60)], + "EventTime": ["2026-01-01T00:00:00Z"] * 60, + "Feature1": [f"value-{i}" for i in range(60)], + }) + + def test_batch_write_max_entries_constant(self): + """Verify batch size constant is 25.""" + assert BATCH_WRITE_MAX_ENTRIES == 25 + + def test_initialization_with_flag(self, feature_definitions): + """Verify use_batch_write_record flag is stored.""" + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=True, + ) + assert mgr.use_batch_write_record is True + + def test_initialization_default_flag(self, feature_definitions): + """Verify use_batch_write_record defaults to False.""" + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + ) + assert mgr.use_batch_write_record is False + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_single_batch(self, mock_fg_class, feature_definitions, sample_dataframe): + """5 rows should result in 1 batch_write_record call.""" + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + mock_response = Mock() + mock_response.unprocessed_entries = [] + mock_response.errors = [] + mock_fg.batch_write_record.return_value = mock_response + + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=True, + ) + mgr.run(data_frame=sample_dataframe) + + mock_fg.batch_write_record.assert_called_once() + assert mgr.failed_rows == [] + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_multiple_batches(self, mock_fg_class, feature_definitions, large_dataframe): + """60 rows should result in 3 batch_write_record calls (25+25+10).""" + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + mock_response = Mock() + mock_response.unprocessed_entries = [] + mock_response.errors = [] + mock_fg.batch_write_record.return_value = mock_response + + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=True, + ) + mgr.run(data_frame=large_dataframe) + + assert mock_fg.batch_write_record.call_count == 3 + assert mgr.failed_rows == [] + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_partial_failure_maps_to_row(self, mock_fg_class, feature_definitions): + """Errors with matching entry should map back to specific row indices.""" + from sagemaker.core.shapes import BatchWriteRecordEntry, FeatureValue + + df = pd.DataFrame({ + "RecordIdentifier": ["good-1", None, "good-3"], + "EventTime": ["2026-01-01T00:00:00Z"] * 3, + "Feature1": ["v1", "v2", "v3"], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + # Build what the entry for row 1 (None RecordIdentifier) would look like + error_entry = BatchWriteRecordEntry( + feature_group_name="test-fg", + record=[ + FeatureValue(feature_name="EventTime", value_as_string="2026-01-01T00:00:00Z"), + FeatureValue(feature_name="Feature1", value_as_string="v2"), + ], + ) + mock_error = Mock() + mock_error.entry = error_entry + mock_error.error_code = "ValidationError" + mock_error.error_message = "Missing RecordIdentifier" + + mock_response = Mock() + mock_response.unprocessed_entries = [] + mock_response.errors = [mock_error] + mock_fg.batch_write_record.return_value = mock_response + + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=True, + ) + + with pytest.raises(IngestionError) as exc_info: + mgr.run(data_frame=df) + + assert 1 in exc_info.value.failed_rows + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_full_failure(self, mock_fg_class, feature_definitions, sample_dataframe): + """If batch_write_record raises an exception, all rows in that batch fail.""" + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + mock_fg.batch_write_record.side_effect = Exception("AccessForbidden") + + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=True, + ) + + with pytest.raises(IngestionError) as exc_info: + mgr.run(data_frame=sample_dataframe) + + assert len(exc_info.value.failed_rows) == 5 + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_skips_null_values(self, mock_fg_class, feature_definitions): + """Null/NaN values should not be included in the record.""" + df = pd.DataFrame({ + "RecordIdentifier": ["id-1"], + "EventTime": ["2026-01-01T00:00:00Z"], + "Feature1": [None], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + mock_response = Mock() + mock_response.unprocessed_entries = [] + mock_response.errors = [] + mock_fg.batch_write_record.return_value = mock_response + + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=True, + ) + mgr.run(data_frame=df) + + call_kwargs = mock_fg.batch_write_record.call_args + entries = call_kwargs.kwargs.get("entries") or call_kwargs[1].get("entries") + record = entries[0].record + feature_names = [fv.feature_name for fv in record] + assert "Feature1" not in feature_names + assert "RecordIdentifier" in feature_names + + def test_use_batch_write_record_false_uses_put_record(self, feature_definitions): + """Verify use_batch_write_record=False still uses PutRecord path.""" + df = pd.DataFrame({"RecordIdentifier": ["id-1"], "EventTime": ["2026-01-01T00:00:00Z"], "Feature1": ["v"]}) + + with patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") as mock_fg_class: + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=False, + ) + mgr.run(data_frame=df) + + mock_fg.put_record.assert_called_once() + mock_fg.batch_write_record.assert_not_called() + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_does_not_pass_none_target_stores(self, mock_fg_class, feature_definitions, sample_dataframe): + """When target_stores=None, entries should NOT have target_stores set (uses Unassigned).""" + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + mock_response = Mock() + mock_response.unprocessed_entries = [] + mock_response.errors = [] + mock_fg.batch_write_record.return_value = mock_response + + mgr = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=True, + ) + mgr.run(data_frame=sample_dataframe, target_stores=None) + + call_kwargs = mock_fg.batch_write_record.call_args + entries = call_kwargs.kwargs.get("entries") or call_kwargs[1].get("entries") + # target_stores should be Unassigned (not None) since we don't pass it + from sagemaker.core.utils.utils import Unassigned + assert isinstance(entries[0].target_stores, Unassigned) + + +class TestBuildRecord: + """Tests for _build_record helper.""" + + def test_build_record_basic(self): + df = pd.DataFrame({ + "RecordIdentifier": ["id-1"], + "EventTime": ["2026-01-01T00:00:00Z"], + "Feature1": ["hello"], + }) + feature_defs = { + "RecordIdentifier": {"FeatureType": "String", "CollectionType": None}, + "EventTime": {"FeatureType": "String", "CollectionType": None}, + "Feature1": {"FeatureType": "String", "CollectionType": None}, + } + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, feature_defs) + assert len(record) == 3 + names = {fv.feature_name for fv in record} + assert names == {"RecordIdentifier", "EventTime", "Feature1"} + + def test_build_record_skips_none(self): + df = pd.DataFrame({ + "RecordIdentifier": ["id-1"], + "EventTime": ["2026-01-01T00:00:00Z"], + "Feature1": [None], + }) + feature_defs = { + "RecordIdentifier": {"FeatureType": "String", "CollectionType": None}, + "EventTime": {"FeatureType": "String", "CollectionType": None}, + "Feature1": {"FeatureType": "String", "CollectionType": None}, + } + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, feature_defs) + assert len(record) == 2 + names = {fv.feature_name for fv in record} + assert "Feature1" not in names + + def test_build_record_skips_nan(self): + df = pd.DataFrame({ + "RecordIdentifier": ["id-1"], + "EventTime": ["2026-01-01T00:00:00Z"], + "Feature1": [np.nan], + }) + feature_defs = { + "RecordIdentifier": {"FeatureType": "String", "CollectionType": None}, + "EventTime": {"FeatureType": "String", "CollectionType": None}, + "Feature1": {"FeatureType": "String", "CollectionType": None}, + } + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, feature_defs) + assert len(record) == 2 + + def test_build_record_collection_type(self): + df = pd.DataFrame({ + "id": ["id-1"], + "tags": [["a", "b", "c"]], + }) + feature_defs = { + "id": {"FeatureType": "String", "CollectionType": None}, + "tags": {"FeatureType": "String", "CollectionType": "List"}, + } + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, feature_defs) + assert len(record) == 2 + tags_fv = [fv for fv in record if fv.feature_name == "tags"][0] + assert tags_fv.value_as_string_list == ["a", "b", "c"] + + +class TestListRecords: + """Tests for list_records function.""" + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_list_records_single_page(self, mock_fg_class): + """list_records returns a single page response.""" + from sagemaker.mlops.feature_store.feature_utils import list_records + + mock_fg = Mock() + mock_fg_class.get.return_value = mock_fg + mock_response = Mock() + mock_response.record_identifiers = ["id-1", "id-2", "id-3"] + mock_response.next_token = None + mock_fg.list_records.return_value = mock_response + + result = list_records("test-fg", region="us-west-2") + + assert result.record_identifiers == ["id-1", "id-2", "id-3"] + assert result.next_token is None + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_list_records_with_params(self, mock_fg_class): + """list_records passes max_results, next_token, include_soft_deleted_records.""" + from sagemaker.mlops.feature_store.feature_utils import list_records + + mock_fg = Mock() + mock_fg_class.get.return_value = mock_fg + mock_response = Mock() + mock_response.record_identifiers = ["id-1"] + mock_response.next_token = "token-abc" + mock_fg.list_records.return_value = mock_response + + list_records( + "test-fg", max_results=1, next_token="prev-token", + include_soft_deleted_records=True, region="us-west-2" + ) + + mock_fg.list_records.assert_called_once_with( + max_results=1, next_token="prev-token", + include_soft_deleted_records=True, region="us-west-2" + ) + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_list_records_no_optional_params(self, mock_fg_class): + """list_records with no optional params only passes region.""" + from sagemaker.mlops.feature_store.feature_utils import list_records + + mock_fg = Mock() + mock_fg_class.get.return_value = mock_fg + mock_response = Mock() + mock_response.record_identifiers = [] + mock_response.next_token = None + mock_fg.list_records.return_value = mock_response + + list_records("test-fg", region="us-east-1") + + mock_fg.list_records.assert_called_once_with(region="us-east-1")