From 5f513fb63992ea4cbd1f8501972e993c54b0447e 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 | 214 ++++- .../test_feature_store_batch_write_record.py | 227 +++++ .../integ/test_feature_store_list_records.py | 121 +++ .../feature_store/test_batch_write_record.py | 795 ++++++++++++++++++ .../mlops/feature_store/test_list_records.py | 78 ++ 7 files changed, 1464 insertions(+), 18 deletions(-) create mode 100644 sagemaker-mlops/tests/integ/test_feature_store_batch_write_record.py create mode 100644 sagemaker-mlops/tests/integ/test_feature_store_list_records.py create mode 100644 sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_batch_write_record.py create mode 100644 sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_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..0762193e28 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,150 @@ 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 entry in 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: + # Can't match entry back — mark all rows in batch as failed + logger.warning( + "Unprocessed entry could not be mapped to a specific row. " + "Marking all rows in this batch as failed." + ) + failed_rows.extend(row_indices) + break + + # 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/integ/test_feature_store_batch_write_record.py b/sagemaker-mlops/tests/integ/test_feature_store_batch_write_record.py new file mode 100644 index 0000000000..54cb4a284d --- /dev/null +++ b/sagemaker-mlops/tests/integ/test_feature_store_batch_write_record.py @@ -0,0 +1,227 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 +"""Integration tests for BatchWriteRecord via ingest_dataframe(use_batch_write_record=True).""" +import time + +import pytest +import pandas as pd +import numpy as np + +from sagemaker.mlops.feature_store.feature_utils import ingest_dataframe +from sagemaker.mlops.feature_store.ingestion_manager_pandas import IngestionError +from sagemaker.core.resources import FeatureGroup + + +FEATURE_GROUP_NAME = "online-offline-fg-chiragvp" +REGION = "us-west-2" + + +@pytest.fixture(scope="module") +def feature_group(): + fg = FeatureGroup.get(feature_group_name=FEATURE_GROUP_NAME, region=REGION) + assert fg.feature_group_status == "Created" + return fg + + +@pytest.fixture +def timestamp(): + from datetime import datetime, timezone + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +class TestBatchWriteRecordIntegration: + """Integration tests calling the real BatchWriteRecord API.""" + + def test_batch_write_basic(self, feature_group, timestamp): + """Write 5 records via BatchWriteRecord, verify no failures.""" + df = pd.DataFrame({ + "RecordIdentifier": [f"integ-bwr-{i}-{int(time.time())}" for i in range(5)], + "EventTime": [timestamp] * 5, + "Feature1": [f"val-{i}" for i in range(5)], + }) + + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=True, + ) + assert mgr.failed_rows == [] + + def test_batch_write_25_boundary(self, feature_group, timestamp): + """Exactly 25 records — single batch boundary.""" + df = pd.DataFrame({ + "RecordIdentifier": [f"integ-b25-{i}-{int(time.time())}" for i in range(25)], + "EventTime": [timestamp] * 25, + "Feature1": [f"val-{i}" for i in range(25)], + }) + + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=True, + ) + assert mgr.failed_rows == [] + + def test_batch_write_26_two_batches(self, feature_group, timestamp): + """26 records — splits into 2 batches (25+1).""" + df = pd.DataFrame({ + "RecordIdentifier": [f"integ-b26-{i}-{int(time.time())}" for i in range(26)], + "EventTime": [timestamp] * 26, + "Feature1": [f"val-{i}" for i in range(26)], + }) + + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=True, + ) + assert mgr.failed_rows == [] + + def test_batch_write_verify_with_get_record(self, feature_group, timestamp): + """Write via BatchWriteRecord, verify with GetRecord.""" + rid = f"integ-bwr-verify-{int(time.time())}" + df = pd.DataFrame({ + "RecordIdentifier": [rid], + "EventTime": [timestamp], + "Feature1": ["verify-value"], + }) + + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=True, + ) + assert mgr.failed_rows == [] + + time.sleep(2) + record = feature_group.get_record(record_identifier_value_as_string=rid, region=REGION) + assert record is not None + val = next(fv.value_as_string for fv in record.record if fv.feature_name == "Feature1") + assert val == "verify-value" + + def test_batch_write_null_skipped(self, feature_group, timestamp): + """None values excluded from record, record still written.""" + rid = f"integ-bwr-null-{int(time.time())}" + df = pd.DataFrame({ + "RecordIdentifier": [rid], + "EventTime": [timestamp], + "Feature1": [None], + }) + + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=True, + ) + assert mgr.failed_rows == [] + + time.sleep(2) + record = feature_group.get_record(record_identifier_value_as_string=rid, region=REGION) + assert record is not None + names = [fv.feature_name for fv in record.record] + assert "Feature1" not in names + + def test_batch_write_partial_failure(self, feature_group, timestamp): + """10 records, row 5 missing RecordIdentifier — only row 5 fails.""" + records = [] + for i in range(10): + if i == 5: + records.append({"RecordIdentifier": None, "EventTime": timestamp, "Feature1": f"v-{i}"}) + else: + records.append({"RecordIdentifier": f"integ-partial-{i}-{int(time.time())}", "EventTime": timestamp, "Feature1": f"v-{i}"}) + df = pd.DataFrame(records) + + try: + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=True, + ) + failed = mgr.failed_rows + except IngestionError as e: + failed = e.failed_rows + + assert 5 in failed + # Valid rows should have been written + time.sleep(2) + rec = feature_group.get_record( + record_identifier_value_as_string=records[0]["RecordIdentifier"], region=REGION + ) + assert rec is not None + + def test_putrecord_same_result(self, feature_group, timestamp): + """PutRecord path produces same outcome for comparison.""" + df = pd.DataFrame({ + "RecordIdentifier": [f"integ-put-{i}-{int(time.time())}" for i in range(5)], + "EventTime": [timestamp] * 5, + "Feature1": [f"val-{i}" for i in range(5)], + }) + + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=False, + ) + assert mgr.failed_rows == [] + + def test_putrecord_partial_failure_same(self, feature_group, timestamp): + """PutRecord partial failure — same row 5 fails.""" + records = [] + for i in range(10): + if i == 5: + records.append({"RecordIdentifier": None, "EventTime": timestamp, "Feature1": f"v-{i}"}) + else: + records.append({"RecordIdentifier": f"integ-put-partial-{i}-{int(time.time())}", "EventTime": timestamp, "Feature1": f"v-{i}"}) + df = pd.DataFrame(records) + + try: + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=False, + ) + failed = mgr.failed_rows + except IngestionError as e: + failed = e.failed_rows + + assert 5 in failed + + def test_batch_write_default_flag_uses_putrecord(self, feature_group, timestamp): + """No flag → defaults to PutRecord.""" + df = pd.DataFrame({ + "RecordIdentifier": [f"integ-default-{int(time.time())}"], + "EventTime": [timestamp], + "Feature1": ["default-test"], + }) + + mgr = ingest_dataframe( + feature_group_name=FEATURE_GROUP_NAME, + data_frame=df, + max_workers=1, max_processes=1, + ) + assert mgr.use_batch_write_record is False + assert mgr.failed_rows == [] + + def test_batch_write_invalid_fg(self, timestamp): + """Invalid feature group fails for both BWR and PutRecord.""" + df = pd.DataFrame({ + "RecordIdentifier": ["x"], + "EventTime": [timestamp], + "Feature1": ["v"], + }) + + with pytest.raises(Exception): + ingest_dataframe( + feature_group_name="non-existent-fg-xyz", + data_frame=df, + max_workers=1, max_processes=1, + use_batch_write_record=True, + ) diff --git a/sagemaker-mlops/tests/integ/test_feature_store_list_records.py b/sagemaker-mlops/tests/integ/test_feature_store_list_records.py new file mode 100644 index 0000000000..3c4bd1d497 --- /dev/null +++ b/sagemaker-mlops/tests/integ/test_feature_store_list_records.py @@ -0,0 +1,121 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 +"""Integration tests for ListRecords via list_records().""" +import pytest + +from sagemaker.mlops.feature_store.feature_utils import list_records +from sagemaker.core.resources import FeatureGroup + + +FEATURE_GROUP_NAME = "online-offline-fg-chiragvp" +REGION = "us-west-2" + + +@pytest.fixture(scope="module") +def feature_group(): + fg = FeatureGroup.get(feature_group_name=FEATURE_GROUP_NAME, region=REGION) + assert fg.feature_group_status == "Created" + return fg + + +class TestListRecordsIntegration: + """Integration tests calling the real ListRecords API.""" + + def test_list_records_returns_identifiers(self, feature_group): + """Basic list_records returns record identifiers.""" + response = list_records(feature_group_name=FEATURE_GROUP_NAME, region=REGION) + + assert response is not None + assert hasattr(response, "record_identifiers") + assert hasattr(response, "next_token") + assert len(response.record_identifiers) > 0 + + def test_list_records_max_results(self, feature_group): + """max_results limits page size.""" + response = list_records( + feature_group_name=FEATURE_GROUP_NAME, + max_results=2, + region=REGION, + ) + + assert response is not None + assert len(response.record_identifiers) <= 100 # API hint, not strict + + def test_list_records_pagination(self, feature_group): + """Pagination with next_token returns different records.""" + page1 = list_records( + feature_group_name=FEATURE_GROUP_NAME, + max_results=2, + region=REGION, + ) + assert page1 is not None + assert len(page1.record_identifiers) > 0 + + if page1.next_token: + page2 = list_records( + feature_group_name=FEATURE_GROUP_NAME, + next_token=page1.next_token, + max_results=2, + region=REGION, + ) + assert page2 is not None + assert len(page2.record_identifiers) > 0 + + def test_list_records_full_pagination_loop(self, feature_group): + """Manual pagination loop (standard AWS pattern).""" + all_ids = [] + next_token = None + page_count = 0 + + while page_count < 3: # Cap at 3 pages for test speed + kwargs = {"feature_group_name": FEATURE_GROUP_NAME, "max_results": 10, "region": REGION} + if next_token: + kwargs["next_token"] = next_token + + response = list_records(**kwargs) + all_ids.extend(response.record_identifiers) + page_count += 1 + + if not response.next_token: + break + next_token = response.next_token + + assert len(all_ids) > 0 + + def test_list_records_include_soft_deleted(self, feature_group): + """include_soft_deleted_records=True doesn't error.""" + response = list_records( + feature_group_name=FEATURE_GROUP_NAME, + include_soft_deleted_records=True, + region=REGION, + ) + + assert response is not None + assert len(response.record_identifiers) > 0 + + def test_list_records_invalid_feature_group(self): + """Invalid feature group raises exception.""" + with pytest.raises(Exception): + list_records( + feature_group_name="non-existent-fg-xyz", + region=REGION, + ) + + def test_list_records_via_feature_group_object(self, feature_group): + """Call list_records directly on FeatureGroup object.""" + response = feature_group.list_records(region=REGION) + + assert response is not None + assert len(response.record_identifiers) > 0 + + def test_list_records_via_feature_group_with_pagination(self, feature_group): + """FeatureGroup.list_records() with pagination.""" + page1 = feature_group.list_records(max_results=3, region=REGION) + assert page1 is not None + + if page1.next_token: + page2 = feature_group.list_records( + next_token=page1.next_token, max_results=3, region=REGION + ) + assert page2 is not None + assert page1.record_identifiers != page2.record_identifiers diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_batch_write_record.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_batch_write_record.py new file mode 100644 index 0000000000..40f30ebfdb --- /dev/null +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_batch_write_record.py @@ -0,0 +1,795 @@ +# 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.""" + 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): + assert BATCH_WRITE_MAX_ENTRIES == 25 + + def test_initialization_with_flag(self, feature_definitions): + 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): + 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 → 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 → 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_exactly_25(self, mock_fg_class, feature_definitions): + """Exactly 25 rows → 1 call (boundary).""" + df = pd.DataFrame({ + "RecordIdentifier": [f"id-{i}" for i in range(25)], + "EventTime": ["2026-01-01T00:00:00Z"] * 25, + "Feature1": [f"v-{i}" for i in range(25)], + }) + 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) + assert mock_fg.batch_write_record.call_count == 1 + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_exactly_26(self, mock_fg_class, feature_definitions): + """Exactly 26 rows → 2 calls (25+1 boundary).""" + df = pd.DataFrame({ + "RecordIdentifier": [f"id-{i}" for i in range(26)], + "EventTime": ["2026-01-01T00:00:00Z"] * 26, + "Feature1": [f"v-{i}" for i in range(26)], + }) + 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) + assert mock_fg.batch_write_record.call_count == 2 + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_empty_dataframe(self, mock_fg_class, feature_definitions): + """Empty DataFrame → no batch_write_record calls.""" + df = pd.DataFrame({ + "RecordIdentifier": [], + "EventTime": [], + "Feature1": [], + }) + 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=True, + ) + mgr.run(data_frame=df) + mock_fg.batch_write_record.assert_not_called() + 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): + """Error with matching entry maps to specific row index.""" + 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 + + # Entry for row 1 (None RecordIdentifier → only EventTime + Feature1) + 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 + assert 0 not in exc_info.value.failed_rows + assert 2 not in exc_info.value.failed_rows + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_multiple_errors_in_batch(self, mock_fg_class, feature_definitions): + """Multiple errors map to correct row indices.""" + from sagemaker.core.shapes import BatchWriteRecordEntry, FeatureValue + + df = pd.DataFrame({ + "RecordIdentifier": ["good-0", None, "good-2", None, "good-4"], + "EventTime": ["2026-01-01T00:00:00Z"] * 5, + "Feature1": ["v0", "v1", "v2", "v3", "v4"], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + # Entries for rows 1 and 3 (missing RecordIdentifier) + error_entry_1 = 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="v1"), + ], + ) + error_entry_3 = 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="v3"), + ], + ) + + mock_error_1 = Mock() + mock_error_1.entry = error_entry_1 + mock_error_1.error_code = "ValidationError" + mock_error_1.error_message = "Missing RecordIdentifier" + + mock_error_3 = Mock() + mock_error_3.entry = error_entry_3 + mock_error_3.error_code = "ValidationError" + mock_error_3.error_message = "Missing RecordIdentifier" + + mock_response = Mock() + mock_response.unprocessed_entries = [] + mock_response.errors = [mock_error_1, mock_error_3] + 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 + assert 3 in exc_info.value.failed_rows + assert 0 not in exc_info.value.failed_rows + assert 2 not in exc_info.value.failed_rows + assert 4 not 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): + """Exception from batch_write_record marks all rows failed.""" + 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_unprocessed_entries(self, mock_fg_class, feature_definitions): + """Unprocessed entries map back to specific row indices.""" + from sagemaker.core.shapes import BatchWriteRecordEntry, FeatureValue + + df = pd.DataFrame({ + "RecordIdentifier": ["id-0", "id-1", "id-2"], + "EventTime": ["2026-01-01T00:00:00Z"] * 3, + "Feature1": ["v0", "v1", "v2"], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + # Entry for row 2 is returned as unprocessed + unprocessed_entry = BatchWriteRecordEntry( + feature_group_name="test-fg", + record=[ + FeatureValue(feature_name="RecordIdentifier", value_as_string="id-2"), + FeatureValue(feature_name="EventTime", value_as_string="2026-01-01T00:00:00Z"), + FeatureValue(feature_name="Feature1", value_as_string="v2"), + ], + ) + + mock_response = Mock() + mock_response.unprocessed_entries = [unprocessed_entry] + 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, + ) + + with pytest.raises(IngestionError) as exc_info: + mgr.run(data_frame=df) + + assert 2 in exc_info.value.failed_rows + assert 0 not in exc_info.value.failed_rows + assert 1 not in exc_info.value.failed_rows + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_unprocessed_entry_no_match_marks_all(self, mock_fg_class, feature_definitions): + """If unprocessed entry can't be matched, all rows in batch marked failed.""" + from sagemaker.core.shapes import BatchWriteRecordEntry, FeatureValue + + df = pd.DataFrame({ + "RecordIdentifier": ["id-0", "id-1"], + "EventTime": ["2026-01-01T00:00:00Z"] * 2, + "Feature1": ["v0", "v1"], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + # Return an entry that doesn't match anything we sent + unmatched_entry = BatchWriteRecordEntry( + feature_group_name="test-fg", + record=[ + FeatureValue(feature_name="RecordIdentifier", value_as_string="UNKNOWN"), + FeatureValue(feature_name="EventTime", value_as_string="2026-01-01T00:00:00Z"), + ], + ) + + mock_response = Mock() + mock_response.unprocessed_entries = [unmatched_entry] + 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, + ) + + with pytest.raises(IngestionError) as exc_info: + mgr.run(data_frame=df) + + # All rows marked failed since we can't identify which one + assert 0 in exc_info.value.failed_rows + assert 1 in exc_info.value.failed_rows + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_multiple_unprocessed_entries(self, mock_fg_class, feature_definitions): + """Multiple unprocessed entries each map to correct row.""" + from sagemaker.core.shapes import BatchWriteRecordEntry, FeatureValue + + df = pd.DataFrame({ + "RecordIdentifier": ["id-0", "id-1", "id-2", "id-3", "id-4"], + "EventTime": ["2026-01-01T00:00:00Z"] * 5, + "Feature1": ["v0", "v1", "v2", "v3", "v4"], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + # Rows 1 and 3 returned as unprocessed + unprocessed_1 = BatchWriteRecordEntry( + feature_group_name="test-fg", + record=[ + FeatureValue(feature_name="RecordIdentifier", value_as_string="id-1"), + FeatureValue(feature_name="EventTime", value_as_string="2026-01-01T00:00:00Z"), + FeatureValue(feature_name="Feature1", value_as_string="v1"), + ], + ) + unprocessed_3 = BatchWriteRecordEntry( + feature_group_name="test-fg", + record=[ + FeatureValue(feature_name="RecordIdentifier", value_as_string="id-3"), + FeatureValue(feature_name="EventTime", value_as_string="2026-01-01T00:00:00Z"), + FeatureValue(feature_name="Feature1", value_as_string="v3"), + ], + ) + + mock_response = Mock() + mock_response.unprocessed_entries = [unprocessed_1, unprocessed_3] + 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, + ) + + with pytest.raises(IngestionError) as exc_info: + mgr.run(data_frame=df) + + assert 1 in exc_info.value.failed_rows + assert 3 in exc_info.value.failed_rows + assert 0 not in exc_info.value.failed_rows + assert 2 not in exc_info.value.failed_rows + assert 4 not in exc_info.value.failed_rows + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_both_errors_and_unprocessed(self, mock_fg_class, feature_definitions): + """Both errors and unprocessed_entries in same response.""" + from sagemaker.core.shapes import BatchWriteRecordEntry, FeatureValue + + df = pd.DataFrame({ + "RecordIdentifier": ["id-0", None, "id-2", "id-3", "id-4"], + "EventTime": ["2026-01-01T00:00:00Z"] * 5, + "Feature1": ["v0", "v1", "v2", "v3", "v4"], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + # Row 1 has error (missing RecordIdentifier) + 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="v1"), + ], + ) + mock_error = Mock() + mock_error.entry = error_entry + mock_error.error_code = "ValidationError" + mock_error.error_message = "Missing RecordIdentifier" + + # Row 4 is unprocessed (throttled) + unprocessed_entry = BatchWriteRecordEntry( + feature_group_name="test-fg", + record=[ + FeatureValue(feature_name="RecordIdentifier", value_as_string="id-4"), + FeatureValue(feature_name="EventTime", value_as_string="2026-01-01T00:00:00Z"), + FeatureValue(feature_name="Feature1", value_as_string="v4"), + ], + ) + + mock_response = Mock() + mock_response.unprocessed_entries = [unprocessed_entry] + 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) + + # Row 1 (error) and row 4 (unprocessed) should both be in failed_rows + assert 1 in exc_info.value.failed_rows + assert 4 in exc_info.value.failed_rows + # Rows 0, 2, 3 should NOT be failed + assert 0 not in exc_info.value.failed_rows + assert 2 not in exc_info.value.failed_rows + assert 3 not in exc_info.value.failed_rows + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_error_without_entry_marks_all(self, mock_fg_class, feature_definitions): + """Error object without .entry attribute marks all rows failed.""" + df = pd.DataFrame({ + "RecordIdentifier": ["id-0", "id-1"], + "EventTime": ["2026-01-01T00:00:00Z"] * 2, + "Feature1": ["v0", "v1"], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + mock_error = Mock() + mock_error.entry = None # No entry object + + 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 0 in exc_info.value.failed_rows + assert 1 in exc_info.value.failed_rows + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_error_entry_no_match_marks_all(self, mock_fg_class, feature_definitions): + """Error with entry that doesn't match → marks all rows in batch failed.""" + from sagemaker.core.shapes import BatchWriteRecordEntry, FeatureValue + + df = pd.DataFrame({ + "RecordIdentifier": ["id-0", "id-1"], + "EventTime": ["2026-01-01T00:00:00Z"] * 2, + "Feature1": ["v0", "v1"], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + + unmatched_entry = BatchWriteRecordEntry( + feature_group_name="test-fg", + record=[FeatureValue(feature_name="UNKNOWN", value_as_string="x")], + ) + mock_error = Mock() + mock_error.entry = unmatched_entry + mock_error.error_code = "InternalError" + mock_error.error_message = "Something went wrong" + + 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 0 in exc_info.value.failed_rows + assert 1 in exc_info.value.failed_rows + + @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 not included in 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): + """False flag → put_record called, not batch_write_record.""" + 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): + """target_stores=None → entries have Unassigned (not 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=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") + from sagemaker.core.utils.utils import Unassigned + assert isinstance(entries[0].target_stores, Unassigned) + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_passes_target_stores_when_set(self, mock_fg_class, feature_definitions, sample_dataframe): + """target_stores=['OnlineStore'] → entries have target_stores set.""" + 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=["OnlineStore"]) + + call_kwargs = mock_fg.batch_write_record.call_args + entries = call_kwargs.kwargs.get("entries") or call_kwargs[1].get("entries") + assert entries[0].target_stores == ["OnlineStore"] + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_sliced_dataframe_preserves_indices(self, mock_fg_class, feature_definitions): + """Sliced DataFrame (non-zero index) reports correct original indices on failure.""" + df = pd.DataFrame({ + "RecordIdentifier": [f"id-{i}" for i in range(10)], + "EventTime": ["2026-01-01T00:00:00Z"] * 10, + "Feature1": [f"v-{i}" for i in range(10)], + }) + + mock_fg = Mock() + mock_fg_class.return_value = mock_fg + mock_fg.batch_write_record.side_effect = Exception("Failure") + + # Simulate multi-process split: process gets rows 5-9 + failed = IngestionManagerPandas._ingest_batch_write( + data_frame=df, + feature_group_name="test-fg", + feature_definitions=feature_definitions, + start_index=5, + end_index=10, + target_stores=None, + ) + + # Should report original DataFrame indices 5,6,7,8,9 + assert failed == [5, 6, 7, 8, 9] + + +class TestBuildRecord: + """Tests for _build_record helper — verifying indexing correctness.""" + + def test_build_record_single_column(self): + """1 column boundary case.""" + df = pd.DataFrame({"Col1": ["val1"]}) + defs = {"Col1": {"FeatureType": "String", "CollectionType": None}} + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, defs) + assert len(record) == 1 + assert record[0].feature_name == "Col1" + assert record[0].value_as_string == "val1" + + def test_build_record_five_columns(self): + """5 columns — verifies first and last column mapped correctly.""" + df = pd.DataFrame({"A": ["a"], "B": ["b"], "C": ["c"], "D": ["d"], "E": ["e"]}) + defs = {c: {"FeatureType": "String", "CollectionType": None} for c in df.columns} + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, defs) + assert len(record) == 5 + assert record[0].feature_name == "A" + assert record[0].value_as_string == "a" + assert record[4].feature_name == "E" + assert record[4].value_as_string == "e" + + def test_build_record_multiple_rows_correct_values(self): + """Each row maps to correct column values (no cross-row contamination).""" + df = pd.DataFrame({ + "Id": ["id-0", "id-1", "id-2"], + "Val": ["v0", "v1", "v2"], + }) + defs = {c: {"FeatureType": "String", "CollectionType": None} for c in df.columns} + for row in df.itertuples(): + record = IngestionManagerPandas._build_record(df, row, defs) + result = {fv.feature_name: fv.value_as_string for fv in record} + idx = row[0] + assert result["Id"] == f"id-{idx}" + assert result["Val"] == f"v{idx}" + + def test_build_record_sliced_dataframe(self): + """Sliced DataFrame (non-zero index) still maps correctly.""" + df = pd.DataFrame({ + "Id": ["id-0", "id-1", "id-2", "id-3", "id-4"], + "Val": ["v0", "v1", "v2", "v3", "v4"], + }) + sliced = df[2:4] # rows at index 2, 3 + defs = {c: {"FeatureType": "String", "CollectionType": None} for c in df.columns} + for row in sliced.itertuples(): + record = IngestionManagerPandas._build_record(sliced, row, defs) + result = {fv.feature_name: fv.value_as_string for fv in record} + idx = row[0] + assert result["Id"] == f"id-{idx}" + assert result["Val"] == f"v{idx}" + + def test_build_record_skips_none(self): + """None values excluded from record.""" + df = pd.DataFrame({ + "RecordIdentifier": ["id-1"], + "EventTime": ["2026-01-01T00:00:00Z"], + "Feature1": [None], + }) + defs = {c: {"FeatureType": "String", "CollectionType": None} for c in df.columns} + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, 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): + """NaN values excluded from record.""" + df = pd.DataFrame({ + "RecordIdentifier": ["id-1"], + "EventTime": ["2026-01-01T00:00:00Z"], + "Feature1": [np.nan], + }) + defs = {c: {"FeatureType": "String", "CollectionType": None} for c in df.columns} + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, defs) + assert len(record) == 2 + + def test_build_record_collection_type(self): + """Collection type feature uses value_as_string_list.""" + df = pd.DataFrame({ + "id": ["id-1"], + "tags": [["a", "b", "c"]], + }) + defs = { + "id": {"FeatureType": "String", "CollectionType": None}, + "tags": {"FeatureType": "String", "CollectionType": "List"}, + } + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, 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"] + + def test_build_record_integer_value_to_string(self): + """Integer values converted to string.""" + df = pd.DataFrame({"Num": [42]}) + defs = {"Num": {"FeatureType": "Integral", "CollectionType": None}} + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, defs) + assert record[0].value_as_string == "42" + + def test_build_record_float_value_to_string(self): + """Float values converted to string.""" + df = pd.DataFrame({"Frac": [3.14]}) + defs = {"Frac": {"FeatureType": "Fractional", "CollectionType": None}} + row = next(df.itertuples()) + record = IngestionManagerPandas._build_record(df, row, defs) + assert record[0].value_as_string == "3.14" + + def test_build_record_row_index_is_row0(self): + """row[0] is the DataFrame index, used for error mapping.""" + df = pd.DataFrame({"A": ["x", "y", "z"]}) + rows = list(df.itertuples()) + assert rows[0][0] == 0 + assert rows[1][0] == 1 + assert rows[2][0] == 2 + + diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_list_records.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_list_records.py new file mode 100644 index 0000000000..6c8d98554f --- /dev/null +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_list_records.py @@ -0,0 +1,78 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 +"""Unit tests for list_records function.""" +import pytest +from unittest.mock import Mock, patch + + +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): + 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_all_params(self, mock_fg_class): + 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): + 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") + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_list_records_does_not_pass_none_params(self, mock_fg_class): + """None values for optional params should not be passed to the API.""" + 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", max_results=None, next_token=None, region="us-west-2") + + # max_results=None and next_token=None should NOT be in the kwargs + mock_fg.list_records.assert_called_once_with(region="us-west-2")