Tutorial: Run Cortex batch inference over a Snowflake table¶

# scorer.py  — runs inside each sandbox
import sys, json

from snowflake.sandbox import get_snowflake_connection

DB, SCHEMA, start_id, end_id = sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4])

conn = get_snowflake_connection()
cur = conn.cursor()

cur.execute(
    f"SELECT review_id, product, review_text "
    f"FROM {DB}.{SCHEMA}.PRODUCT_REVIEWS "
    f"WHERE review_id BETWEEN %s AND %s ORDER BY review_id",
    (start_id, end_id),
)
rows = cur.fetchall()

results = []
for review_id, product, review_text in rows:
    prompt = (
        "Classify the sentiment of this product review. "
        'Return JSON only: {"sentiment": "positive|neutral|negative", "score": 0.0-1.0}. '
        "Review: " + review_text
    )
    cur.execute("SELECT SNOWFLAKE.CORTEX.COMPLETE('llama3.1-8b', %s)", (prompt,))
    raw = (cur.fetchone() or ["{}"])[0].strip()
    # Strip markdown fences that some models add despite the prompt
    if raw.startswith("```"):
        raw = raw.strip("`").lstrip("json").strip()
    try:
        parsed = json.loads(raw)
    except Exception:
        parsed = {"sentiment": "unknown", "score": 0.0}
    results.append({
        "review_id": review_id,
        "product": product,
        "sentiment": parsed.get("sentiment", "unknown"),
        "score": float(parsed.get("score", 0.0)),
    })
    # Progress line written to stderr — streamed live to the driver without
    # mixing with the JSON result that the driver reads from stdout.
    print(f"scored review {review_id} ({product}): {parsed.get('sentiment', 'unknown')}", file=sys.stderr, flush=True)

conn.close()
print(json.dumps(results))

# main.py  — runs on your machine (the driver)
import json, snowflake.connector
from concurrent.futures import ThreadPoolExecutor
from snowflake.sandbox import Image, Sandbox

CONNECTION = "<your_connection_name>"  # the connection name from your connections.toml
DB     = "<your_database>"
SCHEMA = "<your_schema>"

def score_chunk(label: str, start_id: int, end_id: int) -> list:
    with Sandbox.create(
        image=Image.from_catalog("sandbox-base"), connection=CONNECTION
    ) as sb:
        # Send the scorer into this sandbox. scorer.py needs only
        # snowflake-connector-python, which the base image already provides.
        sb.upload_file("scorer.py", "/tmp/scorer.py")
        stream = sb.exec_stream(
            ["python", "/tmp/scorer.py", DB, SCHEMA, str(start_id), str(end_id)],
            timeout=120.0,
        )
        stdout_lines: list[str] = []
        for line in stream:
            if line.stream == "stderr":  # progress from scorer.py
                print(f"  [{label}] {line.data}")
            else:  # stdout: collect for JSON parsing
                stdout_lines.append(line.data)
        if stream.exit_code != 0:
            raise RuntimeError(f"[{label}] exited {stream.exit_code}")
        # The JSON result is the last NON-EMPTY stdout line. Do not use
        # stdout_lines[-1]: the stream ends with a trailing empty line, so
        # indexing the last element parses "" and raises JSONDecodeError.
        payload = next(s for s in reversed(stdout_lines) if s.strip())
        return json.loads(payload)

print("Starting 3 sandboxes in parallel...")
with ThreadPoolExecutor() as pool:
    futures = [
        pool.submit(score_chunk, "sb-1", 1, 2),
        pool.submit(score_chunk, "sb-2", 3, 4),
        pool.submit(score_chunk, "sb-3", 5, 6),
    ]
    chunks = [f.result() for f in futures]
results = [row for chunk in chunks for row in chunk]

# Write results back to Snowflake using the same connection as the driver
conn = snowflake.connector.connect(connection_name=CONNECTION)
cur = conn.cursor()
for r in results:
    cur.execute(
        f"INSERT INTO {DB}.{SCHEMA}.REVIEW_SENTIMENTS "
        f"(review_id, product, sentiment, sentiment_score) VALUES (%s, %s, %s, %s)",
        (r["review_id"], r["product"], r["sentiment"], r["score"]),
    )
conn.close()

print(f"\n{'ID':<4} {'Product':<22} {'Sentiment':<12} {'Score'}")
print("-" * 55)
for r in sorted(results, key=lambda x: x["review_id"]):
    print(f"{r['review_id']:<4} {r['product']:<22} {r['sentiment']:<12} {r['score']:.2f}")
print(f"\n{len(results)} reviews scored and written to {DB}.{SCHEMA}.REVIEW_SENTIMENTS")