Repository: ecmwf/pdbufr

Issue:
pdbufr can take a very long time to filter certain BUFR files/messages
The following refers to a BUFR file in ECMWF internal file system under my username: `$PERM/public/i0tp_07062020_00.buf`. It contains 364,359 messages and 715,925 subsets.
The following call to pdbufr takes around 41 minutes on the HPC:
```
import pdbufr

DIR = './'
filein = DIR + 'i0tp_07062020_00.buf'

df = pdbufr.read_bufr(filein,
    columns=("year",'month', "hour", "minute","latitude", "longitude", "atmosphericPathDelayInSatelliteSignal"),
                      filters={"stationOrSiteName": 'S3AG-EUME'},
    )
print(df)
```

I compared performance by running a pure eccodes-python script to iterate over the messages and unpack each one. That takes around 5 minutes.

I also modified the pdbufr code so as to avoid performing the actual filtering, and then it took around 5 minutes too.

I could see that certain sets of messages in the BUFR file were taking a lot longer to process than others. After putting a little profiling code into pdbufr and printing the indexes of the messages that were taking longer than 0.01 seconds to process, these were some of the message indexes (0-based) that consistently took longer:
```
count 36467 time 0.015330487862229347
count 82749 time 0.022938511800020933
count 110001 time 0.10010456619784236
count 110002 time 0.0794135918840766
… # most of the messages in between above and below
count 132562 time 1.6417972878552973
count 132563 time 0.033606610260903835
count 160917 time 0.04080211604014039
…
count 311422 time 0.02709425799548626
count 311423 time 0.010877656750380993
…
count 311619 time 1.4403185020200908
count 311620 time 5.845485555008054
…
```

Some standout times to process certain messages, which seemed consistent with multiple runs:
```
count 110774 time 15.490661825053394
count 110804 time 12.531519242096692
count 110878 time 10.05591939855367
count 110905 time 17.22438928298652
count 110909 time 14.911728173028678
count 110910 time 11.457727732136846
count 110911 time 12.165616693906486
count 110912 time 11.780960503034294
count 110913 time 9.73410841403529
count 110959 time 14.38328706100583
count 110960 time 0.0479189301840961
count 110961 time 12.562545870896429
count 110962 time 15.093519839923829
count 110963 time 16.94631726015359
count 110992 time 17.39802599698305
count 110993 time 18.96778273070231
count 110994 time 19.078917557373643
count 110996 time 0.017089318949729204
count 110997 time 13.097122138831764
count 110998 time 15.609233817085624
count 110999 time 12.680324002169073
count 111000 time 12.114437516778708
count 111001 time 19.27112303301692
count 111002 time 11.963714307174087
count 111196 time 25.567047986667603
count 111256 time 30.325322085991502
count 111552 time 39.79854283807799
count 111756 time 58.527211253065616
count 311590 time 16.124957408290356
count 311591 time 13.132242653053254
```

I copied a single message into `$PERM/public/one_msg_110773.bfr` (`bufr_copy -wcount=110773 i0tp_07062020_00.buf one_msg_110773.bfr`, which is the message at my 0-based 110774 index printouts) and it takes 15 seconds to process by itself with the same pdbufr script. This particular message has what looks like 6632 time periods encoded in it, and the single message occupies almost 140KB on disk.

Interestingly, calling bufr_dump on this message took only 0.6 seconds:
```
time bufr_dump -j a one_msg_110773.bfr >b.dump
```


Discussion:

Patch:
diff --git a/pdbufr/bufr_structure.py b/pdbufr/bufr_structure.py
index 6d0d528..834bc5b 100644
--- a/pdbufr/bufr_structure.py
+++ b/pdbufr/bufr_structure.py
@@ -8,7 +8,6 @@
 
 import collections
 import datetime
-import re
 import typing as T
 
 import attr
@@ -90,20 +89,53 @@ class UncompressedBufrKey:
 IS_KEY_COORD = {"subsetNumber": True, "operator": False}
 
 
-def message_structure(message: T.Mapping[str, T.Any]) -> T.Iterator[T.Tuple[int, str]]:
+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
+    types of messages."""
+
+    def __init__(self, d: T.Any):
+        self.d = d
+
+    def __enter__(self) -> T.Any:
+        return self.d
+
+    def __exit__(self, *args) -> None:  # type: ignore
+        pass
+
+    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 __getattr__(self, fname):  # type: ignore
+        def call_func(*args, **kwargs):  # type: ignore
+            return getattr(self.d, fname, *args, **kwargs)
+
+        return call_func
+
+
+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)
     for key in message:
         name = key.rpartition("#")[2]
-
         if name in IS_KEY_COORD:
             is_coord = IS_KEY_COORD[name]
         else:
-            try:
-                code = message[key + "->code"]
-                is_coord = int(code[:3]) < 10
-            except KeyError:
-                is_coord = False
+            is_coord = message.is_coord(key, name)
 
         while is_coord and name in coords:
             _, level = coords.popitem()  # OrderedDict.popitem uses LIFO order
@@ -551,23 +583,6 @@ def test_computed_keys(
     return True
 
 
-class CMWrapper:
-    """Makes it possible to use context manager both with BufrMessage and dict type of messages"""
-
-    def __init__(self, d: T.Any):
-        self.d = d
-
-    def __enter__(self) -> T.Any:
-        if isinstance(self.d, BufrMessage):
-            return self.d.__enter__()  # type: ignore
-        else:
-            return self.d
-
-    def __exit__(self, *args) -> None:  # type: ignore
-        if isinstance(self.d, BufrMessage):
-            self.d.__exit__(*args)  # type: ignore
-
-
 def stream_bufr(
     bufr_file: T.Iterable[T.MutableMapping[str, T.Any]],
     columns: T.Union[T.Sequence[str], str],
@@ -618,7 +633,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 CMWrapper(msg) as message:
+        with wrap_message(msg) as message:
             if "count" in value_filters and not value_filters["count"].match(count):
                 continue
 
@@ -717,7 +732,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 CMWrapper(msg) as message:
+        with wrap_message(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 0eb926a..e0d0c7a 100644
--- a/pdbufr/high_level_bufr/bufr.py
+++ b/pdbufr/high_level_bufr/bufr.py
@@ -36,6 +36,8 @@ class BufrMessage(CodesMessage):
         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
@@ -90,6 +92,31 @@ class BufrMessage(CodesMessage):
         """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:
+            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
+            except:
+                return False
+        return c
+
+    @staticmethod
+    def code_is_coord(code) -> bool:
+        try:
+            return code <= 9999
+        except:
+            return int(code[:3]) < 10
+
 
 class BufrFile(CodesFile):
     __doc__ = "\n".join(CodesFile.__doc__.splitlines()[4:]).format(
diff --git a/pdbufr/high_level_bufr/codesmessage.py b/pdbufr/high_level_bufr/codesmessage.py
index a09d159..9070b63 100644
--- a/pdbufr/high_level_bufr/codesmessage.py
+++ b/pdbufr/high_level_bufr/codesmessage.py
@@ -166,6 +166,12 @@ class CodesMessage(object):
 
     def get(self, key, ktype=None):
         """Get value of a given key as its native or specified type."""
+        # if key.endswith("->code"):
+        #     key = key.rpartition("->")[0]
+        #     name = key.rpartition("#")[2]
+        #     # print(name)
+        #     return self.code(key, name)
+
         with raise_keyerror(key):
             if eccodes.codes_get_size(self.codes_id, key) > 1:
                 ret = eccodes.codes_get_array(self.codes_id, key, ktype)
@@ -173,6 +179,9 @@ class CodesMessage(object):
                 ret = eccodes.codes_get(self.codes_id, key, ktype)
             return ret
 
+    def _get(self, key, ktype=None):
+        return eccodes.codes_get(self.codes_id, key, ktype)
+
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Release message handle and inform host file of release."""
         if self.codes_id != -1:


Test patch:
diff --git a/tests/test_20_bufr_structure.py b/tests/test_20_bufr_structure.py
index 33196e8..32537d2 100644
--- a/tests/test_20_bufr_structure.py
+++ b/tests/test_20_bufr_structure.py
@@ -420,3 +420,24 @@ def test_stream_bufr() -> None:
 
     assert len(res) == 1
     assert res == expected_2
+
+
+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
+
+    data = {
+        "000300": True,
+        "010000": False,
+        "009999": True,
+        "010999": False,
+        "100000": False,
+        "201001": False,
+        "301001": False,
+    }
+
+    for k, v in data.items():
+        assert BufrMessage.code_is_coord(k) == v
+        assert BufrMessage.code_is_coord(int(k)) == v