diff --git a/mesa_frames/concrete/polars/agentset.py b/mesa_frames/concrete/polars/agentset.py index 4bec1ea5..85b2d498 100644 --- a/mesa_frames/concrete/polars/agentset.py +++ b/mesa_frames/concrete/polars/agentset.py @@ -143,7 +143,13 @@ def add( if isinstance(obj._mask, pl.Series): original_active_indices = obj._agents.filter(obj._mask)["unique_id"] - obj._agents = pl.concat([obj._agents, new_agents], how="diagonal_relaxed") + combined_agents = pl.concat([obj._agents, new_agents], how="diagonal_relaxed") + if combined_agents["unique_id"].is_duplicated().any(): + raise ValueError( + "Some ids are duplicated in the AgentSet that are trying to be added together." + ) + + obj._agents = combined_agents if isinstance(obj._mask, pl.Series): obj._update_mask(original_active_indices, new_agents["unique_id"]) diff --git a/tests/polars/test_agentset_polars.py b/tests/polars/test_agentset_polars.py index 9c311727..180a5b9d 100644 --- a/tests/polars/test_agentset_polars.py +++ b/tests/polars/test_agentset_polars.py @@ -32,6 +32,17 @@ def fix1_AgentSetPolars() -> ExampleAgentSetPolars: return agents +@pytest.fixture +def fix4_AgentSetPolars() -> ExampleAgentSetPolars: + model = ModelDF() + agents = ExampleAgentSetPolars(model) + agents.add({"unique_id": [0, 1, 2, 3]}) + agents["wealth"] = agents.starting_wealth + agents["age"] = [10, 20, 30, 40] + model.agents.add(agents) + return agents + + @pytest.fixture def fix2_AgentSetPolars() -> ExampleAgentSetPolars: model = ModelDF() @@ -73,9 +84,17 @@ def test_add( self, fix1_AgentSetPolars: ExampleAgentSetPolars, fix2_AgentSetPolars: ExampleAgentSetPolars, + fix4_AgentSetPolars: ExampleAgentSetPolars, ): agents = fix1_AgentSetPolars agents2 = fix2_AgentSetPolars + agents4 = fix4_AgentSetPolars + + with pytest.raises( + ValueError, + match="Some ids are duplicated in the AgentSet that are trying to be added together.", + ): + result = agents.add(agents4.agents, inplace=False) # Test with a DataFrame result = agents.add(agents2.agents, inplace=False)