Skip to content

Commit

Permalink
Fix issue where dark_storage did not handle empty responses
Browse files Browse the repository at this point in the history
  • Loading branch information
eivindjahren committed Nov 19, 2024
1 parent df2874a commit 71ec78c
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 2 deletions.
3 changes: 2 additions & 1 deletion src/ert/dark_storage/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import pandas as pd
import polars
import xarray as xr
from polars.exceptions import ColumnNotFoundError

from ert.config import GenDataConfig, GenKwConfig
from ert.config.field import Field
Expand Down Expand Up @@ -184,7 +185,7 @@ def data_for_key(
return data.astype(float)
except ValueError:
return data
except (ValueError, KeyError):
except (ValueError, KeyError, ColumnNotFoundError):
return pd.DataFrame()

return pd.DataFrame()
Expand Down
32 changes: 31 additions & 1 deletion tests/ert/unit_tests/dark_storage/test_common.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import pandas as pd
import polars
import pytest

from ert.config import SummaryConfig
from ert.config import GenDataConfig, SummaryConfig
from ert.dark_storage.common import data_for_key
from ert.storage import open_storage
from tests.ert.unit_tests.config.summary_generator import (
Expand Down Expand Up @@ -70,3 +71,32 @@ def test_data_for_key_gives_mean_for_duplicate_values(tmp_path):
assert df[pd.Timestamp("2014-01-16 05:00:00")][0] == pytest.approx(
(value1 + value2) / 2
)


def test_data_for_key_returns_empty_gen_data_config(tmp_path):
with open_storage(tmp_path / "storage", mode="w") as storage:
gen_data_config = GenDataConfig(keys=["response"])
experiment = storage.create_experiment(
observations={},
parameters=[],
responses=[gen_data_config],
)
ensemble = experiment.create_ensemble(name="ensemble", ensemble_size=1)

data = data_for_key(ensemble, "response@0")
assert data.empty

ensemble.save_response(
"gen_data",
polars.DataFrame(
{
"response_key": "response",
"report_step": polars.Series([0], dtype=polars.UInt16),
"index": polars.Series([0], dtype=polars.UInt16),
"values": polars.Series([0.0], dtype=polars.Float32),
}
),
0,
)
data = data_for_key(ensemble, "response@0")
assert not data.empty

0 comments on commit 71ec78c

Please sign in to comment.