Use correct grain

metricflow/dataflow/builder/dataflow_plan_builder.py

@@ -1543,19 +1543,20 @@ def _build_aggregated_measure_from_measure_source_node(
         if non_additive_dimension_spec is not None:
             # Apply semi additive join on the node
             agg_time_dimension = measure_properties.agg_time_dimension
+            non_additive_dimension_grain = self._semantic_model_lookup.get_defined_time_granularity(
+                TimeDimensionReference(non_additive_dimension_spec.name)
+            )
+
             queried_time_dimension_spec: Optional[
                 TimeDimensionSpec
             ] = self._find_non_additive_dimension_in_linkable_specs(
                 agg_time_dimension=agg_time_dimension,
                 linkable_specs=queried_linkable_specs.as_tuple,
                 non_additive_dimension_spec=non_additive_dimension_spec,
             )
-            time_dimension_spec = TimeDimensionSpec.from_name(non_additive_dimension_spec.name)
-            if queried_time_dimension_spec:
-                time_dimension_spec = time_dimension_spec.with_grain_and_date_part(
-                    time_granularity=queried_time_dimension_spec.time_granularity,
-                    date_part=queried_time_dimension_spec.date_part,
-                )
+            time_dimension_spec = TimeDimensionSpec.from_name(non_additive_dimension_spec.name).with_grain(
+                time_granularity=non_additive_dimension_grain
+            )
             window_groupings = tuple(
                 LinklessEntitySpec.from_element_name(name) for name in non_additive_dimension_spec.window_groupings
             )