expr-codegen 0.16.6__tar.gz → 0.17.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/PKG-INFO +28 -18
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/README.md +26 -16
- expr_codegen-0.17.0/expr_codegen/_version.py +1 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/expr.py +11 -7
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/model.py +30 -4
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/polars/code.py +3 -3
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/rust/code.py +4 -4
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/sql/code.py +1 -1
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/tool.py +44 -50
- expr_codegen-0.16.6/expr_codegen/_version.py +0 -1
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/.gitignore +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/LICENSE +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/__init__.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/codes.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/dag.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/latex/__init__.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/latex/printer.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/__init__.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/code.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/helper.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/printer.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/ta.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/template.py.j2 +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/polars/__init__.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/polars/printer.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/polars/template.py.j2 +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/rust/__init__.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/rust/printer.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/rust/template.rs.j2 +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/sql/__init__.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/sql/printer.py +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/sql/template.sql.j2 +0 -0
- {expr_codegen-0.16.6 → expr_codegen-0.17.0}/pyproject.toml +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: expr_codegen
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.17.0
|
|
4
4
|
Summary: symbol expression to polars expression tool
|
|
5
5
|
Author-email: wukan <wu-kan@163.com>
|
|
6
6
|
License: BSD 3-Clause License
|
|
@@ -116,19 +116,19 @@ def _code_block_2():
|
|
|
116
116
|
CPV = cs_zscore(_corr) + cs_zscore(_beta)
|
|
117
117
|
|
|
118
118
|
|
|
119
|
-
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='
|
|
120
|
-
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='
|
|
121
|
-
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='
|
|
119
|
+
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file=sys.stdout) # 打印代码
|
|
120
|
+
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file="output.py") # 保存到文件
|
|
121
|
+
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by') # 只执行,不保存代码
|
|
122
122
|
|
|
123
123
|
code = StringIO()
|
|
124
|
-
codegen_exec(None, _code_block_1, _code_block_2, over_null='
|
|
124
|
+
codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file=code) # 保存到字符串
|
|
125
125
|
code.seek(0)
|
|
126
126
|
code.read() # 读取代码
|
|
127
127
|
|
|
128
128
|
# TODO 替换成合适的数据
|
|
129
129
|
df = pl.DataFrame()
|
|
130
|
-
df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='
|
|
131
|
-
df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='
|
|
130
|
+
df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='order_by').collect() # Lazy CPU
|
|
131
|
+
df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='order_by').collect(engine="gpu") # Lazy GPU
|
|
132
132
|
|
|
133
133
|
```
|
|
134
134
|
|
|
@@ -217,11 +217,11 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
|
|
|
217
217
|
|
|
218
218
|
第2种开头的`null`区域,是否影响结果由算子所决定,特别时是多列输入时`null`区域可能有数据
|
|
219
219
|
|
|
220
|
-
1. `over_null='partition_by'
|
|
221
|
-
2. `over_null='order_by'`。分到一个区域,`null
|
|
220
|
+
1. `over_null='partition_by'`。分到两个区域。多`ts`嵌套时,分区导致数据区变短,输出的有效数据变得更短
|
|
221
|
+
2. `over_null='order_by'`。分到一个区域,`null`排在前面。多`ts`嵌套时,数据区长度不变,输出的有效数据符合算法
|
|
222
222
|
3. `over_null=None`。不处理,直接调用,速度更快。如果确信不会中段产生`null`建议使用此参数
|
|
223
223
|
|
|
224
|
-
`codegen_exec(over_null='
|
|
224
|
+
`codegen_exec(over_null='order_by')`为全局使用`order_by`。但遇到`ts_count_nulls`这类`null`
|
|
225
225
|
函数就得使用`over_null=None`,所以本工具还新添了注释功能来指定单行表达式参数
|
|
226
226
|
|
|
227
227
|
1. `# --over_null partition_by`。单行`over_null='partition_by'`
|
|
@@ -233,18 +233,28 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
|
|
|
233
233
|
|
|
234
234
|
1. `# --over_null`传参注释只能写在单行表达式的后面,不能独立成一行,否则会被忽略
|
|
235
235
|
2. `# --over_null # --over_null=order_by`多个`#`时,只取第一个有效
|
|
236
|
-
3.
|
|
236
|
+
3. 只对`ts`函数有效。并能向上游传播。例如:
|
|
237
237
|
```python
|
|
238
|
-
|
|
239
|
-
|
|
238
|
+
X1 = ts_mean(ts_mean(CLOSE, 2), 2) # --over_null order_by
|
|
239
|
+
X2 = ts_mean(CLOSE, 2) # --over_null=order_by
|
|
240
|
+
X3 = ts_mean(CLOSE, 2) # 取`codegen_exec`参数传入的`over_null`值
|
|
241
|
+
X4 = cs_rank(CLOSE) + ts_mean(CLOSE, 2) # --over_null None
|
|
240
242
|
```
|
|
241
243
|
|
|
242
|
-
|
|
244
|
+
会自动翻译成
|
|
243
245
|
|
|
244
246
|
```python
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
247
|
+
#========================================func_0_cs__date
|
|
248
|
+
_x_1 = cs_rank(CLOSE) # --over_null=None
|
|
249
|
+
#========================================func_0_ts__asset
|
|
250
|
+
_x_0 = ts_mean(CLOSE, 2) # --over_null=order_by
|
|
251
|
+
_x_2 = ts_mean(CLOSE, 2) # --over_null=None
|
|
252
|
+
#========================================func_0_ts__asset
|
|
253
|
+
X1 = ts_mean(_x_0, 2) # --over_null=order_by
|
|
254
|
+
#========================================func_1_cl
|
|
255
|
+
X2 = _x_0 # --over_null=order_by
|
|
256
|
+
X3 = _x_0 # --over_null=order_by
|
|
257
|
+
X4 = _x_1 + _x_2 # --over_null=None
|
|
248
258
|
```
|
|
249
259
|
4. 由于很容易搞错,强烈建议生成`output_file`,检查生成的代码是否正确。
|
|
250
260
|
|
|
@@ -66,19 +66,19 @@ def _code_block_2():
|
|
|
66
66
|
CPV = cs_zscore(_corr) + cs_zscore(_beta)
|
|
67
67
|
|
|
68
68
|
|
|
69
|
-
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='
|
|
70
|
-
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='
|
|
71
|
-
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='
|
|
69
|
+
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file=sys.stdout) # 打印代码
|
|
70
|
+
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file="output.py") # 保存到文件
|
|
71
|
+
code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by') # 只执行,不保存代码
|
|
72
72
|
|
|
73
73
|
code = StringIO()
|
|
74
|
-
codegen_exec(None, _code_block_1, _code_block_2, over_null='
|
|
74
|
+
codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file=code) # 保存到字符串
|
|
75
75
|
code.seek(0)
|
|
76
76
|
code.read() # 读取代码
|
|
77
77
|
|
|
78
78
|
# TODO 替换成合适的数据
|
|
79
79
|
df = pl.DataFrame()
|
|
80
|
-
df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='
|
|
81
|
-
df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='
|
|
80
|
+
df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='order_by').collect() # Lazy CPU
|
|
81
|
+
df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='order_by').collect(engine="gpu") # Lazy GPU
|
|
82
82
|
|
|
83
83
|
```
|
|
84
84
|
|
|
@@ -167,11 +167,11 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
|
|
|
167
167
|
|
|
168
168
|
第2种开头的`null`区域,是否影响结果由算子所决定,特别时是多列输入时`null`区域可能有数据
|
|
169
169
|
|
|
170
|
-
1. `over_null='partition_by'
|
|
171
|
-
2. `over_null='order_by'`。分到一个区域,`null
|
|
170
|
+
1. `over_null='partition_by'`。分到两个区域。多`ts`嵌套时,分区导致数据区变短,输出的有效数据变得更短
|
|
171
|
+
2. `over_null='order_by'`。分到一个区域,`null`排在前面。多`ts`嵌套时,数据区长度不变,输出的有效数据符合算法
|
|
172
172
|
3. `over_null=None`。不处理,直接调用,速度更快。如果确信不会中段产生`null`建议使用此参数
|
|
173
173
|
|
|
174
|
-
`codegen_exec(over_null='
|
|
174
|
+
`codegen_exec(over_null='order_by')`为全局使用`order_by`。但遇到`ts_count_nulls`这类`null`
|
|
175
175
|
函数就得使用`over_null=None`,所以本工具还新添了注释功能来指定单行表达式参数
|
|
176
176
|
|
|
177
177
|
1. `# --over_null partition_by`。单行`over_null='partition_by'`
|
|
@@ -183,18 +183,28 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
|
|
|
183
183
|
|
|
184
184
|
1. `# --over_null`传参注释只能写在单行表达式的后面,不能独立成一行,否则会被忽略
|
|
185
185
|
2. `# --over_null # --over_null=order_by`多个`#`时,只取第一个有效
|
|
186
|
-
3.
|
|
186
|
+
3. 只对`ts`函数有效。并能向上游传播。例如:
|
|
187
187
|
```python
|
|
188
|
-
|
|
189
|
-
|
|
188
|
+
X1 = ts_mean(ts_mean(CLOSE, 2), 2) # --over_null order_by
|
|
189
|
+
X2 = ts_mean(CLOSE, 2) # --over_null=order_by
|
|
190
|
+
X3 = ts_mean(CLOSE, 2) # 取`codegen_exec`参数传入的`over_null`值
|
|
191
|
+
X4 = cs_rank(CLOSE) + ts_mean(CLOSE, 2) # --over_null None
|
|
190
192
|
```
|
|
191
193
|
|
|
192
|
-
|
|
194
|
+
会自动翻译成
|
|
193
195
|
|
|
194
196
|
```python
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
197
|
+
#========================================func_0_cs__date
|
|
198
|
+
_x_1 = cs_rank(CLOSE) # --over_null=None
|
|
199
|
+
#========================================func_0_ts__asset
|
|
200
|
+
_x_0 = ts_mean(CLOSE, 2) # --over_null=order_by
|
|
201
|
+
_x_2 = ts_mean(CLOSE, 2) # --over_null=None
|
|
202
|
+
#========================================func_0_ts__asset
|
|
203
|
+
X1 = ts_mean(_x_0, 2) # --over_null=order_by
|
|
204
|
+
#========================================func_1_cl
|
|
205
|
+
X2 = _x_0 # --over_null=order_by
|
|
206
|
+
X3 = _x_0 # --over_null=order_by
|
|
207
|
+
X4 = _x_1 + _x_2 # --over_null=None
|
|
198
208
|
```
|
|
199
209
|
4. 由于很容易搞错,强烈建议生成`output_file`,检查生成的代码是否正确。
|
|
200
210
|
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.17.0"
|
|
@@ -200,7 +200,7 @@ def get_children(func, func_kwargs, expr, output_exprs, output_symbols, date, as
|
|
|
200
200
|
expr
|
|
201
201
|
表达式
|
|
202
202
|
output_exprs
|
|
203
|
-
|
|
203
|
+
输出分割后的子表达式
|
|
204
204
|
output_symbols
|
|
205
205
|
输出每个子表达式中的符号
|
|
206
206
|
date
|
|
@@ -236,17 +236,21 @@ def get_children(func, func_kwargs, expr, output_exprs, output_symbols, date, as
|
|
|
236
236
|
# 外部与内部不同,需处理
|
|
237
237
|
for i, child in enumerate(children):
|
|
238
238
|
# 不在子中即表示不同
|
|
239
|
-
if curr in child:
|
|
240
|
-
|
|
239
|
+
# if curr in child:
|
|
240
|
+
# continue
|
|
241
241
|
if expr.args[i].is_Atom:
|
|
242
242
|
# print(expr.args[i], 'is_Atom 2')
|
|
243
243
|
continue
|
|
244
244
|
output_exprs = append_node(expr.args[i], output_exprs)
|
|
245
245
|
else:
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
246
|
+
for i, child in enumerate(children):
|
|
247
|
+
if not child: # 跳过空集合
|
|
248
|
+
continue
|
|
249
|
+
if next(iter(child))[0] == CL:
|
|
250
|
+
continue
|
|
251
|
+
if expr.args[i].is_Atom:
|
|
252
|
+
continue
|
|
253
|
+
output_exprs = append_node(expr.args[i], output_exprs)
|
|
250
254
|
|
|
251
255
|
# 按需返回,当前是基础算子就返回下一层信息,否则返回当前
|
|
252
256
|
if curr[0] == CL:
|
|
@@ -349,6 +349,8 @@ def skip_expr_node(G: nx.DiGraph, node, keep_nodes):
|
|
|
349
349
|
e = G.nodes[s]['expr']
|
|
350
350
|
e = e.xreplace({symbol: expr})
|
|
351
351
|
G.nodes[s]['expr'] = e
|
|
352
|
+
if d['comment'] != "#":
|
|
353
|
+
G.nodes[s]['comment'] = d['comment']
|
|
352
354
|
|
|
353
355
|
# 这里用了product生成多个关联边
|
|
354
356
|
G.add_edges_from(product(pred, succ))
|
|
@@ -365,14 +367,11 @@ def dag_start(exprs_list, func, func_kwargs, date, asset):
|
|
|
365
367
|
return G
|
|
366
368
|
|
|
367
369
|
|
|
368
|
-
def dag_middle(G, exprs_names, skip_columns, func, func_kwargs, date, asset
|
|
370
|
+
def dag_middle(G, exprs_names, skip_columns, func, func_kwargs, date, asset):
|
|
369
371
|
"""删除几个没有必要的节点"""
|
|
370
372
|
# 以下划线开头的节点,不保留
|
|
371
373
|
keep_nodes = [k for k in exprs_names if not k.startswith('_')]
|
|
372
374
|
|
|
373
|
-
if not skip_simplify:
|
|
374
|
-
# ts_rank(-RET - -RET, 20),防止替换成ts_rank(0, 20)
|
|
375
|
-
G = merge_nodes_1(G, keep_nodes, *keep_nodes)
|
|
376
375
|
G = merge_nodes_2(G, keep_nodes, *keep_nodes)
|
|
377
376
|
|
|
378
377
|
# 移除0出度的节点,但保留部分
|
|
@@ -405,3 +404,30 @@ def dag_end(G):
|
|
|
405
404
|
exprs_ldl._list = exprs_ldl.values()[1:]
|
|
406
405
|
|
|
407
406
|
return exprs_ldl, G
|
|
407
|
+
|
|
408
|
+
|
|
409
|
+
def propagate_comment_recursive(G):
|
|
410
|
+
"""按拓扑排序反序遍历,确保从终点开始"""
|
|
411
|
+
from collections import deque
|
|
412
|
+
|
|
413
|
+
# Kahn算法求拓扑排序
|
|
414
|
+
in_degree = {node: G.in_degree(node) for node in G.nodes()}
|
|
415
|
+
queue = deque([node for node in G.nodes() if in_degree[node] == 0])
|
|
416
|
+
topo_order = []
|
|
417
|
+
|
|
418
|
+
while queue:
|
|
419
|
+
node = queue.popleft()
|
|
420
|
+
topo_order.append(node)
|
|
421
|
+
for succ in G.successors(node):
|
|
422
|
+
in_degree[succ] -= 1
|
|
423
|
+
if in_degree[succ] == 0:
|
|
424
|
+
queue.append(succ)
|
|
425
|
+
|
|
426
|
+
# 反序遍历拓扑排序(从终点开始)
|
|
427
|
+
for node in reversed(topo_order):
|
|
428
|
+
comment = G.nodes[node]['comment']
|
|
429
|
+
if comment == "#":
|
|
430
|
+
continue
|
|
431
|
+
for pred in G.predecessors(node):
|
|
432
|
+
if G.nodes[pred]['comment'] == "#":
|
|
433
|
+
G.nodes[pred]['comment'] = comment
|
|
@@ -44,7 +44,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
|
|
|
44
44
|
filename,
|
|
45
45
|
date='date', asset='asset',
|
|
46
46
|
extra_codes: Sequence[str] = (),
|
|
47
|
-
over_null: Literal['order_by', 'partition_by', None] = '
|
|
47
|
+
over_null: Literal['order_by', 'partition_by', None] = 'order_by',
|
|
48
48
|
ge_date_idx: int = 0,
|
|
49
49
|
**kwargs):
|
|
50
50
|
"""基于模板的代码生成"""
|
|
@@ -83,14 +83,14 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
|
|
|
83
83
|
exprs_dst.append(f"#" + '=' * 40 + func_name)
|
|
84
84
|
else:
|
|
85
85
|
va, ex, sym, comment = kv
|
|
86
|
-
# 多个#时,只取第一个#后的参数
|
|
87
|
-
args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
|
|
88
86
|
s1 = str(ex)
|
|
89
87
|
s2 = p.doprint(ex)
|
|
90
88
|
if s1 != s2:
|
|
91
89
|
# 不想等,打印注释,显示会更直观察
|
|
92
90
|
func_code.append(f"# {va} = {s1}")
|
|
93
91
|
if k[0] == TS:
|
|
92
|
+
# 多个#时,只取第一个#后的参数
|
|
93
|
+
args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
|
|
94
94
|
ts_func_name = func_name
|
|
95
95
|
# https://github.com/pola-rs/polars/issues/12925#issuecomment-2552764629
|
|
96
96
|
_sym = [f"{s}.is_not_null()" for s in set(sym)]
|
|
@@ -15,7 +15,7 @@ def get_groupby_from_tuple(tup, func_name, drop_cols):
|
|
|
15
15
|
prefix2, *_ = tup
|
|
16
16
|
if len(drop_cols) > 0:
|
|
17
17
|
drop_cols = [f'"{c}".into()' for c in drop_cols]
|
|
18
|
-
drop_str = f'.drop(Selector::ByName {{ names: Arc::new([{
|
|
18
|
+
drop_str = f'.drop(Selector::ByName {{ names: Arc::new([{",".join(drop_cols)}]), strict: true }})'
|
|
19
19
|
else:
|
|
20
20
|
drop_str = ""
|
|
21
21
|
|
|
@@ -44,7 +44,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
|
|
|
44
44
|
filename,
|
|
45
45
|
date='date', asset='asset',
|
|
46
46
|
extra_codes: Sequence[str] = (),
|
|
47
|
-
over_null: Literal['order_by', 'partition_by', None] = '
|
|
47
|
+
over_null: Literal['order_by', 'partition_by', None] = 'order_by',
|
|
48
48
|
ge_date_idx: int = 0,
|
|
49
49
|
**kwargs):
|
|
50
50
|
"""基于模板的代码生成"""
|
|
@@ -83,14 +83,14 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
|
|
|
83
83
|
exprs_dst.append(f"#" + '=' * 40 + func_name)
|
|
84
84
|
else:
|
|
85
85
|
va, ex, sym, comment = kv
|
|
86
|
-
# 多个#时,只取第一个#后的参数
|
|
87
|
-
args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
|
|
88
86
|
s1 = str(ex)
|
|
89
87
|
s2 = p.doprint(ex)
|
|
90
88
|
if s1 != s2:
|
|
91
89
|
# 不想等,打印注释,显示会更直观察
|
|
92
90
|
func_code.append(f"// {va} = {s1}")
|
|
93
91
|
if k[0] == TS:
|
|
92
|
+
# 多个#时,只取第一个#后的参数
|
|
93
|
+
args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
|
|
94
94
|
ts_func_name = func_name
|
|
95
95
|
# https://github.com/pola-rs/polars/issues/12925#issuecomment-2552764629
|
|
96
96
|
_sym = [f'col("{s}").is_not_null()' for s in set(sym)]
|
|
@@ -14,7 +14,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
|
|
|
14
14
|
filename,
|
|
15
15
|
date='date', asset='asset',
|
|
16
16
|
extra_codes: Sequence[str] = (),
|
|
17
|
-
over_null: Literal['order_by', 'partition_by', None] = '
|
|
17
|
+
over_null: Literal['order_by', 'partition_by', None] = 'order_by',
|
|
18
18
|
table_name: str = 'self',
|
|
19
19
|
ge_date_idx: int = 0,
|
|
20
20
|
**kwargs):
|
|
@@ -1,9 +1,10 @@
|
|
|
1
|
+
import argparse
|
|
1
2
|
import inspect
|
|
2
3
|
import os
|
|
3
4
|
import pathlib
|
|
4
5
|
from functools import lru_cache
|
|
5
6
|
from io import TextIOBase
|
|
6
|
-
from typing import Sequence, Union, TypeVar, Optional, Literal, Iterable
|
|
7
|
+
from typing import Sequence, Union, TypeVar, Optional, Literal, Iterable
|
|
7
8
|
|
|
8
9
|
import polars as pl
|
|
9
10
|
from black import Mode, format_str
|
|
@@ -15,7 +16,7 @@ from sympy.simplify import cse_opts
|
|
|
15
16
|
|
|
16
17
|
from expr_codegen.codes import sources_to_exprs
|
|
17
18
|
from expr_codegen.expr import get_current_by_prefix, get_children, replace_exprs
|
|
18
|
-
from expr_codegen.model import dag_start, dag_end, dag_middle, _RESERVED_WORD_
|
|
19
|
+
from expr_codegen.model import dag_start, dag_end, dag_middle, _RESERVED_WORD_, propagate_comment_recursive, create_dag_exprs
|
|
19
20
|
|
|
20
21
|
try:
|
|
21
22
|
from pandas import DataFrame as _pd_DataFrame
|
|
@@ -61,6 +62,20 @@ def simplify2(expr, skip_simplify: bool):
|
|
|
61
62
|
return expr
|
|
62
63
|
|
|
63
64
|
|
|
65
|
+
def comment_adjust(expr_src, over_null):
|
|
66
|
+
parser = argparse.ArgumentParser()
|
|
67
|
+
parser.add_argument("--over_null", type=str, nargs="?", default=over_null)
|
|
68
|
+
|
|
69
|
+
expr_dst = []
|
|
70
|
+
for k, v, c in expr_src:
|
|
71
|
+
args, argv = parser.parse_known_args(args=c.split("#")[1].split(" "))
|
|
72
|
+
c = f"# --over_null={args.over_null}"
|
|
73
|
+
|
|
74
|
+
expr_dst.append((k, v, c))
|
|
75
|
+
|
|
76
|
+
return expr_dst
|
|
77
|
+
|
|
78
|
+
|
|
64
79
|
class ExprTool:
|
|
65
80
|
|
|
66
81
|
def __init__(self):
|
|
@@ -135,21 +150,7 @@ class ExprTool:
|
|
|
135
150
|
syms = [str(s) for s in syms]
|
|
136
151
|
return exprs, syms
|
|
137
152
|
|
|
138
|
-
def
|
|
139
|
-
"""减少中间变量数量,有利用减少内存占用"""
|
|
140
|
-
|
|
141
|
-
exprs_list = []
|
|
142
|
-
|
|
143
|
-
# cse前简化一次,cse后不再简化
|
|
144
|
-
# (~开盘涨停 & 昨收涨停) | (~收盘涨停 & 最高涨停)
|
|
145
|
-
for k, v in repl:
|
|
146
|
-
exprs_list.append((k, v, "#"))
|
|
147
|
-
for k, v, c in redu:
|
|
148
|
-
exprs_list.append((k, v, c))
|
|
149
|
-
|
|
150
|
-
return exprs_list
|
|
151
|
-
|
|
152
|
-
def cse(self, exprs, symbols_repl=None, exprs_src=None, skip_simplify=False):
|
|
153
|
+
def cse(self, exprs, symbols_repl=None, exprs_src=None, skip_simplify=False, date='date', asset='asset'):
|
|
153
154
|
"""多个子公式+长公式,提取公共公式
|
|
154
155
|
|
|
155
156
|
Parameters
|
|
@@ -171,38 +172,26 @@ class ExprTool:
|
|
|
171
172
|
表达式
|
|
172
173
|
|
|
173
174
|
"""
|
|
174
|
-
self.exprs_names = [k for k,
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
if skip_simplify
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
symbols_redu = iter(exprs_src)
|
|
188
|
-
for expr in redu[-outputs_len:]:
|
|
189
|
-
# 可能部分表达式只在之前出现过,后面完全用不到如,ts_rank(ts_decay_linear(x_147, 11.4157), 6.72611)
|
|
190
|
-
variable = next(symbols_redu)
|
|
191
|
-
a = symbols(variable[0])
|
|
192
|
-
new_redu.append((a, expr, variable[2]))
|
|
193
|
-
|
|
194
|
-
self.exprs_list = self.reduce(repl, new_redu)
|
|
195
|
-
|
|
196
|
-
# with open("exprs.pickle", "wb") as file:
|
|
197
|
-
# pickle.dump(exprs_dict, file)
|
|
198
|
-
|
|
175
|
+
self.exprs_names = [k for k, _, _ in exprs_src]
|
|
176
|
+
used = set(self.exprs_names) | {date, asset} | _RESERVED_WORD_
|
|
177
|
+
used.update(str(s) for expr, _ in exprs for s in expr.free_symbols)
|
|
178
|
+
pool = symbols_repl if symbols_repl is not None else numbered_symbols('_x_')
|
|
179
|
+
replacements = (s for s in pool if str(s) not in used)
|
|
180
|
+
optimizations = [] if skip_simplify else [(cse_opts.sub_pre, cse_opts.sub_post)]
|
|
181
|
+
self.exprs_list = []
|
|
182
|
+
for comment in dict.fromkeys(c for _, _, c in exprs_src):
|
|
183
|
+
outputs = [(k, v, c) for k, v, c in exprs_src if c == comment]
|
|
184
|
+
selected = [expr for expr, c in exprs if c == comment]
|
|
185
|
+
repl, redu = cse(selected, replacements, optimizations=optimizations)
|
|
186
|
+
self.exprs_list.extend((k, v, comment) for k, v in repl)
|
|
187
|
+
self.exprs_list.extend((symbols(k), expr, c) for (k, _, c), expr in zip(outputs, redu[-len(outputs):]))
|
|
199
188
|
return self.exprs_list
|
|
200
189
|
|
|
201
|
-
def dag(self,
|
|
190
|
+
def dag(self, skip_columns, date, asset):
|
|
202
191
|
"""生成DAG"""
|
|
203
192
|
G = dag_start(self.exprs_list, self.get_current_func, self.get_current_func_kwargs, date, asset)
|
|
204
|
-
|
|
205
|
-
|
|
193
|
+
G = dag_middle(G, self.exprs_names, skip_columns, self.get_current_func, self.get_current_func_kwargs, date, asset)
|
|
194
|
+
propagate_comment_recursive(G)
|
|
206
195
|
return dag_end(G)
|
|
207
196
|
|
|
208
197
|
def all(self, exprs_src, style: Literal['pandas', 'polars', 'sql', 'rust'] = 'polars',
|
|
@@ -210,7 +199,7 @@ class ExprTool:
|
|
|
210
199
|
replace: bool = True, regroup: bool = False, format: bool = True,
|
|
211
200
|
date='date', asset='asset',
|
|
212
201
|
extra_codes: Sequence[object] = (),
|
|
213
|
-
over_null: Literal['order_by', 'partition_by', None] = '
|
|
202
|
+
over_null: Literal['order_by', 'partition_by', None] = 'order_by',
|
|
214
203
|
table_name: str = 'self',
|
|
215
204
|
ge_date_idx: int = 0,
|
|
216
205
|
skip_simplify: bool = False,
|
|
@@ -254,14 +243,19 @@ class ExprTool:
|
|
|
254
243
|
if replace:
|
|
255
244
|
exprs_src = replace_exprs(exprs_src)
|
|
256
245
|
|
|
246
|
+
source_graph = create_dag_exprs((symbols(k), v, c) for k, v, c in exprs_src)
|
|
247
|
+
propagate_comment_recursive(source_graph)
|
|
248
|
+
exprs_context = [(k, v, source_graph.nodes[k]['comment']) for k, v, _ in exprs_src]
|
|
249
|
+
exprs_context = comment_adjust(exprs_context, over_null)
|
|
250
|
+
|
|
257
251
|
# 子表达式在前,原表式在最后
|
|
258
|
-
exprs_dst, syms_dst = self.merge(date, asset,
|
|
252
|
+
exprs_dst, syms_dst = self.merge(date, asset, exprs_context, skip_simplify=skip_simplify)
|
|
259
253
|
syms_dst = list(set(syms_dst) - _RESERVED_WORD_)
|
|
260
254
|
|
|
261
255
|
# 提取公共表达式
|
|
262
|
-
self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), exprs_src=
|
|
256
|
+
self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), exprs_src=exprs_context, skip_simplify=skip_simplify, date=date, asset=asset)
|
|
263
257
|
# 有向无环图流转
|
|
264
|
-
exprs_ldl, G = self.dag(
|
|
258
|
+
exprs_ldl, G = self.dag(skip_columns, date, asset)
|
|
265
259
|
|
|
266
260
|
if regroup:
|
|
267
261
|
exprs_ldl.optimize(merge=style != 'sql')
|
|
@@ -306,7 +300,7 @@ class ExprTool:
|
|
|
306
300
|
style: Literal['pandas', 'polars', 'sql', 'rust'] = 'polars',
|
|
307
301
|
template_file: Optional[str] = None,
|
|
308
302
|
date: str = 'date', asset: str = 'asset',
|
|
309
|
-
over_null: Literal['order_by', 'partition_by', None] = '
|
|
303
|
+
over_null: Literal['order_by', 'partition_by', None] = 'order_by',
|
|
310
304
|
table_name: str = 'self',
|
|
311
305
|
ge_date_idx: int = 0,
|
|
312
306
|
skip_simplify: bool = False,
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "0.16.6"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|