From cc1b552efc3ee4c48452466b14c61b4a73d731ca Mon Sep 17 00:00:00 2001 From: Yash Parihar Date: Wed, 11 Feb 2026 10:08:13 +0530 Subject: [PATCH 1/2] feat: zero-copy data transfer with rust-numpy (closes #17) Replace pylist_to_vec2 with numpy::PyReadonlyArray2 for zero-copy buffer access at the Python-Rust bridge. Eliminates per-element Python object extraction and .tolist() overhead. Changes: - etna_core/Cargo.toml: added numpy = 0.27 dependency - etna_core/src/lib.rs: replaced pylist_to_vec2/PyList with ndarray_to_vec2/PyReadonlyArray2 - etna/preprocessing.py: return numpy arrays instead of .tolist() - etna/api.py: pass contiguous float32 arrays directly to Rust - tests/test_tqdm_progress.py: fixed mock patching for compiled extension --- Cargo.lock | 75 +++++++++++++++++++++++++++++++++++++ etna/api.py | 13 +++++-- etna/preprocessing.py | 4 +- etna_core/Cargo.toml | 3 +- etna_core/src/lib.rs | 30 ++++++++------- tests/test_tqdm_progress.py | 2 + 6 files changed, 106 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dd9eadd..58ff6ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -752,6 +752,7 @@ name = "etna_rust" version = "0.1.0" dependencies = [ "maturin", + "numpy", "pyo3", "rand 0.9.2", "serde", @@ -1253,6 +1254,16 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "maturin" version = "1.11.3" @@ -1422,6 +1433,21 @@ dependencies = [ "tempfile", ] +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "nom" version = "7.1.3" @@ -1450,12 +1476,30 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -1471,6 +1515,22 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "numpy" +version = "0.27.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7aac2e6a6e4468ffa092ad43c39b81c79196c2bb773b8db4085f695efe3bba17" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + [[package]] name = "object" version = "0.32.2" @@ -1619,6 +1679,15 @@ version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f89776e4d69bb58bc6993e99ffa1d11f228b839984854c7daeb5d37f87cbe950" +[[package]] +name = "portable-atomic-util" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a9db96d7fa8782dd8c15ce32ffe8680bbd1e978a43bf51a34d39483540495f5" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.4" @@ -1833,6 +1902,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rayon" version = "1.11.0" diff --git a/etna/api.py b/etna/api.py index 86a08eb..8025140 100644 --- a/etna/api.py +++ b/etna/api.py @@ -87,10 +87,14 @@ def train(self, epochs: int = 100, lr: float = 0.01, batch_size: int = 32, weigh print("[*] Preprocessing data...") X, y = self.preprocessor.fit_transform(self.df, self.target) + # Ensure contiguous float32 arrays for zero-copy transfer to Rust + X = np.ascontiguousarray(X, dtype=np.float32) + y = np.ascontiguousarray(y, dtype=np.float32) + # Cache training data for predict() without arguments - self._cached_X = np.array(X) + self._cached_X = X - self.input_dim = len(X[0]) + self.input_dim = X.shape[1] self.output_dim = self.preprocessor.output_dim optimizer_lower = optimizer.lower() @@ -155,6 +159,8 @@ def predict(self, data_path: str = None): df = load_data(data_path) print("Transforming input data...") X_new = self.preprocessor.transform(df) + # Ensure contiguous float32 array for zero-copy transfer to Rust + X_new = np.ascontiguousarray(X_new, dtype=np.float32) # Case 2: Predict on cached training data else: @@ -163,8 +169,7 @@ def predict(self, data_path: str = None): "No data available for prediction. " "Pass a CSV path to predict(data_path=...)." ) - # Convert numpy array to list for Rust - X_new = self._cached_X.tolist() if isinstance(self._cached_X, np.ndarray) else self._cached_X + X_new = np.ascontiguousarray(self._cached_X, dtype=np.float32) preds = self.rust_model.predict(X_new) diff --git a/etna/preprocessing.py b/etna/preprocessing.py index 543478d..af2c80c 100644 --- a/etna/preprocessing.py +++ b/etna/preprocessing.py @@ -127,7 +127,7 @@ def fit_transform(self, df: pd.DataFrame, target_col: str): y_scaled = (y_vals - self.target_mean) / self.target_std y_final = y_scaled.reshape(-1, 1) - return X_final.tolist(), y_final.tolist() + return X_final, y_final # ------------------------------------------------- @@ -160,7 +160,7 @@ def transform(self, df: pd.DataFrame): X_processed.append(one_hot) X_final = np.hstack(X_processed) if X_processed else np.empty((len(df), 0)) - return X_final.tolist() + return X_final def get_state(self): return { diff --git a/etna_core/Cargo.toml b/etna_core/Cargo.toml index 9857776..2d26317 100644 --- a/etna_core/Cargo.toml +++ b/etna_core/Cargo.toml @@ -13,7 +13,8 @@ name = "etna_rust" crate-type = ["cdylib"] [dependencies] -pyo3 = { version = "0.27.2", features = ["extension-module", "abi3-py38"] } +pyo3 = { version = "0.27.2", features = ["extension-module"] } +numpy = "0.27" rand = "0.9.2" serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.149" diff --git a/etna_core/src/lib.rs b/etna_core/src/lib.rs index ad821e7..d3d41d6 100644 --- a/etna_core/src/lib.rs +++ b/etna_core/src/lib.rs @@ -6,18 +6,21 @@ mod loss_function; mod optimizer; use pyo3::prelude::*; -use pyo3::types::PyList; use pyo3::Py; +use numpy::PyReadonlyArray2; use crate::layers::Activation; use crate::model::{OptimizerType, SimpleNN}; -/// Safe conversion helper -fn pylist_to_vec2(pylist: &Bound<'_, PyList>) -> PyResult>> { - pylist - .iter() - .map(|item| item.extract::>()) - .collect::, _>>() +/// Zero-copy conversion: reads directly from NumPy's contiguous buffer +/// instead of iterating over Python list objects one element at a time. +fn ndarray_to_vec2(arr: PyReadonlyArray2<'_, f32>) -> Vec> { + let array = arr.as_array(); + array + .rows() + .into_iter() + .map(|row: numpy::ndarray::ArrayView1<'_, f32>| row.to_vec()) + .collect() } /// Python Class Wrapper @@ -53,8 +56,8 @@ impl EtnaModel { fn train( &mut self, py: Python<'_>, - x: &Bound<'_, PyList>, - y: &Bound<'_, PyList>, + x: PyReadonlyArray2<'_, f32>, + y: PyReadonlyArray2<'_, f32>, epochs: usize, lr: f32, batch_size: usize, @@ -62,9 +65,8 @@ impl EtnaModel { optimizer: &str, progress_callback: Option>, ) -> PyResult> { - let x_vec = pylist_to_vec2(x)?; - let y_vec = pylist_to_vec2(y)?; - + let x_vec = ndarray_to_vec2(x); + let y_vec = ndarray_to_vec2(y); let optimizer_type = match optimizer { "adam" => OptimizerType::Adam, _ => OptimizerType::Sgd, @@ -92,8 +94,8 @@ impl EtnaModel { Ok(history) } - fn predict(&mut self, x: &Bound<'_, PyList>) -> PyResult> { - let x_vec = pylist_to_vec2(x)?; + fn predict(&mut self, x: PyReadonlyArray2<'_, f32>) -> PyResult> { + let x_vec = ndarray_to_vec2(x); Ok(self.inner.predict(&x_vec)) } diff --git a/tests/test_tqdm_progress.py b/tests/test_tqdm_progress.py index 2274a2e..d36e719 100644 --- a/tests/test_tqdm_progress.py +++ b/tests/test_tqdm_progress.py @@ -38,6 +38,8 @@ def mock_train(X, y, epochs, lr, batch_size, weight_decay, optimizer, progress_c import importlib import etna.api importlib.reload(etna.api) + # Ensure the reloaded module uses our mock, not the real compiled extension + etna.api._etna_rust = mock_etna_rust # Patch load_data and Preprocessor with patch.object(etna.api, 'load_data') as mock_load: From 611070a4af1e47913b084c5f3e6a53f1a01467f4 Mon Sep 17 00:00:00 2001 From: Yash verdhan parihar Date: Fri, 13 Feb 2026 09:39:27 +0530 Subject: [PATCH 2/2] Update etna_core/src/lib.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- etna_core/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/etna_core/src/lib.rs b/etna_core/src/lib.rs index d3d41d6..364cda2 100644 --- a/etna_core/src/lib.rs +++ b/etna_core/src/lib.rs @@ -19,7 +19,7 @@ fn ndarray_to_vec2(arr: PyReadonlyArray2<'_, f32>) -> Vec> { array .rows() .into_iter() - .map(|row: numpy::ndarray::ArrayView1<'_, f32>| row.to_vec()) + .map(|row| row.to_vec()) .collect() }