Skip to content

[SNOW-3203938] Fix ai_parse_document_basic test #4109

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sfc-gh-mrek wants to merge 4 commits into
base: main
Choose a base branch
from
Open

[SNOW-3203938] Fix ai_parse_document_basic test #4109

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b320e4b
[SNOW-3203938] Fix ai_parse_document_basic test
sfc-gh-mrek Mar 6, 2026
128cda4
Add Changelog entry
sfc-gh-mrek Mar 6, 2026
5fe5853
Adjust comments according to graphite-app suggestions
sfc-gh-mrek Mar 6, 2026
6c6df01
Update CHANGELOG.md
sfc-gh-mrek Mar 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 43 additions & 5 deletions tests/integ/test_dataframe_ai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1070,8 +1070,49 @@ def test_dataframe_ai_transcribe_default_output_column(session, resources_path):
assert all("start" in s and "end" in s for s in data["segments"])


def test_dataframe_ai_parse_document_basic(session, resources_path):
"""Test DataFrame.ai.parse_document OCR on a PDF document."""
def test_dataframe_ai_parse_document_basic_legacy(session, resources_path):
"""Test DataFrame.ai.parse_document OCR on a PDF document.
This test covers legacy behavior (post error handling changes, metadata is present in response)."""
session.sql(
"ALTER SESSION SET AI_SQL_ERROR_HANDLING_USE_FAIL_ON_ERROR = FALSE"
).collect()
Comment on lines +1076 to +1078
Copy link
Copy Markdown
Contributor

@sfc-gh-aling sfc-gh-aling Mar 6, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

we normally do not do this because our tests run in parallel and it might pollute other test

Copy link
Copy Markdown
Author

@sfc-gh-mrek sfc-gh-mrek Mar 9, 2026

Choose a reason for hiding this comment

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

So how can I control the behaviour based on the parameter?
For a bit we'll live in a world where user can have AI_SQL_ERROR_HANDLING_USE_FAIL_ON_ERROR either enabled or disabled. It changes the response format for the default ai_parse_document behavior (as i explained in linked doc, the metadata is moved out of the value object for AI_SQL_ERROR_HANDLING_USE_FAIL_ON_ERROR=false).
Also, as this parameter is currently opt in, based on the account cofiguration tests will either succeed or fail if we leave the metadata check.

Should there be an assumption that moving forward metadata won't be at top level so the check should be simply removed (even though for some accounts it can still be here?).
This doesn't look like a good change, but at the same time is the safest one if we can't use session params to fully control functions behavior :)

try:
stage_name = Utils.random_stage_name()
_ = session.sql(
f"CREATE OR REPLACE TEMP STAGE {stage_name} ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')"
).collect()
file_local = TestFiles(resources_path).test_doc_pdf
_ = session.file.put(file_local, f"@{stage_name}", auto_compress=False)

df = session.create_dataframe(
[[f"@{stage_name}/doc.pdf"]], schema=["file_path"]
)

result_df = df.ai.parse_document(
input_column=to_file(col("file_path")),
output_column="parsed",
mode="OCR",
)

assert result_df.columns == ["FILE_PATH", "PARSED"]

results = result_df.collect()
data = json.loads(results[0]["PARSED"]) if results[0]["PARSED"] else {}
assert isinstance(data, dict)
assert "content" in data and isinstance(data["content"], str)
assert isinstance(data.get("metadata", {}), dict)
if "metadata" in data and "pageCount" in data["metadata"]:
assert data["metadata"].get("pageCount", 0) >= 3
finally:
session.sql(
"ALTER SESSION SET AI_SQL_ERROR_HANDLING_USE_FAIL_ON_ERROR = TRUE"
).collect()


def test_dataframe_ai_parse_document_basic_new_eh(session, resources_path):
"""Test DataFrame.ai.parse_document OCR on a PDF document.
This test covers legacy behavior (post error handling changes, metadata is absent in response)."""

stage_name = Utils.random_stage_name()
_ = session.sql(
f"CREATE OR REPLACE TEMP STAGE {stage_name} ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')"
Expand All @@ -1093,9 +1134,6 @@ def test_dataframe_ai_parse_document_basic(session, resources_path):
data = json.loads(results[0]["PARSED"]) if results[0]["PARSED"] else {}
assert isinstance(data, dict)
assert "content" in data and isinstance(data["content"], str)
assert isinstance(data.get("metadata", {}), dict)
if "metadata" in data and "pageCount" in data["metadata"]:
assert data["metadata"].get("pageCount", 0) >= 3


def test_dataframe_ai_parse_document_default_output_column(session, resources_path):
Expand Down
Loading