Onboarding MaterializationSet from Feature Store to DAB

acceptance/bundle/invariant/configs/materialization_set.yml.tmpl

@@ -0,0 +1,34 @@
+bundle:
+  name: test-bundle-$UNIQUE_NAME
+
+resources:
+  features:
+    foo:
+      full_name: main.matset_inv_$UNIQUE_NAME.test_feature
+      source:
+        delta_table_source:
+          full_name: main.matset_inv_$UNIQUE_NAME.source_table
+      function:
+        aggregation_function:
+          count_function:
+            input: value
+          time_window:
+            sliding:
+              window_duration: 604800s
+              slide_duration: 86400s
+      entities:
+        - name: id
+      timeseries_column:
+        name: timestamp
+
+  materialization_sets:
+    foo:
+      materialized_features:
+        - feature_name: ${resources.features.foo.full_name}
+          offline_store_config:
+            catalog_name: main
+            schema_name: matset_inv_$UNIQUE_NAME
+            table_name_prefix: offline
+          cron_schedule_trigger:
+            mode: MANUAL
+            cron_expression: "0 0 * * * ?"

acceptance/bundle/invariant/configs/materialization_set.yml.tmpl-cleanup.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+# The offline materialization leaves its own tables behind, so the schema is dropped with CASCADE.
+echo "Cleaning up temporary source table"
+$CLI tables delete main.matset_inv_$UNIQUE_NAME.source_table || true
+$CLI schemas delete main.matset_inv_$UNIQUE_NAME --force || true

acceptance/bundle/invariant/configs/materialization_set.yml.tmpl-init.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+# The feature's delta_table_source must already exist, and the invariant harness provisions no
+# tables of its own.
+echo "Creating temporary source table: main.matset_inv_$UNIQUE_NAME.source_table"
+
+$CLI schemas create matset_inv_$UNIQUE_NAME main -o json | jq '{full_name}'
+
+# MSYS_NO_PATHCONV=1 prevents Git Bash on Windows from converting /api/... to a filesystem path.
+MSYS_NO_PATHCONV=1 $CLI api post "/api/2.0/sql/statements/" --json "{
+    \"warehouse_id\": \"$TEST_DEFAULT_WAREHOUSE_ID\",
+    \"statement\": \"CREATE TABLE main.matset_inv_$UNIQUE_NAME.source_table (id INT, value STRING, timestamp TIMESTAMP) USING DELTA\",
+    \"wait_timeout\": \"45s\"
+  }" > /dev/null

acceptance/bundle/python/materialization_sets-support/databricks.yml

@@ -0,0 +1,22 @@
+bundle:
+  name: my_project
+
+sync: {paths: []} # don't need to copy files
+
+python:
+  resources:
+    - "resources:load_resources"
+  mutators:
+    - "mutators:update_materialization_set"
+
+resources:
+  materialization_sets:
+    my_materialization_set_1:
+      materialized_features:
+        - feature_name: main.default.my_feature_1
+          offline_store_config:
+            catalog_name: main
+            schema_name: default
+            table_name_prefix: my_feature_1_off
+          cron_schedule_trigger:
+            mode: DERIVED

acceptance/bundle/python/materialization_sets-support/mutators.py

@@ -0,0 +1,23 @@
+from dataclasses import replace
+
+from databricks.bundles.core import materialization_set_mutator
+from databricks.bundles.materialization_sets import MaterializationSet
+
+
+@materialization_set_mutator
+def update_materialization_set(
+    materialization_set: MaterializationSet,
+) -> MaterializationSet:
+    # A set has no string field of its own, so the member's feature name is what a mutator can
+    # observably change.
+    members = list(materialization_set.materialized_features)
+    first = members[0]
+    assert isinstance(first.feature_name, str)
+
+    return replace(
+        materialization_set,
+        materialized_features=[
+            replace(first, feature_name=f"{first.feature_name}_updated"),
+            *members[1:],
+        ],
+    )

acceptance/bundle/python/materialization_sets-support/out.test.toml

@@ -0,0 +1,3 @@
+Cloud = false
+EnvMatrix.DMS = ["", "true"]
+EnvMatrix.PYDAB_VERSION = ["current"]

acceptance/bundle/python/materialization_sets-support/output.txt

@@ -0,0 +1,48 @@
+
+>>> uv run [UV_ARGS] -q [CLI] bundle validate --output json
+{
+  "experimental": {
+    "python": {
+      "mutators": [
+        "mutators:update_materialization_set"
+      ],
+      "resources": [
+        "resources:load_resources"
+      ]
+    }
+  },
+  "resources": {
+    "materialization_sets": {
+      "my_materialization_set_1": {
+        "materialized_features": [
+          {
+            "cron_schedule_trigger": {
+              "mode": "DERIVED"
+            },
+            "feature_name": "main.default.my_feature_1_updated",
+            "offline_store_config": {
+              "catalog_name": "main",
+              "schema_name": "default",
+              "table_name_prefix": "my_feature_1_off"
+            }
+          }
+        ]
+      },
+      "my_materialization_set_2": {
+        "materialized_features": [
+          {
+            "cron_schedule_trigger": {
+              "mode": "DERIVED"
+            },
+            "feature_name": "main.default.my_feature_2_updated",
+            "offline_store_config": {
+              "catalog_name": "main",
+              "schema_name": "default",
+              "table_name_prefix": "my_feature_2_off"
+            }
+          }
+        ]
+      }
+    }
+  }
+}

acceptance/bundle/python/materialization_sets-support/resources.py

@@ -0,0 +1,24 @@
+from databricks.bundles.core import Resources
+
+
+def load_resources() -> Resources:
+    resources = Resources()
+
+    resources.add_materialization_set(
+        "my_materialization_set_2",
+        {
+            "materialized_features": [
+                {
+                    "feature_name": "main.default.my_feature_2",
+                    "offline_store_config": {
+                        "catalog_name": "main",
+                        "schema_name": "default",
+                        "table_name_prefix": "my_feature_2_off",
+                    },
+                    "cron_schedule_trigger": {"mode": "DERIVED"},
+                },
+            ],
+        },
+    )
+
+    return resources

acceptance/bundle/python/materialization_sets-support/script

@@ -0,0 +1,5 @@
+
+trace uv run $UV_ARGS -q $CLI bundle validate --output json | \
+  jq "pick(.experimental.python, .resources)"
+
+rm -fr .databricks __pycache__

acceptance/bundle/python/materialization_sets-support/test.toml

@@ -0,0 +1,4 @@
+Cloud = false # tests don't interact with APIs
+
+# materialization sets are only supported in the current version of the wheel
+EnvMatrix.PYDAB_VERSION = ["current"]

acceptance/bundle/refschema/output/materialization_sets.txt

@@ -0,0 +1,58 @@
+resources.materialization_sets.*.id	string	INPUT
+resources.materialization_sets.*.lifecycle	resources.Lifecycle	INPUT
+resources.materialization_sets.*.lifecycle.depends_on	[]resources.DependsOn	INPUT
+resources.materialization_sets.*.lifecycle.depends_on[*]	resources.DependsOn	INPUT
+resources.materialization_sets.*.lifecycle.depends_on[*].resource	string	INPUT
+resources.materialization_sets.*.lifecycle.prevent_destroy	bool	INPUT
+resources.materialization_sets.*.materialization_set_id	string	ALL
+resources.materialization_sets.*.materialized_features	[]ml.MaterializedFeature	ALL
+resources.materialization_sets.*.materialized_features[*]	ml.MaterializedFeature	ALL
+resources.materialization_sets.*.materialized_features[*].budget_policy_id	string	ALL
+resources.materialization_sets.*.materialized_features[*].cron_schedule	string	ALL
+resources.materialization_sets.*.materialized_features[*].cron_schedule_trigger	*ml.CronSchedule	ALL
+resources.materialization_sets.*.materialized_features[*].cron_schedule_trigger.cron_expression	string	ALL
+resources.materialization_sets.*.materialized_features[*].cron_schedule_trigger.mode	ml.CronScheduleMode	ALL
+resources.materialization_sets.*.materialized_features[*].cron_schedule_trigger.timezone_id	string	ALL
+resources.materialization_sets.*.materialized_features[*].feature_name	string	ALL
+resources.materialization_sets.*.materialized_features[*].is_online	bool	ALL
+resources.materialization_sets.*.materialized_features[*].job_id	int64	ALL
+resources.materialization_sets.*.materialized_features[*].last_materialization_time	string	ALL
+resources.materialization_sets.*.materialized_features[*].latest_backfill_operation	string	ALL
+resources.materialization_sets.*.materialized_features[*].materialization_set_id	string	ALL
+resources.materialization_sets.*.materialized_features[*].materialized_feature_id	string	ALL
+resources.materialization_sets.*.materialized_features[*].notifications	*ml.MaterializationNotifications	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_failure	*ml.MaterializationFailureNotification	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_failure.destination_ids	[]string	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_failure.destination_ids[*]	string	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_failure.email_addresses	[]string	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_failure.email_addresses[*]	string	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_failure.final_attempt_only	bool	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_success	*ml.MaterializationSuccessNotification	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_success.destination_ids	[]string	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_success.destination_ids[*]	string	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_success.email_addresses	[]string	ALL
+resources.materialization_sets.*.materialized_features[*].notifications.on_success.email_addresses[*]	string	ALL
+resources.materialization_sets.*.materialized_features[*].offline_store_config	*ml.OfflineStoreConfig	ALL
+resources.materialization_sets.*.materialized_features[*].offline_store_config.catalog_name	string	ALL
+resources.materialization_sets.*.materialized_features[*].offline_store_config.schema_name	string	ALL
+resources.materialization_sets.*.materialized_features[*].offline_store_config.table_name_prefix	string	ALL
+resources.materialization_sets.*.materialized_features[*].online_store_config	*ml.OnlineStoreConfig	ALL
+resources.materialization_sets.*.materialized_features[*].online_store_config.catalog_name	string	ALL
+resources.materialization_sets.*.materialized_features[*].online_store_config.online_store_name	string	ALL
+resources.materialization_sets.*.materialized_features[*].online_store_config.schema_name	string	ALL
+resources.materialization_sets.*.materialized_features[*].online_store_config.table_name_prefix	string	ALL
+resources.materialization_sets.*.materialized_features[*].pipeline_id	string	ALL
+resources.materialization_sets.*.materialized_features[*].pipeline_schedule_state	ml.MaterializedFeaturePipelineScheduleState	ALL
+resources.materialization_sets.*.materialized_features[*].streaming_mode	*ml.StreamingMode	ALL
+resources.materialization_sets.*.materialized_features[*].streaming_mode.freshness_target	string	ALL
+resources.materialization_sets.*.materialized_features[*].streaming_mode.mode	ml.StreamingModeStreamingModeType	ALL
+resources.materialization_sets.*.materialized_features[*].streaming_mode.shuffle_partitions	int	ALL
+resources.materialization_sets.*.materialized_features[*].table_name	string	ALL
+resources.materialization_sets.*.materialized_features[*].table_trigger	*ml.TableTrigger	ALL
+resources.materialization_sets.*.materialized_features[*].table_trigger.condition	ml.TableTriggerTableTriggerCondition	ALL
+resources.materialization_sets.*.materialized_features[*].table_trigger.min_time_between_triggers	*duration.Duration	ALL
+resources.materialization_sets.*.materialized_features[*].table_trigger.wait_after_last_change	*duration.Duration	ALL
+resources.materialization_sets.*.materialized_features[*].tags	map[string]string	ALL
+resources.materialization_sets.*.materialized_features[*].tags.*	string	ALL
+resources.materialization_sets.*.modified_status	string	INPUT
+resources.materialization_sets.*.url	string	INPUT