Add asset graph integration for partitioned asset checks

This PR integrates partitioned asset checks with the asset graph system:

- Update AssetGraph to handle partitioned check nodes
- Add subset computation logic in AssetGraphView for check partitions
- Implement compute_subset_with_status for different execution statuses
- Update entity_subset to support asset check partition subsets
- Tests validate subset computation works correctly for partitioned and non-partitioned checks

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: prha

python_modules/dagster/dagster/_core/asset_graph_view/asset_graph_view.py

@@ -18,7 +18,6 @@
 from dagster._core.definitions.asset_key import AssetCheckKey, AssetKey, EntityKey, T_EntityKey
 from dagster._core.definitions.assets.graph.asset_graph_subset import AssetGraphSubset
 from dagster._core.definitions.events import AssetKeyPartitionKey
-from dagster._core.definitions.freshness import FreshnessState
 from dagster._core.definitions.partitions.context import (
     PartitionLoadingContext,
     use_partition_loading_context,
@@ -162,10 +161,7 @@ def get_inner_queryer_for_back_compat(self) -> "CachingInstanceQueryer":
         return self._queryer
 
     def _get_partitions_def(self, key: T_EntityKey) -> Optional["PartitionsDefinition"]:
-        if isinstance(key, AssetKey):
-            return self.asset_graph.get(key).partitions_def
-        else:
-            return None
+        return self.asset_graph.get(key).partitions_def
 
     @cached_method
     @use_partition_loading_context
@@ -550,6 +546,11 @@ async def compute_subset_with_status(
         from dagster._core.storage.event_log.base import AssetCheckSummaryRecord
 
         """Returns the subset of an asset check that matches a given status."""
+        # Handle partitioned asset checks with partition-level granularity
+        if self._get_partitions_def(key):
+            return await self._get_partitioned_check_subset_with_status(key, status)
+
+        # Handle non-partitioned asset checks with existing logic
         summary = await AssetCheckSummaryRecord.gen(self, key)
         latest_record = summary.last_check_execution_record if summary else None
         resolved_status = (
@@ -562,26 +563,6 @@ async def compute_subset_with_status(
         else:
             return self.get_empty_subset(key=key)
 
-    async def compute_subset_with_freshness_state(
-        self, key: AssetKey, state: FreshnessState
-    ) -> EntitySubset[AssetKey]:
-        from dagster._core.definitions.asset_health.asset_freshness_health import (
-            AssetFreshnessHealthState,
-        )
-
-        if not self.asset_graph.has(key) or self.asset_graph.get(key).freshness_policy is None:
-            if state == FreshnessState.NOT_APPLICABLE:
-                return self.get_full_subset(key=key)
-            else:
-                return self.get_empty_subset(key=key)
-
-        asset_freshness_health_state = await AssetFreshnessHealthState.compute_for_asset(key, self)
-
-        if asset_freshness_health_state.freshness_state == state:
-            return self.get_full_subset(key=key)
-        else:
-            return self.get_empty_subset(key=key)
-
     async def _compute_run_in_progress_check_subset(
         self, key: AssetCheckKey
     ) -> EntitySubset[AssetCheckKey]:
@@ -609,6 +590,31 @@ async def _compute_missing_check_subset(
     ) -> EntitySubset[AssetCheckKey]:
         return await self.compute_subset_with_status(key, None)
 
+    async def _get_partitioned_check_subset_with_status(
+        self, key: AssetCheckKey, status: Optional["AssetCheckExecutionResolvedStatus"]
+    ) -> EntitySubset[AssetCheckKey]:
+        """Get subset of partitioned asset check matching specific status."""
+        from dagster._core.storage.asset_check_partition_cache import (
+            get_asset_check_partition_status,
+        )
+
+        check_node = self.asset_graph.get(key)
+        if not check_node or not check_node.partitions_def:
+            check.failed(f"Asset check {key} not found or not partitioned.")
+
+        # Get partition status from our cache service
+        partition_status = get_asset_check_partition_status(
+            self._queryer.instance, key, check_node.partitions_def
+        )
+
+        matching_subset = (
+            partition_status.get_subset_for_status(status) if status else partition_status.missing
+        )
+        if not matching_subset.is_empty:
+            return EntitySubset(self, key=key, value=_ValidatedEntitySubsetValue(matching_subset))
+        else:
+            return self.get_empty_subset(key=key)
+
     async def _compute_run_in_progress_asset_subset(self, key: AssetKey) -> EntitySubset[AssetKey]:
         from dagster._core.storage.partition_status_cache import AssetStatusCacheValue
 

python_modules/dagster/dagster/_core/asset_graph_view/entity_subset.py

@@ -69,10 +69,13 @@ def convert_to_serializable_subset(self) -> SerializableEntitySubset[T_EntityKey
         return SerializableEntitySubset(key=self._key, value=self._value)
 
     def expensively_compute_partition_keys(self) -> AbstractSet[str]:
-        return {
-            check.not_none(akpk.partition_key, "No None partition keys")
-            for akpk in self.expensively_compute_asset_partitions()
-        }
+        internal_value = self.get_internal_value()
+        if isinstance(internal_value, PartitionsSubset):
+            return set(internal_value.get_partition_keys())
+        elif internal_value:
+            check.failed("No None partition keys")
+        else:
+            return set()
 
     def expensively_compute_asset_partitions(self) -> AbstractSet[AssetKeyPartitionKey]:
         if not isinstance(self.key, AssetKey):
@@ -106,13 +109,25 @@ def compute_intersection(self, other: Self) -> Self:
         return self._oper(other, operator.and_)
 
     def compute_intersection_with_partition_keys(
-        self: "EntitySubset[AssetKey]", partition_keys: AbstractSet[str]
-    ) -> "EntitySubset[AssetKey]":
-        key = check.inst(self.key, AssetKey)
-        partition_subset = self._asset_graph_view.get_asset_subset_from_asset_partitions(
-            self.key, {AssetKeyPartitionKey(key, pk) for pk in partition_keys}
+        self: "EntitySubset[T_EntityKey]", partition_keys: AbstractSet[str]
+    ) -> "EntitySubset[T_EntityKey]":
+        # Get partitions definition for the entity
+        if not self.partitions_def:
+            # Non-partitioned entity - return self if partition keys provided, else empty
+            return self if partition_keys else self._asset_graph_view.get_empty_subset(key=self.key)
+
+        # Create partition subset directly from partition keys
+        partition_subset = self.partitions_def.subset_with_partition_keys(partition_keys)
+
+        # Create EntitySubset with the partition subset
+        entity_subset_with_partitions = EntitySubset(
+            self._asset_graph_view,
+            key=self.key,
+            value=_ValidatedEntitySubsetValue(inner=partition_subset),
         )
-        return self.compute_intersection(partition_subset)
+
+        # Return intersection of this subset with the partition-based subset
+        return self.compute_intersection(entity_subset_with_partitions)
 
     @cached_method
     def compute_parent_subset(self, parent_key: AssetKey) -> "EntitySubset[AssetKey]":

python_modules/dagster/dagster/_core/definitions/assets/graph/asset_graph.py

@@ -190,6 +190,7 @@ def __init__(
                 v.get_spec_for_check_key(k).description,
                 v.get_spec_for_check_key(k).automation_condition,
                 v.get_spec_for_check_key(k).metadata,
+                v.get_spec_for_check_key(k).partitions_def,
             )
             for k, v in assets_defs_by_check_key.items()
         }

python_modules/dagster/dagster/_core/definitions/assets/graph/base_asset_graph.py

@@ -217,13 +217,15 @@ def __init__(
         description: Optional[str],
         automation_condition: Optional["AutomationCondition[AssetCheckKey]"],
         metadata: ArbitraryMetadataMapping,
+        partitions_def: Optional[PartitionsDefinition],
     ):
         self.key = key
         self.blocking = blocking
         self._automation_condition = automation_condition
         self._additional_deps = additional_deps
         self._description = description
         self._metadata = metadata
+        self._partitions_def = partitions_def
 
     @property
     def parent_entity_keys(self) -> AbstractSet[AssetKey]:
@@ -235,8 +237,7 @@ def child_entity_keys(self) -> AbstractSet[EntityKey]:
 
     @property
     def partitions_def(self) -> Optional[PartitionsDefinition]:
-        # all checks are unpartitioned
-        return None
+        return self._partitions_def
 
     @property
     def partition_mappings(self) -> Mapping[EntityKey, PartitionMapping]:

python_modules/dagster/dagster/_core/definitions/assets/graph/remote_asset_graph.py

@@ -40,10 +40,9 @@
 from dagster._core.definitions.metadata import ArbitraryMetadataMapping
 from dagster._core.definitions.partitions.definition import PartitionsDefinition
 from dagster._core.definitions.partitions.mapping import PartitionMapping
-from dagster._core.definitions.selector import ScheduleSelector, SensorSelector
 from dagster._core.definitions.utils import DEFAULT_GROUP_NAME
 from dagster._core.remote_representation.external import RemoteRepository
-from dagster._core.remote_representation.handle import RepositoryHandle
+from dagster._core.remote_representation.handle import InstigatorHandle, RepositoryHandle
 from dagster._core.workspace.workspace import CurrentWorkspace
 from dagster._record import ImportFrom, record
 from dagster._utils.cached_method import cached_method
@@ -396,33 +395,31 @@ def resolve_to_singular_repo_scoped_node(self) -> "RemoteRepositoryAssetNode":
             )
         )
 
-    def get_targeting_schedule_selectors(
+    def get_targeting_schedule_handles(
         self,
-    ) -> Sequence[ScheduleSelector]:
+    ) -> Sequence[InstigatorHandle]:
         selectors = []
         for node in self.repo_scoped_asset_infos:
             for schedule_name in node.targeting_schedule_names:
                 selectors.append(
-                    ScheduleSelector(
-                        location_name=node.handle.location_name,
-                        repository_name=node.handle.repository_name,
-                        schedule_name=schedule_name,
+                    InstigatorHandle(
+                        repository_handle=node.handle,
+                        instigator_name=schedule_name,
                     )
                 )
 
         return selectors
 
-    def get_targeting_sensor_selectors(
+    def get_targeting_sensor_handles(
         self,
-    ) -> Sequence[SensorSelector]:
+    ) -> Sequence[InstigatorHandle]:
         selectors = []
         for node in self.repo_scoped_asset_infos:
             for sensor_name in node.targeting_sensor_names:
                 selectors.append(
-                    SensorSelector(
-                        location_name=node.handle.location_name,
-                        repository_name=node.handle.repository_name,
-                        sensor_name=sensor_name,
+                    InstigatorHandle(
+                        repository_handle=node.handle,
+                        instigator_name=sensor_name,
                     )
                 )
         return selectors
@@ -480,6 +477,9 @@ def _get_asset_check_node_from_remote_asset_check_node(
             remote_node.asset_check.description,
             remote_node.asset_check.automation_condition,
             {},  # metadata not yet on AssetCheckNodeSnap
+            remote_node.asset_check.partitions_def_snapshot.get_partitions_definition()
+            if remote_node.asset_check.partitions_def_snapshot
+            else None,
         )
 
     ##### COMMON ASSET GRAPH INTERFACE

python_modules/dagster/dagster_tests/storage_tests/test_asset_checks.py

@@ -83,3 +83,97 @@ def test_partitioned_asset_check_graph_structure():
     # Test: check is linked to asset
     asset_node = asset_graph.get(partitioned_asset.key)
     assert partitioned_asset_check.check_key in asset_node.check_keys
+
+
+@pytest.mark.asyncio
+async def test_partitioned_asset_check_subset_computation_empty():
+    """Test subset computation for partitioned asset checks before any executions."""
+    from dagster._core.asset_graph_view.asset_graph_view import AssetGraphView
+    from dagster._core.storage.asset_check_execution_record import AssetCheckExecutionResolvedStatus
+
+    with dg.instance_for_test() as instance:
+        view = AssetGraphView.for_test(partitioned_defs, instance=instance)
+        partitioned_key = partitioned_asset_check.check_key
+
+        # Before executions, all partitions should be missing
+        missing_subset = await view.compute_subset_with_status(partitioned_key, None)
+        assert not missing_subset.is_empty
+        assert missing_subset.expensively_compute_partition_keys() == {"a", "b", "c"}
+
+        # No partitions should have execution statuses
+        for status in [
+            AssetCheckExecutionResolvedStatus.SUCCEEDED,
+            AssetCheckExecutionResolvedStatus.FAILED,
+            AssetCheckExecutionResolvedStatus.IN_PROGRESS,
+        ]:
+            subset = await view.compute_subset_with_status(partitioned_key, status)
+            assert subset.is_empty
+
+
+@pytest.mark.asyncio
+async def test_partitioned_asset_check_subset_computation_after_execution():
+    """Test subset computation after executing partitions with different outcomes."""
+    from dagster._core.asset_graph_view.asset_graph_view import AssetGraphView
+    from dagster._core.storage.asset_check_execution_record import AssetCheckExecutionResolvedStatus
+
+    with dg.instance_for_test() as instance:
+        # Execute partitions with known outcomes
+        result_a = dg.materialize(
+            [partitioned_asset, partitioned_asset_check], instance=instance, partition_key="a"
+        )
+        assert result_a.success
+
+        result_b = dg.materialize(
+            [partitioned_asset, partitioned_asset_check], instance=instance, partition_key="b"
+        )
+        assert result_b.success  # Run succeeds but check fails
+
+        view = AssetGraphView.for_test(partitioned_defs, instance=instance)
+        partitioned_key = partitioned_asset_check.check_key
+
+        # Test: subsets reflect execution outcomes
+        succeeded_subset = await view.compute_subset_with_status(
+            partitioned_key, AssetCheckExecutionResolvedStatus.SUCCEEDED
+        )
+        assert not succeeded_subset.is_empty
+        assert succeeded_subset.expensively_compute_partition_keys() == {"a"}
+
+        failed_subset = await view.compute_subset_with_status(
+            partitioned_key, AssetCheckExecutionResolvedStatus.FAILED
+        )
+        assert not failed_subset.is_empty  # Should contain 'b'
+        assert failed_subset.expensively_compute_partition_keys() == {"b"}
+
+        missing_subset = await view.compute_subset_with_status(partitioned_key, None)
+        assert not missing_subset.is_empty  # Should contain 'c'
+        assert missing_subset.expensively_compute_partition_keys() == {"c"}
+
+
+@pytest.mark.asyncio
+async def test_non_partitioned_asset_check_compatibility():
+    """Test that non-partitioned asset checks still work with existing logic."""
+    from dagster._core.asset_graph_view.asset_graph_view import AssetGraphView
+    from dagster._core.storage.asset_check_execution_record import AssetCheckExecutionResolvedStatus
+
+    with dg.instance_for_test() as instance:
+        view = AssetGraphView.for_test(defs, instance=instance)
+        non_partitioned_key = the_asset_check.check_key
+
+        missing_subset = await view.compute_subset_with_status(non_partitioned_key, None)
+        assert not missing_subset.is_partitioned
+
+        succeeded_subset = await view.compute_subset_with_status(
+            non_partitioned_key, AssetCheckExecutionResolvedStatus.SUCCEEDED
+        )
+        assert not succeeded_subset.is_partitioned
+
+        # Execute the check
+        result = dg.materialize([the_asset, the_asset_check], instance=instance)
+        assert result.success
+
+        # After execution
+        view_after = AssetGraphView.for_test(defs, instance=instance)
+        succeeded_subset_after = await view_after.compute_subset_with_status(
+            non_partitioned_key, AssetCheckExecutionResolvedStatus.SUCCEEDED
+        )
+        assert not succeeded_subset_after.is_partitioned