Skip to content

Search docs and blog

AI/ML Primitives Benchmark

CLI name: ai_primitives - use benchbox run --benchmark ai_primitives

Query templates: every AI/ML Primitives query

The AI/ML Primitives benchmark tests SQL-based AI functions across cloud data platforms. It enables consistent evaluation of AI capabilities in Snowflake Cortex, BigQuery ML, and Databricks AI Functions using TPC-H data.

Overview

PropertyValue
Total Queries16
Categories4 (generative, nlp, transform, embedding)
Data SourceTPC-H (shared data)
Supported PlatformsSnowflake, BigQuery, Databricks
Cost ControlBuilt-in budget enforcement

Why AI Primitives?

Modern cloud data platforms increasingly offer built-in AI/ML capabilities through SQL functions:

  • Snowflake Cortex: COMPLETE, SUMMARIZE, SENTIMENT, EXTRACT_ANSWER, EMBED_TEXT
  • BigQuery ML: ML.GENERATE_TEXT, ML.UNDERSTAND_TEXT, ML.GENERATE_EMBEDDING
  • Databricks AI: ai_query, ai_summarize, ai_classify, ai_generate

The AI Primitives benchmark provides:

  1. Standardized Testing: Same queries across all platforms
  2. Cost Estimation: Pre-execution cost estimates with budget enforcement
  3. Performance Comparison: Consistent metrics for platform evaluation
  4. Production Data: Uses TPC-H customer/supplier comments as realistic text

Query Categories

Generative (4 queries)

Text generation and completion tasks:

Query IDDescription
generative_complete_simpleBasic text completion with short prompt
generative_complete_customer_profileGenerate customer profile summaries from comments
generative_question_answerAnswer questions about customer data
generative_sql_generationGenerate SQL from natural language

NLP (5 queries)

Natural language processing and understanding:

Query IDDescription
nlp_sentiment_analysisAnalyze sentiment of customer comments
nlp_classificationClassify customer feedback into categories
nlp_entity_extractionExtract entities (products, locations, etc.)
nlp_language_detectionDetect language of supplier comments
nlp_keyword_extractionExtract key terms from order comments

Transform (4 queries)

Text transformation tasks:

Query IDDescription
transform_summarizeSummarize long customer comments
transform_translateTranslate supplier comments
transform_grammar_correctionCorrect grammar in customer feedback
transform_rewrite_professionalRewrite comments in professional tone

Embedding (3 queries)

Vector embedding generation:

Query IDDescription
embedding_customer_commentsGenerate embeddings for customer comments
embedding_supplier_descriptionsGenerate embeddings for supplier descriptions
embedding_similarity_searchFind similar items using vector similarity

Quick Start

CLI Usage

The commands below run the benchmark on Snowflake at a small scale factor, run a dry run to estimate costs first, set a cost budget (the run fails if the budget is exceeded), and run only the nlp and transform categories on Databricks.

Terminal window
benchbox run --platform snowflake --benchmark ai_primitives --scale 0.01


benchbox run --platform snowflake --benchmark ai_primitives --dry-run ./preview


benchbox run --platform snowflake --benchmark ai_primitives \
  --benchmark-option max_cost_usd=1.00


benchbox run --platform databricks --benchmark ai_primitives \
  --benchmark-option categories=nlp,transform

Programmatic Usage

from benchbox.core.ai_primitives import AIPrimitivesBenchmark


benchmark = AIPrimitivesBenchmark(
    scale_factor=0.01,
    max_cost_usd=1.00,
    dry_run=False
)


total_cost, estimates = benchmark.estimate_cost(
    platform="snowflake",
    categories=["nlp", "transform"]
)
print(f"Estimated cost: ${total_cost:.4f}")


result = benchmark.run_benchmark(
    connection=snowflake_connection,
    platform="snowflake",
    categories=["nlp"]
)


print(f"Successful: {result.successful_queries}/{result.total_queries}")
print(f"Total cost: ${result.total_cost_estimated_usd:.4f}")

Platform Support

Snowflake (Cortex)

Snowflake Cortex provides LLM functions directly in SQL. The queries below show text completion, sentiment analysis, and embedding generation, in that order:

SELECT SNOWFLAKE.CORTEX.COMPLETE('llama3-8b', 'Summarize: ' || comment) AS summary
FROM customer LIMIT 10;


SELECT SNOWFLAKE.CORTEX.SENTIMENT(c_comment) AS sentiment
FROM customer WHERE c_nationkey = 1;


SELECT SNOWFLAKE.CORTEX.EMBED_TEXT_768('e5-base-v2', c_comment) AS embedding
FROM customer LIMIT 10;

Supported Models: llama3-8b, llama3-70b, mistral-large, mixtral-8x7b

BigQuery (ML Functions)

BigQuery ML provides AI functions through external model connections. The queries below show text generation and then embedding generation:

SELECT ML.GENERATE_TEXT(
    MODEL `project.dataset.llm_model`,
    (SELECT 'Summarize: ' || comment AS prompt),
    STRUCT(100 AS max_output_tokens)
).ml_generate_text_result AS summary
FROM `project.dataset.customer` LIMIT 10;


SELECT ML.GENERATE_EMBEDDING(
    MODEL `project.dataset.embedding_model`,
    (SELECT c_comment AS content)
).ml_generate_embedding_result AS embedding
FROM `project.dataset.customer` LIMIT 10;

Setup Required: Create external model connection to Vertex AI or other LLM provider.

Databricks (AI Functions)

Databricks provides AI functions through Foundation Model APIs. The queries below show text completion and then classification:

SELECT ai_query(
    'databricks-meta-llama-3-1-70b-instruct',
    CONCAT('Summarize this customer: ', c_comment)
) AS summary
FROM customer LIMIT 10;


SELECT ai_classify(c_comment, ARRAY('positive', 'negative', 'neutral')) AS sentiment
FROM customer LIMIT 10;

Supported Models: Foundation Model APIs including Llama, DBRX, and external endpoints.

Cost Management

AI queries incur costs based on token usage. The benchmark provides built-in cost controls.

Cost Estimation

from benchbox.core.ai_primitives import AIPrimitivesBenchmark


benchmark = AIPrimitivesBenchmark(scale_factor=0.01)


total_cost, estimates = benchmark.estimate_cost(platform="snowflake")


for estimate in estimates:
    print(f"{estimate.query_id}: ${estimate.estimated_cost_usd:.4f}")
    print(f"  Tokens: {estimate.estimated_tokens}")
    print(f"  Model: {estimate.model}")

Budget Enforcement

max_cost_usd sets the maximum cost, and execution fails if it is exceeded. A dry run produces estimates only and makes no AI calls.

Terminal window
benchbox run --platform snowflake --benchmark ai_primitives \
  --benchmark-option max_cost_usd=1.00


benchbox run --platform snowflake --benchmark ai_primitives \
  --dry-run ./cost_preview

Typical Costs (Estimated)

CategoryQueriesEst. Cost (SF=0.01)
Generative4~$0.05-0.15
NLP5~$0.02-0.08
Transform4~$0.03-0.10
Embedding3~$0.01-0.05
Total16~$0.10-0.40

Costs vary by platform, model, and actual token usage.

Configuration Options

OptionTypeDefaultDescription
scale_factorfloat0.01TPC-H scale factor (affects data volume)
max_cost_usdfloat0.0Maximum allowed cost (0 = unlimited)
dry_runboolFalseEstimate costs without executing
categorieslistallFilter by category
querieslistallRun specific query IDs

Result Structure

@dataclass
class AIBenchmarkResult:
    benchmark: str = "AI Primitives"
    platform: str = ""
    scale_factor: float = 1.0
    dry_run: bool = False
    total_queries: int = 0
    successful_queries: int = 0
    failed_queries: int = 0
    skipped_queries: int = 0
    total_execution_time_ms: float = 0.0
    total_cost_estimated_usd: float = 0.0
    cost_tracker: CostTracker | None = None
    query_results: list[AIQueryResult] = field(default_factory=list)

Unsupported Platforms

The following platforms do not have built-in AI functions and will skip AI queries:

  • DuckDB
  • ClickHouse
  • SQLite
  • PostgreSQL
  • DataFusion
  • Trino/Presto
  • Redshift
  • Azure Synapse Analytics
  • Microsoft Fabric

When running on these platforms, queries are skipped with informative messages.

Best Practices

1. Start with Dry Run

Always estimate costs before running AI queries:

Terminal window
benchbox run --platform snowflake --benchmark ai_primitives \
  --dry-run ./preview --scale 0.01

2. Use Small Scale Factors

AI functions process text row-by-row. Start with --scale 0.01:

Terminal window
benchbox run --platform databricks --benchmark ai_primitives --scale 0.01

3. Set Cost Budgets

Prevent runaway costs with budget limits. max_cost_usd is a hard limit:

benchmark = AIPrimitivesBenchmark(
    scale_factor=0.01,
    max_cost_usd=1.00
)

4. Test by Category

Run specific categories to focus testing:

Terminal window
benchbox run --platform snowflake --benchmark ai_primitives \
  --benchmark-option categories=nlp

5. Monitor Actual Costs

Compare estimates vs actuals in results:

result = benchmark.run_benchmark(connection, "snowflake")
print(f"Estimated: ${result.total_cost_estimated_usd:.4f}")
print(f"Tracker: {result.cost_tracker.get_summary()}")

Extending the Benchmark

Adding New Queries

Add queries to benchbox/core/ai_primitives/catalog/queries.yaml:

- id: my_new_query
  category: nlp
  description: My new AI query
  model: llama3-8b
  estimated_tokens: 100
  batch_size: 10
  cost_per_1k_tokens: 0.0003
  sql: |
    SELECT 'placeholder' as result
  variants:
    snowflake: |
      SELECT SNOWFLAKE.CORTEX.SENTIMENT(c_comment) FROM customer LIMIT 10
    bigquery: |
      -- BigQuery variant
    databricks: |
      -- Databricks variant
  skip_on:
    - duckdb
    - clickhouse

Adding New Platforms

  1. Add platform to SUPPORTED_PLATFORMS in benchmark.py
  2. Add query variants in queries.yaml
  3. Implement cost estimation for the platform’s pricing model

BenchBox is open source under the MIT License.