SuperAPI AI Directory
Partner with us中文AI newsFavorites
Skillsdi-agent-query-substrait

di-agent-query-substrait

Agents & models

View repository

Updated 10-16 05:56

About this skill

Generate Substrait DSL

Translate natural language data requests into executable Substrait JSON via the Substrait DSL.

Inputs Required

We need the flow description and the input schema. It is important to get the input schema correct and it should be consistent with the names and types in the asset browser. Use list_project_assets to find the relevant asset(s) by names and inspect_project_asset with the asset ID(s) to retrieve the authoritative column names, types, and nullability. Do not use any other information like fewshot examples or prior knowledge to guess the schema — it must be derived from the project assets to ensure the generated Substrait plan is valid against the actual data.

If the user provides the table schemas (one or more source tables with column names and substrait types), verify them against the actual asset in the project using list_project_assets and inspect_project_asset. If there are inconsistencies, flag them to the user. The schema from the asset tools takes precedence unless the user explicitly overrides it.

If the user does not supply table schemas, derive them directly from the project assets by calling list_project_assets and inspect_project_asset.

Once you have the assets build a NamedStruct by mapping them assets to DSL types using the following mapping:

Asset Type Mapping

When constructing a NamedStruct from an inspected asset, map asset column types to DSL types as follows:

Asset type DSL type
longvarchar, varchar, char, text string
bigint, int8 i64
integer, int, int4 i64
smallint, int2 i16
real, float4 fp32
numeric, decimal fp64
double precision, float8 fp64
boolean, bool boolean
date date
timestamp, timestamp without time zone timestamp

Append ? to the DSL type for any column where nullable: true in the asset metadata (e.g., string?).

Workflow

Follow these steps to generate a Substrait plan from a natural language request.

1. Fetch Few-Shot Examples

Call get_substrait_dsl_examples to retrieve relevant DSL examples.

CRITICAL RULE — the user_request parameter MUST be the user's ENTIRE original message, copied verbatim, with no modifications whatsoever. The retrieval system uses sentence-level semantic embeddings that require the full context to match correctly. Passing keywords or summaries returns irrelevant examples and causes downstream failures.

get_substrait_dsl_examples(user_request="<FULL VERBATIM USER MESSAGE — EVERY WORD, NO SUMMARIZATION>", collection="draft_generation", n=5)

WRONG (keywords/summary): user_request="join two tables aggregate count sum divide cast round NPS score" CORRECT (full verbatim text): user_request="Given a fct_reviews table (listing_id, review_date, reviewer_name, review_text, review_sentiment) and a dim_listings table (listing_id, listing_name, minimum_nights, host_id, room_type, price, created_at, updated_at), produce a per-listing NPS aggregation result set. For each listing: listing_id, listing_name, room_type, nps_total (overall NPS in [-100, 100] across all reviews of that listing, computed as (positive - negative) / total * 100, rounded to 0 decimals), reviews_total (total review count). Inner join on listing_id; only include listings that have at least one review."

The user_request string length should match the user's message length. If your user_request is shorter than the user's message, you are doing it wrong.

Study the returned examples to understand the syntax of the DSL. Never use the few-shot examples to derive any input schema even if seemingly more appropriate — that must come from the project assets as described in the next step.

2. Verify Schema from Assets

It is mandatory to verify the schema using the asset tools. Otherwise, the result will be an invalid plan that fails to execute!

  1. Call list_project_assets to locate the asset(s) by name in the project
  2. Call inspect_project_asset with the asset ID(s) and asset_type="data_asset" to retrieve column metadata
  3. For each column, map the asset type to a DSL type using the Asset Type Mapping table above
  4. It is imperative to append ? to nullable column types (e.g., string?)
  5. Use the resulting schema as the NamedStruct for ReadTable

The schema found in the asset browser takes precedence over any schema mentioned in the few-shot examples. If there are inconsistencies to user provided schemas, flag them and ask the user to confirm which one to use.

3. Generate DSL

Call get_substrait_dsl_spec to retrieve the language specification. Study it carefully before writing any DSL code. Write a DSL query that matches the user's request using the verified table schemas by following the rules in the spec for how to construct the query, how to name and reference columns, and how to use aliases.

  • Work with immutable VirtualTable objects transformed through Constructs
  • Each line: variable = Construct(...)
  • End with return <vtable>
  • Expressions must be inlined — no expression variables
  • Use left. and right. prefixes in Join conditions and only in Join conditions
  • Use alias on aggregation functions and grouping expressions

4. Compile via compile_substrait_dsl

Call the compile_substrait_dsl MCP tool with the generated DSL:

compile_substrait_dsl(substrait_dsl_code="<your_dsl>", read_tables=[...])

You may only proceed if compile_substrait_dsl succeeds. If it succeeds, the Substrait JSON is in the response — proceed to step 6. If it does not succeed, continue with step 5.

5. Self-Correct (up to 3 attempts)

When errors occur:

  1. Read the error message carefully
  2. Optionally call get_substrait_dsl_examples with collection="generation_with_error_correction" for correction examples
  3. Fix the DSL based on the error feedback
  4. Call compile_substrait_dsl again

Common errors and fixes:

  • Column not found: check the vtable schema after each construct — use only columns that exist
  • Type mismatch: use cast() to convert types, or check literal types
  • Scope errors: aliases in the same construct's exprs list are not visible to each other — chain a second Select/Project
  • Join column errors: use left. and right. prefixes

If all 3 attempts fail: report the final error to the user, show the last DSL attempt, and suggest they refine their query or check the table schemas.

6. Execute or Convert

Depending on the user's request, execute or convert the DSL. Only execute the request if the user explicitly asks for it.

  • Execute the Substrait DSL (if the user asks to run it):

    Two execution tools are available. Both compile the DSL and run the resulting flow. Only call either after a successful compile_substrait_dsl validation.

    run_substrait_dsl(substrait_dsl_code="<dsl>", project_id="<project_id>", source_asset_ids=["<id1>", ...])

    run_substrait_dsl_static(substrait_dsl_code="<dsl>", project_id="<project_id>")

    If one fails, retry with the other as a fallback — either can fail in cases where the other succeeds.

  • Convert to Elyra pipeline JSON (if the user requests it):

    convert_substrait_dsl_to_elyra(substrait_dsl_code="<dsl>", read_tables=[...])

    Returns {"success": bool, "elyra_json": dict|null, "errors": str|null}.

7. Return Results

Present the user with:

  • The Substrait JSON from the successful compilation
  • The DSL code generated (for transparency)

Output Format

### Generated DSL

\`\`\`query
<DSL code, excluding ReadTable lines>
\`\`\`

### Substrait JSON

\`\`\`json
<substrait_json from compile_substrait_dsl>
\`\`\`

Important Notes

  • get_substrait_dsl_examples user_request parameter: You MUST pass the user's complete unmodified message as the user_request value. Do not extract keywords, summarize, or shorten it in any way. Also always pass collection="draft_generation" and n=5. Violation of this rule produces irrelevant examples.
  • You must call get_substrait_dsl_spec and study the spec before generating any DSL code. It is the authoritative reference for all constructs, expressions, types, and scoping rules — do not only rely on few-shot examples or prior knowledge for syntax details
  • Select keeps ONLY listed columns; Project keeps ALL columns plus new ones
  • Aggregation measures must NOT be wrapped in cast()
  • Final output column names must not contain dots — use Select with aliases to rename

Runtime Considerations

  • Nullability: Any column marked nullable: true in the asset metadata MUST use the nullable type suffix (e.g., string?, i64?). DataStage will crash at runtime if NULL values are read into non-nullable fields.
  • Case-insensitive filtering: When filtering on text/enum columns, use lower() for case-insensitive comparison (e.g., equal(lower(col("status")), literal("pending", "string"))) rather than assuming a specific casing in the data.
  • Column disambiguation: When joining multiple tables, pre-rename overlapping columns via Select before joining (e.g., rename customer_id to orders_customer_id using a table-derived prefix). Do not rely on left./right. prefixes outside of Join conditions.