feat: Wire BatchWriteRecord and ListRecords into ingest_dataframe#6026
Open
chiragvp-aws wants to merge 1 commit into
Open
feat: Wire BatchWriteRecord and ListRecords into ingest_dataframe#6026chiragvp-aws wants to merge 1 commit into
chiragvp-aws wants to merge 1 commit into
Conversation
chiragvp-aws
temporarily deployed
to
auto-approve
July 15, 2026 17:27 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
chiragvp-aws
temporarily deployed
to
auto-approve
July 15, 2026 17:27 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
chiragvp-aws
temporarily deployed
to
auto-approve
July 15, 2026 22:35 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
chiragvp-aws
temporarily deployed
to
auto-approve
July 15, 2026 22:35 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
grraman
reviewed
Jul 17, 2026
Comment on lines
+384
to
+385
| for index in range(1, len(row)): | ||
| feature_name = data_frame.columns[index - 1] |
There was a problem hiding this comment.
Not sure I'm following the 1-based indexing used here
grraman
reviewed
Jul 17, 2026
| return record | ||
|
|
||
| @staticmethod | ||
| def _ingest_batch_write( |
There was a problem hiding this comment.
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?
grraman
reviewed
Jul 17, 2026
| 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 |
There was a problem hiding this comment.
Should we find their position in the original entries list, or in the original data frame?
grraman
reviewed
Jul 17, 2026
| idx = entries.index(entry) | ||
| failed_rows.append(row_indices[idx]) | ||
| except ValueError: | ||
| # If we can't find the exact entry, mark by position |
There was a problem hiding this comment.
When would we run into this except block?
grraman
reviewed
Jul 17, 2026
| @@ -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.""" | |||
There was a problem hiding this comment.
Why couple them into the same test module? Should we separate them into 2 modules?
grraman
reviewed
Jul 17, 2026
| 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): |
There was a problem hiding this comment.
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?
chiragvp-aws
force-pushed
the
wiring-batchwrite-list
branch
from
July 17, 2026 17:35
d0b514f to
9eda65d
Compare
chiragvp-aws
temporarily deployed
to
auto-approve
July 17, 2026 17:35 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
chiragvp-aws
temporarily deployed
to
auto-approve
July 17, 2026 17:35 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
chiragvp-aws
force-pushed
the
wiring-batchwrite-list
branch
from
July 17, 2026 19:17
9eda65d to
f79f93a
Compare
chiragvp-aws
temporarily deployed
to
auto-approve
July 17, 2026 19:19 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
chiragvp-aws
temporarily deployed
to
auto-approve
July 17, 2026 19:19 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
- 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
chiragvp-aws
force-pushed
the
wiring-batchwrite-list
branch
from
July 17, 2026 21:32
f79f93a to
5f513fb
Compare
chiragvp-aws
temporarily deployed
to
auto-approve
July 17, 2026 21:33 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
chiragvp-aws
temporarily deployed
to
auto-approve
July 17, 2026 21:33 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.