Skip to content

refactor: migrate from FeatureStoreClient to FeatureEngineeringClient#222

Open
veenaramesh wants to merge 4 commits into
databricks:testfrom
veenaramesh:refactor/upgrade-feature-engineering
Open

refactor: migrate from FeatureStoreClient to FeatureEngineeringClient#222
veenaramesh wants to merge 4 commits into
databricks:testfrom
veenaramesh:refactor/upgrade-feature-engineering

Conversation

@veenaramesh

Copy link
Copy Markdown
Collaborator

Summary

  • Replace all FeatureStoreClient with FeatureEngineeringClient (the current recommended API)
  • Replace from databricks.feature_store with from databricks.feature_engineering
  • Fix score_batch call signatures to use keyword args (model_uri=, df=)
  • Fix dynamic feature module import path for flattened structure

Depends on #219 (simplify-project-structure).

Test plan

  • Feature engineering job (write_feature_table) passes on e2-demo-field-eng
  • Model training pipeline (Train, ModelValidation, ModelDeployment) passes with feature store
  • Batch inference with feature store (pre-existing timestamp column issue, not related to this change)

Move all Databricks notebooks out of nested notebooks/ subdirectories
to sit alongside their sibling helper modules. This eliminates the
%cd .. hacks that were needed to import from parent directories and
simplifies the project layout for a data science workflow.

Moves:
- training/notebooks/*.py → training/*.py
- validation/notebooks/ModelValidation.py → validation/ModelValidation.py
- deployment/*/notebooks/*.py → deployment/*/*.py
- feature_engineering/notebooks/*.py → feature_engineering/*.py
- monitoring/notebooks/*.py → monitoring/*.py

Updates all resource YAMLs, update_layout.tmpl, and README references.

Co-authored-by: Isaac
Databricks bundle deployments set the notebook CWD correctly, so
the notebook_path/%cd $notebook_path pattern is unnecessary. Sibling
imports work naturally after the flatten-notebooks change.

Also inlines the trivial get_deployed_model_stage_for_env helper
(a dict lookup) to eliminate cross-directory sys.path.append in
deploy.py and BatchInference.py.

Co-authored-by: Isaac
No longer needed with databricks bundle init.

Co-authored-by: Isaac
Migrate all feature store code from the legacy databricks.feature_store
(FeatureStoreClient) to databricks.feature_engineering
(FeatureEngineeringClient). Remove UC conditional branches that
duplicated the same logic with different imports.

Co-authored-by: Isaac
@veenaramesh veenaramesh force-pushed the refactor/upgrade-feature-engineering branch from d60f471 to 37013e8 Compare June 12, 2026 17:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant