Skip to content

Commit ad04889

Browse files
sfc-gh-jdusfc-gh-jdu
committed
SNOW-2301878: Support DataFrame.ai.extract (#3719)
1 parent d8970d9 commit ad04889

3 files changed

Lines changed: 263 additions & 0 deletions

File tree

‎docs/source/snowpark/dataframe.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ DataFrame
124124
DataFrameAIFunctions.classify
125125
DataFrameAIFunctions.complete
126126
DataFrameAIFunctions.embed
127+
DataFrameAIFunctions.extract
127128
DataFrameAIFunctions.filter
128129
DataFrameAIFunctions.parse_document
129130
DataFrameAIFunctions.sentiment

‎src/snowflake/snowpark/dataframe_ai_functions.py‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
ai_filter,
1616
ai_agg,
1717
ai_classify,
18+
ai_extract,
1819
ai_similarity,
1920
ai_sentiment,
2021
ai_embed,
@@ -1095,3 +1096,162 @@ def parse_document(
10951096
"DataFrame.ai.parse_document",
10961097
)
10971098
return df
1099+
1100+
@experimental(version="1.37.0")
1101+
def extract(
1102+
self,
1103+
input_column: ColumnOrName,
1104+
*,
1105+
response_format: Optional[Union[Dict[str, str], List]] = None,
1106+
output_column: Optional[str] = None,
1107+
_emit_ast: bool = True,
1108+
) -> "snowflake.snowpark.DataFrame":
1109+
"""Extract structured information from text or files using a response schema.
1110+
1111+
Args:
1112+
input_column: The column (Column object or column name as string) containing the text
1113+
or FILE data to extract information from. Use ``to_file`` for staged file paths.
1114+
response_format: The schema describing information to extract. Supports:
1115+
1116+
- Simple object schema (dict) mapping feature names to extraction prompts:
1117+
``{'name': 'What is the last name of the employee?', 'address': 'What is the address of the employee?'}``
1118+
- Array of strings containing the information to be extracted:
1119+
``['What is the last name of the employee?', 'What is the address of the employee?']``
1120+
- Array of arrays containing two strings (feature name and extraction prompt):
1121+
``[['name', 'What is the last name of the employee?'], ['address', 'What is the address of the employee?']]``
1122+
- Array of strings with colon-separated feature names and extraction prompts:
1123+
``['name: What is the last name of the employee?', 'address: What is the address of the employee?']``
1124+
1125+
output_column: The name of the output column to be appended.
1126+
If not provided, a column named ``AI_EXTRACT_OUTPUT`` is appended.
1127+
1128+
Returns:
1129+
A new DataFrame with an appended JSON object containing the extracted fields
1130+
under ``response``.
1131+
1132+
Examples::
1133+
1134+
>>> # Extract from text string
1135+
>>> from snowflake.snowpark.functions import col
1136+
>>> df = session.create_dataframe([
1137+
... ["John Smith lives in San Francisco and works for Snowflake"],
1138+
... ], schema=["text"])
1139+
>>> result_df = df.ai.extract(
1140+
... input_column="text",
1141+
... response_format={'name': 'What is the first name of the employee?', 'city': 'What is the address of the employee?'},
1142+
... output_column="extracted",
1143+
... )
1144+
>>> result_df.select("EXTRACTED").show()
1145+
--------------------------------
1146+
|"EXTRACTED" |
1147+
--------------------------------
1148+
|{ |
1149+
| "response": { |
1150+
| "city": "San Francisco", |
1151+
| "name": "John" |
1152+
| } |
1153+
|} |
1154+
--------------------------------
1155+
<BLANKLINE>
1156+
1157+
>>> # Extract using array format
1158+
>>> df = session.create_dataframe(
1159+
... [
1160+
... ["Alice Johnson works in Seattle"],
1161+
... ["Bob Williams works in Portland"],
1162+
... ],
1163+
... schema=["text"]
1164+
... )
1165+
>>> result_df = df.ai.extract(
1166+
... input_column=col("text"),
1167+
... response_format=[["name", "What is the first name?"], ["city", "What city do they work in?"]],
1168+
... output_column="info",
1169+
... )
1170+
>>> result_df.show()
1171+
------------------------------------------------------------
1172+
|"TEXT" |"INFO" |
1173+
------------------------------------------------------------
1174+
|Alice Johnson works in Seattle |{ |
1175+
| | "response": { |
1176+
| | "city": "Seattle", |
1177+
| | "name": "Alice" |
1178+
| | } |
1179+
| |} |
1180+
|Bob Williams works in Portland |{ |
1181+
| | "response": { |
1182+
| | "city": "Portland", |
1183+
| | "name": "Bob" |
1184+
| | } |
1185+
| |} |
1186+
------------------------------------------------------------
1187+
<BLANKLINE>
1188+
1189+
>>> # Extract lists using List: prefix
1190+
>>> df = session.create_dataframe(
1191+
... [["Python, Java, and JavaScript are popular programming languages"]],
1192+
... schema=["text"]
1193+
... )
1194+
>>> result_df = df.ai.extract(
1195+
... input_column="text",
1196+
... response_format=[["languages", "List: What programming languages are mentioned?"]],
1197+
... output_column="extracted",
1198+
... )
1199+
>>> result_df.select("EXTRACTED").show()
1200+
----------------------
1201+
|"EXTRACTED" |
1202+
----------------------
1203+
|{ |
1204+
| "response": { |
1205+
| "languages": [ |
1206+
| "Python", |
1207+
| "Java", |
1208+
| "JavaScript" |
1209+
| ] |
1210+
| } |
1211+
|} |
1212+
----------------------
1213+
<BLANKLINE>
1214+
1215+
>>> # Extract from file
1216+
>>> from snowflake.snowpark.functions import to_file
1217+
>>> _ = session.sql("CREATE OR REPLACE TEMP STAGE mystage ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')").collect()
1218+
>>> _ = session.file.put("tests/resources/invoice.pdf", "@mystage", auto_compress=False)
1219+
>>> df = session.create_dataframe([["@mystage/invoice.pdf"]], schema=["file_path"])
1220+
>>> result_df = df.ai.extract(
1221+
... input_column=to_file(col("file_path")),
1222+
... response_format=[["date", "What is the invoice date?"], ["amount", "What is the amount?"]],
1223+
... output_column="info",
1224+
... )
1225+
>>> result_df.select("INFO").show()
1226+
--------------------------------
1227+
|"INFO" |
1228+
--------------------------------
1229+
|{ |
1230+
| "response": { |
1231+
| "amount": "USD $950.00", |
1232+
| "date": "Nov 26, 2016" |
1233+
| } |
1234+
|} |
1235+
--------------------------------
1236+
<BLANKLINE>
1237+
1238+
"""
1239+
1240+
input_col = _to_col_if_str(input_column, "DataFrame.ai.extract")
1241+
1242+
result_col = ai_extract(
1243+
input=input_col,
1244+
response_format=response_format,
1245+
_emit_ast=False,
1246+
)
1247+
1248+
output_column_name = output_column or "AI_EXTRACT_OUTPUT"
1249+
df = self._dataframe.with_column(
1250+
output_column_name, result_col, _emit_ast=False
1251+
)
1252+
1253+
add_api_call(
1254+
df,
1255+
"DataFrame.ai.extract",
1256+
)
1257+
return df

‎tests/integ/test_dataframe_ai.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1125,3 +1125,105 @@ def test_dataframe_ai_parse_document_default_output_column(session, resources_pa
11251125
and "index" in first_page
11261126
and "content" in first_page
11271127
)
1128+
1129+
1130+
def test_dataframe_ai_extract_text_basic(session):
1131+
"""Test DataFrame.ai.extract with text input and dict response format."""
1132+
df = session.create_dataframe(
1133+
[
1134+
["John Smith lives in San Francisco"],
1135+
["Alice Johnson works in Seattle"],
1136+
],
1137+
schema=["text"],
1138+
)
1139+
1140+
result_df = df.ai.extract(
1141+
input_column="text",
1142+
response_format={
1143+
"name": "What is the first name?",
1144+
"city": "What city is mentioned?",
1145+
},
1146+
output_column="extracted",
1147+
)
1148+
1149+
assert result_df.columns == ["TEXT", "EXTRACTED"]
1150+
1151+
results = result_df.collect(_emit_ast=False)
1152+
assert len(results) == 2
1153+
text_to_response = {}
1154+
for row in results:
1155+
assert row["EXTRACTED"] is not None
1156+
data = json.loads(row["EXTRACTED"]) if row["EXTRACTED"] else {}
1157+
assert isinstance(data, dict) and "response" in data
1158+
text_to_response[row["TEXT"]] = data["response"]
1159+
1160+
assert text_to_response["John Smith lives in San Francisco"]["name"] == "John"
1161+
assert (
1162+
text_to_response["John Smith lives in San Francisco"]["city"] == "San Francisco"
1163+
)
1164+
assert text_to_response["Alice Johnson works in Seattle"]["name"] == "Alice"
1165+
assert text_to_response["Alice Johnson works in Seattle"]["city"] == "Seattle"
1166+
1167+
1168+
def test_dataframe_ai_extract_default_output_column(session):
1169+
"""Test DataFrame.ai.extract uses default output column name."""
1170+
df = session.create_dataframe([["Bob lives in Denver"]], schema=["text"])
1171+
1172+
result_df = df.ai.extract(
1173+
input_column="text",
1174+
response_format=[
1175+
["person", "What is the first name?"],
1176+
["location", "What city is mentioned?"],
1177+
],
1178+
)
1179+
1180+
assert "AI_EXTRACT_OUTPUT" in result_df.columns
1181+
results = result_df.collect(_emit_ast=False)
1182+
data = (
1183+
json.loads(results[0]["AI_EXTRACT_OUTPUT"])
1184+
if results[0]["AI_EXTRACT_OUTPUT"]
1185+
else {}
1186+
)
1187+
assert isinstance(data, dict) and isinstance(data.get("response", {}), dict)
1188+
assert data["response"]["person"] == "Bob"
1189+
assert data["response"]["location"] == "Denver"
1190+
1191+
1192+
def test_dataframe_ai_extract_file(session, resources_path):
1193+
"""Test DataFrame.ai.extract on a staged PDF file."""
1194+
stage_name = Utils.random_stage_name()
1195+
_ = session.sql(
1196+
f"CREATE OR REPLACE TEMP STAGE {stage_name} ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')"
1197+
).collect()
1198+
file_local = TestFiles(resources_path).test_invoice_pdf
1199+
_ = session.file.put(file_local, f"@{stage_name}", auto_compress=False)
1200+
1201+
df = session.create_dataframe(
1202+
[[f"@{stage_name}/invoice.pdf"]], schema=["file_path"]
1203+
)
1204+
1205+
result_df = df.ai.extract(
1206+
input_column=to_file(col("file_path")),
1207+
response_format=[
1208+
["date", "What is the invoice date?"],
1209+
["amount", "What is the amount?"],
1210+
],
1211+
output_column="info",
1212+
)
1213+
1214+
assert result_df.columns == ["FILE_PATH", "INFO"]
1215+
results = result_df.collect(_emit_ast=False)
1216+
data = json.loads(results[0]["INFO"]) if results[0]["INFO"] else {}
1217+
assert isinstance(data, dict) and isinstance(data.get("response", {}), dict)
1218+
assert data["response"]["date"] == "Nov 26, 2016"
1219+
assert data["response"]["amount"] == "USD $950.00"
1220+
1221+
1222+
def test_dataframe_ai_extract_error_handling(session):
1223+
"""Test error handling in DataFrame.ai.extract."""
1224+
df = session.create_dataframe([["text"]], schema=["col"])
1225+
with pytest.raises(TypeError, match="expected Column or str"):
1226+
df.ai.extract(
1227+
input_column=123, # Invalid type
1228+
response_format={"a": "What is a?"},
1229+
)

0 commit comments

Comments
 (0)