diff --git a/design_notes.md b/design_notes.md new file mode 100644 index 0000000..9bba04b --- /dev/null +++ b/design_notes.md @@ -0,0 +1,31 @@ +# Questions + +* How are you supposed to interact with the Qube + + * I think the compressed nodes and structure should be hidden from the user as much as possible + + * Select/filter to return similar Qubes makes sense + + * Would a "path" make sense? + qube[("class", &["od"])][("type", &["fc"])] + Not sure how it would work with coordinates. These don't point to a whole node. Would need to store coordinates along the path. + It would prevent nodes from being visible to the user. + + is the path really any different to a filter though? + it stores the route separately to the result, which is quite nice perhaps + + a path suggests you can only navigate down the tree, whereas a filter does not care about the structure of the tree + in other words a path reveals the hierarchy, whereas the qube does not really need to show that + + A path does not make sense. Treat the Qube not as a tree from outside. + + * Methods should include: + sel + dims() + coords(dim) -> list of coordinates along that dim + + transpose (change internal ordering of axes) + align (align two Qubes to have the same axes?) + squeeze() removes length-1 dims. + drop_dim() + group_by() to group along a dimension by some function diff --git a/py_qubed/src/lib.rs b/py_qubed/src/lib.rs index 8ca3150..854e31c 100644 --- a/py_qubed/src/lib.rs +++ b/py_qubed/src/lib.rs @@ -1,16 +1,12 @@ use ::qubed::Coordinates; -use ::qubed::Datacube; use ::qubed::Qube; use ::qubed::select::SelectMode; use pyo3::exceptions::PyTypeError; -use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList, PyModule, PySet}; +use pyo3::types::{PyDict, PyList, PyModule}; use serde_json::Value as JsonValue; -mod helpers; - -#[pyclass(name = "Qube", unsendable)] +#[pyclass(unsendable)] pub struct PyQube { inner: Qube, } @@ -75,30 +71,6 @@ impl PyQube { } } - #[staticmethod] - #[pyo3(signature = (datacube, order=None))] - pub fn from_datacube( - datacube: Bound<'_, PyDict>, - order: Option>, - ) -> PyResult { - let dc = pydict_to_datacube(datacube)?; - let order_slices: Option<&[String]> = order.as_deref(); - Ok(Self { inner: Qube::from_datacube(&dc, order_slices) }) - } - - #[pyo3(signature = (datacube, order=None, accept_existing_order=false))] - pub fn append_datacube( - &mut self, - datacube: Bound<'_, PyDict>, - order: Option>, - accept_existing_order: bool, - ) -> PyResult<()> { - let dc = pydict_to_datacube(datacube)?; - let order_slices: Option<&[String]> = order.as_deref(); - self.inner.append_datacube(dc, order_slices, accept_existing_order); - Ok(()) - } - pub fn select( &self, request: Bound<'_, PyDict>, @@ -200,7 +172,7 @@ impl PyQube { let mut validated_qubes = Vec::with_capacity(others.len()); for item in others.iter() { let other_cell = - item.cast::().map_err(|_| PyTypeError::new_err("expected Qube"))?; + item.cast::().map_err(|_| PyTypeError::new_err("expected PyQube"))?; validated_qubes.push(other_cell.clone().unbind()); } @@ -216,290 +188,6 @@ impl PyQube { pub fn __repr__(&self) -> PyResult { Ok(format!("PyQube(root_id={:?})", self.inner.root())) } - - // ----------------------------------------------------------------------- - // New utility methods (mirrors of the @support_qubed_output functions - // in forecast-in-a-box/fiab_plugin_ecmwf/qubed_utils.py) - // ----------------------------------------------------------------------- - - /// Return all dimension names and their union of coordinate values. - /// - /// This is an alias for `all_unique_dim_coords()` and returns the same - /// ``dict[str, list[str]]`` mapping. The name `axes` matches the - /// terminology used in the qubed-utils helper library. - /// - /// ```python - /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) - /// ax = q.axes() - /// assert set(ax["param"]) == {"2t", "tp"} - /// assert set(ax["time"]) == {"0", "1", "2"} - /// ``` - pub fn axes(&self, py: Python<'_>) -> PyResult> { - let dim_coords = self.inner.all_unique_dim_coords(); - let py_dict = PyDict::new(py); - - for (dimension, coordinates) in dim_coords { - let coord_str = coordinates.to_string(); - let values: Vec<&str> = if coord_str.is_empty() { - vec![] - } else if coord_str.contains('/') { - coord_str.split('/').collect() - } else { - vec![&coord_str] - }; - - let py_list = PyList::empty(py); - for value in values { - py_list.append(value)?; - } - py_dict.set_item(dimension, py_list)?; - } - - Ok(py_dict.into_any().unbind()) - } - - /// Return the set of all dimension names present anywhere in the Qube. - /// - /// ```python - /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) - /// assert q.dimensions() == {"param", "time"} - /// ``` - pub fn dimensions(&self, py: Python<'_>) -> PyResult> { - let dims = self.inner.dimensions(); - let py_set = PySet::new(py, dims.iter())?; - Ok(py_set.into_any().unbind()) - } - - /// Return the set of dimension names present in **every** leaf path. - /// - /// For a Qube with uniform depth this equals `dimensions()`. When - /// branches differ in depth, only dimensions that appear in *all* - /// branches are returned. - /// - /// ```python - /// q1 = Qube.from_datacube({"param": "2t", "time": "0/1"}, ["param", "time"]) - /// q2 = Qube.from_datacube({"param": "msl"}, ["param"]) - /// q1.append(q2) - /// assert q1.common_dimensions() == {"param"} - /// ``` - pub fn common_dimensions(&self, py: Python<'_>) -> PyResult> { - let common = self.inner.common_dimensions(); - let py_set = PySet::new(py, common.iter())?; - Ok(py_set.into_any().unbind()) - } - - /// Wrap the entire Qube under one or more new outer dimensions. - /// - /// `dimension` is a ``dict[str, list]`` where each key becomes a new - /// dimension name and the corresponding list supplies its coordinate - /// values. Dimensions are applied in insertion order: the *last* entry - /// in the dict becomes the outermost dimension. - /// - /// The operation mutates the Qube in place and does **not** return a new - /// object (unlike the Python-reference implementation which follows an - /// immutable style). - /// - /// ```python - /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) - /// q.expand({"ensemble": ["ens1", "ens2"]}) - /// assert "ensemble" in q.dimensions() - /// assert "param" in q.dimensions() - /// assert "time" in q.dimensions() - /// ``` - pub fn expand(&mut self, dimension: Bound<'_, PyDict>) -> PyResult<()> { - for (k, v) in dimension.iter() { - let key: String = k - .extract() - .map_err(|_| PyTypeError::new_err("expand: dimension keys must be strings"))?; - - let coords = if v.is_instance_of::() { - let lst = - v.downcast::().map_err(|e| PyTypeError::new_err(e.to_string()))?; - let joined = join_pylist_as_path(lst)?; - Coordinates::from_string(&joined) - } else { - let py_str = v.str()?; - let s: String = py_str.extract()?; - Coordinates::from_string(&s) - }; - - self.inner.expand(&key, coords).map_err(PyTypeError::new_err)?; - } - Ok(()) - } - - /// Remove one or more dimensions from the Qube. - /// - /// `axis` may be a single dimension name (``str``) or a list of names. - /// All children of the removed nodes are re-parented to their - /// grandparent, preserving the rest of the structure. The result is - /// automatically compressed. - /// - /// Raises ``ValueError`` if any of the specified dimensions do not exist - /// in the Qube. - /// - /// ```python - /// q = Qube.from_datacube( - /// {"param": "2t/tp", "time": "0/1/2", "level": "1000/850"}, - /// ["param", "time", "level"], - /// ) - /// q.collapse("level") - /// assert "level" not in q.dimensions() - /// assert "param" in q.dimensions() - /// ``` - pub fn collapse(&mut self, axis: Bound<'_, PyAny>) -> PyResult<()> { - // Accept either a single string or a list of strings. - let axes: Vec = if axis.is_instance_of::() { - let lst = axis.downcast::().map_err(|e| PyTypeError::new_err(e.to_string()))?; - lst.iter() - .map(|item| { - item.str() - .and_then(|s| s.extract::()) - .map_err(|_| PyTypeError::new_err("collapse: axis names must be strings")) - }) - .collect::>()? - } else { - let s: String = axis - .extract() - .map_err(|_| PyTypeError::new_err("collapse: axis must be a str or list[str]"))?; - vec![s] - }; - - // Validate that every requested dimension actually exists. - let existing_dims = self.inner.dimensions(); - let missing: Vec<&str> = axes - .iter() - .filter(|a| !existing_dims.contains(a.as_str())) - .map(String::as_str) - .collect(); - - if !missing.is_empty() { - return Err(PyValueError::new_err(format!( - "Dimension(s) '{}' not found in qube (existing: {:?})", - missing.join(", "), - { - let mut sorted: Vec<&str> = existing_dims.iter().map(String::as_str).collect(); - sorted.sort(); - sorted - } - ))); - } - - self.inner.drop(axes).map_err(PyTypeError::new_err) - } - - /// Collapse then expand in a single operation. - /// - /// Equivalent to calling `collapse(axis)` followed by `expand(dimension)` - /// on the same Qube. Useful for replacing one dimension with another - /// (e.g. replacing "level" with "step"). - /// - /// ```python - /// q = Qube.from_datacube( - /// {"param": "2t/tp", "time": "0/1/2"}, - /// ["param", "time"], - /// ) - /// q.coxpand("time", {"step": ["s1", "s2"]}) - /// assert "time" not in q.dimensions() - /// assert "step" in q.dimensions() - /// ``` - pub fn coxpand( - &mut self, - axis: Bound<'_, PyAny>, - dimension: Bound<'_, PyDict>, - ) -> PyResult<()> { - self.collapse(axis)?; - self.expand(dimension) - } - - /// Check whether the Qube contains a given dimension or set of values. - /// - /// * **`str`** — returns ``True`` if the named dimension exists anywhere - /// in the Qube. - /// * **`dict[str, list]`** — returns ``True`` if every key in the dict is - /// a dimension that exists *and* every value listed for that key is - /// present in that dimension's coordinate set. - /// * **`Qube`** — checks that every dimension/value present in the other - /// Qube also exists in this one (subset check). - /// - /// ```python - /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) - /// - /// assert q.contains("param") is True - /// assert q.contains("level") is False - /// assert q.contains({"param": ["2t"]}) is True - /// assert q.contains({"param": ["xyz"]}) is False - /// assert q.contains({"time": ["0", "1"]}) is True - /// ``` - pub fn contains(&self, py: Python<'_>, item: Bound<'_, PyAny>) -> PyResult { - // Helper: convert a Coordinates value to a Vec of individual values. - let coords_to_strings = |coords: &::qubed::Coordinates| -> Vec { - let s = coords.to_string(); - if s.is_empty() { - vec![] - } else if s.contains('/') { - s.split('/').map(|v| v.to_owned()).collect() - } else { - vec![s] - } - }; - - // Pre-compute current axes as Vec per dimension. - let current_axes: std::collections::BTreeMap> = self - .inner - .all_unique_dim_coords() - .into_iter() - .map(|(k, v)| (k, coords_to_strings(&v))) - .collect(); - - // --- branch 1: item is a plain string (dimension name check) --- - if let Ok(dim_str) = item.extract::() { - return Ok(current_axes.contains_key(&dim_str)); - } - - // --- branch 2: item is a dict --- - if let Ok(dict) = item.downcast::() { - return helpers::check_dict_against_axes(dict, ¤t_axes, py); - } - - // --- branch 3: item is a Qube --- - if let Ok(other_cell) = item.downcast::() { - let other_ref = other_cell.borrow(); - let other_axes: std::collections::BTreeMap> = other_ref - .inner - .all_unique_dim_coords() - .into_iter() - .map(|(k, v)| (k, coords_to_strings(&v))) - .collect(); - - for (key, other_vals) in &other_axes { - match current_axes.get(key) { - None => return Ok(false), - Some(cur_vals) => { - for oval in other_vals { - if !cur_vals.contains(oval) { - return Ok(false); - } - } - } - } - } - return Ok(true); - } - - Err(PyTypeError::new_err("contains: item must be a str, dict[str, list], or Qube")) - } -} - -fn pydict_to_datacube(datacube: Bound<'_, PyDict>) -> PyResult { - let mut dc = Datacube::new(); - for (k, v) in datacube.iter() { - let key: String = - k.extract().map_err(|_| PyTypeError::new_err("datacube keys must be strings"))?; - let val_str: String = v.str()?.extract()?; - dc.add_coordinate(&key, Coordinates::from_string(&val_str)); - } - Ok(dc) } pub(crate) fn join_pylist_as_path(lst: &Bound<'_, PyList>) -> PyResult { diff --git a/py_qubed/tests/test_qubed_api.py b/py_qubed/tests/test_qubed_api.py index 9ee429d..c12d4e3 100644 --- a/py_qubed/tests/test_qubed_api.py +++ b/py_qubed/tests/test_qubed_api.py @@ -1,4 +1,4 @@ -from qubed import Qube +from qubed import PyQube import pytest ASCII_INPUT = """root @@ -13,7 +13,7 @@ def test_ascii_roundtrip_contains_expected_nodes() -> None: - qube = Qube.from_ascii(ASCII_INPUT) + qube = PyQube.from_ascii(ASCII_INPUT) output = qube.to_ascii() assert output == ASCII_INPUT @@ -23,15 +23,15 @@ def test_ascii_roundtrip_contains_expected_nodes() -> None: def test_append_and_append_many_smoke() -> None: # Each source has a distinct class value, so all three should survive the merge. - left = Qube.from_ascii("""root + left = PyQube.from_ascii("""root └── class=1 └── param=10 """) - right = Qube.from_ascii("""root + right = PyQube.from_ascii("""root └── class=2 └── param=20 """) - third = Qube.from_ascii("""root + third = PyQube.from_ascii("""root └── class=3 └── param=30 """) @@ -46,14 +46,14 @@ def test_append_and_append_many_smoke() -> None: def test_append_many_rejects_non_qube_items() -> None: - target = Qube() + target = PyQube() - with pytest.raises(TypeError, match="expected Qube"): + with pytest.raises(TypeError, match="expected PyQube"): target.append_many(["not-a-qube"]) def test_to_datacubes_shape() -> None: - qube = Qube.from_ascii("""root + qube = PyQube.from_ascii("""root └── class=5 └── param=42 """) @@ -66,7 +66,7 @@ def test_to_datacubes_shape() -> None: def test_str_and_len_dunder_methods() -> None: - qube = Qube.from_ascii("""root + qube = PyQube.from_ascii("""root └── class=5 └── param=42 """) @@ -75,7 +75,7 @@ def test_str_and_len_dunder_methods() -> None: def test_to_from_arena_json_roundtrip() -> None: - qube = Qube.from_ascii("""root + qube = PyQube.from_ascii("""root └── class=5 └── param=42 """) @@ -94,107 +94,12 @@ def test_to_from_arena_json_roundtrip() -> None: assert any(isinstance(item, dict) and "dim" in item and "coords" in item for item in qube_list) # Reconstruct and verify ascii equality - reconstructed = Qube.from_arena_json(arena_json) + reconstructed = PyQube.from_arena_json(arena_json) assert reconstructed.to_ascii() == qube.to_ascii() -def test_from_datacube_basic() -> None: - """Build a Qube from a single datacube dict and verify all dimensions appear.""" - dc = {"class": "od", "expver": "0001", "param": "1"} - q = Qube.from_datacube(dc, ["class", "expver", "param"]) - output = q.to_ascii() - - assert "class=od" in output - assert "expver=0001" in output - assert "param=1" in output - - -def test_from_datacube_order_controls_dimension_levels() -> None: - """The `order` parameter determines the nesting order of dimensions in the tree.""" - dc = {"class": "od", "expver": "0001", "param": "1"} - q = Qube.from_datacube(dc, ["class", "expver", "param"]) - output = q.to_ascii() - - # class must appear before expver, expver before param in the rendered tree - assert output.index("class=od") < output.index("expver=0001") - assert output.index("expver=0001") < output.index("param=1") - - -def test_from_datacube_no_order() -> None: - """When order is None all dimensions should still be present.""" - dc = {"class": "od", "param": "1"} - q = Qube.from_datacube(dc, None) - output = q.to_ascii() - - assert "class=od" in output - assert "param=1" in output - - -def test_from_datacube_roundtrip_via_to_datacubes() -> None: - """from_datacube + to_datacubes should recover the original mapping.""" - dc = {"class": "od", "expver": "0001", "param": "1"} - q = Qube.from_datacube(dc, ["class", "expver", "param"]) - datacubes = q.to_datacubes() - - assert len(datacubes) == 1 - assert datacubes[0]["class"] == "od" - assert datacubes[0]["expver"] == "0001" - assert datacubes[0]["param"] == "1" - - -def test_from_datacube_multi_value_coords() -> None: - """Coordinates with multiple values (slash-separated) should all be preserved.""" - dc = {"class": "od", "param": "1/2/3"} - q = Qube.from_datacube(dc, ["class", "param"]) - coords = q.all_unique_dim_coords() - - assert set(coords["param"]) == {"1", "2", "3"} - assert coords["class"] == ["od"] - - -def test_append_datacube_merges_new_dimension_values() -> None: - """append_datacube with a new dimension value should add it to the Qube.""" - q = Qube.from_ascii("""root -└── class=od - └── param=1 -""") - q.append_datacube({"class": "rd", "param": "1"}, ["class", "param"]) - coords = q.all_unique_dim_coords() - - assert set(coords["class"]) == {"od", "rd"} - assert set(coords["param"]) == {"1"} - - -def test_append_datacube_merges_into_existing_structure() -> None: - """append_datacube should extend an existing branch rather than duplicate it.""" - q = Qube.from_ascii("""root -└── class=od - └── expver=0001 - └── param=1 -""") - q.append_datacube( - {"class": "od", "expver": "0002", "param": "1"}, - ["class", "expver", "param"], - ) - coords = q.all_unique_dim_coords() - - assert set(coords["expver"]) == {"0001", "0002"} - assert coords["class"] == ["od"] - assert coords["param"] == ["1"] - - -def test_append_datacube_multiple_times_builds_correct_tree() -> None: - """Calling append_datacube repeatedly should produce the same result as append_many.""" - q = Qube() - for cls in ("a", "b", "c"): - q.append_datacube({"class": cls, "param": "1"}, ["class", "param"]) - - coords = q.all_unique_dim_coords() - assert set(coords["class"]) == {"a", "b", "c"} - - def test_arena_preserves_leading_zeros() -> None: - qube = Qube.from_ascii("""root + qube = PyQube.from_ascii("""root └── class=od └── expver=0001 └── param=1 @@ -203,5 +108,5 @@ def test_arena_preserves_leading_zeros() -> None: arena_json = qube.to_arena_json() assert "0001" in arena_json - reconstructed = Qube.from_arena_json(arena_json) + reconstructed = PyQube.from_arena_json(arena_json) assert "expver=0001" in reconstructed.to_ascii() diff --git a/py_qubed/tests/test_select_api.py b/py_qubed/tests/test_select_api.py index 4d64d05..a172ea3 100644 --- a/py_qubed/tests/test_select_api.py +++ b/py_qubed/tests/test_select_api.py @@ -19,7 +19,7 @@ def test_select_1(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) selected = q.select({"class": [1]}, None, None) @@ -32,7 +32,7 @@ def test_select_1(): ├── param=1 └── param=2""" - assert selected.to_ascii() == qubed.Qube.from_ascii(expected).to_ascii() + assert selected.to_ascii() == qubed.PyQube.from_ascii(expected).to_ascii() def test_select_2(): @@ -53,7 +53,7 @@ def test_select_2(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) selected = q.select({"class": [1], "param": [1]}, None, None) @@ -64,7 +64,7 @@ def test_select_2(): └── expver=0002 └── param=1""" - assert selected.to_ascii() == qubed.Qube.from_ascii(expected).to_ascii() + assert selected.to_ascii() == qubed.PyQube.from_ascii(expected).to_ascii() def test_select_3(): @@ -85,7 +85,7 @@ def test_select_3(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) selected = q.select({"expver": ["0001"]}, None, None) @@ -100,7 +100,7 @@ def test_select_3(): ├── param=2 └── param=3""" - assert selected.to_ascii() == qubed.Qube.from_ascii(expected).to_ascii() + assert selected.to_ascii() == qubed.PyQube.from_ascii(expected).to_ascii() def test_all_unique_dim_coords(): @@ -121,7 +121,7 @@ def test_all_unique_dim_coords(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) dim_coords = q.all_unique_dim_coords() @@ -166,7 +166,7 @@ def test_compress(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) # Get the ASCII representation before compression ascii_before = q.to_ascii() @@ -191,7 +191,7 @@ def test_compress_2(): └── expver=0002 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) # Get the ASCII representation before compression ascii_before = q.to_ascii() @@ -228,7 +228,7 @@ def test_select_multiple_values(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) # Select multiple values for the same key selected = q.select({"param": [1, 3]}, None, None) @@ -246,7 +246,7 @@ def test_select_multiple_values(): └── expver=0002 └── param=1""" - assert selected.to_ascii() == qubed.Qube.from_ascii(expected).to_ascii() + assert selected.to_ascii() == qubed.PyQube.from_ascii(expected).to_ascii() def test_default(): @@ -268,7 +268,7 @@ def test_default(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) # Default mode: shows full subtree default_result = q.select({"class": [1]}, None, None) @@ -284,7 +284,7 @@ def test_default(): assert ( default_result.to_ascii() - == qubed.Qube.from_ascii(default_expected).to_ascii() + == qubed.PyQube.from_ascii(default_expected).to_ascii() ) @@ -298,14 +298,14 @@ def test_drop(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) q.drop(["expver"]) expected = r"""root └── class=1 └── param=1/2""" - assert q.to_ascii() == qubed.Qube.from_ascii(expected).to_ascii() + assert q.to_ascii() == qubed.PyQube.from_ascii(expected).to_ascii() def test_squeeze(): @@ -318,7 +318,7 @@ def test_squeeze(): ├── param=1 └── param=2""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) q.squeeze() # class has only one value (1), so it gets squeezed out @@ -326,7 +326,7 @@ def test_squeeze(): └── expver=0001/0002 └── param=1/2""" - assert q.to_ascii() == qubed.Qube.from_ascii(expected).to_ascii() + assert q.to_ascii() == qubed.PyQube.from_ascii(expected).to_ascii() def test_select_drops_branches_without_matching_deep_key(): @@ -339,14 +339,14 @@ def test_select_drops_branches_without_matching_deep_key(): ├── param=3 └── param=4""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) selected = q.select({"param": [1]}, None, None) expected = r"""root └── expver=0001 └── param=1""" - assert selected.to_ascii() == qubed.Qube.from_ascii(expected).to_ascii(), ( + assert selected.to_ascii() == qubed.PyQube.from_ascii(expected).to_ascii(), ( "expver=0002 (no param=1 descendants) should be absent from the result" ) @@ -366,7 +366,7 @@ def test_select_deep_key_multi_level_unselected_prefix(): ├── param=5 └── param=6""" - q = qubed.Qube.from_ascii(input_qube) + q = qubed.PyQube.from_ascii(input_qube) selected = q.select({"param": [1]}, None, None) expected = r"""root @@ -374,6 +374,6 @@ def test_select_deep_key_multi_level_unselected_prefix(): └── expver=0001 └── param=1""" - assert selected.to_ascii() == qubed.Qube.from_ascii(expected).to_ascii(), ( + assert selected.to_ascii() == qubed.PyQube.from_ascii(expected).to_ascii(), ( "only class=1/expver=0001 contains param=1; all other branches must be pruned" ) diff --git a/py_qubed_meteo/Cargo.toml b/py_qubed_meteo/Cargo.toml new file mode 100644 index 0000000..b62a329 --- /dev/null +++ b/py_qubed_meteo/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "py_qubed_meteo" +version = "0.1.0" +edition = "2024" + +[dependencies] +qubed = { path = "../qubed" } +qubed_meteo = { path = "../qubed_meteo" } +serde_json = "1.0.145" +pyo3 = { version = "0.28", features = ["extension-module", "abi3-py38"] } + +[lib] +name = "qubed_meteo" +path = "src/lib.rs" +crate-type = ["cdylib"] diff --git a/py_qubed_meteo/pyproject.toml b/py_qubed_meteo/pyproject.toml new file mode 100644 index 0000000..6a81668 --- /dev/null +++ b/py_qubed_meteo/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["maturin>=1.7,<2.0"] +build-backend = "maturin" + +[project] +name = "qubed-meteo" +version = "0.1.0" +description = "Python bindings for meteorological adapters built on qubed" +requires-python = ">=3.8" +classifiers = [ + "Programming Language :: Python :: 3", + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] + +[tool.maturin] +manifest-path = "Cargo.toml" +features = [] diff --git a/py_qubed_meteo/src/lib.rs b/py_qubed_meteo/src/lib.rs new file mode 100644 index 0000000..0ce5a5a --- /dev/null +++ b/py_qubed_meteo/src/lib.rs @@ -0,0 +1,43 @@ +use ::qubed::Qube; +use ::qubed_meteo::adapters::fdb::FromFDBList; +use ::qubed_meteo::adapters::mars_list::FromMARSList; +use ::qubed_meteo::adapters::to_constraints::ToDssConstraints; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyModule; +use pyo3::wrap_pyfunction; + +#[pyfunction] +pub fn from_mars_list_py(text: &str) -> PyResult { + match Qube::from_mars_list(text) { + // ASCII is our stable bridge format so callers can pipe into PyQube.from_ascii(). + Ok(qube) => Ok(qube.to_ascii()), + Err(e) => Err(PyValueError::new_err(e)), + } +} + +#[pymodule] +#[pyo3(name = "qubed_meteo")] +fn py_qubed_meteo_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(from_mars_list_py, m)?)?; + m.add_function(wrap_pyfunction!(from_fdb_list_py, m)?)?; + m.add_function(wrap_pyfunction!(to_dss_constraints_py, m)?)?; + Ok(()) +} + +#[pyfunction] +pub fn from_fdb_list_py(request_json: &str) -> PyResult { + let v: serde_json::Value = + serde_json::from_str(request_json).map_err(|e| PyValueError::new_err(e.to_string()))?; + match Qube::from_fdb_list(&v) { + Ok(qube) => Ok(qube.to_ascii()), + Err(e) => Err(PyValueError::new_err(e)), + } +} + +#[pyfunction] +pub fn to_dss_constraints_py(ascii: &str) -> PyResult { + let qube = Qube::from_ascii(ascii).map_err(|e| PyValueError::new_err(e))?; + let v = qube.to_dss_constraints(); + serde_json::to_string(&v).map_err(|e| PyValueError::new_err(e.to_string())) +} diff --git a/py_qubed_meteo/tests/test_api.py b/py_qubed_meteo/tests/test_api.py new file mode 100644 index 0000000..c9cc394 --- /dev/null +++ b/py_qubed_meteo/tests/test_api.py @@ -0,0 +1,55 @@ +import json + +import qubed_meteo + + +def test_from_fdb_list_py_basic() -> None: + items = [ + "class=od,expver=0001,param=1/2", + "class=rd,expver=0003,param=3/4", + ] + + ascii_out = qubed_meteo.from_fdb_list_py(items) + assert isinstance(ascii_out, str) + # basic tokens + assert "class=od" in ascii_out + assert "expver=0001" in ascii_out + assert "param=1/2" in ascii_out + + +def test_to_dss_constraints_py_basic() -> None: + items = [ + "class=od,expver=0001,param=1/2", + "class=rd,expver=0003,param=3/4", + "class=rd,expver=0002,param=5/6", + ] + + ascii_out = qubed_meteo.from_fdb_list_py(items) + json_out = qubed_meteo.to_dss_constraints_py(ascii_out) + parsed = json.loads(json_out) + + assert isinstance(parsed, list) + # Expect three objects (one per leaf path) + assert len(parsed) == 3 + + found_od = False + found_rd_34 = False + found_rd_56 = False + + for obj in parsed: + assert "class" in obj and "expver" in obj and "param" in obj + class_vals = obj["class"] + exp_vals = obj["expver"] + param_vals = obj["param"] + + if "od" in class_vals: + found_od = True + assert "0001" in exp_vals + assert "1" in param_vals and "2" in param_vals + if "rd" in class_vals: + if set(param_vals) >= {"3", "4"}: + found_rd_34 = True + if set(param_vals) >= {"5", "6"}: + found_rd_56 = True + + assert found_od and found_rd_34 and found_rd_56 diff --git a/py_qubed_meteo/tests/test_qubed_meteo_api.py b/py_qubed_meteo/tests/test_qubed_meteo_api.py new file mode 100644 index 0000000..05bccd4 --- /dev/null +++ b/py_qubed_meteo/tests/test_qubed_meteo_api.py @@ -0,0 +1,23 @@ +from qubed import PyQube +import qubed_meteo + + +MARS_LIST_SAMPLE = """class=od,expver=1,param=2 +class=rd,expver=2,param=3 +""" + + +def test_from_mars_list_py_produces_parseable_ascii() -> None: + ascii_tree = qubed_meteo.from_mars_list_py(MARS_LIST_SAMPLE) + + # Parsing the adapter output through PyQube verifies end-to-end contract compatibility. + parsed = PyQube.from_ascii(ascii_tree) + datacubes = parsed.to_datacubes() + + assert ascii_tree.startswith("root") + assert len(datacubes) >= 1 + + +def test_from_mars_list_py_handles_empty_input() -> None: + ascii_tree = qubed_meteo.from_mars_list_py("\n\n") + assert ascii_tree == "root\n" diff --git a/qubed/src/compress.rs b/qubed/src/compress.rs index 4b35d52..e697935 100644 --- a/qubed/src/compress.rs +++ b/qubed/src/compress.rs @@ -200,7 +200,8 @@ impl Qube { } } - /// Merges the coordinates of a group of nodes into the first node in the group. + /// Merges the coordinates of a group of nodes into the first node in the group, + /// then attempts to compress the merged coordinates into a tighter range representation. fn merge_coords(&mut self, group: Vec) { // The coordinates of all other nodes in the group are set to `Coordinates::Empty`. @@ -213,6 +214,9 @@ impl Qube { merged.extend(coords); } + // Try to compress consecutive integer or datetime values into ranges. + merged.try_compress(); + { let node = self.node_mut(group[0]).unwrap(); *node.coords_mut() = merged; diff --git a/qubed/src/coordinates/datetime.rs b/qubed/src/coordinates/datetime.rs index 7ed02a7..e3ecda2 100644 --- a/qubed/src/coordinates/datetime.rs +++ b/qubed/src/coordinates/datetime.rs @@ -1,13 +1,118 @@ use std::hash::Hash; -use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; +use chrono::{DateTime, Duration, NaiveDate, NaiveDateTime, NaiveTime, Utc}; use tiny_vec::TinyVec; use crate::coordinates::{Coordinates, IntersectionResult}; +/// An inclusive datetime range `[start, end]` with a given step (as a `Duration`). +/// All values `start + k * step` where `start + k * step <= end` are members. +#[derive(Debug, Clone, PartialEq)] +pub struct DateTimeRange { + pub start: NaiveDateTime, + pub end: NaiveDateTime, + /// Step between members. Must be positive (> 0 duration). + pub step: Duration, +} + +impl DateTimeRange { + /// Create a range with an explicit step. + pub fn new(start: NaiveDateTime, end: NaiveDateTime, step: Duration) -> Self { + assert!(start <= end, "DateTimeRange: start must be <= end"); + assert!(step > Duration::zero(), "DateTimeRange: step must be positive"); + DateTimeRange { start, end, step } + } + + /// Create a range with a 1-day step. + pub fn daily(start: NaiveDateTime, end: NaiveDateTime) -> Self { + Self::new(start, end, Duration::days(1)) + } + + /// Create a range with a 1-hour step. + pub fn hourly(start: NaiveDateTime, end: NaiveDateTime) -> Self { + Self::new(start, end, Duration::hours(1)) + } + + /// Number of elements in this range. + pub fn len(&self) -> usize { + let step_ns = self.step.num_nanoseconds().unwrap_or(i64::MAX); + let total_ns = (self.end - self.start).num_nanoseconds().unwrap_or(0); + if step_ns <= 0 || total_ns < 0 { + return 0; + } + (total_ns / step_ns + 1) as usize + } + + pub fn contains(&self, value: NaiveDateTime) -> bool { + if value < self.start || value > self.end { + return false; + } + let elapsed_ns = (value - self.start).num_nanoseconds().unwrap_or(-1); + let step_ns = self.step.num_nanoseconds().unwrap_or(0); + if step_ns <= 0 { + return false; + } + elapsed_ns % step_ns == 0 + } + + /// Iterate over all values in this range. + pub fn iter(&self) -> impl Iterator + '_ { + let mut current = self.start; + let end = self.end; + let step = self.step; + std::iter::from_fn(move || { + if current <= end { + let val = current; + current = current + step; + Some(val) + } else { + None + } + }) + } + + /// Intersect two ranges that share the same step. Returns None if no overlap. + pub fn intersect_range(&self, other: &DateTimeRange) -> Option { + if self.step != other.step { + return None; + } + let new_start = self.start.max(other.start); + let new_end = self.end.min(other.end); + if new_start > new_end { + return None; + } + // Align new_start to a grid point of self + let step_ns = self.step.num_nanoseconds()?; + let offset_ns = (new_start - self.start).num_nanoseconds()?; + let remainder = offset_ns % step_ns; + let aligned_start = if remainder == 0 { + new_start + } else { + new_start + Duration::nanoseconds(step_ns - remainder) + }; + if aligned_start > new_end { + return None; + } + Some(DateTimeRange::new(aligned_start, new_end, self.step)) + } + + pub fn to_string(&self) -> String { + let step_secs = self.step.num_seconds(); + format!( + "{}:{}s:{}", + self.start.format("%Y-%m-%dT%H:%M:%S"), + step_secs, + self.end.format("%Y-%m-%dT%H:%M:%S") + ) + } +} + +// ---- DateTimeCoordinates ---- + #[derive(Debug, Clone, PartialEq)] pub enum DateTimeCoordinates { List(TinyVec), + RangeSet(TinyVec), } impl DateTimeCoordinates { @@ -18,24 +123,51 @@ impl DateTimeCoordinates { list.push(*v); } } + (DateTimeCoordinates::RangeSet(ranges), DateTimeCoordinates::RangeSet(new_ranges)) => { + for r in new_ranges.iter() { + ranges.push(r.clone()); + } + } + (self_coords, other) => { + // Mixed List/RangeSet: materialise both into a List + let all_vals: Vec = match &*self_coords { + DateTimeCoordinates::List(l) => l.iter().copied().collect(), + DateTimeCoordinates::RangeSet(rs) => rs.iter().flat_map(|r| r.iter()).collect(), + }; + let other_vals: Vec = match other { + DateTimeCoordinates::List(l) => l.iter().copied().collect(), + DateTimeCoordinates::RangeSet(rs) => rs.iter().flat_map(|r| r.iter()).collect(), + }; + let mut merged: TinyVec = TinyVec::new(); + for v in all_vals.into_iter().chain(other_vals) { + merged.push(v); + } + *self_coords = DateTimeCoordinates::List(merged); + } } } pub(crate) fn append(&mut self, new_coord: NaiveDateTime) { match self { DateTimeCoordinates::List(list) => list.push(new_coord), + DateTimeCoordinates::RangeSet(ranges) => { + // Append as a single-element "range" + ranges.push(DateTimeRange::new(new_coord, new_coord, Duration::seconds(1))); + } } } pub(crate) fn len(&self) -> usize { match self { DateTimeCoordinates::List(list) => list.len(), + DateTimeCoordinates::RangeSet(ranges) => ranges.iter().map(|r| r.len()).sum(), } } pub(crate) fn contains(&self, value: NaiveDateTime) -> bool { match self { DateTimeCoordinates::List(list) => list.iter().any(|&v| v == value), + DateTimeCoordinates::RangeSet(ranges) => ranges.iter().any(|r| r.contains(value)), } } @@ -46,6 +178,48 @@ impl DateTimeCoordinates { .map(|dt| dt.format("%Y-%m-%dT%H:%M:%S").to_string()) .collect::>() .join("/"), + DateTimeCoordinates::RangeSet(ranges) => { + ranges.iter().map(|r| r.to_string()).collect::>().join("/") + } + } + } + + /// Human-readable ASCII representation. + /// + /// Each range is rendered as: + /// - singleton `v` → `v` (using `%Y-%m-%dT%H:%M:%S`) + /// - daily range → `start/to/end` + /// - range with step → `start/to/end/by/s` + /// + /// Multiple ranges / singletons are separated by `|`. + /// Plain `List` values use the standard `/`-joined format. + pub(crate) fn to_ascii_string(&self) -> String { + let fmt = "%Y-%m-%dT%H:%M:%S"; + match self { + DateTimeCoordinates::List(list) => { + list.iter().map(|dt| dt.format(fmt).to_string()).collect::>().join("/") + } + DateTimeCoordinates::RangeSet(ranges) => { + let day_secs = Duration::days(1).num_seconds(); + ranges + .iter() + .map(|r| { + if r.start == r.end { + r.start.format(fmt).to_string() + } else if r.step.num_seconds() == day_secs { + format!("{}/to/{}", r.start.format(fmt), r.end.format(fmt)) + } else { + format!( + "{}/to/{}/by/{}s", + r.start.format(fmt), + r.end.format(fmt), + r.step.num_seconds() + ) + } + }) + .collect::>() + .join("|") + } } } @@ -53,12 +227,20 @@ impl DateTimeCoordinates { "datetime".hash(hasher); match self { DateTimeCoordinates::List(list) => { + "list".hash(hasher); for dt in list.iter() { - // use seconds and nanoseconds for stable hashing dt.and_utc().timestamp().hash(hasher); dt.and_utc().timestamp_subsec_nanos().hash(hasher); } } + DateTimeCoordinates::RangeSet(ranges) => { + "range_set".hash(hasher); + for r in ranges.iter() { + r.start.and_utc().timestamp().hash(hasher); + r.end.and_utc().timestamp().hash(hasher); + r.step.num_seconds().hash(hasher); + } + } } } @@ -67,23 +249,17 @@ impl DateTimeCoordinates { other: &DateTimeCoordinates, ) -> IntersectionResult { match (self, other) { + // List ∩ List (DateTimeCoordinates::List(list_a), DateTimeCoordinates::List(list_b)) => { use std::collections::HashSet; - let mut set_b: HashSet = HashSet::new(); - for v in list_b.iter() { - set_b.insert(*v); - } - - let mut set_a: HashSet = HashSet::new(); - for v in list_a.iter() { - set_a.insert(*v); - } + let set_b: HashSet = list_b.iter().copied().collect(); + let set_a: HashSet = list_a.iter().copied().collect(); let mut intersection = TinyVec::new(); let mut only_a = TinyVec::new(); - let mut added: HashSet = HashSet::new(); + for v in list_a.iter() { if set_b.contains(v) { if !added.contains(v) { @@ -108,16 +284,48 @@ impl DateTimeCoordinates { only_b: DateTimeCoordinates::List(only_b), } } + + // RangeSet ∩ RangeSet + (DateTimeCoordinates::RangeSet(ranges_a), DateTimeCoordinates::RangeSet(ranges_b)) => { + intersect_dt_range_sets(ranges_a, ranges_b) + } + + // RangeSet ∩ List (range is `self`) + (DateTimeCoordinates::RangeSet(ranges), DateTimeCoordinates::List(list)) => { + intersect_dt_rangeset_with_list(ranges, list, false) + } + + // List ∩ RangeSet (list is `self`) + (DateTimeCoordinates::List(list), DateTimeCoordinates::RangeSet(ranges)) => { + intersect_dt_rangeset_with_list(ranges, list, true) + } } } + /// Try to parse a `DateTimeRange` from its textual `to_string()` representation. + /// Format: `":s:"` e.g. `"2020-01-01T00:00:00:86400s:2020-01-10T00:00:00"`. + /// Multiple ranges separated by `/` are also supported. + pub(crate) fn parse_range_from_str(s: &str) -> Option { + let mut ranges: TinyVec = TinyVec::new(); + for part in s.split('/') { + let r = parse_single_dt_range(part)?; + ranges.push(r); + } + if ranges.is_empty() { None } else { Some(DateTimeCoordinates::RangeSet(ranges)) } + } + /// Try to parse a string into `NaiveDateTime` using common formats. pub(crate) fn parse_from_str(s: &str) -> Option { - // Try RFC3339 / ISO 8601 + // Try RFC3339 / ISO 8601 with timezone if let Ok(dt) = DateTime::parse_from_rfc3339(s) { return Some(dt.with_timezone(&Utc).naive_utc()); } + // Try YYYY-MM-DDTHH:MM:SS (no timezone — the canonical to_ascii_string format) + if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S") { + return Some(ndt); + } + // Try YYYY-MM-DD HH:MM:SS if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") { return Some(ndt); @@ -146,12 +354,246 @@ impl DateTimeCoordinates { } } +// ---- Range intersection helpers ---- + +fn intersect_dt_range_sets( + ranges_a: &TinyVec, + ranges_b: &TinyVec, +) -> IntersectionResult { + let mut intersection: TinyVec = TinyVec::new(); + let mut a_consumed = vec![false; ranges_a.len()]; + let mut b_consumed = vec![false; ranges_b.len()]; + + for (ia, ra) in ranges_a.iter().enumerate() { + for (ib, rb) in ranges_b.iter().enumerate() { + if ra.step == rb.step { + if let Some(inter) = ra.intersect_range(rb) { + intersection.push(inter); + a_consumed[ia] = true; + b_consumed[ib] = true; + } + } else { + // Different steps: materialise + for v in ra.iter().filter(|&v| rb.contains(v)) { + intersection.push(DateTimeRange::new(v, v, ra.step)); + a_consumed[ia] = true; + b_consumed[ib] = true; + } + } + } + } + + let mut only_a: TinyVec = TinyVec::new(); + let mut only_b: TinyVec = TinyVec::new(); + + for (ia, ra) in ranges_a.iter().enumerate() { + if !a_consumed[ia] { + only_a.push(ra.clone()); + } else { + for v in ra.iter() { + if !ranges_b.iter().any(|rb| rb.contains(v)) { + only_a.push(DateTimeRange::new(v, v, ra.step)); + } + } + } + } + + for (ib, rb) in ranges_b.iter().enumerate() { + if !b_consumed[ib] { + only_b.push(rb.clone()); + } else { + for v in rb.iter() { + if !ranges_a.iter().any(|ra| ra.contains(v)) { + only_b.push(DateTimeRange::new(v, v, rb.step)); + } + } + } + } + + IntersectionResult { + intersection: DateTimeCoordinates::RangeSet(intersection), + only_a: DateTimeCoordinates::RangeSet(only_a), + only_b: DateTimeCoordinates::RangeSet(only_b), + } +} + +/// Intersect a RangeSet with a List. +/// `swapped` = true when original call was (List, RangeSet), so we swap only_a/only_b. +fn intersect_dt_rangeset_with_list( + ranges: &TinyVec, + list: &TinyVec, + swapped: bool, +) -> IntersectionResult { + use std::collections::HashSet; + + let list_set: HashSet = list.iter().copied().collect(); + + let mut intersection: TinyVec = TinyVec::new(); + let mut only_list: TinyVec = TinyVec::new(); + let mut only_ranges: TinyVec = TinyVec::new(); + + let mut added: HashSet = HashSet::new(); + for &v in list.iter() { + if ranges.iter().any(|r| r.contains(v)) { + if !added.contains(&v) { + intersection.push(v); + added.insert(v); + } + } else { + only_list.push(v); + } + } + + for r in ranges.iter() { + for v in r.iter() { + if !list_set.contains(&v) { + only_ranges.push(DateTimeRange::new(v, v, r.step)); + } + } + } + + let intersection_coord = DateTimeCoordinates::List(intersection); + let (only_a, only_b) = if swapped { + // original: (List, RangeSet) → only_a = list side, only_b = range side + (DateTimeCoordinates::List(only_list), DateTimeCoordinates::RangeSet(only_ranges)) + } else { + // original: (RangeSet, List) → only_a = range side, only_b = list side + (DateTimeCoordinates::RangeSet(only_ranges), DateTimeCoordinates::List(only_list)) + }; + + IntersectionResult { intersection: intersection_coord, only_a, only_b } +} + +// ---- Default ---- + impl Default for DateTimeCoordinates { fn default() -> Self { DateTimeCoordinates::List(TinyVec::new()) } } +impl DateTimeCoordinates { + /// Attempt to compress the coordinate list into a tighter `RangeSet` representation. + /// + /// Works analogously to `IntegerCoordinates::try_compress_to_ranges`: + /// - Collect all values (sorted). + /// - Greedy scan for runs with a uniform `Duration` step (≥ 3 elements per run to save space). + /// - Only replace `self` with a `RangeSet` when the resulting number of ranges is strictly + /// less than the original element count. + pub fn try_compress_to_ranges(&mut self) { + let mut values: Vec = match self { + DateTimeCoordinates::List(list) => list.iter().copied().collect(), + DateTimeCoordinates::RangeSet(ranges) => { + let mut v: Vec = ranges.iter().flat_map(|r| r.iter()).collect(); + v.sort_unstable(); + v.dedup(); + v + } + }; + + if values.len() < 3 { + return; + } + + values.sort_unstable(); + values.dedup(); + + let ranges = compress_datetimes_to_ranges(&values); + + let range_count = ranges.len(); + if range_count < values.len() { + let mut rv: TinyVec = TinyVec::new(); + for r in ranges { + rv.push(r); + } + *self = DateTimeCoordinates::RangeSet(rv); + } + } +} + +/// Partition a sorted, deduplicated slice of `NaiveDateTime` values into the +/// minimum set of uniform-step ranges. +/// +/// Uses the same "emit one singleton, retry" strategy as the integer equivalent: +/// if the run starting at position `i` using step `values[i+1]-values[i]` is +/// fewer than 3 elements, only a singleton is emitted and the scan retries from +/// `i+1`. This prevents a short 2-element run from consuming a value that +/// could start a longer run immediately after. +pub(crate) fn compress_datetimes_to_ranges(values: &[NaiveDateTime]) -> Vec { + if values.is_empty() { + return vec![]; + } + + let singleton = |v: NaiveDateTime| DateTimeRange::new(v, v, Duration::seconds(1)); + + let mut result: Vec = Vec::new(); + let mut i = 0; + + while i < values.len() { + let run_len = if i + 1 < values.len() { + let step = values[i + 1] - values[i]; + if step > Duration::zero() { + let mut j = i; + while j + 1 < values.len() && values[j + 1] - values[j] == step { + j += 1; + } + j - i + 1 + } else { + 1 + } + } else { + 1 + }; + + if run_len >= 3 { + let step = values[i + 1] - values[i]; + result.push(DateTimeRange::new(values[i], values[i + run_len - 1], step)); + i += run_len; + } else { + result.push(singleton(values[i])); + i += 1; + } + } + + result +} + +/// Parse a single datetime range string of the form `":s:"`. +fn parse_single_dt_range(s: &str) -> Option { + // The format is: YYYY-MM-DDTHH:MM:SS:s:YYYY-MM-DDTHH:MM:SS + // We locate the step token by finding `:s:` in the middle. + // Start and end datetimes are 19 chars each in "%Y-%m-%dT%H:%M:%S" format. + // So layout is: [19 chars]:[step]s:[19 chars] + if s.len() < 19 + 1 + 1 + 1 + 19 { + return None; + } + let start_str = &s[..19]; + let rest = &s[19..]; + // rest starts with ":s:" + if !rest.starts_with(':') { + return None; + } + let rest = &rest[1..]; // skip leading ':' + // Find the 's:' separator + let sep_pos = rest.find("s:")?; + let step_str = &rest[..sep_pos]; + let end_str = &rest[sep_pos + 2..]; + if end_str.len() < 19 { + return None; + } + let end_str = &end_str[..19]; + + let start = NaiveDateTime::parse_from_str(start_str, "%Y-%m-%dT%H:%M:%S").ok()?; + let end = NaiveDateTime::parse_from_str(end_str, "%Y-%m-%dT%H:%M:%S").ok()?; + let step_secs: i64 = step_str.parse().ok()?; + if step_secs <= 0 { + return None; + } + Some(DateTimeRange::new(start, end, Duration::seconds(step_secs))) +} + +// ---- From impls ---- + impl From for Coordinates { fn from(value: NaiveDateTime) -> Self { let mut vec = TinyVec::new(); @@ -198,15 +640,47 @@ impl From<&[NaiveDateTime; N]> for Coordinates { } } +/// Construct `Coordinates` from a single `DateTimeRange`. +impl From for Coordinates { + fn from(value: DateTimeRange) -> Self { + let mut ranges: TinyVec = TinyVec::new(); + ranges.push(value); + Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) + } +} + +/// Construct `Coordinates` from a slice of `DateTimeRange`. +impl From<&[DateTimeRange]> for Coordinates { + fn from(value: &[DateTimeRange]) -> Self { + let mut ranges: TinyVec = TinyVec::new(); + for r in value { + ranges.push(r.clone()); + } + Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) + } +} + +// ---- Tests ---- + #[cfg(test)] mod tests { use super::*; + fn dt(year: i32, month: u32, day: u32) -> NaiveDateTime { + NaiveDate::from_ymd_opt(year, month, day).unwrap().and_hms_opt(0, 0, 0).unwrap() + } + + fn dt_h(year: i32, month: u32, day: u32, hour: u32) -> NaiveDateTime { + NaiveDate::from_ymd_opt(year, month, day).unwrap().and_hms_opt(hour, 0, 0).unwrap() + } + + // ---- Existing List tests (preserved) ---- + #[test] fn test_datetime_append_and_len() { let mut coords = DateTimeCoordinates::default(); - let d1 = NaiveDate::from_ymd(2020, 1, 1).and_hms(0, 0, 0); - let d2 = NaiveDate::from_ymd(2020, 1, 2).and_hms(12, 30, 0); + let d1 = dt(2020, 1, 1); + let d2 = NaiveDate::from_ymd_opt(2020, 1, 2).unwrap().and_hms_opt(12, 30, 0).unwrap(); coords.append(d1); coords.append(d2); @@ -216,14 +690,15 @@ mod tests { assert_eq!(list[0], d1); assert_eq!(list[1], d2); } + _ => panic!("Expected List"), } } #[test] fn test_datetime_extend() { + let d1 = dt(2020, 1, 1); + let d2 = dt(2020, 1, 2); let mut a = DateTimeCoordinates::default(); - let d1 = NaiveDate::from_ymd(2020, 1, 1).and_hms(0, 0, 0); - let d2 = NaiveDate::from_ymd(2020, 1, 2).and_hms(0, 0, 0); a.append(d1); let mut b = DateTimeCoordinates::default(); @@ -237,29 +712,31 @@ mod tests { assert_eq!(list[0], d1); assert_eq!(list[1], d2); } + _ => panic!("Expected List"), } } #[test] fn test_datetime_to_string_and_parse() { - let d = NaiveDate::from_ymd(2021, 5, 4).and_hms(6, 7, 8); + let d = NaiveDate::from_ymd_opt(2021, 5, 4).unwrap().and_hms_opt(6, 7, 8).unwrap(); let mut c = DateTimeCoordinates::default(); c.append(d); let s = c.to_string(); assert!(s.contains("2021-05-04T06:07:08")); - // parse from iso string let parsed = DateTimeCoordinates::parse_from_str("2021-05-04T06:07:08Z"); assert!(parsed.is_some()); assert_eq!(parsed.unwrap(), d); } #[test] - fn test_datetime_intersect() { + fn test_datetime_intersect_list_list() { + let d1 = dt(2020, 1, 1); + let d2 = dt(2020, 1, 2); + let d3 = dt(2020, 1, 3); + let d4 = dt(2020, 1, 4); + let mut a = DateTimeCoordinates::default(); - let d1 = NaiveDate::from_ymd(2020, 1, 1).and_hms(0, 0, 0); - let d2 = NaiveDate::from_ymd(2020, 1, 2).and_hms(0, 0, 0); - let d3 = NaiveDate::from_ymd(2020, 1, 3).and_hms(0, 0, 0); a.append(d1); a.append(d2); a.append(d3); @@ -267,7 +744,7 @@ mod tests { let mut b = DateTimeCoordinates::default(); b.append(d2); b.append(d3); - b.append(NaiveDate::from_ymd(2020, 1, 4).and_hms(0, 0, 0)); + b.append(d4); let res = a.intersect(&b); @@ -277,6 +754,369 @@ mod tests { assert_eq!(list[0], d2); assert_eq!(list[1], d3); } + _ => panic!("Expected List"), + } + } + + // ---- DateTimeRange unit tests ---- + + #[test] + fn test_datetime_range_contains_daily() { + let r = DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 5)); + assert!(r.contains(dt(2020, 1, 1))); + assert!(r.contains(dt(2020, 1, 3))); + assert!(r.contains(dt(2020, 1, 5))); + assert!(!r.contains(dt(2019, 12, 31))); + assert!(!r.contains(dt(2020, 1, 6))); + // Midnight check — hours not aligned + assert!(!r.contains(dt_h(2020, 1, 2, 6))); + } + + #[test] + fn test_datetime_range_len_daily() { + let r = DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 5)); + assert_eq!(r.len(), 5); // Jan 1,2,3,4,5 + } + + #[test] + fn test_datetime_range_iter() { + let r = DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 3)); + let vals: Vec = r.iter().collect(); + assert_eq!(vals, vec![dt(2020, 1, 1), dt(2020, 1, 2), dt(2020, 1, 3)]); + } + + #[test] + fn test_datetime_range_to_string() { + let r = DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 5)); + let s = r.to_string(); + assert!(s.contains("2020-01-01T00:00:00")); + assert!(s.contains("2020-01-05T00:00:00")); + assert!(s.contains("86400s")); // 1 day in seconds + } + + // ---- RangeSet ∩ RangeSet ---- + + #[test] + fn test_rangeset_intersect_rangeset_overlapping() { + // [Jan 1..Jan 10 daily] ∩ [Jan 5..Jan 15 daily] = [Jan 5..Jan 10] + let a = Coordinates::from(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 10))); + let b = Coordinates::from(DateTimeRange::daily(dt(2020, 1, 5), dt(2020, 1, 15))); + + let result = a.intersect(&b); + + if let Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) = &result.intersection + { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].start, dt(2020, 1, 5)); + assert_eq!(ranges[0].end, dt(2020, 1, 10)); + } else { + panic!("Expected RangeSet intersection, got {:?}", result.intersection); + } + + // only_a: Jan 1..4 + let only_a_vals: Vec = match &result.only_a { + Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) => { + ranges.iter().flat_map(|r| r.iter()).collect() + } + other => panic!("Expected RangeSet only_a, got {:?}", other), + }; + assert_eq!( + only_a_vals, + vec![dt(2020, 1, 1), dt(2020, 1, 2), dt(2020, 1, 3), dt(2020, 1, 4)] + ); + + // only_b: Jan 11..15 + let only_b_vals: Vec = match &result.only_b { + Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) => { + ranges.iter().flat_map(|r| r.iter()).collect() + } + other => panic!("Expected RangeSet only_b, got {:?}", other), + }; + assert_eq!( + only_b_vals, + vec![ + dt(2020, 1, 11), + dt(2020, 1, 12), + dt(2020, 1, 13), + dt(2020, 1, 14), + dt(2020, 1, 15) + ] + ); + } + + #[test] + fn test_rangeset_intersect_rangeset_no_overlap() { + let a = Coordinates::from(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 5))); + let b = Coordinates::from(DateTimeRange::daily(dt(2020, 2, 1), dt(2020, 2, 5))); + + let result = a.intersect(&b); + + if let Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) = &result.intersection + { + assert_eq!(ranges.len(), 0); + } else { + panic!("Expected empty RangeSet intersection"); + } + } + + #[test] + fn test_rangeset_intersect_rangeset_different_steps() { + // [Jan 1..Jan 6 daily] = {1,2,3,4,5,6} + // [Jan 1..Jan 6 every 2 days] = {1,3,5} + // intersection = {1,3,5} + let a = Coordinates::from(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 6))); + let b = Coordinates::from(DateTimeRange::new( + dt(2020, 1, 1), + dt(2020, 1, 6), + Duration::days(2), + )); + + let result = a.intersect(&b); + + let inter_vals: Vec = match &result.intersection { + Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) => { + ranges.iter().flat_map(|r| r.iter()).collect() + } + other => panic!("Expected RangeSet, got {:?}", other), + }; + assert_eq!(inter_vals, vec![dt(2020, 1, 1), dt(2020, 1, 3), dt(2020, 1, 5)]); + } + + // ---- RangeSet ∩ List ---- + + #[test] + fn test_rangeset_intersect_list() { + // [Jan 1..Jan 10 daily] ∩ {Jan 3, Jan 7, Jan 15, Jan 20} + // intersection = {Jan 3, Jan 7} + // only_range = Jan 1,2,4,5,6,8,9,10 + // only_list = {Jan 15, Jan 20} + let range_coords = Coordinates::from(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 10))); + + let list_coords = Coordinates::from( + [dt(2020, 1, 3), dt(2020, 1, 7), dt(2020, 1, 15), dt(2020, 1, 20)].as_slice(), + ); + + let result = range_coords.intersect(&list_coords); + + // intersection: {Jan 3, Jan 7} + if let Coordinates::DateTimes(DateTimeCoordinates::List(list)) = &result.intersection { + assert_eq!(list.len(), 2); + assert_eq!(list[0], dt(2020, 1, 3)); + assert_eq!(list[1], dt(2020, 1, 7)); + } else { + panic!("Expected List intersection, got {:?}", result.intersection); + } + + // only_a (range side): all range values not in list + let only_a_vals: Vec = match &result.only_a { + Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) => { + let mut v: Vec = ranges.iter().flat_map(|r| r.iter()).collect(); + v.sort(); + v + } + other => panic!("Expected RangeSet only_a, got {:?}", other), + }; + let expected_only_a: Vec = + (1..=10).filter(|&d| d != 3 && d != 7).map(|d| dt(2020, 1, d)).collect(); + assert_eq!(only_a_vals, expected_only_a); + + // only_b (list side): {Jan 15, Jan 20} + if let Coordinates::DateTimes(DateTimeCoordinates::List(list)) = &result.only_b { + assert_eq!(list.len(), 2); + assert_eq!(list[0], dt(2020, 1, 15)); + assert_eq!(list[1], dt(2020, 1, 20)); + } else { + panic!("Expected List only_b, got {:?}", result.only_b); + } + } + + #[test] + fn test_list_intersect_rangeset_symmetry() { + // Swapped: {Jan 3, Jan 7, Jan 15} ∩ [Jan 1..Jan 10 daily] + // only_a should be the list side = {Jan 15} + // only_b should be the range side + let list_coords = + Coordinates::from([dt(2020, 1, 3), dt(2020, 1, 7), dt(2020, 1, 15)].as_slice()); + let range_coords = Coordinates::from(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 10))); + + let result = list_coords.intersect(&range_coords); + + // intersection: {Jan 3, Jan 7} + if let Coordinates::DateTimes(DateTimeCoordinates::List(list)) = &result.intersection { + assert_eq!(list.len(), 2); + } else { + panic!("Expected List intersection, got {:?}", result.intersection); + } + + // only_a (list side): {Jan 15} + if let Coordinates::DateTimes(DateTimeCoordinates::List(list)) = &result.only_a { + assert_eq!(list.len(), 1); + assert_eq!(list[0], dt(2020, 1, 15)); + } else { + panic!("Expected List only_a, got {:?}", result.only_a); + } + } + + #[test] + fn test_rangeset_contains() { + let coords = Coordinates::from(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 5))); + assert!(coords.contains(dt(2020, 1, 1))); + assert!(coords.contains(dt(2020, 1, 3))); + assert!(coords.contains(dt(2020, 1, 5))); + assert!(!coords.contains(dt(2020, 1, 6))); + assert!(!coords.contains(dt(2019, 12, 31))); + } + + #[test] + fn test_rangeset_len() { + let coords = Coordinates::from(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 5))); + assert_eq!(coords.len(), 5); + } + + #[test] + fn test_rangeset_to_string() { + let r = DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 3)); + let coords = Coordinates::from(r); + let s = coords.to_string(); + assert!(s.contains("2020-01-01T00:00:00")); + assert!(s.contains("2020-01-03T00:00:00")); + } + + #[test] + fn test_hourly_range_contains() { + let r = DateTimeRange::hourly(dt_h(2020, 1, 1, 0), dt_h(2020, 1, 1, 6)); // 0,1,2,3,4,5,6h + assert!(r.contains(dt_h(2020, 1, 1, 0))); + assert!(r.contains(dt_h(2020, 1, 1, 3))); + assert!(r.contains(dt_h(2020, 1, 1, 6))); + assert!(!r.contains(dt_h(2020, 1, 1, 7))); + // Midpoint check — not on hourly grid + let mid = dt_h(2020, 1, 1, 0) + Duration::minutes(30); + assert!(!r.contains(mid)); + } + + // ---- try_compress_to_ranges tests ---- + + #[test] + fn test_compress_list_daily_to_range() { + // List of 5 consecutive days → RangeSet([Jan1..Jan5 daily]) + let mut c = DateTimeCoordinates::default(); + for day in 1..=5 { + c.append(dt(2020, 1, day)); } + c.try_compress_to_ranges(); + match &c { + DateTimeCoordinates::RangeSet(ranges) => { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].start, dt(2020, 1, 1)); + assert_eq!(ranges[0].end, dt(2020, 1, 5)); + assert_eq!(ranges[0].step, Duration::days(1)); + } + other => panic!("Expected RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_list_hourly_to_range() { + // 6 consecutive hours → RangeSet([0h..5h hourly]) + let mut c = DateTimeCoordinates::default(); + for h in 0..=5 { + c.append(dt_h(2020, 1, 1, h)); + } + c.try_compress_to_ranges(); + match &c { + DateTimeCoordinates::RangeSet(ranges) => { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].step, Duration::hours(1)); + assert_eq!(ranges[0].len(), 6); + } + other => panic!("Expected RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_list_two_runs_to_two_ranges() { + // Jan 1..3 (daily) + Jan 10..12 (daily) → two ranges + let mut c = DateTimeCoordinates::default(); + for day in [1u32, 2, 3, 10, 11, 12] { + c.append(dt(2020, 1, day)); + } + c.try_compress_to_ranges(); + match &c { + DateTimeCoordinates::RangeSet(ranges) => { + assert_eq!(ranges.len(), 2, "Expected 2 ranges, got {:?}", ranges); + assert_eq!(ranges[0].start, dt(2020, 1, 1)); + assert_eq!(ranges[0].end, dt(2020, 1, 3)); + assert_eq!(ranges[1].start, dt(2020, 1, 10)); + assert_eq!(ranges[1].end, dt(2020, 1, 12)); + } + other => panic!("Expected RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_two_elements_does_not_compress() { + // Only 2 datetimes — not worth compressing + let mut c = DateTimeCoordinates::default(); + c.append(dt(2020, 1, 1)); + c.append(dt(2020, 1, 2)); + c.try_compress_to_ranges(); + assert!(matches!(c, DateTimeCoordinates::List(_))); + } + + #[test] + fn test_compress_already_rangeset_merges_adjacent_ranges() { + // Two adjacent daily ranges: [Jan 1..Jan 3] and [Jan 4..Jan 6] → merge to [Jan 1..Jan 6] + use tiny_vec::TinyVec; + let mut ranges: TinyVec = TinyVec::new(); + ranges.push(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 3))); + ranges.push(DateTimeRange::daily(dt(2020, 1, 4), dt(2020, 1, 6))); + let mut c = DateTimeCoordinates::RangeSet(ranges); + c.try_compress_to_ranges(); + match &c { + DateTimeCoordinates::RangeSet(rs) => { + assert_eq!(rs.len(), 1, "Expected merged into 1 range, got {:?}", rs); + assert_eq!(rs[0].start, dt(2020, 1, 1)); + assert_eq!(rs[0].end, dt(2020, 1, 6)); + } + other => panic!("Expected RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_via_coordinates_try_compress() { + // Via top-level Coordinates::try_compress() + let mut coords = Coordinates::from( + [dt(2020, 6, 1), dt(2020, 6, 2), dt(2020, 6, 3), dt(2020, 6, 4)].as_slice(), + ); + coords.try_compress(); + match &coords { + Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)) => { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].start, dt(2020, 6, 1)); + assert_eq!(ranges[0].end, dt(2020, 6, 4)); + } + other => panic!("Expected compressed DateTimes RangeSet, got {:?}", other), + } + } + + #[test] + fn test_datetime_rangeset_to_string_from_string_roundtrip() { + let r = DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 10)); + let coords = Coordinates::from(r); + let s = coords.to_string(); + let parsed = Coordinates::from_string(&s); + assert_eq!(parsed, coords, "Roundtrip failed: {:?} → {:?} → {:?}", coords, s, parsed); + } + + #[test] + fn test_datetime_two_range_from_string_roundtrip() { + use tiny_vec::TinyVec; + let mut ranges: TinyVec = TinyVec::new(); + ranges.push(DateTimeRange::daily(dt(2020, 1, 1), dt(2020, 1, 5))); + ranges.push(DateTimeRange::daily(dt(2020, 2, 1), dt(2020, 2, 5))); + let c = Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)); + let s = c.to_string(); + let parsed = Coordinates::from_string(&s); + assert_eq!(parsed, c, "Roundtrip failed: {:?} → {:?} → {:?}", c, s, parsed); } } diff --git a/qubed/src/coordinates/integers.rs b/qubed/src/coordinates/integers.rs index 5ad704e..0cfca15 100644 --- a/qubed/src/coordinates/integers.rs +++ b/qubed/src/coordinates/integers.rs @@ -10,13 +10,101 @@ pub enum IntegerCoordinates { RangeSet(TinyVec), } +/// An inclusive integer range `[start, end]` with a given step size. +/// All values `start + k * step` where `start + k * step <= end` are members. #[derive(Debug, Clone, PartialEq)] pub struct IntegerRange { - start: i32, - end: i32, - step: std::num::NonZeroU16, + pub start: i32, + pub end: i32, + pub step: std::num::NonZeroU16, } +impl IntegerRange { + /// Create a new range. Panics if start > end. + pub fn new(start: i32, end: i32, step: u16) -> Self { + assert!(start <= end, "IntegerRange: start ({}) must be <= end ({})", start, end); + IntegerRange { + start, + end, + step: std::num::NonZeroU16::new(step).expect("step must be non-zero"), + } + } + + /// Create a unit-step range `[start, end]`. + pub fn new_step1(start: i32, end: i32) -> Self { + Self::new(start, end, 1) + } + + pub fn step_size(&self) -> i32 { + self.step.get() as i32 + } + + /// Number of elements in this range. + pub fn len(&self) -> usize { + if self.start > self.end { + return 0; + } + ((self.end - self.start) / self.step_size() + 1) as usize + } + + pub fn contains(&self, value: i32) -> bool { + if value < self.start || value > self.end { + return false; + } + (value - self.start) % self.step_size() == 0 + } + + /// Iterate over all values in this range. + pub fn iter(&self) -> impl Iterator + '_ { + let step = self.step_size(); + let end = self.end; + let mut current = self.start; + std::iter::from_fn(move || { + if current <= end { + let val = current; + current += step; + Some(val) + } else { + None + } + }) + } + + /// Intersect two ranges. Returns `(intersection_range, only_a_ranges, only_b_ranges)`. + /// Both ranges must have the same step (or step 1) for a clean range result; + /// otherwise the result is materialised as a Set. + /// + /// Returns `None` if the ranges do not overlap. + pub fn intersect_range(&self, other: &IntegerRange) -> Option { + let step = self.step_size(); + if step != other.step_size() { + // Different steps – can still compute overlapping anchor, but only if + // they share common elements. For simplicity we materialise below via + // the Set path instead. + return None; + } + // Same step. Compute overlap window. + let new_start = self.start.max(other.start); + let new_end = self.end.min(other.end); + if new_start > new_end { + return None; + } + // Align new_start to a multiple of step from self.start + let offset = (new_start - self.start).rem_euclid(step); + let aligned_start = if offset == 0 { new_start } else { new_start + (step - offset) }; + if aligned_start > new_end { + return None; + } + Some(IntegerRange::new(aligned_start, new_end, step as u16)) + } + + pub fn to_string(&self) -> String { + format!("{}:{}:{}", self.start, self.step, self.end) + } +} + +// ---- IntegerCoordinates methods ---- + impl IntegerCoordinates { pub(crate) fn extend(&mut self, new_coords: &IntegerCoordinates) { match new_coords { @@ -25,9 +113,30 @@ impl IntegerCoordinates { self.append(*val); } } - IntegerCoordinates::RangeSet(_) => { - unimplemented!("Integer Range compression not currently supported"); - } + IntegerCoordinates::RangeSet(ranges) => match self { + IntegerCoordinates::RangeSet(self_ranges) => { + for r in ranges.iter() { + self_ranges.push(r.clone()); + } + } + IntegerCoordinates::Set(_) => { + // Promote self to RangeSet, materialising current set members as individual ranges + let materialized: Vec = match self { + IntegerCoordinates::Set(s) => { + s.iter().map(|&v| IntegerRange::new_step1(v, v)).collect() + } + _ => unreachable!(), + }; + let mut new_ranges: TinyVec = TinyVec::new(); + for r in materialized { + new_ranges.push(r); + } + for r in ranges.iter() { + new_ranges.push(r.clone()); + } + *self = IntegerCoordinates::RangeSet(new_ranges); + } + }, } } @@ -36,8 +145,9 @@ impl IntegerCoordinates { IntegerCoordinates::Set(set) => { set.insert(new_coord); } - IntegerCoordinates::RangeSet(_) => { - unimplemented!("Integer Range compression not currently supported"); + IntegerCoordinates::RangeSet(ranges) => { + // Append as a single-element range + ranges.push(IntegerRange::new_step1(new_coord, new_coord)); } } } @@ -45,22 +155,48 @@ impl IntegerCoordinates { pub(crate) fn len(&self) -> usize { match self { IntegerCoordinates::Set(list) => list.len(), - IntegerCoordinates::RangeSet(_) => { - unimplemented!("Integer Range compression not currently supported") - } + IntegerCoordinates::RangeSet(ranges) => ranges.iter().map(|r| r.len()).sum(), } } pub(crate) fn to_string(&self) -> String { + match self { + IntegerCoordinates::Set(set) => { + set.iter().map(|v| v.to_string()).collect::>().join("/") + } + IntegerCoordinates::RangeSet(ranges) => { + ranges.iter().map(|v| v.to_string()).collect::>().join("/") + } + } + } + + /// Human-readable ASCII representation. + /// + /// Each range is rendered as: + /// - singleton `v` → `v` + /// - step-1 range → `start/to/end` + /// - stepped range → `start/to/end/by/step` + /// + /// Multiple ranges / singletons are separated by `|`. + /// Plain `Set` values use the standard `/`-joined format (no ranges present). + pub(crate) fn to_ascii_string(&self) -> String { match self { IntegerCoordinates::Set(set) => { set.iter().map(|v| v.to_string()).collect::>().join("/") } IntegerCoordinates::RangeSet(ranges) => ranges .iter() - .map(|v| format!("{}:{}:{}", v.start, v.step, v.end)) + .map(|r| { + if r.start == r.end { + r.start.to_string() + } else if r.step_size() == 1 { + format!("{}/to/{}", r.start, r.end) + } else { + format!("{}/to/{}/by/{}", r.start, r.end, r.step_size()) + } + }) .collect::>() - .join("/"), + .join("|"), } } @@ -69,6 +205,7 @@ impl IntegerCoordinates { other: &IntegerCoordinates, ) -> IntersectionResult { match (self, other) { + // Set ∩ Set — fast merge-walk on sorted sets (IntegerCoordinates::Set(set_a), IntegerCoordinates::Set(set_b)) => { let result = set_a.intersect(set_b); IntersectionResult { @@ -77,8 +214,18 @@ impl IntegerCoordinates { only_b: IntegerCoordinates::Set(result.only_b), } } - _ => { - unimplemented!("Integer Range compression not currently supported"); + + // RangeSet ∩ RangeSet + (IntegerCoordinates::RangeSet(ranges_a), IntegerCoordinates::RangeSet(ranges_b)) => { + intersect_range_sets(ranges_a, ranges_b) + } + + // RangeSet ∩ Set (or Set ∩ RangeSet — handled symmetrically) + (IntegerCoordinates::RangeSet(ranges), IntegerCoordinates::Set(set)) => { + intersect_rangeset_with_set(ranges, set, false) + } + (IntegerCoordinates::Set(set), IntegerCoordinates::RangeSet(ranges)) => { + intersect_rangeset_with_set(ranges, set, true) } } } @@ -102,6 +249,126 @@ impl IntegerCoordinates { } } +/// Intersect two range sets. Returns an IntersectionResult where each part is a RangeSet. +fn intersect_range_sets( + ranges_a: &TinyVec, + ranges_b: &TinyVec, +) -> IntersectionResult { + let mut intersection: TinyVec = TinyVec::new(); + // Track which ranges in a/b were fully consumed by intersections + // We use a materialised set approach: collect all values from each side, + // then do set intersect, then try to re-compress into ranges. + // For same-step ranges we can do range arithmetic; for mixed steps we materialise. + + // Collect which (range_a_idx, range_b_idx) pairs overlap + let mut a_consumed: Vec = vec![false; ranges_a.len()]; + let mut b_consumed: Vec = vec![false; ranges_b.len()]; + + for (ia, ra) in ranges_a.iter().enumerate() { + for (ib, rb) in ranges_b.iter().enumerate() { + if ra.step == rb.step { + if let Some(inter) = ra.intersect_range(rb) { + intersection.push(inter); + a_consumed[ia] = true; + b_consumed[ib] = true; + } + } else { + // Different steps: materialise overlap as individual values + let start = ra.start.max(rb.start); + let end = ra.end.min(rb.end); + if start <= end { + for v in ra.iter().filter(|&v| v >= start && v <= end && rb.contains(v)) { + intersection.push(IntegerRange::new_step1(v, v)); + a_consumed[ia] = true; + b_consumed[ib] = true; + } + } + } + } + } + + // only_a: ranges in a not covered by any intersection + // For simplicity we keep unconsumed full ranges in only_a / only_b + // (partially consumed ranges are a complex case; we emit the whole range minus intersection + // which requires range subtraction — we keep it simple and store the full unconsumed ranges + // plus materialise the partial ones) + let mut only_a: TinyVec = TinyVec::new(); + let mut only_b: TinyVec = TinyVec::new(); + + for (ia, ra) in ranges_a.iter().enumerate() { + if !a_consumed[ia] { + only_a.push(ra.clone()); + } else { + // Emit values from ra not in any rb + for v in ra.iter() { + if !ranges_b.iter().any(|rb| rb.contains(v)) { + only_a.push(IntegerRange::new_step1(v, v)); + } + } + } + } + + for (ib, rb) in ranges_b.iter().enumerate() { + if !b_consumed[ib] { + only_b.push(rb.clone()); + } else { + for v in rb.iter() { + if !ranges_a.iter().any(|ra| ra.contains(v)) { + only_b.push(IntegerRange::new_step1(v, v)); + } + } + } + } + + IntersectionResult { + intersection: IntegerCoordinates::RangeSet(intersection), + only_a: IntegerCoordinates::RangeSet(only_a), + only_b: IntegerCoordinates::RangeSet(only_b), + } +} + +/// Intersect a RangeSet (ranges) with a Set of individual values. +/// `swapped` indicates whether the original call had (Set, RangeSet) — used to swap only_a/only_b. +fn intersect_rangeset_with_set( + ranges: &TinyVec, + set: &TinyOrderedSet, + swapped: bool, +) -> IntersectionResult { + let mut intersection_set: TinyOrderedSet = TinyOrderedSet::new(); + let mut only_set: TinyOrderedSet = TinyOrderedSet::new(); + let mut only_ranges_vals: TinyVec = TinyVec::new(); + + for &v in set.iter() { + if ranges.iter().any(|r| r.contains(v)) { + intersection_set.insert(v); + } else { + only_set.insert(v); + } + } + + // For values in ranges not covered by set, keep the ranges (minus matched values) + for r in ranges.iter() { + for v in r.iter() { + if !set.contains(&v) { + only_ranges_vals.push(IntegerRange::new_step1(v, v)); + } + } + } + + let intersection = IntegerCoordinates::Set(intersection_set); + let (only_a, only_b) = if swapped { + // original was (Set, RangeSet): only_a = set side, only_b = range side + (IntegerCoordinates::Set(only_set), IntegerCoordinates::RangeSet(only_ranges_vals)) + } else { + // original was (RangeSet, Set): only_a = range side, only_b = set side + (IntegerCoordinates::RangeSet(only_ranges_vals), IntegerCoordinates::Set(only_set)) + }; + + IntersectionResult { intersection, only_a, only_b } +} + +// ---- From impls ---- + impl From for Coordinates { fn from(value: IntegerCoordinates) -> Self { Coordinates::Integers(value) @@ -136,17 +403,135 @@ impl From<&[i32; N]> for Coordinates { } } +/// Construct `Coordinates` from a single `IntegerRange`. +impl From for Coordinates { + fn from(value: IntegerRange) -> Self { + let mut ranges: TinyVec = TinyVec::new(); + ranges.push(value); + Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) + } +} + +/// Construct `Coordinates` from a slice of `IntegerRange`. +impl From<&[IntegerRange]> for Coordinates { + fn from(value: &[IntegerRange]) -> Self { + let mut ranges: TinyVec = TinyVec::new(); + for r in value { + ranges.push(r.clone()); + } + Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) + } +} + impl Default for IntegerCoordinates { fn default() -> Self { IntegerCoordinates::Set(TinyOrderedSet::new()) } } +impl IntegerCoordinates { + /// Attempt to compress the coordinate set into a tighter `RangeSet` representation. + /// + /// The algorithm: + /// 1. Collect all values (sorted). + /// 2. Greedy scan: extend the current run as long as the step is consistent. + /// A run is only emitted as a `Range` if it has ≥ 3 elements (otherwise storing + /// as individual values is equally compact or smaller). + /// 3. Only replace `self` with a `RangeSet` when the resulting number of + /// ranges (where singletons also count as one range each) is strictly less + /// than the original element count — i.e. it actually saves space. + pub fn try_compress_to_ranges(&mut self) { + // Collect sorted values (works for both variants) + let mut values: Vec = match self { + IntegerCoordinates::Set(set) => set.iter().copied().collect(), + IntegerCoordinates::RangeSet(ranges) => { + let mut v: Vec = ranges.iter().flat_map(|r| r.iter()).collect(); + v.sort_unstable(); + v.dedup(); + v + } + }; + + if values.len() < 3 { + // Nothing to compress — a range only helps at 3+ elements. + return; + } + + values.sort_unstable(); + values.dedup(); + + let ranges = compress_integers_to_ranges(&values); + + // Only replace if the number of ranges is strictly smaller than the + // original element count (a range with N elements saves N-1 "slots"). + let range_count = ranges.len(); + if range_count < values.len() { + let mut rv: TinyVec = TinyVec::new(); + for r in ranges { + rv.push(r); + } + *self = IntegerCoordinates::RangeSet(rv); + } + } +} + +/// Partition a sorted, deduplicated slice of integers into the minimum set of +/// uniform-step ranges. +/// +/// At each position the algorithm measures the run length achievable with step +/// `values[i+1] - values[i]`. If the run is ≥ 3 elements it is emitted as a +/// range; otherwise only a single singleton is emitted and the position advances +/// by one. This "emit one, retry" behaviour means the algorithm never locks a +/// value into a short run that would prevent a longer run starting one position +/// later. For example [1,2,3,7,10,11,12] produces: +/// [1:1:3] (run of 3) +/// 7 (singleton — run [7,10] step 3 is only length 2) +/// [10:1:12] (run of 3 detected once 10 is re-evaluated as the start) +pub(crate) fn compress_integers_to_ranges(values: &[i32]) -> Vec { + if values.is_empty() { + return vec![]; + } + + let mut result: Vec = Vec::new(); + let mut i = 0; + + while i < values.len() { + // Measure the run starting at i using the step to the very next element. + let run_len = if i + 1 < values.len() { + let step = values[i + 1] - values[i]; + if step > 0 && step <= u16::MAX as i32 { + let mut j = i; + while j + 1 < values.len() && values[j + 1] - values[j] == step { + j += 1; + } + j - i + 1 + } else { + 1 + } + } else { + 1 + }; + + if run_len >= 3 { + let step = values[i + 1] - values[i]; + result.push(IntegerRange::new(values[i], values[i + run_len - 1], step as u16)); + i += run_len; + } else { + // Run too short — emit just this element as a singleton and retry + // from the next position so a better run can be found. + result.push(IntegerRange::new_step1(values[i], values[i])); + i += 1; + } + } + + result +} + impl IntegerCoordinates { pub fn contains(&self, value: i32) -> bool { match self { IntegerCoordinates::Set(set) => set.contains(&value), - IntegerCoordinates::RangeSet(_) => unimplemented!("RangeSet contains not implemented"), + IntegerCoordinates::RangeSet(ranges) => ranges.iter().any(|r| r.contains(value)), } } } @@ -155,6 +540,8 @@ impl IntegerCoordinates { mod tests { use super::*; + // ---- Set tests (existing) ---- + #[test] fn test_integer_coordinates_intersect_tiny_ordered_tiny_ordered() { let mut coords_a = Coordinates::Empty; @@ -181,10 +568,391 @@ mod tests { let mut expected_only_b = Coordinates::Empty; expected_only_b.append(4); - println!("Result: {:?}", result); - assert_eq!(result.intersection, expected_intersection); assert_eq!(result.only_a, expected_only_a); assert_eq!(result.only_b, expected_only_b); } + + // ---- IntegerRange unit tests ---- + + #[test] + fn test_integer_range_contains() { + let r = IntegerRange::new(0, 10, 2); // 0, 2, 4, 6, 8, 10 + assert!(r.contains(0)); + assert!(r.contains(4)); + assert!(r.contains(10)); + assert!(!r.contains(1)); + assert!(!r.contains(3)); + assert!(!r.contains(11)); + assert!(!r.contains(-2)); + } + + #[test] + fn test_integer_range_len() { + let r = IntegerRange::new(1, 9, 2); // 1, 3, 5, 7, 9 → 5 elements + assert_eq!(r.len(), 5); + + let r2 = IntegerRange::new_step1(1, 5); // 1,2,3,4,5 → 5 + assert_eq!(r2.len(), 5); + + let r3 = IntegerRange::new(0, 0, 1); // single element + assert_eq!(r3.len(), 1); + } + + #[test] + fn test_integer_range_iter() { + let r = IntegerRange::new(0, 6, 3); // 0, 3, 6 + let vals: Vec = r.iter().collect(); + assert_eq!(vals, vec![0, 3, 6]); + } + + #[test] + fn test_integer_range_to_string() { + let r = IntegerRange::new(1, 10, 2); + assert_eq!(r.to_string(), "1:2:10"); + } + + // ---- RangeSet intersect tests ---- + + #[test] + fn test_rangeset_intersect_rangeset_overlapping() { + // [1..10 step 1] ∩ [5..15 step 1] = [5..10], only_a=[1..4], only_b=[11..15] + let a = Coordinates::from(IntegerRange::new_step1(1, 10)); + let b = Coordinates::from(IntegerRange::new_step1(5, 15)); + + let result = a.intersect(&b); + + // intersection should contain 5..10 + if let Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) = &result.intersection { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].start, 5); + assert_eq!(ranges[0].end, 10); + } else { + panic!("Expected RangeSet intersection, got {:?}", result.intersection); + } + + // only_a: values 1..4 (each as singleton range) + let only_a_vals: Vec = match &result.only_a { + Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) => { + ranges.iter().flat_map(|r| r.iter()).collect() + } + other => panic!("Expected RangeSet only_a, got {:?}", other), + }; + assert_eq!(only_a_vals, vec![1, 2, 3, 4]); + + // only_b: values 11..15 + let only_b_vals: Vec = match &result.only_b { + Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) => { + ranges.iter().flat_map(|r| r.iter()).collect() + } + other => panic!("Expected RangeSet only_b, got {:?}", other), + }; + assert_eq!(only_b_vals, vec![11, 12, 13, 14, 15]); + } + + #[test] + fn test_rangeset_intersect_rangeset_no_overlap() { + let a = Coordinates::from(IntegerRange::new_step1(1, 5)); + let b = Coordinates::from(IntegerRange::new_step1(10, 20)); + + let result = a.intersect(&b); + + // intersection empty + if let Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) = &result.intersection { + assert_eq!(ranges.len(), 0); + } else { + panic!("Expected empty RangeSet intersection"); + } + } + + #[test] + fn test_rangeset_intersect_rangeset_different_steps() { + // [0..6 step 2] = {0,2,4,6} ∩ [0..9 step 3] = {0,3,6} → intersection = {0,6} + let mut a_ranges: TinyVec = TinyVec::new(); + a_ranges.push(IntegerRange::new(0, 6, 2)); + let a = Coordinates::Integers(IntegerCoordinates::RangeSet(a_ranges)); + + let mut b_ranges: TinyVec = TinyVec::new(); + b_ranges.push(IntegerRange::new(0, 9, 3)); + let b = Coordinates::Integers(IntegerCoordinates::RangeSet(b_ranges)); + + let result = a.intersect(&b); + + let inter_vals: Vec = match &result.intersection { + Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) => { + ranges.iter().flat_map(|r| r.iter()).collect() + } + other => panic!("Expected RangeSet, got {:?}", other), + }; + assert_eq!(inter_vals, vec![0, 6]); + } + + #[test] + fn test_rangeset_intersect_set() { + // [1..10 step 1] ∩ {3, 7, 11, 20} → intersection = {3, 7}, only_range = {1,2,4,5,6,8,9,10}, only_set = {11, 20} + let range_coords = Coordinates::from(IntegerRange::new_step1(1, 10)); + + let mut set_coords = Coordinates::Empty; + set_coords.append(3); + set_coords.append(7); + set_coords.append(11); + set_coords.append(20); + + let result = range_coords.intersect(&set_coords); + + // intersection is a Set with {3, 7} + if let Coordinates::Integers(IntegerCoordinates::Set(set)) = &result.intersection { + let vals: Vec = set.iter().copied().collect(); + assert_eq!(vals, vec![3, 7]); + } else { + panic!("Expected Set intersection, got {:?}", result.intersection); + } + + // only_a (range side): all range values not in {3,7,11,20} + let only_a_vals: Vec = match &result.only_a { + Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) => { + let mut v: Vec = ranges.iter().flat_map(|r| r.iter()).collect(); + v.sort(); + v + } + other => panic!("Expected RangeSet only_a, got {:?}", other), + }; + assert_eq!(only_a_vals, vec![1, 2, 4, 5, 6, 8, 9, 10]); + + // only_b (set side): {11, 20} + if let Coordinates::Integers(IntegerCoordinates::Set(set)) = &result.only_b { + let vals: Vec = set.iter().copied().collect(); + assert_eq!(vals, vec![11, 20]); + } else { + panic!("Expected Set only_b, got {:?}", result.only_b); + } + } + + #[test] + fn test_set_intersect_rangeset_symmetry() { + // Swapped: {3, 7, 11} ∩ [1..10] — only_a/only_b should be swapped relative to above + let mut set_coords = Coordinates::Empty; + set_coords.append(3); + set_coords.append(7); + set_coords.append(11); + + let range_coords = Coordinates::from(IntegerRange::new_step1(1, 10)); + + let result = set_coords.intersect(&range_coords); + + // intersection: {3, 7} + if let Coordinates::Integers(IntegerCoordinates::Set(set)) = &result.intersection { + let vals: Vec = set.iter().copied().collect(); + assert_eq!(vals, vec![3, 7]); + } else { + panic!("Expected Set intersection, got {:?}", result.intersection); + } + + // only_a (set side): {11} + if let Coordinates::Integers(IntegerCoordinates::Set(set)) = &result.only_a { + let vals: Vec = set.iter().copied().collect(); + assert_eq!(vals, vec![11]); + } else { + panic!("Expected Set only_a, got {:?}", result.only_a); + } + } + + #[test] + fn test_rangeset_contains() { + let coords = Coordinates::from(IntegerRange::new(0, 10, 2)); // 0,2,4,6,8,10 + assert!(coords.contains(0i32)); + assert!(coords.contains(4i32)); + assert!(coords.contains(10i32)); + assert!(!coords.contains(1i32)); + assert!(!coords.contains(11i32)); + } + + #[test] + fn test_rangeset_len() { + let coords = Coordinates::from(IntegerRange::new_step1(1, 5)); // 5 elements + assert_eq!(coords.len(), 5); + } + + #[test] + fn test_rangeset_to_string() { + let coords = Coordinates::from(IntegerRange::new(1, 10, 2)); + assert_eq!(coords.to_string(), "1:2:10"); + } + + #[test] + fn test_rangeset_extend_with_set() { + // A RangeSet extended with a Set should keep the range and add individual points + let range_coords = IntegerRange::new_step1(1, 5); + let mut coords = Coordinates::from(range_coords); + coords.append(10i32); + // Should still be a RangeSet with 6 elements + assert_eq!(coords.len(), 6); + assert!(coords.contains(3i32)); + assert!(coords.contains(10i32)); + } + + // ---- try_compress_to_ranges tests ---- + + #[test] + fn test_compress_set_consecutive_step1_to_range() { + // {1,2,3,4,5} → RangeSet([1:1:5]) + let mut c = IntegerCoordinates::default(); + for v in [1, 2, 3, 4, 5] { + c.append(v); + } + c.try_compress_to_ranges(); + match &c { + IntegerCoordinates::RangeSet(ranges) => { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].start, 1); + assert_eq!(ranges[0].end, 5); + assert_eq!(ranges[0].step_size(), 1); + } + other => panic!("Expected RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_set_even_numbers_to_range() { + // {0,2,4,6,8,10} → RangeSet([0:2:10]) + let mut c = IntegerCoordinates::default(); + for v in [0, 2, 4, 6, 8, 10] { + c.append(v); + } + c.try_compress_to_ranges(); + match &c { + IntegerCoordinates::RangeSet(ranges) => { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].start, 0); + assert_eq!(ranges[0].end, 10); + assert_eq!(ranges[0].step_size(), 2); + } + other => panic!("Expected RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_set_two_runs_to_two_ranges() { + // {1,2,3, 10,12,14} → two ranges + let mut c = IntegerCoordinates::default(); + for v in [1, 2, 3, 10, 12, 14] { + c.append(v); + } + c.try_compress_to_ranges(); + match &c { + IntegerCoordinates::RangeSet(ranges) => { + assert_eq!(ranges.len(), 2, "Expected 2 ranges, got {:?}", ranges); + // First range: 1..3 step 1 + assert_eq!(ranges[0].start, 1); + assert_eq!(ranges[0].end, 3); + assert_eq!(ranges[0].step_size(), 1); + // Second range: 10..14 step 2 + assert_eq!(ranges[1].start, 10); + assert_eq!(ranges[1].end, 14); + assert_eq!(ranges[1].step_size(), 2); + } + other => panic!("Expected RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_set_with_isolated_singletons() { + // {1,2,3, 7, 10,11,12} → [1:1:3], singleton 7, [10:1:12] + // The algorithm emits singleton 7 (because the run 7→10 is only length 2) + // and then re-evaluates from 10, finding the run [10,11,12]. + let mut c = IntegerCoordinates::default(); + for v in [1, 2, 3, 7, 10, 11, 12] { + c.append(v); + } + c.try_compress_to_ranges(); + match &c { + IntegerCoordinates::RangeSet(ranges) => { + assert_eq!(ranges.len(), 3, "Expected 3 ranges, got {:?}", ranges); + let vals: Vec = ranges.iter().flat_map(|r| r.iter()).collect(); + assert_eq!(vals, vec![1, 2, 3, 7, 10, 11, 12]); + // First range: 1..3 step 1 + assert_eq!(ranges[0].start, 1); + assert_eq!(ranges[0].end, 3); + assert_eq!(ranges[0].step_size(), 1); + // Second: singleton 7 + assert_eq!(ranges[1].start, 7); + assert_eq!(ranges[1].end, 7); + // Third range: 10..12 step 1 + assert_eq!(ranges[2].start, 10); + assert_eq!(ranges[2].end, 12); + assert_eq!(ranges[2].step_size(), 1); + } + other => panic!("Expected RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_two_elements_does_not_compress() { + // {5, 6} — only 2 elements, not worth compressing + let mut c = IntegerCoordinates::default(); + c.append(5); + c.append(6); + c.try_compress_to_ranges(); + // Should remain a Set (compression threshold is 3+ elements per run) + assert!(matches!(c, IntegerCoordinates::Set(_))); + } + + #[test] + fn test_compress_already_rangeset_merges_adjacent_ranges() { + // Two adjacent step-1 ranges: [1..3] and [4..6] → should merge to [1..6] + use tiny_vec::TinyVec; + let mut ranges: TinyVec = TinyVec::new(); + ranges.push(IntegerRange::new_step1(1, 3)); + ranges.push(IntegerRange::new_step1(4, 6)); + let mut c = IntegerCoordinates::RangeSet(ranges); + c.try_compress_to_ranges(); + match &c { + IntegerCoordinates::RangeSet(rs) => { + assert_eq!(rs.len(), 1); + assert_eq!(rs[0].start, 1); + assert_eq!(rs[0].end, 6); + } + other => panic!("Expected merged RangeSet, got {:?}", other), + } + } + + #[test] + fn test_compress_via_coordinates_try_compress() { + // Using the top-level Coordinates::try_compress() API + let mut coords = Coordinates::Empty; + for v in [10, 11, 12, 13, 14, 15] { + coords.append(v); + } + coords.try_compress(); + match &coords { + Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)) => { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].start, 10); + assert_eq!(ranges[0].end, 15); + } + other => panic!("Expected compressed Integers RangeSet, got {:?}", other), + } + } + + #[test] + fn test_integer_rangeset_from_string_roundtrip() { + // to_string → from_string roundtrip for integer RangeSet + let coords = Coordinates::from(IntegerRange::new(1, 10, 2)); + let s = coords.to_string(); + assert_eq!(s, "1:2:10"); + let parsed = Coordinates::from_string(&s); + assert_eq!(parsed, coords); + } + + #[test] + fn test_integer_two_range_from_string_roundtrip() { + use tiny_vec::TinyVec; + let mut ranges: TinyVec = TinyVec::new(); + ranges.push(IntegerRange::new_step1(1, 5)); + ranges.push(IntegerRange::new_step1(10, 15)); + let c = Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)); + let s = c.to_string(); + let parsed = Coordinates::from_string(&s); + assert_eq!(parsed, c); + } } diff --git a/qubed/src/coordinates/mod.rs b/qubed/src/coordinates/mod.rs index 83fddcb..ae439c6 100644 --- a/qubed/src/coordinates/mod.rs +++ b/qubed/src/coordinates/mod.rs @@ -51,6 +51,32 @@ impl Coordinates { if s.is_empty() { return Coordinates::Empty; } + + // ── Range notations (checked before the generic '/' split) ────────── + + // ASCII range notation: `start/to/end[/by/step]` for integers/datetimes, + // or multiple such ranges joined by `|`. + // The presence of `/to/` is the distinguishing marker. + if s.contains("/to/") || s.contains('|') { + if let Some(coords) = parse_ascii_range_notation(s) { + return coords; + } + } + + // Machine range notation for datetime RangeSets + // (format: "start:s:end/..." — contains 's:' after digits). + if let Some(dt_ranges) = datetime::DateTimeCoordinates::parse_range_from_str(s) { + return Coordinates::DateTimes(dt_ranges); + } + + // Machine range notation for integer RangeSets ("start:step:end/..."). + if looks_like_integer_rangeset(s) { + if let Some(coords) = parse_integer_rangeset_str(s) { + return coords; + } + } + + // ── Generic '/' separated values ──────────────────────────────────── let mut coords = Coordinates::Empty; let split: Vec<&str> = s.split('/').collect(); @@ -61,7 +87,6 @@ impl Coordinates { && part.chars().nth(1).map_or(false, |c| c.is_ascii_digit()); if has_leading_zero { - // Preserve as string to keep formatting coords.append(part.to_string()); } else if let Ok(int_val) = part.parse::() { coords.append(int_val); @@ -87,6 +112,19 @@ impl Coordinates { } } + /// Serialize to a human-readable string for ASCII tree output. + /// + /// Ranges are written in the `start/to/end` (step-1 or daily) or + /// `start/to/end/by/step` notation. Multiple ranges are joined by `|`. + /// All other variants fall back to `to_string()`. + pub fn to_ascii_string(&self) -> String { + match self { + Coordinates::Integers(ints) => ints.to_ascii_string(), + Coordinates::DateTimes(dts) => dts.to_ascii_string(), + other => other.to_string(), + } + } + pub fn len(&self) -> usize { match self { Coordinates::Empty => 0, @@ -107,6 +145,20 @@ impl Coordinates { self.len() == 0 } + /// Attempt to compress coordinate values into a more compact range representation. + /// + /// - `Integers(Set)` or `Integers(RangeSet)` → may become `RangeSet` if + /// consecutive runs are found that reduce the total number of stored items. + /// - `DateTimes(List)` or `DateTimes(RangeSet)` → same treatment. + /// - All other variants are unchanged. + pub fn try_compress(&mut self) { + match self { + Coordinates::Integers(ints) => ints.try_compress_to_ranges(), + Coordinates::DateTimes(dts) => dts.try_compress_to_ranges(), + _ => {} + } + } + pub fn contains(&self, value: T) -> bool where T: Into, @@ -175,6 +227,14 @@ impl Coordinates { only_b: Coordinates::Strings(result.only_b), } } + (Coordinates::DateTimes(dts_a), Coordinates::DateTimes(dts_b)) => { + let result = dts_a.intersect(dts_b); + IntersectionResult { + intersection: Coordinates::DateTimes(result.intersection), + only_a: Coordinates::DateTimes(result.only_a), + only_b: Coordinates::DateTimes(result.only_b), + } + } _ => { unimplemented!("Intersection not implemented for these coordinate types"); } @@ -216,6 +276,159 @@ impl Default for Coordinates { } } +/// Parse the human-readable ASCII range notation produced by `to_ascii_string`. +/// +/// Supported forms (one or more items joined by `|`): +/// - Integer singleton: `7` +/// - Integer step-1 range: `1/to/10` +/// - Integer stepped range: `0/to/10/by/2` +/// - DateTime singleton: `2020-01-01T00:00:00` +/// - DateTime daily range: `2020-01-01T00:00:00/to/2020-01-10T00:00:00` +/// - DateTime range with step: `2020-01-01T00:00:00/to/2020-01-10T00:00:00/by/3600s` +/// +/// Returns `None` if the input doesn't match this grammar at all (so the +/// caller can fall through to other parsers). +fn parse_ascii_range_notation(s: &str) -> Option { + use chrono::Duration; + use datetime::{DateTimeCoordinates, DateTimeRange}; + use integers::{IntegerCoordinates, IntegerRange}; + use tiny_vec::TinyVec; + + // Each `|`-separated segment is one range or singleton. + let segments: Vec<&str> = s.split('|').collect(); + + enum Item { + Int(IntegerRange), + Dt(DateTimeRange), + } + + let mut items: Vec = Vec::new(); + + for seg in &segments { + // Does this segment contain `/to/`? + if let Some(to_pos) = seg.find("/to/") { + let start_str = &seg[..to_pos]; + let rest = &seg[to_pos + 4..]; // after "/to/" + + // Check for optional `/by/` suffix + let (end_str, step_str) = if let Some(by_pos) = rest.find("/by/") { + (&rest[..by_pos], Some(&rest[by_pos + 4..])) + } else { + (rest, None) + }; + + // Try integer range + if let (Ok(start), Ok(end)) = (start_str.parse::(), end_str.parse::()) { + let step: i32 = if let Some(ss) = step_str { ss.parse().ok()? } else { 1 }; + if step <= 0 || step > u16::MAX as i32 || start > end { + return None; + } + items.push(Item::Int(IntegerRange::new(start, end, step as u16))); + continue; + } + + // Try datetime range + let start_dt = DateTimeCoordinates::parse_from_str(start_str)?; + let end_dt = DateTimeCoordinates::parse_from_str(end_str)?; + let step_dur = if let Some(ss) = step_str { + parse_duration_str(ss)? + } else { + Duration::days(1) // default step for datetime ranges + }; + if start_dt > end_dt || step_dur <= Duration::zero() { + return None; + } + items.push(Item::Dt(DateTimeRange::new(start_dt, end_dt, step_dur))); + } else { + // No `/to/` — must be a singleton. + // Try integer first + if let Ok(v) = seg.parse::() { + items.push(Item::Int(IntegerRange::new_step1(v, v))); + continue; + } + // Try datetime + if let Some(dt) = DateTimeCoordinates::parse_from_str(seg) { + items.push(Item::Dt(DateTimeRange::new(dt, dt, Duration::seconds(1)))); + continue; + } + // Unknown token — fall through to generic parser + return None; + } + } + + if items.is_empty() { + return None; + } + + // All items must be the same type (Int or Dt). + let all_int = items.iter().all(|it| matches!(it, Item::Int(_))); + let all_dt = items.iter().all(|it| matches!(it, Item::Dt(_))); + + if all_int { + let mut ranges: TinyVec = TinyVec::new(); + for it in items { + if let Item::Int(r) = it { + ranges.push(r); + } + } + Some(Coordinates::Integers(IntegerCoordinates::RangeSet(ranges))) + } else if all_dt { + let mut ranges: TinyVec = TinyVec::new(); + for it in items { + if let Item::Dt(r) = it { + ranges.push(r); + } + } + Some(Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges))) + } else { + // Mixed int/datetime — shouldn't occur in practice; fall through. + None + } +} + +/// Parse a duration string of the form `s` (e.g. `"86400s"`, `"3600s"`). +fn parse_duration_str(s: &str) -> Option { + let s = s.trim(); + if let Some(num) = s.strip_suffix('s') { + let secs: i64 = num.parse().ok()?; + Some(chrono::Duration::seconds(secs)) + } else { + None + } +} +fn looks_like_integer_rangeset(s: &str) -> bool { + s.split('/').all(|part| { + let segs: Vec<&str> = part.splitn(3, ':').collect(); + segs.len() == 3 && segs.iter().all(|seg| seg.parse::().is_ok()) + }) +} + +/// Parse a string of the form `"start:step:end/start:step:end/..."` into +/// `Coordinates::Integers(RangeSet(...))`. +fn parse_integer_rangeset_str(s: &str) -> Option { + use integers::IntegerRange; + use tiny_vec::TinyVec; + + let mut ranges: TinyVec = TinyVec::new(); + for part in s.split('/') { + let segs: Vec<&str> = part.splitn(3, ':').collect(); + if segs.len() != 3 { + return None; + } + let start: i32 = segs[0].parse().ok()?; + let step: i32 = segs[1].parse().ok()?; + let end: i32 = segs[2].parse().ok()?; + if step <= 0 || step > u16::MAX as i32 || start > end { + return None; + } + ranges.push(IntegerRange::new(start, end, step as u16)); + } + if ranges.is_empty() { + return None; + } + Some(Coordinates::Integers(integers::IntegerCoordinates::RangeSet(ranges))) +} + // ------------- Intersection ------------------ #[derive(Debug, Clone, PartialEq)] @@ -374,6 +587,15 @@ impl Coordinates { map.insert("datetimes".to_string(), Value::Array(vals)); } } + datetime::DateTimeCoordinates::RangeSet(_) => { + // fallback to textual form + if boxed.datetimes.len() > 0 { + map.insert( + "datetimes".to_string(), + Value::String(boxed.datetimes.to_string()), + ); + } + } } Value::Object(map) @@ -388,6 +610,7 @@ impl Coordinates { .collect(); Value::Array(vals) } + datetime::DateTimeCoordinates::RangeSet(_) => Value::String(coords.to_string()), }, } } diff --git a/qubed/src/qube.rs b/qubed/src/qube.rs index e5c4d74..c909495 100644 --- a/qubed/src/qube.rs +++ b/qubed/src/qube.rs @@ -194,14 +194,15 @@ impl Qube { Ok(node_id) } - pub fn all_unique_dim_coords(&self) -> BTreeMap { + pub fn all_unique_dim_coords(&mut self) -> BTreeMap { + // TODO let mut map: BTreeMap = BTreeMap::new(); for (_id, node) in self.nodes.iter() { if let Some(dim_str) = self.dimension_str(&node.dim) { let coords = node.coords.clone(); if coords.is_empty() { - continue; // Skip empty coordinates (incl. the virtual root node) + continue; // Skip empty coordinates } // if there is no entry in map for this dimension, just fill it in with coords, otherwise extend the current entry with coords map.entry(dim_str.to_string()) @@ -212,135 +213,6 @@ impl Qube { map } - /// Returns the set of all dimension names present anywhere in the Qube. - /// - /// This is the set of keys from [`all_unique_dim_coords`]. - /// - /// # Examples - /// ``` - /// use qubed::Qube; - /// let q = Qube::from_ascii("root\n└── class=od\n └── param=1/2").unwrap(); - /// let dims = q.dimensions(); - /// assert!(dims.contains("class")); - /// assert!(dims.contains("param")); - /// assert!(!dims.contains("root")); - /// ``` - pub fn dimensions(&self) -> HashSet { - self.all_unique_dim_coords().into_keys().collect() - } - - /// Returns the set of dimension names present in **every** leaf path (datacube). - /// - /// For a Qube with uniform depth this equals [`dimensions`]. For an - /// irregular Qube some branches may be missing a dimension; only those - /// that appear in *all* branches are returned. - /// - /// # Examples - /// ``` - /// use qubed::Qube; - /// use qubed::Datacube; - /// use qubed::Coordinates; - /// - /// // Both datacubes share "param"; only one has "time". - /// let mut dc1 = Datacube::new(); - /// dc1.add_coordinate("param", Coordinates::from_string("2t/tp")); - /// dc1.add_coordinate("time", Coordinates::from_string("0/1/2")); - /// let mut qube = Qube::from_datacube(&dc1, Some(&["param".to_string(), "time".to_string()])); - /// - /// let mut dc2 = Datacube::new(); - /// dc2.add_coordinate("param", Coordinates::from_string("msl")); - /// let mut other = Qube::from_datacube(&dc2, None); - /// - /// qube.append(&mut other); - /// - /// let common = qube.common_dimensions(); - /// assert!(common.contains("param")); - /// assert!(!common.contains("time")); - /// ``` - pub fn common_dimensions(&self) -> HashSet { - let datacubes = self.to_datacubes(); - if datacubes.is_empty() { - return HashSet::new(); - } - - let mut iter = datacubes.iter().map(|dc| { - dc.coordinates() - .iter() - .filter(|(_, v)| !v.is_empty()) // exclude the virtual root node - .map(|(k, _)| k.clone()) - .collect::>() - }); - - let first = match iter.next() { - Some(s) => s, - None => return HashSet::new(), - }; - - iter.fold(first, |acc, keys| acc.intersection(&keys).cloned().collect()) - } - - /// Wraps the entire existing Qube tree under a new outer dimension. - /// - /// All current children of the root are re-parented to a new node that - /// has the given `key` and `values` as its dimension and coordinates. - /// The result is that `key` becomes the outermost dimension of the Qube. - /// - /// Calling `expand` multiple times nests the dimensions from the inside - /// out: each call wraps the *current* tree, so the last call produces - /// the outermost dimension. - /// - /// # Examples - /// ``` - /// use qubed::{Qube, Coordinates}; - /// - /// let mut q = Qube::from_ascii("root\n└── param=2t/tp\n └── time=0/1/2").unwrap(); - /// q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); - /// - /// let dims = q.dimensions(); - /// assert!(dims.contains("ensemble")); - /// assert!(dims.contains("param")); - /// assert!(dims.contains("time")); - /// ``` - pub fn expand(&mut self, key: &str, values: Coordinates) -> Result<(), String> { - let root_id = self.root_id; - - // 1. Clone root's current children before any mutation. - let old_root_children: BTreeMap> = self - .nodes - .get(root_id) - .ok_or_else(|| "Root node not found".to_string())? - .children - .clone(); - - // 2. Clear root's children so get_or_create_child starts with a clean slate. - if let Some(root) = self.nodes.get_mut(root_id) { - root.children.clear(); - root.structural_hash.store(0, Ordering::Release); - } - - // 3. Create the new dimension node as the sole child of root. - let new_node_id = self.get_or_create_child(key, root_id, Some(values))?; - - // 4. Move the saved children into the new node. - if let Some(new_node) = self.nodes.get_mut(new_node_id) { - new_node.children = old_root_children.clone(); - } - - // 5. Fix parent pointers for the moved subtree roots. - for child_ids in old_root_children.values() { - for &child_id in child_ids.iter() { - if let Some(child) = self.nodes.get_mut(child_id) { - child.parent = Some(new_node_id); - } - } - } - - // 6. Invalidate cached structural hashes up to (and including) root. - self.invalidate_ancestors(new_node_id); - - Ok(()) - } - pub fn remove_node(&mut self, id: NodeIdx) -> Result<(), String> { let node = self.nodes.remove(id).ok_or_else(|| format!("Node {:?} not found", id))?; @@ -927,99 +799,6 @@ mod tests { assert!(class1_node.children(qube.dimension("expver").unwrap()).is_some()); } - #[test] - fn test_dimensions_returns_dim_names() { - let q = Qube::from_ascii( - "root\n└── class=od\n ├── expver=0001\n │ └── param=1\n └── expver=0002\n └── param=2", - ) - .unwrap(); - let dims = q.dimensions(); - assert!(dims.contains("class")); - assert!(dims.contains("expver")); - assert!(dims.contains("param")); - assert!(!dims.contains("root"), "root should not appear as a dimension"); - } - - #[test] - fn test_common_dimensions_uniform_depth() { - let q = Qube::from_ascii("root\n└── class=od\n └── param=1/2").unwrap(); - let common = q.common_dimensions(); - assert!(common.contains("class")); - assert!(common.contains("param")); - } - - #[test] - fn test_common_dimensions_irregular_depth() { - use crate::Datacube; - // Branch 1: param + time - let mut dc1 = Datacube::new(); - dc1.add_coordinate("param", Coordinates::from_string("2t/tp")); - dc1.add_coordinate("time", Coordinates::from_string("0/1/2")); - let mut qube = Qube::from_datacube(&dc1, Some(&["param".to_string(), "time".to_string()])); - - // Branch 2: only param - let mut dc2 = Datacube::new(); - dc2.add_coordinate("param", Coordinates::from_string("msl")); - let mut other = Qube::from_datacube(&dc2, None); - - qube.append(&mut other); - - let common = qube.common_dimensions(); - assert!(common.contains("param"), "'param' should be common"); - assert!(!common.contains("time"), "'time' is absent in one branch"); - } - - #[test] - fn test_expand_wraps_tree_under_new_outer_dimension() { - let mut q = Qube::from_ascii("root\n└── param=2t/tp\n └── time=0/1/2").unwrap(); - q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); - - let dims = q.dimensions(); - assert!(dims.contains("ensemble")); - assert!(dims.contains("param")); - assert!(dims.contains("time")); - - let ascii = q.to_ascii(); - assert!(ascii.contains("ensemble=ens1/ens2"), "new dimension should appear in ascii"); - } - - #[test] - fn test_expand_on_empty_qube() { - let mut q = Qube::new(); - q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); - - let dims = q.dimensions(); - assert!(dims.contains("ensemble")); - } - - #[test] - fn test_expand_twice_nests_outermost_last() { - let mut q = Qube::from_ascii("root\n└── param=2t").unwrap(); - q.expand("ensemble", Coordinates::from_string("ens1")).unwrap(); - q.expand("member", Coordinates::from_string("m1/m2")).unwrap(); - - let ascii = q.to_ascii(); - // "member" was added last so it must appear higher (earlier) in the tree - let member_pos = ascii.find("member").expect("member not found"); - let ensemble_pos = ascii.find("ensemble").expect("ensemble not found"); - let param_pos = ascii.find("param").expect("param not found"); - assert!(member_pos < ensemble_pos, "member should be outer of ensemble"); - assert!(ensemble_pos < param_pos, "ensemble should be outer of param"); - } - - #[test] - fn test_expand_preserves_original_coords() { - let mut q = Qube::from_ascii("root\n└── param=2t/tp\n └── time=0/1/2").unwrap(); - q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); - - let all = q.all_unique_dim_coords(); - // Original coords must still be present - let param_str = all.get("param").unwrap().to_string(); - assert!(param_str.contains("2t") && param_str.contains("tp")); - let ens_str = all.get("ensemble").unwrap().to_string(); - assert!(ens_str.contains("ens1") && ens_str.contains("ens2")); - } - #[test] fn test_squeeze() -> Result<(), String> { let input = r#"root diff --git a/qubed/src/serde/ascii.rs b/qubed/src/serde/ascii.rs index 2e18d67..391dde5 100644 --- a/qubed/src/serde/ascii.rs +++ b/qubed/src/serde/ascii.rs @@ -134,7 +134,7 @@ fn serialize_children(qube: &Qube, parent_id: NodeIdx, prefix: &str, output: &mu let key = child_node.dimension().unwrap_or("unknown"); let values = child_node.coordinates(); - let values_str = values.to_string(); + let values_str = values.to_ascii_string(); output.push_str(prefix); output.push_str(branch); @@ -225,4 +225,86 @@ mod tests { assert_eq!(input, serialized); assert_eq!(serialized, re_serialized); } + + #[test] + fn test_ascii_integer_range_step1_roundtrip() { + // ASCII format for a step-1 integer range + let input = "root\n└── param=1/to/10\n"; + let qube = Qube::from_ascii(input).unwrap(); + let out = qube.to_ascii(); + assert_eq!(input, out, "Step-1 integer range ASCII roundtrip failed:\n{}", out); + } + + #[test] + fn test_ascii_integer_range_stepped_roundtrip() { + // ASCII format for a stepped integer range + let input = "root\n└── param=0/to/10/by/2\n"; + let qube = Qube::from_ascii(input).unwrap(); + let out = qube.to_ascii(); + assert_eq!(input, out, "Stepped integer range ASCII roundtrip failed:\n{}", out); + } + + #[test] + fn test_ascii_integer_multi_range_roundtrip() { + // Two ranges joined by `|`, plus a singleton + let input = "root\n└── param=1/to/5|7|10/to/15\n"; + let qube = Qube::from_ascii(input).unwrap(); + let out = qube.to_ascii(); + assert_eq!(input, out, "Multi-range ASCII roundtrip failed:\n{}", out); + } + + #[test] + fn test_ascii_datetime_range_daily_roundtrip() { + // Daily datetime range (no /by/ suffix — daily is the default) + let input = "root\n└── date=2020-01-01T00:00:00/to/2020-01-10T00:00:00\n"; + let qube = Qube::from_ascii(input).unwrap(); + let out = qube.to_ascii(); + assert_eq!(input, out, "Daily datetime range ASCII roundtrip failed:\n{}", out); + } + + #[test] + fn test_ascii_datetime_range_hourly_roundtrip() { + // Hourly datetime range (3600s step) + let input = "root\n└── date=2020-01-01T00:00:00/to/2020-01-01T06:00:00/by/3600s\n"; + let qube = Qube::from_ascii(input).unwrap(); + let out = qube.to_ascii(); + assert_eq!(input, out, "Hourly datetime range ASCII roundtrip failed:\n{}", out); + } + + #[test] + fn test_ascii_datetime_multi_range_roundtrip() { + // Two daily ranges joined by `|` + let input = "root\n└── date=2020-01-01T00:00:00/to/2020-01-05T00:00:00|2020-02-01T00:00:00/to/2020-02-05T00:00:00\n"; + let qube = Qube::from_ascii(input).unwrap(); + let out = qube.to_ascii(); + assert_eq!(input, out, "Multi datetime range ASCII roundtrip failed:\n{}", out); + } + + #[test] + fn test_compress_then_ascii_roundtrip() { + // Build a Qube with 10 individual integer params, compress, then verify + // the ASCII output uses range notation and round-trips correctly. + let mut qube = Qube::new(); + let root = qube.root(); + let class = { + let mut c = Coordinates::Empty; + c.append("od".to_string()); + qube.get_or_create_child("class", root, Some(c)).unwrap() + }; + for v in 1..=10i32 { + let mut c = Coordinates::Empty; + c.append(v); + qube.get_or_create_child("param", class, Some(c)).unwrap(); + } + qube.compress(); + + let ascii = qube.to_ascii(); + println!("Compressed ASCII:\n{}", ascii); + assert!(ascii.contains("1/to/10"), "Expected range notation in ASCII: {}", ascii); + + // Roundtrip: parse back and re-serialize + let reparsed = Qube::from_ascii(&ascii).unwrap(); + let re_ascii = reparsed.to_ascii(); + assert_eq!(ascii, re_ascii, "ASCII roundtrip after compress failed"); + } } diff --git a/qubed/src/serde/json.rs b/qubed/src/serde/json.rs index fa3741e..05e0835 100644 --- a/qubed/src/serde/json.rs +++ b/qubed/src/serde/json.rs @@ -129,14 +129,18 @@ impl Qube { } }, crate::Coordinates::DateTimes(_) => { - // Fallback to whatever the generic serializer produces (not implemented elsewhere yet) let v = nref.coordinates().to_json_value(); match v { Value::Array(arr) => { map.insert("datetimes".to_string(), Value::Array(arr)); Value::Object(map) } - other => Value::Object(map), + Value::String(s) => { + // RangeSet serialised as a textual string + map.insert("datetimes_range".to_string(), Value::String(s)); + Value::Object(map) + } + _ => Value::Object(map), } } crate::Coordinates::Mixed(_) => { @@ -253,12 +257,19 @@ impl Qube { let has_strings = map.get("strings").is_some(); let has_floats = map.get("floats").is_some(); let has_datetimes = map.get("datetimes").is_some(); - - let typed_key_count = - [has_ints, has_ints_text, has_strings, has_floats, has_datetimes] - .iter() - .filter(|&&b| b) - .count(); + let has_datetimes_range = map.get("datetimes_range").is_some(); + + let typed_key_count = [ + has_ints, + has_ints_text, + has_strings, + has_floats, + has_datetimes, + has_datetimes_range, + ] + .iter() + .filter(|&&b| b) + .count(); if has_ints_text && typed_key_count == 1 { // textual integer representation -> parse as string @@ -273,6 +284,9 @@ impl Qube { map.get("floats").cloned().unwrap_or(Value::Null) } else if has_datetimes && typed_key_count == 1 { map.get("datetimes").cloned().unwrap_or(Value::Null) + } else if has_datetimes_range && typed_key_count == 1 { + // datetime range stored as textual string — parse via from_string + map.get("datetimes_range").cloned().unwrap_or(Value::Null) } else { // Mixed or unknown: pass the whole object so // `from_json_value` can create a MixedCoordinates @@ -473,4 +487,184 @@ mod json_tests { let reconstructed = Qube::from_arena_json(arena).expect("from_arena_json"); assert_eq!(qube.to_json(), reconstructed.to_json()); } + + #[test] + fn test_arena_roundtrip_integer_rangeset() { + use crate::coordinates::integers::{IntegerCoordinates, IntegerRange}; + use tiny_vec::TinyVec; + + let mut qube = Qube::new(); + let root = qube.root(); + + // param = 1:1:10 (range 1..10 step 1) + let mut ranges: TinyVec = TinyVec::new(); + ranges.push(IntegerRange::new_step1(1, 10)); + let coords = Coordinates::Integers(IntegerCoordinates::RangeSet(ranges)); + qube.get_or_create_child("param", root, Some(coords)).unwrap(); + + let arena = qube.to_arena_json(); + println!("Integer RangeSet arena JSON:\n{}", serde_json::to_string_pretty(&arena).unwrap()); + + let reconstructed = Qube::from_arena_json(arena).expect("from_arena_json"); + assert_eq!( + qube.to_json(), + reconstructed.to_json(), + "Integer RangeSet arena roundtrip failed" + ); + } + + #[test] + fn test_arena_roundtrip_datetime_rangeset() { + use crate::coordinates::datetime::{DateTimeCoordinates, DateTimeRange}; + use chrono::NaiveDate; + use tiny_vec::TinyVec; + + let mut qube = Qube::new(); + let root = qube.root(); + + let d_start = NaiveDate::from_ymd_opt(2020, 1, 1).unwrap().and_hms_opt(0, 0, 0).unwrap(); + let d_end = NaiveDate::from_ymd_opt(2020, 1, 10).unwrap().and_hms_opt(0, 0, 0).unwrap(); + + let mut ranges: TinyVec = TinyVec::new(); + ranges.push(DateTimeRange::daily(d_start, d_end)); + let coords = Coordinates::DateTimes(DateTimeCoordinates::RangeSet(ranges)); + qube.get_or_create_child("date", root, Some(coords)).unwrap(); + + let arena = qube.to_arena_json(); + println!( + "DateTime RangeSet arena JSON:\n{}", + serde_json::to_string_pretty(&arena).unwrap() + ); + + let reconstructed = Qube::from_arena_json(arena).expect("from_arena_json"); + assert_eq!( + qube.to_json(), + reconstructed.to_json(), + "DateTime RangeSet arena roundtrip failed" + ); + } + + #[test] + fn test_qube_compress_integers_to_rangeset() { + // Building a Qube with many single-integer param nodes, then compressing + // should merge them and then compress into a range. + let mut qube = Qube::new(); + let root = qube.root(); + + // class=od branch + let class = { + let mut c = Coordinates::Empty; + c.append("od".to_string()); + qube.get_or_create_child("class", root, Some(c)).unwrap() + }; + + // Add params 1..10 as individual nodes under the same parent + for v in 1..=10i32 { + let mut c = Coordinates::Empty; + c.append(v); + qube.get_or_create_child("param", class, Some(c)).unwrap(); + } + + qube.compress(); + + // After compression the param node should have a single RangeSet coord + let class_node = qube.node(class).unwrap(); + let param_kids: Vec<_> = class_node.all_children().collect(); + assert_eq!(param_kids.len(), 1, "Expected params merged into 1 node"); + + let param_node = qube.node(param_kids[0]).unwrap(); + match param_node.coordinates() { + Coordinates::Integers(crate::coordinates::integers::IntegerCoordinates::RangeSet( + ranges, + )) => { + // Should have compressed to a single 1..10 range + assert_eq!(ranges.len(), 1, "Expected single range, got {:?}", ranges); + assert_eq!(ranges[0].start, 1); + assert_eq!(ranges[0].end, 10); + } + other => panic!("Expected IntegerCoordinates::RangeSet, got {:?}", other), + } + } + + #[test] + fn test_qube_compress_datetimes_to_rangeset() { + use chrono::NaiveDate; + + let mut qube = Qube::new(); + let root = qube.root(); + + let class = { + let mut c = Coordinates::Empty; + c.append("od".to_string()); + qube.get_or_create_child("class", root, Some(c)).unwrap() + }; + + // Add dates Jan 1..10 as individual nodes + for day in 1..=10u32 { + let d = NaiveDate::from_ymd_opt(2020, 1, day).unwrap().and_hms_opt(0, 0, 0).unwrap(); + let mut c = Coordinates::Empty; + c.append(d); + qube.get_or_create_child("date", class, Some(c)).unwrap(); + } + + qube.compress(); + + let class_node = qube.node(class).unwrap(); + let date_kids: Vec<_> = class_node.all_children().collect(); + assert_eq!(date_kids.len(), 1, "Expected dates merged into 1 node"); + + let date_node = qube.node(date_kids[0]).unwrap(); + match date_node.coordinates() { + Coordinates::DateTimes( + crate::coordinates::datetime::DateTimeCoordinates::RangeSet(ranges), + ) => { + assert_eq!(ranges.len(), 1, "Expected single daily range, got {:?}", ranges); + assert_eq!( + ranges[0].start, + NaiveDate::from_ymd_opt(2020, 1, 1).unwrap().and_hms_opt(0, 0, 0).unwrap() + ); + assert_eq!( + ranges[0].end, + NaiveDate::from_ymd_opt(2020, 1, 10).unwrap().and_hms_opt(0, 0, 0).unwrap() + ); + } + other => panic!("Expected DateTimeCoordinates::RangeSet, got {:?}", other), + } + } + + #[test] + fn test_qube_compress_two_integer_ranges() { + // Two disjoint ranges of integers should both be preserved as ranges + let mut qube = Qube::new(); + let root = qube.root(); + + let class = { + let mut c = Coordinates::Empty; + c.append("od".to_string()); + qube.get_or_create_child("class", root, Some(c)).unwrap() + }; + + // Add 1..5 and 10..15 as individual integer nodes + for v in (1..=5i32).chain(10..=15) { + let mut c = Coordinates::Empty; + c.append(v); + qube.get_or_create_child("param", class, Some(c)).unwrap(); + } + + qube.compress(); + + let class_node = qube.node(class).unwrap(); + let param_kids: Vec<_> = class_node.all_children().collect(); + assert_eq!(param_kids.len(), 1); + + let param_node = qube.node(param_kids[0]).unwrap(); + match param_node.coordinates() { + Coordinates::Integers(crate::coordinates::integers::IntegerCoordinates::RangeSet( + ranges, + )) => { + assert_eq!(ranges.len(), 2, "Expected 2 ranges, got {:?}", ranges); + } + other => panic!("Expected IntegerCoordinates::RangeSet with 2 ranges, got {:?}", other), + } + } } diff --git a/qubed/tests/test_compression.rs b/qubed/tests/test_compression.rs index deb0818..4154b29 100644 --- a/qubed/tests/test_compression.rs +++ b/qubed/tests/test_compression.rs @@ -37,7 +37,7 @@ fn compress_uncompressed_tree() { │ └── param=1/2 └── class=2 ├── expver=0001 - │ └── param=1/2/3 + │ └── param=1/to/3 └── expver=0002 └── param=1/2 "#; diff --git a/qubed_meteo/bin/fdb_db_reader.rs b/qubed_meteo/bin/fdb_db_reader.rs new file mode 100644 index 0000000..014c676 --- /dev/null +++ b/qubed_meteo/bin/fdb_db_reader.rs @@ -0,0 +1,50 @@ +use qubed::Qube; +use qubed_meteo::adapters::fdb::FromFDBList; +use rsfdb::{FDB, request::Request}; +use serde_json::json; +use std::env; +use std::time::Instant; +use std::fs::File; + +fn main() -> Result<(), Box> { + // Ensure FDB config is set so the internal listing can open the DB + use std::path::PathBuf; + + let config_path = PathBuf::from("xxx"); // Adjust this path to point to your local FDB config.yaml + unsafe { + std::env::set_var("FDB5_CONFIG_FILE", config_path.to_str().expect("Invalid config path")); + } + + let lib_path = PathBuf::from("xxx"); // Adjust this path to point to the directory containing FDB shared libraries + + unsafe { + std::env::set_var("DYLD_LIBRARY_PATH", lib_path.to_str().expect("Invalid path to shared libraries")); + } + + let request_map = json!({ + "class" : "d1", + "dataset": "extremes-dt", + "expver" : "0001", + "stream" : "oper", + "date": "20260303", + "time" : "0000", + "domain" : "g", + "levtype" : "sfc", + }); + let start_time = Instant::now(); + + // Build the Qube directly from the request; the adapter will open FDB and list. + let qube = Qube::from_fdb_list(&request_map).expect("Failed to build Qube from FDB list"); + + // Stop the timer + let duration = start_time.elapsed(); + + // Print the time taken + println!("Time taken to construct Qube: {:?}", duration); + + let file = File::create("extremes_eg.json")?; + serde_json::to_writer(file, &qube.to_arena_json())?; + + Ok(()) + +} \ No newline at end of file diff --git a/qubed_meteo/examples/read_from_dss_constraints.rs b/qubed_meteo/examples/read_from_dss_constraints.rs new file mode 100644 index 0000000..de5314f --- /dev/null +++ b/qubed_meteo/examples/read_from_dss_constraints.rs @@ -0,0 +1,119 @@ +use qubed::Qube; +use qubed_meteo::adapters::dss_constraints::FromDssConstraints; +use std::time::Instant; + +fn main() { + let path = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/data/medium2_era5_constraints.json"); + + let dss_json = std::fs::read_to_string(path).expect("Failed to read DSS constraints JSON file"); + + let dss_json = + serde_json::from_str::(&dss_json).expect("Failed to parse JSON file"); + + let start_time = Instant::now(); + + let qube = Qube::from_dss_constraints(&dss_json); + + // Stop the timer + let duration = start_time.elapsed(); + + // Print the time taken + println!("Time taken to construct Qube: {:?}", duration); + + println!("Constructed Qube: {:?}", qube.unwrap().to_ascii()); +} + +// #[cfg(test)] +// mod tests { +// use super::*; +// use serde_json::json; + +// #[test] +// fn test_merge_datacubes() { +// // Define the first datacube as JSON +// let datacube1 = json!({ +// "variable": [ +// "10m_u_component_of_wind", +// "10m_v_component_of_wind", +// "convective_precipitation", +// "mean_sea_level_pressure", +// "surface_pressure" +// ], +// "day": ["01"], +// "year": ["2017", "2018", "2019", "2020"], +// "month": ["01"], +// "leadtime_hour": [ +// "1008", "1032", "1056", "1080", "1104", "1128", "1152", "1176", "120", "1200", +// "1224", "1248", "1272", "1296", "1320", "1344", "1368", "1392", "1416", "144", +// "1440", "168", "192", "216", "24", "240", "264", "288", "312", "336", "360", +// "384", "408", "432", "456", "48", "480", "504", "528", "552", "576", "600", +// "624", "648", "672", "696", "72", "720", "744", "768", "792", "816", "840", +// "864", "888", "912", "936", "96", "960", "984" +// ], +// "forecast_type": ["control_forecast", "perturbed_forecast"], +// "hday": ["01"], +// "hmonth": ["01"], +// "hyear": [ +// "1991", "1992", "1993", "1994", "1995", "1996", "1997", "1998", "1999", "2000", +// "2001", "2002", "2003", "2004", "2005", "2006", "2007", "2008", "2009", "2010" +// ], +// "level_type": ["single_level"], +// "origin": ["kma"], +// "time": ["00:00"] +// }); + +// // Define the second datacube as JSON +// let datacube2 = json!({ +// "origin": ["kma"], +// "hday": ["01"], +// "forecast_type": ["control_forecast", "perturbed_forecast"], +// "leadtime_hour": [ +// "0_24", "1008_1032", "1032_1056", "1056_1080", "1080_1104", "1104_1128", +// "1128_1152", "1152_1176", "1176_1200", "1200_1224", "120_144", "1224_1248", +// "1248_1272", "1272_1296", "1296_1320", "1320_1344", "1344_1368", "1368_1392", +// "1392_1416", "1416_1440", "144_168", "168_192", "192_216", "216_240", "240_264", +// "24_48", "264_288", "288_312", "312_336", "336_360", "360_384", "384_408", +// "408_432", "432_456", "456_480", "480_504", "48_72", "504_528", "528_552", +// "552_576", "576_600", "600_624", "624_648", "648_672", "672_696", "696_720", +// "720_744", "72_96", "744_768", "768_792", "792_816", "816_840", "840_864", +// "864_888", "888_912", "912_936", "936_960", "960_984", "96_120", "984_1008" +// ], +// "hmonth": ["01"], +// "variable": [ +// "2m_dewpoint_temperature", +// "2m_temperature", +// "sea_ice_area_fraction", +// "skin_temperature" +// ], +// "day": ["01"], +// "hyear": [ +// "1991", "1992", "1993", "1994", "1995", "1996", "1997", "1998", "1999", "2000", +// "2001", "2002", "2003", "2004", "2005", "2006", "2007", "2008", "2009", "2010" +// ], +// "level_type": ["single_level"], +// "month": ["01"], +// "time": ["00:00"] +// }); + +// // Combine the two datacubes into a JSON array +// let dss_constraints = json!([datacube1, datacube2]); + +// // Call the `from_dss_constraints` method +// let result = Qube::from_dss_constraints(&dss_constraints); + +// // Assert that the result is Ok +// assert!(result.is_ok()); + +// // Get the resulting Qube +// let qube = result.unwrap(); + +// // Print the resulting Qube for debugging +// println!("Merged Qube: {:?}", qube.to_ascii()); + +// // Add assertions to verify the merged structure +// // For example, check that certain dimensions exist in the merged Qube +// assert!(qube.contains_dimension("variable")); +// assert!(qube.contains_dimension("leadtime_hour")); +// assert!(qube.contains_dimension("origin")); +// } +// } diff --git a/qubed_meteo/examples/read_from_fdb_list.rs b/qubed_meteo/examples/read_from_fdb_list.rs new file mode 100644 index 0000000..a370516 --- /dev/null +++ b/qubed_meteo/examples/read_from_fdb_list.rs @@ -0,0 +1,47 @@ +use qubed::Qube; +use qubed_meteo::adapters::fdb::FromFDBList; +use rsfdb::{FDB, request::Request}; +use serde_json::json; +use std::env; +use std::time::Instant; + +fn main() { + // Ensure FDB config is set so the internal listing can open the DB + let config_path = + env::current_dir().unwrap().join("/Users/male/git/fdb-home/etc/fdb/config.yaml"); + unsafe { + std::env::set_var("FDB5_CONFIG_FILE", config_path.to_str().expect("Invalid config path")); + } + + let request_map = json!({ + "class" : "od", + "expver" : "0001", + "stream" : "oper", + "time" : "0000", + "domain" : "g", + "levtype" : "sfc", + }); + let start_time = Instant::now(); + + // Build the Qube directly from the request; the adapter will open FDB and list. + let qube = Qube::from_fdb_list(&request_map).expect("Failed to build Qube from FDB list"); + + // println!("Qube structure:\n{}", qube.to_ascii()); + + println!("Qube in arena json format:\n{}", qube.to_arena_json()); + + // Stop the timer + let duration = start_time.elapsed(); + + // Print the time taken + println!("Time taken to construct Qube: {:?}", duration); + + // let start_time2 = Instant::now(); + + // let datacubes = qube.unwrap().to_datacubes(); + + // let duration2 = start_time2.elapsed(); + + // println!("Time taken to convert Qube to datacubes: {:?}", duration2); + // println!("Number of datacubes: {}", datacubes.len()); +} diff --git a/qubed_meteo/examples/read_from_mars_list.rs b/qubed_meteo/examples/read_from_mars_list.rs new file mode 100644 index 0000000..145a3e2 --- /dev/null +++ b/qubed_meteo/examples/read_from_mars_list.rs @@ -0,0 +1,28 @@ +use qubed::Qube; +use qubed_meteo::adapters::mars_list::FromMARSList; +use std::time::Instant; + +fn main() { + let path = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/data/large_mars.list"); + + let mars_list = std::fs::read_to_string(path).expect("Failed to read MARS list file"); + + let start_time = Instant::now(); + + let qube = Qube::from_mars_list(&mars_list); + + // Stop the timer + let duration = start_time.elapsed(); + + // Print the time taken + println!("Time taken to construct Qube: {:?}", duration); + + let start_time2 = Instant::now(); + + let datacubes = qube.unwrap().to_datacubes(); + + let duration2 = start_time2.elapsed(); + + println!("Time taken to convert Qube to datacubes: {:?}", duration2); + println!("Number of datacubes: {}", datacubes.len()); +} diff --git a/qubed_meteo/qubed_meteo/examples/read_from_fdb_list.rs b/qubed_meteo/qubed_meteo/examples/read_from_fdb_list.rs index cb01c98..2b58f15 100644 --- a/qubed_meteo/qubed_meteo/examples/read_from_fdb_list.rs +++ b/qubed_meteo/qubed_meteo/examples/read_from_fdb_list.rs @@ -11,10 +11,8 @@ use std::time::Instant; fn main() { // Ensure FDB config is set so the internal listing can open the DB let config_path = - env::current_dir().unwrap().join("/Users/male/git/fdb-home/etc/fdb/config.yaml"); - unsafe { - std::env::set_var("FDB5_CONFIG_FILE", config_path.to_str().expect("Invalid config path")); - } + env::var("FDB5_CONFIG_FILE").expect("Set FDB5_CONFIG_FILE to your local FDB config path"); + unsafe { std::env::set_var("FDB5_CONFIG_FILE", config_path) }; let request_map = json!({ "class" : "od", diff --git a/qubed_meteo/qubed_meteo/src/adapters/fdb.rs b/qubed_meteo/qubed_meteo/src/adapters/fdb.rs index 783592f..3b592e3 100644 --- a/qubed_meteo/qubed_meteo/src/adapters/fdb.rs +++ b/qubed_meteo/qubed_meteo/src/adapters/fdb.rs @@ -108,15 +108,12 @@ mod tests { #[test] #[cfg(feature = "rsfdb-support")] fn test_from_fdb_list_basic() { - // Ensure FDB config is set (adjust path for local environment if needed) - let config_path = - env::current_dir().unwrap().join("/Users/male/git/fdb-home/etc/fdb/config.yaml"); - unsafe { - std::env::set_var( - "FDB5_CONFIG_FILE", - config_path.to_str().expect("Invalid config path"), - ); - } + // Skip when FDB is unavailable in the test environment. + let Ok(config_path) = env::var("FDB5_CONFIG_FILE") else { + eprintln!("Skipping test_from_fdb_list_basic: FDB5_CONFIG_FILE is not set"); + return; + }; + unsafe { std::env::set_var("FDB5_CONFIG_FILE", config_path) }; let request_map = json!({ "class" : "od", diff --git a/qubed_meteo/src/adapters/dss_constraints.rs b/qubed_meteo/src/adapters/dss_constraints.rs new file mode 100644 index 0000000..892a372 --- /dev/null +++ b/qubed_meteo/src/adapters/dss_constraints.rs @@ -0,0 +1,84 @@ +use qubed::{Datacube, Qube}; +use serde_json::Value; + +pub trait FromDssConstraints { + fn from_dss_constraints(dss_constraints: &Value) -> Result; +} + +impl FromDssConstraints for Qube { + fn from_dss_constraints(dss_constraints: &Value) -> Result { + let datacubes: &Vec = + dss_constraints.as_array().expect("DSS constraints should be a JSON array"); + + let order = vec![ + "origin".to_string(), + "forecast_type".to_string(), + "hday".to_string(), + "day".to_string(), + "hmonth".to_string(), + "hyear".to_string(), + "year".to_string(), + "month".to_string(), + "time".to_string(), + "leadtime_hour".to_string(), + "level_type".to_string(), + "variable".to_string(), + ]; + + // Parse the first datacube and initialize the main Qube + let first_datacube = parse_datacube(&datacubes[0])?; + let mut qube = Qube::from_datacube(&first_datacube, Some(&order)); + + // Collect all other Qubes into a Vec + let mut other_qubes: Vec = Vec::new(); + for datacube in &datacubes[1..] { + let qube_part = parse_datacube(datacube); + + let qube_part = match qube_part { + Ok(dc) => Qube::from_datacube(&dc, Some(&order)), + Err(e) => return Err(format!("Failed to parse datacube: {}", e)), + }; + + other_qubes.push(qube_part); + } + + // Use append_many to merge all Qubes together + qube.append_many(&mut other_qubes); + + Ok(qube) + } +} + +fn parse_datacube(dss_datacube: &Value) -> Result { + let mut datacube = Datacube::new(); + + let dimensions = dss_datacube.as_object().expect("DSS datacube should be a JSON object"); + + for (dimension_name, coordinates) in dimensions { + let coord_array = coordinates.as_array().expect( + format!("Datacube dimension {} should be a JSON array", dimension_name).as_str(), + ); + + let mut coords = qubed::Coordinates::new(); + + for coord in coord_array { + match coord { + Value::Number(num) => { + if num.is_i64() { + coords.append(num.as_i64().unwrap() as i32); + } else if num.is_f64() { + coords.append(num.as_f64().unwrap()); + } + } + Value::String(s) => { + coords.append(s.clone()); + } + _ => panic!("Unsupported coordinate type in dimension {}", dimension_name), + } + } + + datacube.add_coordinate(dimension_name, coords); + } + + Ok(datacube) +} diff --git a/qubed_meteo/src/adapters/fdb.rs b/qubed_meteo/src/adapters/fdb.rs new file mode 100644 index 0000000..63fdb18 --- /dev/null +++ b/qubed_meteo/src/adapters/fdb.rs @@ -0,0 +1,133 @@ +use qubed::{Coordinates, NodeIdx, Qube}; +use rsfdb::{FDB, request::Request}; +use serde_json::Value as JsonValue; + +pub trait FromFDBList { + /// Build a `Qube` from a JSON request map by performing an internal list. + /// + /// The `request_map` should be a JSON object (the same structure accepted + /// by `rsfdb::request::Request::from_json`). The implementation will build + /// an `rsfdb::request::Request` internally, call `list` with + /// `splitkey=true` and iterate the results. + fn from_fdb_list(request_map: &JsonValue) -> Result; +} + +impl FromFDBList for Qube { + fn from_fdb_list(request_map: &JsonValue) -> Result { + // Build Request from provided JSON map + let request = Request::from_json(request_map.clone()) + .map_err(|e| format!("Failed to build Request from JSON: {:?}", e))?; + + let fdb = FDB::new(None).map_err(|e| format!("Failed to open FDB: {:?}", e))?; + let list_iter = + fdb.list(&request, true, false).map_err(|e| format!("FDB list failed: {:?}", e))?; + + let mut qube = Qube::new(); + let root = qube.root(); + + fn make_coords(vals: &[&str]) -> Option { + let mut coords = Coordinates::new(); + for v in vals { + let s = v.trim(); + if s.is_empty() { + continue; + } + // Check for leading zeros to preserve formatting (e.g., "0001") + let has_leading_zero = s.len() > 1 + && s.starts_with('0') + && s.chars().nth(1).map_or(false, |c| c.is_ascii_digit()); + + if has_leading_zero { + // Preserve as string to keep formatting + coords.append(s.to_string()); + } else if let Ok(i) = s.parse::() { + coords.append(i); + } else if let Ok(f) = s.parse::() { + coords.append(f); + } else { + coords.append(s.to_string()); + } + } + if coords.is_empty() { None } else { Some(coords) } + } + + for item in list_iter { + // Each item may contain a splitkey metadata (request-like key/value pairs) + // Build a comma-separated path string from the splitkey metadata similar + // to the previous external representation. + let mut parts_vec: Vec = Vec::new(); + + if let Some(metadata) = item.request { + for kv in metadata.iter() { + parts_vec.push(format!("{}={}", kv.key, kv.value)); + } + } + + if parts_vec.is_empty() { + continue; + } + + let mut parent = root; + for part in parts_vec.iter() { + if let Some((key, val)) = part.split_once('=') { + let vals: Vec<&str> = + val.split('/').map(|s| s.trim()).filter(|s| !s.is_empty()).collect(); + + // If there are no value parts (e.g. "key=") skip creating an empty child + if vals.is_empty() { + continue; + } + let coords = make_coords(&vals); + let child = qube + .get_or_create_child(key.trim(), parent, coords) + .map_err(|e| format!("create_child failed: {:?}", e))?; + parent = child; + } else { + let child = qube + .get_or_create_child(part.trim(), parent, None) + .map_err(|e| format!("get_or_create_child failed: {:?}", e))?; + parent = child; + } + } + } + + qube.compress(); + Ok(qube) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use std::env; + + #[test] + fn test_from_fdb_list_basic() { + // Ensure FDB config is set (adjust path for local environment if needed) + let config_path = + env::current_dir().unwrap().join("/Users/male/git/fdb-home/etc/fdb/config.yaml"); + unsafe { + std::env::set_var( + "FDB5_CONFIG_FILE", + config_path.to_str().expect("Invalid config path"), + ); + } + + let request_map = json!({ + "class" : "od", + "expver" : "0001", + "stream" : "oper", + "time" : "0000", + "domain" : "g", + "levtype" : "sfc", + }); + + let qube = + ::from_fdb_list(&request_map).expect("failed to build qube"); + println!("Qube structure:\n{}", qube.to_ascii()); + + let serialized = qube.to_ascii(); + assert!(!serialized.is_empty()); + } +} diff --git a/qubed_meteo/src/adapters/mars_list.rs b/qubed_meteo/src/adapters/mars_list.rs new file mode 100644 index 0000000..f6eb492 --- /dev/null +++ b/qubed_meteo/src/adapters/mars_list.rs @@ -0,0 +1,246 @@ +use qubed::{Coordinates, NodeIdx, Qube}; + +pub trait FromMARSList { + fn from_mars_list(mars_list: &str) -> Result; +} + +impl FromMARSList for Qube { + fn from_mars_list(mars_list: &str) -> Result { + let mut qube = Qube::new(); + let root = qube.root(); + + // Stack of (indent, node). Start with the root sentinel (indent 0). + let mut stack: Vec<(usize, NodeIdx)> = vec![(0, root)]; + + let mut prev_indent = 0usize; + // remember the last created node from the previous non-empty line + let mut last_line_last_created: Option = None; + + fn make_coords(vals: &[&str]) -> Option { + let mut coords = Coordinates::new(); + for v in vals { + let s = v.trim(); + if s.is_empty() { + continue; + } + // Check for leading zeros to preserve formatting (e.g., "0001") + let has_leading_zero = s.len() > 1 + && s.starts_with('0') + && s.chars().nth(1).map_or(false, |c| c.is_ascii_digit()); + + if has_leading_zero { + // Preserve as string to keep formatting + coords.append(s.to_string()); + } else if let Ok(i) = s.parse::() { + coords.append(i); + } else if let Ok(f) = s.parse::() { + coords.append(f); + } else { + coords.append(s.to_string()); + } + } + if coords.is_empty() { None } else { Some(coords) } + } + + for raw_line in mars_list.lines() { + // normalize and keep leading whitespace to compute indent + let raw = raw_line.replace('\r', ""); + let indent = raw.chars().take_while(|c| c.is_whitespace()).count(); + let trimmed = raw.trim(); + + if trimmed.is_empty() { + prev_indent = indent; + continue; + } + + // tokens split by commas on this line + let tokens: Vec<&str> = + trimmed.split(',').map(|t| t.trim()).filter(|t| !t.is_empty()).collect(); + + if tokens.is_empty() { + prev_indent = indent; + continue; + } + + // find nearest shallower parent for this indent (stack fallback) + while let Some(&(top_indent, _)) = stack.last() { + if top_indent >= indent { + stack.pop(); + } else { + break; + } + } + let stack_parent = stack.last().map(|&(_, id)| id).unwrap_or(root); + + // choose parent and creation strategy: + // - if deeper than previous line and we have last_line_last_created -> build a chain under that node + // - otherwise use stack_parent + let mut last_created: Option = None; + + if indent > prev_indent { + if let Some(mut parent) = last_line_last_created { + // Build a chain under `parent`: for each token create a child and + // make that child the new parent for the next token. + for tok in tokens.iter() { + if let Some((key, val)) = tok.split_once('=') { + let vals: Vec<&str> = val + .split('/') + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .collect(); + let coords = make_coords(&vals); + let child = qube + .get_or_create_child(key.trim(), parent, coords) + .map_err(|e| format!("create_child failed: {:?}", e))?; + parent = child; + last_created = Some(child); + } else { + let child = qube + .get_or_create_child(tok, parent, None) + .map_err(|e| format!("get_or_create_child failed: {:?}", e))?; + parent = child; + last_created = Some(child); + } + } + // ensure last_created references the final node in the chain + last_created = Some(parent); + } else { + // no previous node: create siblings under stack_parent (unchanged fallback) + for tok in tokens.iter() { + if let Some((key, val)) = tok.split_once('=') { + let vals: Vec<&str> = val + .split('/') + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .collect(); + let coords = make_coords(&vals); + let child = qube + .get_or_create_child(key.trim(), stack_parent, coords) + .map_err(|e| format!("get_or_create_child failed: {:?}", e))?; + last_created = Some(child); + } else { + let child = qube + .get_or_create_child(tok, stack_parent, None) + .map_err(|e| format!("get_or_create_child failed: {:?}", e))?; + last_created = Some(child); + } + } + } + } else { + // same or shallower indent -> create a chain under stack_parent + let mut current_parent = stack_parent; + for tok in tokens.iter() { + if let Some((key, val)) = tok.split_once('=') { + let vals: Vec<&str> = + val.split('/').map(|s| s.trim()).filter(|s| !s.is_empty()).collect(); + let coords = make_coords(&vals); + let child = qube + .get_or_create_child(key.trim(), current_parent, coords) + .map_err(|e| format!("get_or_create_child failed: {:?}", e))?; + current_parent = child; + last_created = Some(child); + } else { + let child = qube + .get_or_create_child(tok, current_parent, None) + .map_err(|e| format!("get_or_create_child failed: {:?}", e))?; + current_parent = child; + last_created = Some(child); + } + } + } + + // Update stack: remove entries with indent >= current, then push current indent -> last_created (if any) + while let Some(&(top_indent, _)) = stack.last() { + if top_indent >= indent { + stack.pop(); + } else { + break; + } + } + if let Some(node) = last_created { + stack.push((indent, node)); + } + + // save last-created from this line for possible attachment by next indented line + last_line_last_created = last_created; + prev_indent = indent; + } + + qube.compress(); + Ok(qube) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_from_mars_list_basic_structure() { + // Construct a small MARS list that exercises indent/chain behavior + let mars = "alpha=0,beta=1/2\n gamma=3\ndelta=4"; + + let qube = ::from_mars_list(mars).expect("failed to parse"); + let root = qube.root(); + let root_ref = qube.node(root).expect("root node missing"); + + // root should have two top-level children (alpha and delta) + assert_eq!(root_ref.children_count(), 2); + + // find `alpha` under root + let mut alpha_id = None; + for child in root_ref.all_children() { + if let Some(nr) = qube.node(child) { + if nr.dimension() == Some("alpha") { + alpha_id = Some(child); + break; + } + } + } + let alpha_id = alpha_id.expect("alpha child not found"); + + // alpha should have one child (beta) + let alpha_ref = qube.node(alpha_id).unwrap(); + assert_eq!(alpha_ref.children_count(), 1); + + // find `beta` under alpha and assert its coordinates length (1/2 -> two entries) + let mut beta_id = None; + for child in alpha_ref.all_children() { + if let Some(nr) = qube.node(child) { + if nr.dimension() == Some("beta") { + beta_id = Some(child); + break; + } + } + } + let beta_id = beta_id.expect("beta child not found"); + let beta_ref = qube.node(beta_id).unwrap(); + assert_eq!(beta_ref.coordinates().len(), 2); + + // gamma should be a child of beta with one coordinate + let mut gamma_id = None; + for child in beta_ref.all_children() { + if let Some(nr) = qube.node(child) { + if nr.dimension() == Some("gamma") { + gamma_id = Some(child); + break; + } + } + } + let gamma_id = gamma_id.expect("gamma child not found"); + let gamma_ref = qube.node(gamma_id).unwrap(); + assert_eq!(gamma_ref.coordinates().len(), 1); + + // delta should exist under root with one coordinate + let mut delta_found = false; + for child in root_ref.all_children() { + if let Some(nr) = qube.node(child) { + if nr.dimension() == Some("delta") { + delta_found = true; + assert_eq!(nr.coordinates().len(), 1); + } + } + } + assert!(delta_found, "delta child not found under root"); + } +} diff --git a/qubed_meteo/src/adapters/mod.rs b/qubed_meteo/src/adapters/mod.rs new file mode 100644 index 0000000..1038e33 --- /dev/null +++ b/qubed_meteo/src/adapters/mod.rs @@ -0,0 +1,4 @@ +pub mod dss_constraints; +pub mod fdb; +pub mod mars_list; +pub mod to_constraints; diff --git a/qubed_meteo/src/adapters/to_constraints.rs b/qubed_meteo/src/adapters/to_constraints.rs new file mode 100644 index 0000000..b9ddae5 --- /dev/null +++ b/qubed_meteo/src/adapters/to_constraints.rs @@ -0,0 +1,206 @@ +use qubed::{Coordinates, Qube}; +use serde_json::{Map, Value}; + +pub trait ToDssConstraints { + fn to_dss_constraints(&self) -> Value; +} + +impl ToDssConstraints for Qube { + /// Encode the Qube as an array of constraint-like objects. + /// + /// Each array entry corresponds to a leaf-path in the Qube. For every dimension + /// seen anywhere in the Qube a key is present in the object; if the current + /// leaf-path contains that dimension its coordinates are serialized as an + /// array of strings, otherwise an empty array is emitted. This mirrors the + /// sample "array of maps" format used in the example constraints file. + fn to_dss_constraints(&self) -> Value { + // Use existing Datacube conversion then map each Datacube to a JSON object + let datacubes = self.to_datacubes(); + + // Collect union of all dimension keys so every object has the same keys + let mut all_dims: Vec = Vec::new(); + { + let mut dims_set = std::collections::BTreeSet::new(); + for dc in &datacubes { + for k in dc.coordinates().keys() { + // Exclude the internal root dimension from output + if k == "root" { + continue; + } + dims_set.insert(k.clone()); + } + } + all_dims = dims_set.into_iter().collect(); + } + + let mut out: Vec = Vec::new(); + for dc in datacubes { + let mut map = Map::new(); + for dim in all_dims.iter() { + // defensive: skip root if present in list + if dim == "root" { + continue; + } + if let Some(coords) = dc.coordinates().get(dim) { + let coord_str = coords.to_string(); + if coord_str.is_empty() { + map.insert(dim.clone(), Value::Array(Vec::new())); + } else { + let arr = + coord_str.split('/').map(|s| Value::String(s.to_string())).collect(); + map.insert(dim.clone(), Value::Array(arr)); + } + } else { + map.insert(dim.clone(), Value::Array(Vec::new())); + } + } + out.push(Value::Object(map)); + } + + Value::Array(out) + } +} + +#[cfg(test)] +mod tests_constraints_json { + use super::*; + use serde_json::json; + + #[test] + fn test_to_constraints_json_multiple_coordinates() { + let mut qube = Qube::new(); + let root = qube.root(); + + // class=od + let class_coords = { + let mut c = Coordinates::Empty; + c.append("od".to_string()); + c + }; + let class = qube.get_or_create_child("class", root, Some(class_coords)).unwrap(); + + // expver=0001/0002 (multiple coordinates on same node) + let exp_coords = { + let mut c = Coordinates::Empty; + c.append("0001".to_string()); + c.append("0002".to_string()); + c + }; + let expver = qube.get_or_create_child("expver", class, Some(exp_coords)).unwrap(); + + // param=1/2 + let param_coords = { + let mut c = Coordinates::Empty; + c.append("1".to_string()); + c.append("2".to_string()); + c + }; + let _param = qube.get_or_create_child("param", expver, Some(param_coords)).unwrap(); + + // Add a second branch so we have two leaf paths + // Second branch: class=rd / expver=0003 / param=3/4 + let class2_coords = { + let mut c = Coordinates::Empty; + c.append("rd".to_string()); + c + }; + let class2 = qube.get_or_create_child("class", root, Some(class2_coords)).unwrap(); + + let exp2_coords = { + let mut c = Coordinates::Empty; + c.append("0003".to_string()); + c + }; + let expver2 = qube.get_or_create_child("expver", class2, Some(exp2_coords)).unwrap(); + + let param2_coords = { + let mut c = Coordinates::Empty; + c.append("3".to_string()); + c.append("4".to_string()); + c + }; + let _param2 = qube.get_or_create_child("param", expver2, Some(param2_coords)).unwrap(); + + let param3_coords = { + let mut c = Coordinates::Empty; + c.append("5".to_string()); + c.append("6".to_string()); + c + }; + let _param3 = qube.get_or_create_child("param", expver2, Some(param3_coords)).unwrap(); + + let json_out = qube.to_dss_constraints(); + + println!("{}", json_out); + assert!(json_out.is_array()); + let arr = json_out.as_array().unwrap(); + // Three leaf paths -> three objects + assert_eq!(arr.len(), 3); + + // Find and validate both objects + let mut found_od = false; + let mut found_rd_34 = false; + let mut found_rd_56 = false; + for v in arr.iter() { + let obj = v.as_object().unwrap(); + // Keys present + assert!(obj.contains_key("class")); + assert!(obj.contains_key("expver")); + assert!(obj.contains_key("param")); + + let class_vals: Vec = obj["class"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + if class_vals.contains(&"od".to_string()) { + found_od = true; + let exp_vals: Vec = obj["expver"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + assert!(exp_vals.contains(&"0001".to_string())); + let param_vals: Vec = obj["param"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + assert!(param_vals.contains(&"1".to_string())); + assert!(param_vals.contains(&"2".to_string())); + } + if class_vals.contains(&"rd".to_string()) { + let exp_vals: Vec = obj["expver"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + assert!(exp_vals.contains(&"0003".to_string())); + let param_vals: Vec = obj["param"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + // rd branch may appear twice with different params; record which one we saw + if param_vals.contains(&"3".to_string()) && param_vals.contains(&"4".to_string()) { + found_rd_34 = true; + } else if param_vals.contains(&"5".to_string()) + && param_vals.contains(&"6".to_string()) + { + found_rd_56 = true; + } else { + panic!("Unexpected rd param values: {:?}", param_vals); + } + } + } + assert!( + found_od && found_rd_34 && found_rd_56, + "Both branches should be present in output" + ); + } +} diff --git a/qubed_meteo/src/lib.rs b/qubed_meteo/src/lib.rs new file mode 100644 index 0000000..4e9af61 --- /dev/null +++ b/qubed_meteo/src/lib.rs @@ -0,0 +1 @@ +pub mod adapters; diff --git a/todo.md b/todo.md new file mode 100644 index 0000000..fa8ac26 --- /dev/null +++ b/todo.md @@ -0,0 +1,22 @@ +# TODO: Feature Parity with Python + +[x] Selection +[x] Implement merge/union +[x] Compression +[x] Implement serializers and deserializers for different format +[ ] Implement more benchmarks and compare with Python implementation + +[ ] Implement metadata +[ ] Python API +[ ] FDB adapter + +[ ] Implement user-provided formatting options for serialization (e.g. to turn expver=1 to "0001") +[ ] Evaluate different select modes +[ ] Implement formatters (maybe in adapters/serde) +[ ] Implement stac_server changes to new API +[ ] Check with FIAB and PPROC to see if we broke anything +[ ] Documentation + +# Improvements + +[ ] Make a macro to make expressing selections easier \ No newline at end of file