-
Notifications
You must be signed in to change notification settings - Fork 128
Add AI assistant skills for DQX #1125
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
mwojtyczka
merged 11 commits into
databrickslabs:main
from
ghanse:feature/dqx-skills-1088
May 4, 2026
Merged
Changes from all commits
Commits
Show all changes
11 commits
Select commit
Hold shift + click to select a range
f0eb1c3
feat: AI assistant skills for DQX (Genie Code, Claude Code)
ghanse 34adc92
Merge branch 'main' into feature/dqx-skills-1088
mwojtyczka fb3552a
Merge branch 'main' into feature/dqx-skills-1088
mwojtyczka 00b418e
Update docs and add plugin.json
ghanse 86a2253
Remove stale docs page
ghanse 4a39c9b
Merge branch 'main' into feature/dqx-skills-1088
mwojtyczka d766d04
Merge branch 'main' into feature/dqx-skills-1088
mwojtyczka d8c6f5d
Update docs and add unit test
ghanse e8b8aec
Apply suggestion from @mwojtyczka
mwojtyczka 2c9df02
refactor
mwojtyczka e3fa934
fix the repo url
mwojtyczka File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,15 @@ | ||
| { | ||
| "name": "databrickslabs-dqx", | ||
| "owner": { | ||
| "name": "Databricks Labs", | ||
| "email": "labs@databricks.com", | ||
| "url": "https://databrickslabs.github.io/" | ||
| }, | ||
| "plugins": [ | ||
| { | ||
| "name": "dqx", | ||
| "source": "./skills", | ||
| "description": "DQX by Databricks Labs — AI assistant skills for defining, applying, storing, and profiling data quality rules in PySpark." | ||
| } | ||
| ] | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -20,3 +20,4 @@ Resolves #.. | |
|
|
||
| - [ ] added/updated demos | ||
| - [ ] added/updated docs | ||
| - [ ] added/updated agent skills | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,118 @@ | ||
| --- | ||
| sidebar_position: 314 | ||
| --- | ||
|
|
||
| import Admonition from '@theme/Admonition'; | ||
| import Tabs from '@theme/Tabs'; | ||
| import TabItem from '@theme/TabItem'; | ||
|
|
||
| # AI Tools & Skills | ||
|
|
||
| DQX ships [Agent Skills](https://agentskills.io/) that teach AI assistants how to use the library correctly. They're small, focused Markdown files with YAML frontmatter — the open format supported by [Databricks Genie Code](https://docs.databricks.com/aws/en/genie-code/skills), [Claude Code](https://docs.claude.com/en/docs/claude-code/plugins), and any other tool that follows the standard. | ||
|
|
||
| The skills live in the [`skills/`](https://github.com/databrickslabs/dqx/tree/main/skills) directory of the DQX repo. | ||
|
|
||
| ## What's included | ||
|
|
||
| | Skill | What it covers | Canonical docs | | ||
| |----------------------------|------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------| | ||
| | `dqx-define-checks` | Creating quality rules — `DQRowRule`, `DQDatasetRule`, `DQForEachColRule`, YAML / dict metadata form | [Quality Checks Definition](/docs/guide/quality_checks_definition) | | ||
| | `dqx-apply-checks` | Validating a DataFrame or table against a set of rules | [Applying Quality Checks](/docs/guide/quality_checks_apply) | | ||
| | `dqx-end-to-end` | Read → check → write in one call with `apply_checks_and_save_in_table` | [End-to-end apply](/docs/guide/quality_checks_apply#applying-checks-defined-with-dqx-classes) | | ||
| | `dqx-profile-and-generate` | Profiling and generating rule candidates with `DQProfiler` / `DQGenerator` | [Data Profiling](/docs/guide/data_profiling) | | ||
| | `dqx-storage` | Loading / saving checks across file, workspace, volume, table, installation, and Lakebase backends | [Loading and Storing Quality Checks](/docs/guide/quality_checks_storage) | | ||
|
|
||
| Each skill is a folder containing a single `SKILL.md` with the standard `name` + `description` frontmatter. Tools auto-load skills based on the description; users can also invoke them explicitly (`@dqx-define-checks`, slash commands, etc.) depending on the tool. | ||
|
|
||
| ## Installing DQX skills | ||
|
|
||
| <Tabs> | ||
| <TabItem value="genie" label="Databricks Genie Code" default> | ||
|
|
||
| Clone the DQX repository in your workspace using a [Git folder](https://docs.databricks.com/aws/en/repos/git-operations-with-repos), | ||
| then copy the `skills/` folder into a workspace-level or user-level skills directory: | ||
|
|
||
| ```bash | ||
| # Option A — workspace-wide (all Genie Code users see these) | ||
| databricks workspace import-dir skills /Workspace/.assistant/skills | ||
|
mwojtyczka marked this conversation as resolved.
|
||
|
|
||
| # Option B — current user only — substitute your workspace email below | ||
| databricks workspace import-dir skills /Users/<your-email>/.assistant/skills | ||
| ``` | ||
|
|
||
| Genie Code picks up skills from these directories automatically. Use the following prompt with Agent mode to confirm: | ||
|
|
||
| > *"List the DQX skills you can use."* | ||
|
|
||
| The skill fires when its description matches your request. Use `@` to manually invoke a specific skill (e.g. `@dqx-define-checks add a uniqueness check on order_id`). | ||
|
|
||
| See the [Databricks Genie Code Documentation](https://docs.databricks.com/aws/en/genie-code/skills) for more details. | ||
|
|
||
| </TabItem> | ||
| <TabItem value="claude" label="Claude Code"> | ||
|
|
||
| The DQX repo is a Claude Code plugin marketplace ([`.claude-plugin/marketplace.json`](https://github.com/databrickslabs/dqx/blob/v0.13.0/.claude-plugin/marketplace.json) at the repo root) containing the `dqx` plugin under [`skills/`](https://github.com/databrickslabs/dqx/tree/main/skills). Add the marketplace once, then install: | ||
|
|
||
| ```bash | ||
| # Inside Claude Code | ||
| /plugin marketplace add databrickslabs/dqx | ||
|
mwojtyczka marked this conversation as resolved.
|
||
| /plugin install dqx@databrickslabs-dqx | ||
| ``` | ||
|
|
||
| Or drop the folder into your personal skills directory for a one-off install: | ||
|
|
||
| ```bash | ||
| cp -R skills/dqx-* ~/.claude/skills/ | ||
| ``` | ||
|
|
||
| Each skill becomes available as `/dqx:dqx-define-checks` (plugin-installed) or as a normal skill (user-level). Claude auto-invokes on description match. | ||
|
mwojtyczka marked this conversation as resolved.
|
||
|
|
||
| See the [Claude Code Plugins documentation](https://docs.claude.com/en/docs/claude-code/plugins) for more details. | ||
|
|
||
| </TabItem> | ||
| <TabItem value="other" label="Other AI dev tools"> | ||
|
|
||
| Any tool that follows the [Agent Skills](https://agentskills.io/) open standard can install and load DQX skills. To install DQX skills with your tool: | ||
|
|
||
| 1. Copy the per-skill folder into the tool's configured skills directory. | ||
| 2. Make sure the path contains `SKILL.md` directly (not nested one level deeper). | ||
| 3. Restart the tool if it caches on startup. | ||
|
|
||
| <Admonition type="note" title="Flattening skill files"> | ||
| If your tool expects a single flat skills file instead of a directory with multiple files, concatenate the files: | ||
|
|
||
| ```bash | ||
| cat skills/dqx-*/SKILL.md > dqx-skills.md | ||
| ``` | ||
| </Admonition> | ||
| </TabItem> | ||
| </Tabs> | ||
|
|
||
| ## Using DQX skills | ||
|
|
||
| After installing DQX skills, you can either let your tool use skills automatically or invoke them by name. Typical prompts: | ||
|
|
||
| ``` | ||
| Add a DQX uniqueness check on (order_id, line_item_id) to my pipeline. | ||
| Split my bronze table into valid and quarantine outputs using these rules: … | ||
| Profile catalog.schema.orders and suggest quality checks. | ||
| Load DQX checks from a Delta table and apply them to a streaming DataFrame. | ||
| ``` | ||
|
|
||
| Agents will load the relevant skill into context, follow its patterns, and link back to the canonical documentation for anything outside the skill's scope. | ||
|
|
||
| ## Extending DQX skills | ||
|
mwojtyczka marked this conversation as resolved.
|
||
|
|
||
| DQX's agent skills are scoped to DQX's public APIs. Follow these guidelines when extending them: | ||
|
|
||
| - Skills should stay short; the full `SKILL.md` is loaded every time the skill fires, so every line costs tokens on every invocation. | ||
| - Prefer linking to `/docs/...` over duplicating content; the skill's job is to tell the model **when** and **how** to use the API, not to reprint the reference. | ||
| - Always import from `databricks.labs.dqx.*` — never guess module paths. | ||
| - Point to the canonical documentation for any topic outside the skill's core responsibility. | ||
| - Changes to the public DQX API should be reflected in the matching skill in the same PR. See [Contributing](/docs/dev/contributing) for the full workflow. | ||
|
|
||
| ## Source | ||
|
|
||
| - Skills: [`skills/`](https://github.com/databrickslabs/dqx/tree/main/skills) | ||
| - Plugin manifest (Claude Code / generic): [`skills/.claude-plugin/plugin.json`](https://github.com/databrickslabs/dqx/blob/v0.13.0/skills/.claude-plugin/plugin.json) | ||
| - Marketplace manifest (Claude Code): [`.claude-plugin/marketplace.json`](https://github.com/databrickslabs/dqx/blob/v0.13.0/.claude-plugin/marketplace.json) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| { | ||
| "name": "dqx", | ||
| "version": "0.1.0", | ||
| "description": "DQX by Databricks Labs — AI assistant skills for defining, applying, storing, and profiling data quality rules in PySpark.", | ||
| "author": { | ||
| "name": "Databricks Labs", | ||
| "email": "labs@databricks.com" | ||
| }, | ||
| "homepage": "https://databrickslabs.github.io/dqx/", | ||
| "repository": "https://github.com/databrickslabs/dqx.git", | ||
| "license": "SEE LICENSE IN LICENSE", | ||
|
mwojtyczka marked this conversation as resolved.
|
||
| "keywords": [ | ||
| "databricks", | ||
| "dqx", | ||
| "data-quality", | ||
| "pyspark", | ||
| "delta", | ||
| "unity-catalog" | ||
| ], | ||
| "skills": [ | ||
| "./dqx-define-checks", | ||
| "./dqx-apply-checks", | ||
|
mwojtyczka marked this conversation as resolved.
|
||
| "./dqx-end-to-end", | ||
| "./dqx-profile-and-generate", | ||
| "./dqx-storage" | ||
| ] | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,23 @@ | ||
| # DQX AI Assistant Skills | ||
|
|
||
| DQX provides agent skills that can be used with [Databricks Genie Code](https://docs.databricks.com/aws/en/genie-code/skills) or any tool that follows the [Agent Skills](https://agentskills.io/) open standard. Skills teach AI assistants how to use DQX. | ||
|
|
||
| ## Skills | ||
|
|
||
| Each skill is a folder with a `SKILL.md` file that documents usage patterns. The following skills are available: | ||
|
|
||
| | Skill | What it covers | | ||
| |-----------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------| | ||
| | [`dqx-define-checks`](./dqx-define-checks/SKILL.md) | Creating quality rules (`DQRowRule`, `DQDatasetRule`, `DQForEachColRule`) and the equivalent YAML / dict metadata form. | | ||
| | [`dqx-apply-checks`](./dqx-apply-checks/SKILL.md) | Validating data with a set of rules (checks loadable from YAML / JSON / Delta). | | ||
| | [`dqx-end-to-end`](./dqx-end-to-end/SKILL.md) | Running check → save cycles in one call with `apply_checks_and_save_in_table`. | | ||
| | [`dqx-profile-and-generate`](./dqx-profile-and-generate/SKILL.md) | Profiling a dataset and generating rule candidates with `DQProfiler` / `DQGenerator`. | | ||
| | [`dqx-storage`](./dqx-storage/SKILL.md) | Loading and saving checks across file / workspace / volume / table / installation / Lakebase backends. | | ||
|
|
||
| ## Install | ||
|
|
||
| See the [installation and usage guide](https://databrickslabs.github.io/dqx/docs/guide/ai_tools_skills) for Databricks Genie Code and other AI tools. | ||
|
|
||
| ## Scope and guardrails | ||
|
|
||
| DQX's agent skills are scoped to DQX's public APIs. See [Extending DQX skills](https://databrickslabs.github.io/dqx/docs/guide/ai_tools_skills#extending-dqx-skills) on the docs site for the full list of design guidelines. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,103 @@ | ||
| --- | ||
|
mwojtyczka marked this conversation as resolved.
|
||
| name: dqx-apply-checks | ||
| description: > | ||
| Validate a PySpark DataFrame or Delta table against a set of DQX quality rules using | ||
| DQEngine. Use when the user asks to "run data quality checks", "apply DQX rules to a | ||
| DataFrame/table", "split valid and invalid rows", "quarantine bad records", or | ||
| "integrate DQX into a streaming pipeline". Covers apply_checks, apply_checks_and_split, | ||
| the by_metadata variants, and the shape of the result columns. | ||
| --- | ||
|
|
||
| # DQX — Apply quality checks | ||
|
|
||
| Apply checks with `DQEngine`. There are six method families — pick by **(1) rule form** and **(2) what output you want**. | ||
|
|
||
| | Rule form | Append results as columns | Split valid / invalid | End-to-end (read + check + write) | | ||
| |--------------------------------------------------------------------------------|----------------------------|---------------------------------------|----------------------------------------------| | ||
| | Classes (`DQRowRule`, …) | `apply_checks` | `apply_checks_and_split` | `apply_checks_and_save_in_table` | | ||
| | Metadata (`list[dict]`, loadable from YAML / JSON / Delta — see `dqx-storage`) | `apply_checks_by_metadata` | `apply_checks_by_metadata_and_split` | `apply_checks_by_metadata_and_save_in_table` | | ||
|
|
||
| For the end-to-end methods see the `dqx-end-to-end` skill. Streaming uses the same methods as batch. | ||
|
|
||
| ## Examples | ||
|
|
||
| ### Using DQX Classes | ||
|
|
||
| ```python | ||
| from databricks.labs.dqx import check_funcs | ||
| from databricks.labs.dqx.engine import DQEngine | ||
| from databricks.labs.dqx.rule import DQRowRule, DQDatasetRule | ||
| from databricks.sdk import WorkspaceClient | ||
| # `spark` is available in Databricks notebooks / jobs. Locally, create it with | ||
| # `from pyspark.sql import SparkSession; spark = SparkSession.builder.getOrCreate()`. | ||
|
|
||
| dq = DQEngine(WorkspaceClient()) | ||
|
|
||
| checks = [ | ||
| DQRowRule(criticality="warn", check_func=check_funcs.is_not_null, column="col3"), | ||
| DQDatasetRule(criticality="error", check_func=check_funcs.is_unique, columns=["col1", "col2"]), | ||
| ] | ||
|
|
||
| df = spark.read.table("catalog.schema.input") | ||
|
|
||
| # Option A — one output DataFrame with _warnings / _errors columns appended. | ||
| annotated = dq.apply_checks(df, checks) | ||
|
|
||
| # Option B — split on error severity. Warnings ride along with valid_df. | ||
| valid_df, invalid_df = dq.apply_checks_and_split(df, checks) | ||
| ``` | ||
|
|
||
| ### Using Metadata (loaded from storage) | ||
|
|
||
| ```python | ||
| dq = DQEngine(WorkspaceClient()) | ||
|
|
||
| checks_metadata = """ | ||
| - criticality: warn | ||
| check: | ||
| function: is_not_null | ||
| arguments: | ||
| column: col3 | ||
|
|
||
| - criticality: error | ||
| check: | ||
| function: is_unique | ||
| arguments: | ||
| columns: [col1, col2] | ||
| """ | ||
|
|
||
| # Option A — one output DataFrame with _warnings / _errors columns appended. | ||
| annotated = dq.apply_checks_by_metadata(df, checks_metadata) | ||
|
|
||
| # Option B — split on error severity. Warnings ride along with valid_df. | ||
| valid_df, invalid_df = dq.apply_checks_by_metadata_and_split(df, checks_metadata) | ||
| ``` | ||
|
|
||
| ## Result shape | ||
|
|
||
| By default two struct columns are appended: | ||
|
|
||
| - **`_errors`** — an array of result structs, one per failing error-criticality rule. | ||
| - **`_warnings`** — same, for warn-criticality rules. | ||
|
|
||
| Each result struct contains: `name`, `message`, `columns`, `filter`, `function`, `run_time`, `run_id`, `user_metadata`, `rule_fingerprint`, `rule_set_fingerprint`, and (for rules that couldn't be evaluated) `skipped=true` with a reason. | ||
|
|
||
| Column names and the skipped-entry behavior are configurable — see [Additional Configuration](https://databrickslabs.github.io/dqx/docs/guide/additional_configuration) for `ExtraParams`, `result_column_names`, and `suppress_skipped`. | ||
|
|
||
| ## Streaming | ||
|
|
||
| Use `spark.readStream` / `writeStream` as normal; the same `apply_checks*` calls work on streaming DataFrames. For incremental batch-style execution out of the box, use `apply_checks_and_save_in_table` with the default `AvailableNow` trigger (see `dqx-end-to-end`). | ||
|
|
||
| ## Multiple inputs | ||
|
|
||
| - `apply_checks_and_save_in_tables(run_configs)` — one call processes many `(input, output)` pairs defined by `RunConfig` entries. | ||
| - `apply_checks_and_save_in_tables_for_patterns(patterns=["catalog.schema.*"], checks_location="catalog.schema.checks", run_config_template=RunConfig(...))` — expand wildcards against Unity Catalog. `patterns` is a `list[str]`; `checks_location` is required and is used as a template for per-table file lookups or as the shared Delta table. See [Applying Checks](https://databrickslabs.github.io/dqx/docs/guide/quality_checks_apply) for the full signature. | ||
|
|
||
| ## Do / Don't | ||
|
|
||
| - **Do** fail fast: `DQEngine.validate_checks(checks)` raises on bad definitions before you touch data. | ||
| - **Do** persist the split result (`invalid_df`) to a quarantine table — not just the annotated DataFrame — so downstream consumers see only good rows. | ||
| - **Don't** mix DQX classes and metadata in one call; pick one form per invocation. | ||
| - **Don't** drop the `_errors` / `_warnings` columns unless you've materialised the result elsewhere — they're how downstream tooling sees what failed. | ||
|
|
||
| Canonical docs: <https://databrickslabs.github.io/dqx/docs/guide/quality_checks_apply>. | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.