snowflake.snowpark.DataFrameAIFunctions.extract¶

DataFrameAIFunctions.extract(input_column: Union[snowflake.snowpark.column.Column, str], *, response_format: Optional[Union[Dict[str, Any], List]] = None, scores: Optional[bool] = None, config: Optional[Dict] = None, output_column: Optional[str] = None) → snowflake.snowpark.DataFrame[source]¶

Extract structured information from text or files using a response schema.

Parameters:

• input_column – The column (Column object or column name as string) containing the text or FILE data to extract information from. Use to_file for staged file paths.

• response_format –

  The schema describing information to extract. Supports:

  • Simple object schema (dict) mapping feature names to extraction prompts: {'name': 'What is the last name of the employee?', 'address': 'What is the address of the employee?'}

  • Array of strings containing the information to be extracted: ['What is the last name of the employee?', 'What is the address of the employee?']

  • Array of arrays containing two strings (feature name and extraction prompt): [['name', 'What is the last name of the employee?'], ['address', 'What is the address of the employee?']]

  • Array of strings with colon-separated feature names and extraction prompts: ['name: What is the last name of the employee?', 'address: What is the address of the employee?']

  • JSON Schema format: a dict with a top-level schema key whose value is a JSON Schema object ('type': 'object'). Each entry under properties describes one field to extract, where description is the extraction prompt and type is one of 'string' (a single value), 'array' (a list), or 'object' (a table, which must also specify column_ordering and its own properties), e.g.:

    CopyExpand

    {
        "schema": {
            "type": "object",
            "properties": {
                "name": {"description": "What is the last name of the employee?", "type": "string"},
                "skills": {"description": "What skills are listed?", "type": "array"},
                "salary_table": {
                    "description": "Salary by year",
                    "type": "object",
                    "column_ordering": ["year", "salary"],
                    "properties": {
                        "year": {"description": "Year", "type": "array"},
                        "salary": {"description": "Salary", "type": "array"},
                    },
                },
            },
        }
    }
    

    Show lessSee more

    Scroll to top

    You can’t combine the JSON Schema format with the other formats above: if response_format contains a schema key, every field to extract must be defined within that JSON Schema.

• scores – When True, the output JSON includes a scoring object alongside response with per-field confidence scores between 0 and 1. Default is None (no scores).

• config –

  Optional dict of configuration settings for file inputs. Supported key:

  • scale_factor: Numeric value 1.0–4.0 to scale pages before processing, improving OCR quality for dense or small-text documents.

• output_column – The name of the output column to be appended. If not provided, a column named AI_EXTRACT_OUTPUT is appended.

Returns:

A new DataFrame with an appended JSON object containing the extracted fields under response.

Examples:

CopyExpand

>>> # Extract from text string
>>> from snowflake.snowpark.functions import col
>>> df = session.create_dataframe([
...     ["John Smith lives in San Francisco and works for Snowflake"],
... ], schema=["text"])
>>> result_df = df.ai.extract(
...     input_column="text",
...     response_format={'name': 'What is the first name of the employee?', 'city': 'What is the address of the employee?'},
...     output_column="extracted",
... )
>>> result_df.select("EXTRACTED").show()
--------------------------------
|"EXTRACTED"                   |
--------------------------------
|{                             |
|  "error": null,              |
|  "response": {               |
|    "city": "San Francisco",  |
|    "name": "John"            |
|  }                           |
|}                             |
--------------------------------


>>> # Extract using array format
>>> df = session.create_dataframe(
...     [
...         ["Alice Johnson works in Seattle"],
...         ["Bob Williams works in Portland"],
...     ],
...     schema=["text"]
... )
>>> result_df = df.ai.extract(
...     input_column=col("text"),
...     response_format=[["name", "What is the first name?"], ["city", "What city do they work in?"]],
...     output_column="info",
... )
>>> result_df.show()
------------------------------------------------------------
|"TEXT"                          |"INFO"                   |
------------------------------------------------------------
|Alice Johnson works in Seattle  |{                        |
|                                |  "error": null,         |
|                                |  "response": {          |
|                                |    "city": "Seattle",   |
|                                |    "name": "Alice"      |
|                                |  }                      |
|                                |}                        |
|Bob Williams works in Portland  |{                        |
|                                |  "error": null,         |
|                                |  "response": {          |
|                                |    "city": "Portland",  |
|                                |    "name": "Bob"        |
|                                |  }                      |
|                                |}                        |
------------------------------------------------------------


>>> # Extract lists using List: prefix
>>> df = session.create_dataframe(
...     [["Python, Java, and JavaScript are popular programming languages"]],
...     schema=["text"]
... )
>>> result_df = df.ai.extract(
...     input_column="text",
...     response_format=[["languages", "List: What programming languages are mentioned?"]],
...     output_column="extracted",
... )
>>> result_df.select("EXTRACTED").show()
----------------------
|"EXTRACTED"         |
----------------------
|{                   |
|  "error": null,    |
|  "response": {     |
|    "languages": [  |
|      "Python",     |
|      "Java",       |
|      "JavaScript"  |
|    ]               |
|  }                 |
|}                   |
----------------------


>>> # Extract using JSON Schema format (required for combining entities, lists, and
>>> # tables in a single call; can't be combined with the other response_format shapes)
>>> df = session.create_dataframe(
...     ["John Smith lives in San Francisco and lists Python and SQL as skills"],
...     schema=["text"]
... )
>>> result_df = df.ai.extract(
...     input_column="text",
...     response_format={
...         "schema": {
...             "type": "object",
...             "properties": {
...                 "name": {"description": "What is the first name?", "type": "string"},
...                 "skills": {"description": "What skills are listed?", "type": "array"},
...             },
...         }
...     },
...     output_column="extracted",
... )
>>> result_df.select("EXTRACTED").show()
-----------------------
|"EXTRACTED"          |
-----------------------
|{                    |
|  "error": null,     |
|  "response": {      |
|    "name": "John",  |
|    "skills": [      |
|      "Python",      |
|      "SQL"          |
|    ]                |
|  }                  |
|}                    |
-----------------------


>>> # Extract from file
>>> from snowflake.snowpark.functions import to_file
>>> _ = session.sql("CREATE OR REPLACE TEMP STAGE mystage ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')").collect()
>>> _ = session.file.put("tests/resources/invoice.pdf", "@mystage", auto_compress=False)
>>> df = session.create_dataframe([["@mystage/invoice.pdf"]], schema=["file_path"])
>>> result_df = df.ai.extract(
...     input_column=to_file(col("file_path")),
...     response_format=[["date", "What is the invoice date?"], ["amount", "What is the amount?"]],
...     output_column="info",
... )
>>> result_df.select("INFO").show()
--------------------------------
|"INFO"                        |
--------------------------------
|{                             |
|  "error": null,              |
|  "response": {               |
|    "amount": "USD $950.00",  |
|    "date": "Nov 26, 2016"    |
|  }                           |
|}                             |
--------------------------------

Show lessSee more

Scroll to top