From 3627a5782cfa01840753ba11968e171d05139d44 Mon Sep 17 00:00:00 2001
From: nkaz001 <nkaz001@protonmail.com>
Date: Fri, 22 Mar 2024 10:40:18 -0400
Subject: [PATCH] feat: add `compress` and `structured_array` options.

---
 hftbacktest/data/utils/binancefutures.py | 15 +++++++++++----
 hftbacktest/data/utils/snapshot.py       | 16 ++++++++++++++--
 hftbacktest/data/utils/tardis.py         | 16 ++++++++++++----
 3 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/hftbacktest/data/utils/binancefutures.py b/hftbacktest/data/utils/binancefutures.py
index 5821eeb..f4c9769 100644
--- a/hftbacktest/data/utils/binancefutures.py
+++ b/hftbacktest/data/utils/binancefutures.py
@@ -5,10 +5,17 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from .. import correct, validate_data
+from .. import validate_data
 from ..validation import correct_event_order, convert_to_struct_arr
-from ... import DEPTH_EVENT, DEPTH_CLEAR_EVENT, DEPTH_SNAPSHOT_EVENT, TRADE_EVENT, correct_local_timestamp, \
-    COL_EXCH_TIMESTAMP, COL_LOCAL_TIMESTAMP
+from ... import (
+    DEPTH_EVENT,
+    DEPTH_CLEAR_EVENT,
+    DEPTH_SNAPSHOT_EVENT,
+    TRADE_EVENT,
+    COL_EXCH_TIMESTAMP,
+    COL_LOCAL_TIMESTAMP,
+    correct_local_timestamp
+)
 
 
 def convert(
@@ -57,7 +64,7 @@ def convert(
         base_latency: The value to be added to the feed latency.
                       See :func:`.correct_local_timestamp`.
         compress: If this is set to True, the output file will be compressed.
-        structured_array: If this is set to True, the output is converted into the new format.
+        structured_array: If this is set to True, the output is converted into the new format(currently only Rust impl).
 
     Returns:
         Converted data compatible with HftBacktest.
diff --git a/hftbacktest/data/utils/snapshot.py b/hftbacktest/data/utils/snapshot.py
index 83abf28..b15efcd 100644
--- a/hftbacktest/data/utils/snapshot.py
+++ b/hftbacktest/data/utils/snapshot.py
@@ -3,6 +3,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from ..validation import convert_to_struct_arr
 from ... import HftBacktest, Linear, ConstantLatency
 from ...typing import DataCollection, Data
 from ...reader import UNTIL_END_OF_DATA, DEPTH_SNAPSHOT_EVENT
@@ -13,7 +14,9 @@ def create_last_snapshot(
         tick_size: float,
         lot_size: float,
         initial_snapshot: Optional[Data] = None,
-        output_snapshot_filename: Optional[str] = None
+        output_snapshot_filename: Optional[str] = None,
+        compress: bool = False,
+        structured_array: bool = False
 ) -> NDArray:
     r"""
     Creates a snapshot of the last market depth for the specified data, which can be used as the initial snapshot data
@@ -26,6 +29,9 @@ def create_last_snapshot(
          initial_snapshot: The initial market depth snapshot.
          output_snapshot_filename: If provided, the snapshot data will be saved to the specified filename in ``npz``
                                    format.
+         compress: If this is set to True, the output file will be compressed.
+         structured_array: If this is set to True, the output is converted into the new
+                           format(currently only Rust impl).
 
     Returns:
         Snapshot of the last market depth compatible with HftBacktest.
@@ -65,7 +71,13 @@ def create_last_snapshot(
 
     snapshot = np.asarray(snapshot, np.float64)
 
+    if structured_array:
+        snapshot = convert_to_struct_arr(snapshot)
+
     if output_snapshot_filename is not None:
-        np.savez(output_snapshot_filename, data=snapshot)
+        if compress:
+            np.savez_compressed(output_snapshot_filename, data=snapshot)
+        else:
+            np.savez(output_snapshot_filename, data=snapshot)
 
     return snapshot
diff --git a/hftbacktest/data/utils/tardis.py b/hftbacktest/data/utils/tardis.py
index 22898b2..21db6b9 100644
--- a/hftbacktest/data/utils/tardis.py
+++ b/hftbacktest/data/utils/tardis.py
@@ -6,15 +6,22 @@
 
 from .. import validate_data
 from ..validation import correct_event_order, convert_to_struct_arr
-from ... import DEPTH_CLEAR_EVENT, DEPTH_SNAPSHOT_EVENT, TRADE_EVENT, DEPTH_EVENT, COL_LOCAL_TIMESTAMP, \
-    correct_local_timestamp, COL_EXCH_TIMESTAMP
+from ... import (
+    DEPTH_CLEAR_EVENT,
+    DEPTH_SNAPSHOT_EVENT,
+    TRADE_EVENT,
+    DEPTH_EVENT,
+    COL_LOCAL_TIMESTAMP,
+    COL_EXCH_TIMESTAMP,
+    correct_local_timestamp
+)
 
 
 def convert(
         input_files: List[str],
         output_filename: Optional[str] = None,
         buffer_size: int = 100_000_000,
-        ss_buffer_size: int = 10_000,
+        ss_buffer_size: int = 1_000_000,
         base_latency: float = 0,
         snapshot_mode: Literal['process', 'ignore_sod', 'ignore'] = 'process',
         compress: bool = False,
@@ -28,6 +35,7 @@ def convert(
                      e.g. ['incremental_book.csv', 'trades.csv'].
         output_filename: If provided, the converted data will be saved to the specified filename in ``npz`` format.
         buffer_size: Sets a preallocated row size for the buffer.
+        ss_buffer_size: Sets a preallocated row size for the snapshot.
         base_latency: The value to be added to the feed latency.
                       See :func:`.correct_local_timestamp`.
         snapshot_mode: - If this is set to 'ignore', all snapshots are ignored. The order book will converge to a
@@ -39,7 +47,7 @@ def convert(
                          for more details.
                        - Otherwise, all snapshot events will be processed.
         compress: If this is set to True, the output file will be compressed.
-        structured_array: If this is set to True, the output is converted into the new format.
+        structured_array: If this is set to True, the output is converted into the new format(currently only Rust impl).
 
     Returns:
         Converted data compatible with HftBacktest.