rust-pyfunc 0.78.0__tar.gz → 0.80.0__tar.gz

{rust_pyfunc-0.78.0 → rust_pyfunc-0.80.0}/.gitignore

@@ -239,3 +239,4 @@ docs/theme_feature_expansion_explanation.md
 .mindfs/sessions/session-list.db
 .maturin_zig/*
 dist_versions/*
+python/rust_pyfunc/rust_pyfunc_worker

{rust_pyfunc-0.78.0 → rust_pyfunc-0.80.0}/Cargo.lock

@@ -550,6 +550,17 @@ version = "0.8.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
 
+[[package]]
+name = "core_affinity"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342"
+dependencies = [
+ "libc",
+ "num_cpus",
+ "winapi",
+]
+
 [[package]]
 name = "cpufeatures"
 version = "0.2.17"
@@ -2181,12 +2192,13 @@ checksum = "082f11ffa03bbef6c2c6ea6bea1acafaade2fd9050ae0234ab44a2153742b058"
 
 [[package]]
 name = "rust_pyfunc"
-version = "0.78.0"
+version = "0.80.0"
 dependencies = [
  "arrow",
  "base64 0.21.7",
  "bincode",
  "chrono",
+ "core_affinity",
  "criterion",
  "crossbeam",
  "csv",
@@ -2671,6 +2683,22 @@ dependencies = [
  "safe_arch",
 ]
 
+[[package]]
+name = "winapi"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
+dependencies = [
+ "winapi-i686-pc-windows-gnu",
+ "winapi-x86_64-pc-windows-gnu",
+]
+
+[[package]]
+name = "winapi-i686-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
+
 [[package]]
 name = "winapi-util"
 version = "0.1.9"
@@ -2680,6 +2708,12 @@ dependencies = [
  "windows-sys 0.59.0",
 ]
 
+[[package]]
+name = "winapi-x86_64-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
+
 [[package]]
 name = "windows-core"
 version = "0.61.2"

{rust_pyfunc-0.78.0 → rust_pyfunc-0.80.0}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rust_pyfunc"
-version = "0.78.0"
+version = "0.80.0"
 edition = "2021"
 description = "A collection of high-performance Python functions implemented in Rust"
 readme = "README.md"
@@ -46,9 +46,15 @@ log = "0.4"
 once_cell = "1.19"
 csv = "1.3"
 ndarray-npy = "0.8"
+core_affinity = "0.8"
 hdf5-metno = { version = "0.12", features = ["static", "zlib"], optional = true }
 zstd = "0.13"
 
+# Worker 二进制（多进程因子流水线）
+[[bin]]
+name = "rust_pyfunc_worker"
+path = "src/bin/worker_pipeline.rs"
+
 [features]
 default = ["hdf5"]
 hdf5 = ["hdf5-metno"]
@@ -65,8 +71,9 @@ criterion = "0.5"
 [profile.release]
 lto = true
 codegen-units = 1
-panic = "abort"
 opt-level = 3
+# 注意：不用 panic = "abort"，因为 worker bin 需要 catch_unwind 捕获计算 panic，
+# 避免整个 worker 进程崩溃重启。pyo3 0.18 在 unwind 下正常工作。
 
 [build-dependencies]
 pkg-config = "0.3"

{rust_pyfunc-0.78.0 → rust_pyfunc-0.80.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rust_pyfunc
-Version: 0.78.0
+Version: 0.80.0
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{rust_pyfunc-0.78.0 → rust_pyfunc-0.80.0}/alter.sh

@@ -26,5 +26,10 @@ conda activate chenzongwei311
 export PATH="/home/chenzongwei/.local/bin:$PATH"
 /home/chenzongwei/.local/bin/mold -run maturin develop --release
 
+# 5. 编译并部署 worker 二进制（多进程因子流水线 mode="multiprocess" 需要）
+/home/chenzongwei/.local/bin/mold -run cargo build --release --bin rust_pyfunc_worker
+cp target/release/rust_pyfunc_worker python/rust_pyfunc/rust_pyfunc_worker
+echo "✅ worker 二进制已部署到 python/rust_pyfunc/rust_pyfunc_worker"
+
 # 如果需要，你也可以在脚本结束时停用环境
 # conda deactivate

{rust_pyfunc-0.78.0 → rust_pyfunc-0.80.0}/python/rust_pyfunc/core_functions.pyi

@@ -2288,3 +2288,104 @@ def compute_ts_stats(series: List[float]) -> List[float]:
         数据不足5个时返回空列表
     """
     ...
+
+
+def read_trade_fast(
+    code: str,
+    date: int,
+    with_retreat: int = 0,
+    with_afternoon_adjust: bool = False,
+) -> NDArray[np.float64]:
+    """高速读取逐笔成交数据（Rust 多线程实现，替代 read_trade 的加速版）。
+
+    用 memmap2 内存映射文件，大文件按字节块切分并行解析（rayon），
+    读取时即完成全部预处理：flag!=32 过滤、exchtime 微秒整数→epoch 秒、
+    可选的下午时段平移（adjust_afternoon）。
+
+    与 read_trade 的差异：
+    - 直接返回 (n, 7) float64 numpy 数组，跳过 DataFrame 构造与 Timestamp 往返转换
+    - with_afternoon_adjust=True 时在读时一步完成下午时段平移并过滤非盘中数据
+    - time_sec 列为 epoch 秒（与 prepare() 中 exchtime.astype(int64)//1e9 一致）
+
+    列顺序（n×7 矩阵）：
+        [0] time_sec    成交时间，epoch 秒
+        [1] price       成交价
+        [2] volume      成交量
+        [3] turnover    成交额
+        [4] flag        成交标志（66=主买, 83=主卖；32=撤单，已被过滤）
+        [5] bid_order   买方订单号
+        [6] ask_order   卖方订单号
+
+    Parameters
+    ----------
+    code : str
+        股票代码，如 "000001"
+    date : int
+        交易日，如 20260605
+    with_retreat : int, 默认 0
+        1 时保留 flag==32 的撤单记录，0 时过滤（与 read_trade 默认一致）
+    with_afternoon_adjust : bool, 默认 False
+        True 时做 adjust_afternoon 平移：保留 [09:30,11:30]，下午 [13:00,14:57]
+        整体前移 90 分钟，过滤集合竞价前/收盘后的数据
+
+    Returns
+    -------
+    numpy.ndarray
+        (n, 7) float64 数组
+    """
+    ...
+
+
+def run_factor_pipeline(
+    pipeline: str,
+    tasks: List,
+    n_jobs: int,
+    backup_file: str,
+    expected_result_length: int,
+    trading_days: List[int],
+    params: object = None,
+    update_mode: bool = False,
+    bind_cores: bool = True,
+    backup_batch_size: int = 2000,
+    progress_log: bool = False,
+) -> None:
+    """纯 Rust 因子流水线引擎（run_factor_pipeline 优化方案 Phase 3）。
+
+    把 go 函数（如 hm90.go）的整条链路在 Rust 内部一气呵成，消除 Python worker、
+    pandas、pyo3 往返、msgpack 等所有流转浪费。并行控制采用自适应嵌套并行 + 核绑定，
+    总线程恒定 = n_jobs。
+
+    backup 格式与现有 run_pools_queue 完全兼容（复用 backup_writer），
+    后续 query_backup / query_backup_factor_only_ultra_fast 无需改动。
+
+    Parameters
+    ----------
+    pipeline : str
+        流水线标识，目前支持 "order_pair_hm90"（hm90.go 的 Rust 翻译）
+    tasks : List[[int, str]]
+        任务列表，每个元素 [date, code]
+    n_jobs : int
+        并行线程数（总 CPU 占用 ≈ n_jobs × 100%，默认生产用 200）
+    backup_file : str
+        备份文件路径（与 run_pools_queue 的 backup_file 格式兼容）
+    expected_result_length : int
+        每个任务的预期结果长度（因子数）
+    trading_days : List[int]
+        交易日历数组（由 rp.td.trading_days.tolist() 提供，用于 last_trading_day 查找）
+    params : dict, 可选
+        流水线参数，如 {"tolerance_v1": 0.001, "tolerance_v2": 0.00001}
+    update_mode : bool, 默认 False
+        断点续算：True 时跳过 backup 中已完成的任务
+    bind_cores : bool, 默认 True
+        核绑定：把 n_jobs 个线程各绑定到一个物理核，CPU 占用精确稳定
+    backup_batch_size : int, 默认 2000
+        每攒满 N 个结果写一次 backup（批量 zstd 压缩）
+    progress_log : bool, 默认 False
+        每 500 个任务打印一次进度
+
+    Notes
+    -----
+    "order_pair_hm90" 流水线关闭了 lyapunov 特征（原 get_features_factors 默认开启，
+    但花 69% 时间只产生 2.3% 特征）。如需 lyapunov，仍可用原 run_pools_queue。
+    """
+    ...

rust_pyfunc-0.80.0/src/bin/worker_pipeline.rs

@@ -0,0 +1,101 @@
+//! rust_pyfunc_worker —— 多进程因子流水线的 worker 二进制入口。
+//!
+//! 由主进程（run_factor_pipeline mode="multiprocess"）通过 Command::new spawn。
+//! 通信协议：stdin/stdout pipe + 4字节长度前缀 + bincode。
+//!
+//! 生命周期：
+//!   1. 启动时从环境变量 RUST_PYFUNC_CORE_AFFINITY_IDX 读取核绑定索引（可选）
+//!   2. 读 Init 消息（获取 params、trading_days、expected_len），回复 Ready
+//!   3. while loop { 读 Task → 调 pipeline_order_pair_hm90 → 写 Result }
+//!   4. 收到 Shutdown（长度0）则退出
+//!   5. 计算错误不 panic，回传 Error 消息（避免进程崩溃重启开销）
+use rust_pyfunc::factor_pipeline::{
+    ipc_read_result, ipc_read_task, ipc_write, ipc_write_result,
+    pipeline_order_pair_hm90, Hm90Params, ResultMessage, TaskMessage,
+};
+use std::io::{BufReader, BufWriter};
+
+fn main() {
+    // 核绑定（可选）
+    if let Ok(idx_str) = std::env::var("RUST_PYFUNC_CORE_AFFINITY_IDX") {
+        if let Ok(idx) = idx_str.parse::<usize>() {
+            if let Some(cores) = core_affinity::get_core_ids() {
+                if idx < cores.len() {
+                    let _ = core_affinity::set_for_current(cores[idx].clone());
+                }
+            }
+        }
+    }
+
+    let stdin = std::io::stdin();
+    let stdout = std::io::stdout();
+    // stdin 用锁避免竞争，BufReader 缓冲
+    let mut reader = BufReader::new(stdin.lock());
+    let mut writer = BufWriter::new(stdout.lock());
+
+    // 1. 读 Init 消息
+    let (params, trading_days, _expected_len) = match ipc_read_task(&mut reader) {
+        Ok(Some(TaskMessage::Init { params, trading_days, expected_len })) => {
+            (params, trading_days, expected_len)
+        }
+        _ => {
+            eprintln!("worker: 未收到 Init 消息，退出");
+            std::process::exit(1);
+        }
+    };
+
+    // 回复 Ready
+    if ipc_write_result(&mut writer, &ResultMessage::Ready).is_err() {
+        eprintln!("worker: 无法回复 Ready，退出");
+        std::process::exit(1);
+    }
+
+    // 2. 主循环：读任务 → 计算 → 写结果
+    loop {
+        let msg = match ipc_read_task(&mut reader) {
+            Ok(Some(TaskMessage::Task { date, code })) => (date, code),
+            Ok(Some(TaskMessage::Shutdown)) | Ok(None) => {
+                // 优雅退出
+                break;
+            }
+            Ok(Some(TaskMessage::Init { .. })) => {
+                // 忽略重复 Init
+                continue;
+            }
+            Err(e) => {
+                eprintln!("worker: 读取任务失败: {}, 退出", e);
+                break;
+            }
+        };
+
+        let (date, code) = msg;
+
+        // 调用流水线计算
+        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+            pipeline_order_pair_hm90(date, &code, &params, &trading_days, 0)
+        }));
+
+        match result {
+            Ok(facs) => {
+                let task_result = rust_pyfunc::backup_reader::TaskResult {
+                    date,
+                    code,
+                    timestamp: 0,
+                    facs,
+                };
+                if ipc_write_result(&mut writer, &ResultMessage::Result(task_result)).is_err() {
+                    eprintln!("worker: 写结果失败，退出");
+                    break;
+                }
+            }
+            Err(_) => {
+                // panic 被捕获，回传 Error
+                let _ = ipc_write_result(&mut writer, &ResultMessage::Error {
+                    date,
+                    code,
+                    msg: "计算 panic".to_string(),
+                });
+            }
+        }
+    }
+}

{rust_pyfunc-0.78.0 → rust_pyfunc-0.80.0}/src/factor_neutralization_io_optimized.rs

@@ -1,4 +1,4 @@
-use arrow::array::{Array, Float64Array, Int32Array, Int64Array, LargeStringArray, StringArray};
+use arrow::array::{Array, Float32Array, Float64Array, Int32Array, Int64Array, LargeStringArray, StringArray};
 use chrono::Local;
 use nalgebra::{DMatrix, DVector};
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
@@ -262,17 +262,38 @@ impl IOOptimizedStyleData {
             }
         };
 
-        // 批量提取风格因子列引用
-        let style_columns: Result<Vec<&Float64Array>, _> = (2..13)
+        // 批量提取风格因子列取值闭包: value_0~value_9(Float64) + ind_1~ind_31(Float32) = 41列
+        let style_getters: Vec<Box<dyn Fn(usize) -> f64 + Send + Sync>> = (2..43)
             .map(|i| {
-                batch
-                    .column(i)
-                    .as_any()
-                    .downcast_ref::<Float64Array>()
-                    .ok_or_else(|| PyRuntimeError::new_err(format!("风格因子列{}类型错误", i - 2)))
+                let col = batch.column(i);
+                let as_any = col.as_any();
+                let getter: Box<dyn Fn(usize) -> f64 + Send + Sync> = if let Some(arr) =
+                    as_any.downcast_ref::<Float64Array>()
+                {
+                    Box::new(move |row_idx: usize| {
+                        if arr.is_null(row_idx) {
+                            f64::NAN
+                        } else {
+                            arr.value(row_idx)
+                        }
+                    })
+                } else if let Some(arr) = as_any.downcast_ref::<Float32Array>() {
+                    Box::new(move |row_idx: usize| {
+                        if arr.is_null(row_idx) {
+                            f64::NAN
+                        } else {
+                            arr.value(row_idx) as f64
+                        }
+                    })
+                } else {
+                    return Err(PyRuntimeError::new_err(format!(
+                        "风格因子列{}类型错误：期望Float64或Float32",
+                        i - 2
+                    )));
+                };
+                Ok(getter)
             })
-            .collect();
-        let style_columns = style_columns?;
+            .collect::<PyResult<Vec<_>>>()?;
 
         // 向量化处理每一行
         for row_idx in 0..batch.num_rows() {
@@ -280,16 +301,7 @@ impl IOOptimizedStyleData {
             let stock = get_stock_value(row_idx);
 
             // 使用迭代器和collect优化风格值提取
-            let style_values: Vec<f64> = style_columns
-                .iter()
-                .map(|col| {
-                    if col.is_null(row_idx) {
-                        f64::NAN
-                    } else {
-                        col.value(row_idx)
-                    }
-                })
-                .collect();
+            let style_values: Vec<f64> = style_getters.iter().map(|g| g(row_idx)).collect();
 
             data_by_date
                 .entry(date)
@@ -310,19 +322,18 @@ impl IOOptimizedStyleData {
         // 预分配所有数据结构
         let mut stocks = Vec::with_capacity(n_stocks);
         let mut stock_index_map = HashMap::with_capacity(n_stocks);
-        let mut style_matrix = DMatrix::zeros(n_stocks, 12);
+        let mut style_matrix = DMatrix::zeros(n_stocks, 41);
 
-        // 单次遍历填充所有数据结构
+        // 单次遍历填充所有数据结构（41维风格: 10 barra + 31 行业哑变量, 行业和=1已含截距, 不再加常数项）
         for (i, (stock, style_values)) in stock_data.into_iter().enumerate() {
             stock_index_map.insert(stock.clone(), i);
             stocks.push(stock);
 
             // 直接写入矩阵（避免边界检查）
             unsafe {
-                for j in 0..11 {
+                for j in 0..41 {
                     *style_matrix.get_unchecked_mut((i, j)) = style_values[j];
                 }
-                *style_matrix.get_unchecked_mut((i, 11)) = 1.0;
             }
         }
 
@@ -939,12 +950,13 @@ fn process_single_date_io_optimized(
 
         let beta = &selected_regression_matrix * &aligned_y_vector;
 
+        let n_features = day_data.style_matrix.ncols();
         let mut result_values = Vec::new();
         for (i, &union_idx) in valid_union_indices.iter().enumerate() {
             let style_idx = valid_style_indices[i];
 
             let mut predicted_value = 0.0;
-            for j in 0..12 {
+            for j in 0..n_features {
                 predicted_value += day_data.style_matrix[(style_idx, j)] * beta[j];
             }