Skip to content

Commit 6e867df

Browse files
committed
fix(py): modify the Tardis CSV schema definition to handle string IDs in trade data.
1 parent 11d0067 commit 6e867df

File tree

1 file changed

+5
-26
lines changed
  • py-hftbacktest/hftbacktest/data/utils

1 file changed

+5
-26
lines changed

py-hftbacktest/hftbacktest/data/utils/tardis.py

+5-26
Original file line numberDiff line numberDiff line change
@@ -18,34 +18,13 @@
1818
event_dtype
1919
)
2020

21-
trade_cols = [
22-
'exchange',
23-
'symbol',
24-
'timestamp',
25-
'local_timestamp',
26-
'id',
27-
'side',
28-
'price',
29-
'amount'
30-
]
31-
32-
depth_cols = [
33-
'exchange',
34-
'symbol',
35-
'timestamp',
36-
'local_timestamp',
37-
'is_snapshot',
38-
'side',
39-
'price',
40-
'amount'
41-
]
4221

4322
trade_schema = {
4423
'exchange': pl.String,
4524
'symbol': pl.String,
4625
'timestamp': pl.Int64,
4726
'local_timestamp': pl.Int64,
48-
'id': pl.UInt64,
27+
'id': pl.String,
4928
'side': pl.String,
5029
'price': pl.Float64,
5130
'amount': pl.Float64,
@@ -125,16 +104,16 @@ def convert(
125104
with open(file) as f:
126105
line = f.readline()
127106
header = line.strip().split(',')
128-
if header == trade_cols:
107+
if header == list(trade_schema.keys()):
129108
schema = trade_schema
130-
elif header == depth_cols:
109+
elif header == list(depth_schema.keys()):
131110
schema = depth_schema
132111
except:
133112
# Fails to infer the file type; let Polars infer the schema.
134113
pass
135114

136115
df = pl.read_csv(file, schema=schema)
137-
if df.columns == trade_cols:
116+
if df.columns == list(trade_schema.keys()):
138117
arr = (
139118
df.with_columns(
140119
pl.when(pl.col('side') == 'buy')
@@ -171,7 +150,7 @@ def convert(
171150
)
172151
tmp[row_num:row_num + len(arr)] = arr[:]
173152
row_num += len(arr)
174-
elif df.columns == depth_cols:
153+
elif df.columns == list(depth_schema.keys()):
175154
arr = (
176155
df.with_columns(
177156
(pl.col('timestamp') * 1000)

0 commit comments

Comments
 (0)