fix: validate pyarrow columns in stage input checks#2156
Conversation
Greptile SummaryThis PR fixes a regression where
Confidence Score: 5/5Safe to merge — the fix is narrow, directly targets the reported false-negative, and does not touch any execution or scheduling paths. The change is a self-contained guard function replacing a single hasattr call. Both PyArrow and pandas branches are tested, the logic is straightforward, and there are no side effects outside of validate_input. No files require special attention; the two observations are minor and do not affect correctness of the fix. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["_has_data_attr(data, attr)"] --> B{"hasattr(data, attr)?"}
B -- Yes --> C["return True"]
B -- No --> D{"isinstance(data, pa.Table)?"}
D -- Yes --> E["attr in data.column_names"]
E --> F["return result"]
D -- No --> G{"isinstance(data, pd.DataFrame)?"}
G -- Yes --> H["attr in data.columns"]
H --> I["return result"]
G -- No --> J["return False"]
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
A["_has_data_attr(data, attr)"] --> B{"hasattr(data, attr)?"}
B -- Yes --> C["return True"]
B -- No --> D{"isinstance(data, pa.Table)?"}
D -- Yes --> E["attr in data.column_names"]
E --> F["return result"]
D -- No --> G{"isinstance(data, pd.DataFrame)?"}
G -- Yes --> H["attr in data.columns"]
H --> I["return result"]
G -- No --> J["return False"]
Reviews (6): Last reviewed commit: "fix: remove unreachable pyarrow schema f..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if hasattr(data, "columns"): | ||
| return attr in data.columns |
There was a problem hiding this comment.
columns branch is ambiguous for PyArrow-like types
For pa.Table, data.columns returns a list of ChunkedArray objects, not strings — so attr in data.columns would perform object-identity comparison and always return False for a plain string attribute name. This branch works today only because pa.Table also exposes column_names, which is checked first. If another table type (e.g. a future columnar format) has columns returning arrays but lacks column_names, this branch would silently misreport all columns as missing. A brief inline comment or guard (e.g. checking that the first element is a string) would make the intent and the ordering dependency explicit.
| class TestProcessingStageValidateInput: | ||
| def test_validate_input_accepts_pyarrow_columns(self): | ||
| stage = RequiredColumnStage() | ||
| batch = DocumentBatch( | ||
| task_id="batch", | ||
| dataset_name="dataset", | ||
| data=pa.table({"text": ["hello"]}), | ||
| ) | ||
|
|
||
| assert stage.validate_input(batch) is True | ||
|
|
||
| def test_validate_input_rejects_missing_pyarrow_columns(self): | ||
| stage = RequiredColumnStage() | ||
| batch = DocumentBatch( | ||
| task_id="batch", | ||
| dataset_name="dataset", | ||
| data=pa.table({"other": ["hello"]}), | ||
| ) | ||
|
|
||
| assert stage.validate_input(batch) is False |
There was a problem hiding this comment.
DocumentBatch.data is typed as pa.Table | pd.DataFrame, and pd.DataFrame takes a different code path through _has_data_attr (it lacks column_names, so it falls through to the columns branch). The new test class only exercises the PyArrow path. Adding parallel test_validate_input_accepts_pandas_columns / test_validate_input_rejects_missing_pandas_columns cases would ensure that both union branches stay correct and guard against a future regression in the columns fallback.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
+1 let's try to make sure the test coverage here reaches every possible block of _has_data_attr please.
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
|
/ok to test 0a69c05 |
|
/ok to test 5a2e951 |
| if hasattr(data, attr): | ||
| return True | ||
|
|
||
| if hasattr(data, "column_names"): |
There was a problem hiding this comment.
Should it look more like
| if hasattr(data, "column_names"): | |
| if isinstance(data, pa.Table) and hasattr(data, "column_names"): |
?
| if hasattr(data, "column_names"): | ||
| return attr in data.column_names | ||
|
|
||
| if hasattr(data, "columns"): |
There was a problem hiding this comment.
| if hasattr(data, "columns"): | |
| if isinstance(data, pd.DataFrame) and hasattr(data, "columns"): |
?
| if hasattr(data, "columns"): | ||
| return attr in data.columns | ||
|
|
||
| if hasattr(data, "schema") and hasattr(data.schema, "names"): |
There was a problem hiding this comment.
| if hasattr(data, "schema") and hasattr(data.schema, "names"): | |
| if isinstance(data, pa.Table) and hasattr(data, "schema") and hasattr(data.schema, "names"): |
?
| return attr in data.columns | ||
|
|
||
| if hasattr(data, "schema") and hasattr(data.schema, "names"): | ||
| return attr in data.schema.names |
There was a problem hiding this comment.
Is this block supposed to cover the empty table case? Like what conditions would you expect for it to get all the way down here?
There was a problem hiding this comment.
Bumping this comment. @nightcityblade should this condition be removed or is there a case where the function could enter this block?
| def test_validate_input_accepts_pyarrow_columns(self): | ||
| stage = RequiredColumnStage() | ||
| batch = DocumentBatch( | ||
| task_id="batch", |
There was a problem hiding this comment.
Make sure you have the latest version of main:
| task_id="batch", |
we no longer pass task_id directly to create a DocumentBatch.
| def test_validate_input_rejects_missing_pyarrow_columns(self): | ||
| stage = RequiredColumnStage() | ||
| batch = DocumentBatch( | ||
| task_id="batch", |
There was a problem hiding this comment.
| task_id="batch", |
| class TestProcessingStageValidateInput: | ||
| def test_validate_input_accepts_pyarrow_columns(self): | ||
| stage = RequiredColumnStage() | ||
| batch = DocumentBatch( | ||
| task_id="batch", | ||
| dataset_name="dataset", | ||
| data=pa.table({"text": ["hello"]}), | ||
| ) | ||
|
|
||
| assert stage.validate_input(batch) is True | ||
|
|
||
| def test_validate_input_rejects_missing_pyarrow_columns(self): | ||
| stage = RequiredColumnStage() | ||
| batch = DocumentBatch( | ||
| task_id="batch", | ||
| dataset_name="dataset", | ||
| data=pa.table({"other": ["hello"]}), | ||
| ) | ||
|
|
||
| assert stage.validate_input(batch) is False |
There was a problem hiding this comment.
+1 let's try to make sure the test coverage here reaches every possible block of _has_data_attr please.
|
Pushed a follow-up commit for the review notes. I narrowed the fallback checks to concrete / cases, updated the construction to match current , and added coverage for the pandas path plus the explicit-attribute/unknown-attribute branches.\n\nLocal verification: All checks passed! passed. I also attempted the focused test module, but the local environment here is missing repo test dependencies ( from , plus package deps needed to import end-to-end), so I wasn’t able to run that pytest target locally before pushing. |
|
Pushed a follow-up commit for the review notes. I narrowed the fallback checks to concrete Local verification: |
|
Addressed the latest review note: I removed the unreachable fallback and pushed a follow-up commit, since for we already return through the branch.\n\nLocal verification: passed. I also retried |
|
Addressed the latest review note: I removed the unreachable schema.names fallback and pushed a follow-up commit, since for pa.Table we already return through the column_names branch. Local verification: python3 -m ruff check nemo_curator/stages/base.py tests/stages/common/test_base.py passed. I also retried pytest tests/stages/common/test_base.py -q, but the local environment here still fails at import time because ray is not installed via tests/conftest.py. |
Signed-off-by: nightcityblade <nightcityblade@gmail.com>
|
Pushed a follow-up force-push to add the missing DCO sign-off to the latest commit, so the DCO check should be able to clear on the updated head. |
|
Too many files changed for review. ( |
3 participants
Description
closes #2151
ProcessingStage.validate_input()to recognize column names on table-like task data instead of relying only onhasattr()DocumentBatchtests using PyArrow tablesUsage
Checklist
Validation run locally:
python3 -m ruff check nemo_curator/stages/base.py tests/stages/common/test_base.pypython3 -m pytest --noconftest tests/stages/common/test_base.py -q(blocked in this environment:pyarrowis not installed; the full test setup also requiresrayviatests/conftest.py)