Repository: ecmwf/pdbufr (language: python)

Issue:
Feature/performance
Fixes #53 

The problem in #53 was caused by calling the following code too many times to determine whether  a BUFR key is a coordinate:

```
code = eccodes.codes_get(key + "->code")
try:
    is_coord = code[:3] < 10
```

This PR improves the performance by caching the `is_coord` values within a given message and providing access to them via the newly added `BufrMessage.is_coord()` method. To make it work for messages defined as a mapping (e.g. dict) the message wrapper had to be modified. Please note that the coordinate test in the code snippet above was also optimised and  implemented roughly like this:

```
code = self._get(name + "->code", int)
is_coord = code <= 9999
```

The results are encouraging. For the full BUFR file from #53 this code:

```
import pdbufr

f = "i0tp_07062020_00.buf"
df = pdbufr.read_bufr(f,
    columns=("year",'month', "hour", "minute","latitude", "longitude", "atmosphericPathDelayInSatelliteSignal"),
                      filters={"stationOrSiteName": 'S3AG-EUME'},
    )
print(len(df))
```
runs in 8 min 24 sec on my MacBook.

As a reference, iterating through the messages with eccodes and unpacking each with the code below takes 7:26.

```
import eccodes

f = open("i0tp_07062020_00.buf", "rb")
    
while 1:
    bufr = eccodes.codes_bufr_new_from_file(f)
    if bufr is None:
        break
 
    eccodes.codes_set(bufr, "unpack", 1)
    eccodes.codes_release(bufr)
 
f.close()
```


Discussion / hints:
## [Codecov](https://app.codecov.io/gh/ecmwf/pdbufr/pull/54?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) Report
Patch coverage: **`87.93`**% and project coverage change: **`-0.33`** :warning:
> Comparison is base [(`b77c1cd`)](https://app.codecov.io/gh/ecmwf/pdbufr/commit/b77c1cd5985f83c9a0ce81690359a8a6a5b08527?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) 95.46% compared to head [(`7c70ea4`)](https://app.codecov.io/gh/ecmwf/pdbufr/pull/54?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) 95.14%.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master      #54      +/-   ##
==========================================
- Coverage   95.46%   95.14%   -0.33%     
==========================================
  Files          15       15              
  Lines        1587     1626      +39     
  Branches      210      212       +2     
==========================================
+ Hits         1515     1547      +32     
- Misses         55       61       +6     
- Partials       17       18       +1     
```


| [Impacted Files](https://app.codecov.io/gh/ecmwf/pdbufr/pull/54?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) | Coverage Δ | |
|---|---|---|
| [pdbufr/high\_level\_bufr/bufr.py](https://app.codecov.io/gh/ecmwf/pdbufr/pull/54?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf#diff-cGRidWZyL2hpZ2hfbGV2ZWxfYnVmci9idWZyLnB5) | `85.41% <82.60%> (-2.59%)` | :arrow_down: |
| [pdbufr/bufr\_structure.py](https://app.codecov.io/gh/ecmwf/pdbufr/pull/54?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf#diff-cGRidWZyL2J1ZnJfc3RydWN0dXJlLnB5) | `93.67% <88.88%> (-0.65%)` | :arrow_down: |
| [pdbufr/high\_level\_bufr/codesmessage.py](https://app.codecov.io/gh/ecmwf/pdbufr/pull/54?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf#diff-cGRidWZyL2hpZ2hfbGV2ZWxfYnVmci9jb2Rlc21lc3NhZ2UucHk=) | `63.15% <100.00%> (+0.99%)` | :arrow_up: |
| [tests/test\_20\_bufr\_structure.py](https://app.codecov.io/gh/ecmwf/pdbufr/pull/54?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf#diff-dGVzdHMvdGVzdF8yMF9idWZyX3N0cnVjdHVyZS5weQ==) | `100.00% <100.00%> (ø)` | |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ecmwf/pdbufr/pull/54?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).

That's wonderful @sandorkertesz! I've just run the code on an interactive node of our HPC and it takes 7m39s, a fantastic improvement from 41 minutes!

Fix patch:
diff --git a/pdbufr/bufr_structure.py b/pdbufr/bufr_structure.py
index 834bc5b..eb3210a 100644
--- a/pdbufr/bufr_structure.py
+++ b/pdbufr/bufr_structure.py
@@ -14,7 +14,7 @@
 import eccodes  # type: ignore
 import numpy as np
 
-from pdbufr.high_level_bufr.bufr import BufrMessage
+from pdbufr.high_level_bufr.bufr import bufr_code_is_coord
 
 from . import bufr_filters
 
@@ -89,34 +89,55 @@ def relative_key(self) -> str:
 IS_KEY_COORD = {"subsetNumber": True, "operator": False}
 
 
-def wrap_message(m: T.Any) -> T.Any:
-    if isinstance(m, BufrMessage):
-        return m
-    else:
-        return MessageWrapper(m)
-
-
 class MessageWrapper:
-    """Makes it possible to use context manager and is_coord methods for all
+    """Makes it possible to use context manager and is_coord method for all
     types of messages."""
 
+    WRAP: T.Dict[T.Any, T.Any] = {}
+
     def __init__(self, d: T.Any):
         self.d = d
+        self.wrap_enter = not hasattr(d, "__enter__")
+        self.wrap_exit = not hasattr(d, "__exit__")
+        self.wrap_is_coord = not hasattr(d, "is_coord")
+
+    @staticmethod
+    def wrap(m: T.Any) -> T.Any:
+        t = type(m)
+        w = MessageWrapper.WRAP.get(t, None)
+        if w is None:
+            w = not all(
+                [
+                    hasattr(m, "__enter__"),
+                    hasattr(m, "__exit__"),
+                    hasattr(m, "is_coord"),
+                ]
+            )
+
+            MessageWrapper.WRAP[t] = w
+
+        if not w:
+            return m
+        else:
+            return MessageWrapper(m)
 
     def __enter__(self) -> T.Any:
-        return self.d
+        if self.wrap_enter:
+            return self.d
+        else:
+            return self.d.__enter__()
 
     def __exit__(self, *args) -> None:  # type: ignore
-        pass
+        if self.wrap_exit:
+            pass
+        else:
+            self.d.__exit__(*args)
 
     def __iter__(self):  # type: ignore
         return self.d.__iter__()
 
-    def is_coord(self, key, name=None):  # type: ignore
-        try:
-            return BufrMessage.code_is_coord(self.d[key + "->code"])
-        except Exception:
-            return False
+    def __getitem__(self, key: str) -> T.Any:
+        return self.d[key]
 
     def __getattr__(self, fname):  # type: ignore
         def call_func(*args, **kwargs):  # type: ignore
@@ -124,19 +145,51 @@ def call_func(*args, **kwargs):  # type: ignore
 
         return call_func
 
+    def is_coord(self, key: str) -> bool:
+        if self.wrap_is_coord:
+            try:
+                return bufr_code_is_coord(self.d[key + "->code"])
+            except Exception:
+                return False
+        else:
+            return self.d.is_coord(key)  # type: ignore
+
+
+class IsCoordCache:
+    """Caches if a BUFR key is a coordinate descriptor"""
+
+    def __init__(self, message: T.Any) -> None:
+        self.message = message
+        self.cache: T.Dict[str, bool] = {}
+
+    def check(self, key: str, name: str) -> bool:
+        c = self.cache.get(name, None)
+        if c is None:
+            if name in IS_KEY_COORD:
+                c = IS_KEY_COORD[name]
+                self.cache[name] = c
+                return c
+            else:
+                try:
+                    c = self.message.is_coord(key)
+                    self.cache[name] = c
+                    return c
+                except:
+                    return False
+        return c
+
 
 def message_structure(message: T.Any) -> T.Iterator[T.Tuple[int, str]]:
     level = 0
     coords: T.Dict[str, int] = collections.OrderedDict()
 
-    message = wrap_message(message)
+    message = MessageWrapper.wrap(message)
+    is_coord_cache = IsCoordCache(message)
+
     for key in message:
         name = key.rpartition("#")[2]
-        if name in IS_KEY_COORD:
-            is_coord = IS_KEY_COORD[name]
-        else:
-            is_coord = message.is_coord(key, name)
 
+        is_coord = is_coord_cache.check(key, name)
         while is_coord and name in coords:
             _, level = coords.popitem()  # OrderedDict.popitem uses LIFO order
 
@@ -633,7 +686,7 @@ def stream_bufr(
     for count, msg in enumerate(bufr_file, 1):
         # we use a context manager to automatically delete the handle of the BufrMessage.
         # We have to use a wrapper object here because a message can also be a dict
-        with wrap_message(msg) as message:
+        with MessageWrapper.wrap(msg) as message:
             if "count" in value_filters and not value_filters["count"].match(count):
                 continue
 
@@ -732,7 +785,7 @@ def stream_bufr_flat(
     for count, msg in enumerate(bufr_file, 1):
         # we use a context manager to automatically delete the handle of the BufrMessage.
         # We have to use a wrapper object here because a message can also be a dict
-        with wrap_message(msg) as message:
+        with MessageWrapper.wrap(msg) as message:
             # count filter
             if count_filter is not None and not count_filter.match(count):
                 continue
diff --git a/pdbufr/high_level_bufr/bufr.py b/pdbufr/high_level_bufr/bufr.py
index e0d0c7a..4145216 100644
--- a/pdbufr/high_level_bufr/bufr.py
+++ b/pdbufr/high_level_bufr/bufr.py
@@ -17,6 +17,13 @@
 from .codesmessage import CodesMessage
 
 
+def bufr_code_is_coord(code) -> bool:
+    try:
+        return code <= 9999
+    except:
+        return int(code[:3]) < 10
+
+
 class BufrMessage(CodesMessage):
     __doc__ = "\n".join(CodesMessage.__doc__.splitlines()[4:]).format(
         prod_type="BUFR", classname="BufrMessage", parent="BufrFile", alias="bufr"
@@ -36,8 +43,6 @@ def __init__(self, codes_file=None, clone=None, sample=None, headers_only=False)
         super(self.__class__, self).__init__(codes_file, clone, sample, headers_only)
         # self._unpacked = False
 
-        self._is_coord_cache = dict()
-
     # def get(self, key, ktype=None):
     #    """Return requested value, unpacking data values if necessary."""
     #    # TODO: Only do this if accessing arrays that need unpacking
@@ -92,30 +97,15 @@ def copy_data(self, destMsg):
         """Copy data values from this message to another message"""
         return eccodes.codes_bufr_copy_data(self.codes_id, destMsg.codes_id)
 
-    def is_coord(self, key, name=None):
-        if name is None:
-            name = key
-
-        c = self._is_coord_cache.get(name, None)
-        if c is None:
+    def is_coord(self, key):
+        try:
+            c = self._get(key + "->code", int)
             try:
-                c = self._get(name + "->code", int)
-                try:
-                    c = self.code_is_coord(c)
-                    self._is_coord_cache[name] = c
-                    return c
-                except:
-                    return False
+                return bufr_code_is_coord(c)
             except:
                 return False
-        return c
-
-    @staticmethod
-    def code_is_coord(code) -> bool:
-        try:
-            return code <= 9999
         except:
-            return int(code[:3]) < 10
+            return False
 
 
 class BufrFile(CodesFile):


Test patch:
diff --git a/tests/test_20_bufr_structure.py b/tests/test_20_bufr_structure.py
index 32537d2..8d7f1ed 100644
--- a/tests/test_20_bufr_structure.py
+++ b/tests/test_20_bufr_structure.py
@@ -426,7 +426,7 @@ def test_code_is_coord() -> None:
     """Ensures that the different ways of identifying a coordinate key from
     the BUFR descriptor/code agree"""
 
-    from pdbufr.high_level_bufr.bufr import BufrMessage
+    from pdbufr.high_level_bufr.bufr import bufr_code_is_coord
 
     data = {
         "000300": True,
@@ -439,5 +439,5 @@ def test_code_is_coord() -> None:
     }
 
     for k, v in data.items():
-        assert BufrMessage.code_is_coord(k) == v
-        assert BufrMessage.code_is_coord(int(k)) == v
+        assert bufr_code_is_coord(k) == v
+        assert bufr_code_is_coord(int(k)) == v