FIX: Fix DBNStore.to_df() when empty

Author: nmacholl

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## 0.14.1 - TBD
+- Fixed issue where `DBNStore.to_df()` would raise an exception if no records were present
+
 ## 0.14.0 - 2023-06-14
 - Added support for reusing a `Live` client to reconnect
 - Added `metadata` property to `Live`

databento/common/data.py

@@ -28,15 +28,6 @@ def get_deriv_ba_types(level: int) -> list[tuple[str, type | str]]:
     Schema.TRADES,
 )
 
-
-OHLCV_SCHEMAS = (
-    Schema.OHLCV_1S,
-    Schema.OHLCV_1M,
-    Schema.OHLCV_1H,
-    Schema.OHLCV_1D,
-)
-
-
 RECORD_HEADER: list[tuple[str, type | str]] = [
     ("length", np.uint8),
     ("rtype", np.uint8),
@@ -265,6 +256,7 @@ def get_deriv_ba_fields(level: int) -> list[str]:
 
 
 DERIV_HEADER_COLUMNS = [
+    "ts_recv",
     "ts_event",
     "ts_in_delta",
     "publisher_id",
@@ -279,6 +271,7 @@ def get_deriv_ba_fields(level: int) -> list[str]:
 ]
 
 OHLCV_HEADER_COLUMNS = [
+    "ts_event",
     "publisher_id",
     "instrument_id",
     "open",
@@ -289,7 +282,6 @@ def get_deriv_ba_fields(level: int) -> list[str]:
 ]
 
 DEFINITION_DROP_COLUMNS = [
-    "ts_recv",
     "length",
     "rtype",
     "reserved1",
@@ -299,14 +291,12 @@ def get_deriv_ba_fields(level: int) -> list[str]:
 ]
 
 IMBALANCE_DROP_COLUMNS = [
-    "ts_recv",
     "length",
     "rtype",
     "dummy",
 ]
 
 STATISTICS_DROP_COLUMNS = [
-    "ts_recv",
     "length",
     "rtype",
     "dummy",
@@ -330,6 +320,7 @@ def get_deriv_ba_fields(level: int) -> list[str]:
 
 COLUMNS = {
     Schema.MBO: [
+        "ts_recv",
         "ts_event",
         "ts_in_delta",
         "publisher_id",

databento/common/dbnstore.py

@@ -2,6 +2,7 @@
 
 import abc
 import datetime as dt
+import functools
 import logging
 from collections.abc import Generator
 from io import BytesIO
@@ -453,10 +454,6 @@ def _prepare_dataframe(
         df: pd.DataFrame,
         schema: Schema,
     ) -> pd.DataFrame:
-        # Setup column ordering and index
-        df.set_index(self._get_index_column(schema), inplace=True)
-        df = df.reindex(columns=COLUMNS[schema])
-
         if schema == Schema.MBO or schema in DERIV_SCHEMAS:
             df["flags"] = df["flags"] & 0xFF  # Apply bitmask
             df["side"] = df["side"].str.decode("utf-8")
@@ -941,7 +938,12 @@ def to_df(
                 raise ValueError("a schema must be specified for mixed DBN data")
             schema = self.schema
 
-        df = pd.DataFrame(self.to_ndarray(schema=schema))
+        df = pd.DataFrame(
+            self.to_ndarray(schema),
+            columns=COLUMNS[schema],
+        )
+        df.set_index(self._get_index_column(schema), inplace=True)
+
         df = self._prepare_dataframe(df, schema)
 
         if pretty_ts:
@@ -1049,12 +1051,10 @@ def to_ndarray(
             self,
         )
 
-        result = []
-        for record in schema_records:
-            np_rec = np.frombuffer(
-                bytes(record),
-                dtype=STRUCT_MAP[schema],
-            )
-            result.append(np_rec[0])
+        decoder = functools.partial(np.frombuffer, dtype=STRUCT_MAP[schema])
+        result = tuple(map(decoder, map(bytes, schema_records)))
+
+        if not result:
+            return np.empty(shape=(0, 1), dtype=STRUCT_MAP[schema])
 
-        return np.asarray(result)
+        return np.ravel(result)

databento/version.py

@@ -1 +1 @@
-__version__ = "0.14.0"
+__version__ = "0.14.1"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databento"
-version = "0.14.0"
+version = "0.14.1"
 description = "Official Python client library for Databento"
 authors = [
     "Databento <[email protected]>",

tests/test_historical_bento.py

@@ -413,17 +413,7 @@ def test_to_df_with_pretty_px_with_various_schemas_converts_prices_as_expected(
 
 @pytest.mark.parametrize(
     "expected_schema",
-    [
-        Schema.MBO,
-        Schema.MBP_1,
-        Schema.MBP_10,
-        Schema.TBBO,
-        Schema.TRADES,
-        Schema.OHLCV_1S,
-        Schema.OHLCV_1M,
-        Schema.OHLCV_1H,
-        Schema.OHLCV_1D,
-    ],
+    [pytest.param(schema, id=str(schema)) for schema in Schema],
 )
 def test_from_file_given_various_paths_returns_expected_metadata(
     test_data_path: Callable[[Schema], Path],
@@ -474,7 +464,7 @@ def test_mbo_to_csv_writes_expected_file_to_disk(
     written = open(path, mode="rb").read()
     assert path.exists()
     expected = (
-        b"ts_recv,ts_event,ts_in_delta,publisher_id,channel_id,instrument_id,order_id,act" # noqa
+        b"ts_recv,ts_event,ts_in_delta,publisher_id,channel_id,instrument_id,order_id,act"  # noqa
         b"ion,side,flags,price,size,sequence\n1609160400000704060,16091604000004298"  # noqa
         b"31,22993,1,0,5482,647784973705,C,A,128,3722750000000,1,1170352\n160916040"  # noqa
         b"0000711344,1609160400000431665,19621,1,0,5482,647784973631,C,A,128,372300000"  # noqa
@@ -718,16 +708,7 @@ def test_mbp_1_to_json_with_all_options_writes_expected_file_to_disk(
 
 @pytest.mark.parametrize(
     "schema",
-    [
-        s
-        for s in Schema
-        if s
-        not in (
-            Schema.OHLCV_1H,
-            Schema.OHLCV_1D,
-            Schema.DEFINITION,
-        )
-    ],
+    [pytest.param(schema, id=str(schema)) for schema in Schema],
 )
 def test_dbnstore_repr(
     test_data: Callable[[Schema], bytes],
@@ -820,17 +801,7 @@ def test_dbnstore_iterable_parallel(
 
 @pytest.mark.parametrize(
     "schema",
-    [
-        Schema.MBO,
-        Schema.MBP_1,
-        Schema.MBP_10,
-        Schema.OHLCV_1D,
-        Schema.OHLCV_1H,
-        Schema.OHLCV_1M,
-        Schema.OHLCV_1S,
-        Schema.TBBO,
-        Schema.TRADES,
-    ],
+    [pytest.param(schema, id=str(schema)) for schema in Schema],
 )
 def test_dbnstore_compression_equality(
     test_data: Callable[[Schema], bytes],
@@ -923,17 +894,7 @@ def test_dbnstore_buffer_long(
 
 @pytest.mark.parametrize(
     "schema",
-    [
-        Schema.MBO,
-        Schema.MBP_1,
-        Schema.MBP_10,
-        Schema.OHLCV_1D,
-        Schema.OHLCV_1H,
-        Schema.OHLCV_1M,
-        Schema.OHLCV_1S,
-        Schema.TBBO,
-        Schema.TRADES,
-    ],
+    [pytest.param(schema, id=str(schema)) for schema in Schema],
 )
 def test_dbnstore_to_ndarray_with_schema(
     schema: Schema,
@@ -957,19 +918,30 @@ def test_dbnstore_to_ndarray_with_schema(
         assert row == expected[i]
 
 
+def test_dbnstore_to_ndarray_with_schema_empty(
+    test_data: Callable[[Schema], bytes],
+) -> None:
+    """
+    Test that calling to_ndarray on a DBNStore that contains no data of the
+    specified schema returns an empty DataFrame.
+    """
+    # Arrange
+    dbn_stub_data = (
+        zstandard.ZstdDecompressor().stream_reader(test_data(Schema.TRADES)).read()
+    )
+
+    # Act
+    dbnstore = DBNStore.from_bytes(data=dbn_stub_data)
+
+    array = dbnstore.to_ndarray(schema=Schema.MBO)
+
+    # Assert
+    assert len(array) == 0
+
+
 @pytest.mark.parametrize(
     "schema",
-    [
-        Schema.MBO,
-        Schema.MBP_1,
-        Schema.MBP_10,
-        Schema.OHLCV_1D,
-        Schema.OHLCV_1H,
-        Schema.OHLCV_1M,
-        Schema.OHLCV_1S,
-        Schema.TBBO,
-        Schema.TRADES,
-    ],
+    [pytest.param(schema, id=str(schema)) for schema in Schema],
 )
 def test_dbnstore_to_df_with_schema(
     schema: Schema,
@@ -990,3 +962,24 @@ def test_dbnstore_to_df_with_schema(
 
     # Assert
     assert actual.equals(expected)
+
+
+def test_dbnstore_to_df_with_schema_empty(
+    test_data: Callable[[Schema], bytes],
+) -> None:
+    """
+    Test that calling to_df on a DBNStore that contains no data of the
+    specified schema returns an empty DataFrame.
+    """
+    # Arrange
+    dbn_stub_data = (
+        zstandard.ZstdDecompressor().stream_reader(test_data(Schema.TRADES)).read()
+    )
+
+    # Act
+    dbnstore = DBNStore.from_bytes(data=dbn_stub_data)
+
+    df = dbnstore.to_df(schema=Schema.MBO)
+
+    # Assert
+    assert df.empty