AI-Assisted Primary Key Detection and Uniqueness Rules Generation

[jominjohny]

Changes

• Added AI-Assisted (LLM) Primary Key (PK) detection. Added detect_primary_keys_with_llm method to the DQProfiler.
• Added AI-Assisted (LLM) uniqueness rules generation. Extended the DQGenerator to leverage uniqueness profiles from the profiler.

BREAKING CHANGES!

• Renamed level parameter to criticality in generate_dq_rules method in DQGenerator for consistency.
• Replaced table: str parameter with input_config: InputConfig in profile_table method of DQProfiler for greater flexibility.

Linked issues

Resolves #484

Tests

• manually tested
• added unit tests
• added integration tests
• added end-to-end tests
• added performance tests

[github-actions]

✅ 443/443 passed, 1 flaky, 41 skipped, 3h4m20s total

Flaky tests:

• 🤪 test_e2e_workflow_serverless (9m9.07s)

_{Running from acceptance #3280}

[Copilot]

Pull Request Overview

This PR introduces LLM-based primary key detection functionality for the dqx library, enabling automatic identification of primary key columns in database tables using AI analysis.

Key changes:

• Added DatabricksPrimaryKeyDetector class for LLM-powered primary key detection with duplicate validation and retry logic
• Integrated PK detection into the profiler workflow via ProfilerRunner
• Created compare_datasets_with_llm wrapper function that auto-detects primary keys for dataset comparison

Reviewed Changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
src/databricks/labs/dqx/llm/llm_pk_detector.py	Core implementation of DSPy-based primary key detector with Databricks integration
src/databricks/labs/dqx/profiler/generator.py	Added detect_primary_keys_with_llm method to DQGenerator
src/databricks/labs/dqx/profiler/profiler.py	Added convenience method for PK detection via DQProfiler
src/databricks/labs/dqx/profiler/profiler_runner.py	Integrated automatic PK detection into profiler workflow
src/databricks/labs/dqx/check_funcs.py	Added wrapper function for dataset comparison with auto PK detection
src/databricks/labs/dqx/llm/init.py	Updated LLM dependencies to include langchain packages
tests/unit/test_llm_pk_integration.py	Unit tests for LLM PK detection functionality
tests/integration/test_llm_pk_detection.py	Integration tests for end-to-end PK detection
pyproject.toml	Updated dependencies and pylint configuration
demos/dqx_demo_llm_pk_detection.py	Demo notebook showcasing PK detection capabilities

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[mwojtyczka]

LGTM

[Copilot]

Pull request overview

Copilot reviewed 45 out of 45 changed files in this pull request and generated 3 comments.

Comments suppressed due to low confidence (1)

tests/integration/test_dlt_rules_generator.py:1

• Import path tests.integration.test_generator is incorrect. The module should be tests.integration.test_rules_generator based on the file structure visible in the diffs.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review for a chance to win a $100 gift card. Take the survey.

[Copilot]

Pull request overview

Copilot reviewed 45 out of 45 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (2)

tests/integration/test_dlt_rules_generator.py:1

• Import path is incorrect. The module test_generator does not exist - it should be test_rules_generator based on the file being renamed from.
  tests/unit/test_llm_pk_detector.py:1
• These helper functions _create_mock_detector_result and _create_mock_spark_results are defined but never used in any of the tests. Consider removing them to reduce code clutter, or add tests that utilize these helpers.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review for a chance to win a $100 gift card. Take the survey.