Skip to content

feat: Wire BatchWriteRecord and ListRecords into ingest_dataframe#6026

Open
chiragvp-aws wants to merge 1 commit into
aws:masterfrom
chiragvp-aws:wiring-batchwrite-list
Open

feat: Wire BatchWriteRecord and ListRecords into ingest_dataframe#6026
chiragvp-aws wants to merge 1 commit into
aws:masterfrom
chiragvp-aws:wiring-batchwrite-list

Conversation

@chiragvp-aws

Copy link
Copy Markdown
Collaborator
  • 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

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Comment on lines +384 to +385
for index in range(1, len(row)):
feature_name = data_frame.columns[index - 1]

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure I'm following the 1-based indexing used here

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In dataframe, column 0 is the index of the row which should be discarded. We have similar logic for _ingest_data using putRecord.

return record

@staticmethod
def _ingest_batch_write(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lots of chances to accidentally author off-by-1 errors as we're doing a lot of indexing/slicing arrays here. Do we have sufficient unit tests to cover a wide range of scenarios?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have good number of unit tests. I will check if we have some missing scenarios here.

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we find their position in the original entries list, or in the original data frame?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not required as we are building entries for request body at line 451.

idx = entries.index(entry)
failed_rows.append(row_indices[idx])
except ValueError:
# If we can't find the exact entry, mark by position

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When would we run into this except block?

@@ -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."""

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why couple them into the same test module? Should we separate them into 2 modules?

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):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are there any other kinds of partial failures we need to test? Perhaps a scenario with only unprocessed_entries, one with only errors, one with both?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll see if we can add more tests.

- 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants