feat: Zero-Copy Data Transfer with rust-numpy

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)
 

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 {

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"

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<Vec<Vec<f32>>> {
-    pylist
-        .iter()
-        .map(|item| item.extract::<Vec<f32>>())
-        .collect::<Result<Vec<_>, _>>()
+/// 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<Vec<f32>> {
+    let array = arr.as_array();
+    array
+        .rows()
+        .into_iter()
+        .map(|row: numpy::ndarray::ArrayView1<'_, f32>| row.to_vec())
+        .collect()
 }
 
 /// Python Class Wrapper
@@ -53,18 +56,17 @@ 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,
         weight_decay: f32,
         optimizer: &str,
         progress_callback: Option<Py<PyAny>>,
     ) -> PyResult<Vec<f32>> {
-        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<Vec<f32>> {
-        let x_vec = pylist_to_vec2(x)?;
+    fn predict(&mut self, x: PyReadonlyArray2<'_, f32>) -> PyResult<Vec<f32>> {
+        let x_vec = ndarray_to_vec2(x);
         Ok(self.inner.predict(&x_vec))
     }
 

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: