qdp/qdp-python/src/loader.rs

//
// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Loader bindings (Linux only; qdp-core pipeline types only built on Linux)
#[cfg(target_os = "linux")]
mod loader_impl {
    use crate::tensor::QuantumTensor;
    use pyo3::exceptions::PyRuntimeError;
    use pyo3::prelude::*;
    use qdp_core::reader::NullHandling;
    use qdp_core::{Dtype, Encoding, PipelineConfig, PipelineIterator, QdpEngine as CoreEngine};

    /// Rust-backed iterator yielding one QuantumTensor per batch; used by QuantumDataLoader.
    #[pyclass]
    pub struct PyQuantumLoader {
        pub inner: Option<PipelineIterator>,
    }

    impl PyQuantumLoader {
        pub fn new(inner: Option<PipelineIterator>) -> Self {
            Self { inner }
        }
    }

    #[pymethods]
    impl PyQuantumLoader {
        fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
            slf
        }

        fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult<QuantumTensor> {
            let mut iter: PipelineIterator = match slf.inner.take() {
                Some(i) => i,
                None => return Err(pyo3::exceptions::PyStopIteration::new_err("")),
            };
            // Call next_batch without releasing GIL (return type *mut DLManagedTensor is !Send).
            let result = iter.next_batch();
            match result {
                Ok(Some(ptr)) => {
                    slf.inner = Some(iter);
                    Ok(QuantumTensor {
                        ptr,
                        consumed: false,
                    })
                }
                Ok(None) => {
                    // Exhausted; do not put iterator back
                    Err(pyo3::exceptions::PyStopIteration::new_err(""))
                }
                Err(e) => {
                    slf.inner = Some(iter);
                    Err(PyRuntimeError::new_err(format!(
                        "Pipeline next_batch failed: {}",
                        e
                    )))
                }
            }
        }
    }

    /// Parse a Python null-handling string into the Rust enum.
    pub fn parse_null_handling(s: Option<&str>) -> PyResult<NullHandling> {
        match s {
            None | Some("fill_zero") => Ok(NullHandling::FillZero),
            Some("reject") => Ok(NullHandling::Reject),
            Some(other) => Err(pyo3::exceptions::PyValueError::new_err(format!(
                "Invalid null_handling policy '{}'. Expected 'fill_zero' or 'reject'.",
                other
            ))),
        }
    }

    /// Parse an optional Python dtype string into the Rust enum. Defaults to f64 so
    /// file loads are lossless unless the caller explicitly opts into f32 narrowing.
    pub fn parse_dtype(s: Option<&str>) -> PyResult<Dtype> {
        match s {
            None => Ok(Dtype::Float64),
            Some(v) => Dtype::from_str_ci(v).map_err(|e| {
                pyo3::exceptions::PyValueError::new_err(format!(
                    "Invalid dtype '{v}': {e}. Expected 'float32'/'f32' or 'float64'/'f64'."
                ))
            }),
        }
    }

    /// Build PipelineConfig from Python args. device_id is 0 (engine does not expose it); iterator uses engine clone with correct device.
    #[allow(clippy::too_many_arguments)]
    pub fn config_from_args(
        _engine: &CoreEngine,
        batch_size: usize,
        num_qubits: u32,
        encoding_method: &str,
        total_batches: usize,
        seed: Option<u64>,
        null_handling: NullHandling,
        dtype: Dtype,
    ) -> PyResult<PipelineConfig> {
        let encoding = Encoding::from_str_ci(encoding_method)
            .map_err(|e| PyRuntimeError::new_err(format!("Invalid encoding: {e}")))?;
        Ok(PipelineConfig {
            device_id: 0,
            num_qubits,
            batch_size,
            total_batches,
            encoding,
            seed,
            warmup_batches: 0,
            null_handling,
            dtype,
            prefetch_depth: 16,
        })
    }

    /// Resolve path from Python str or pathlib.Path (__fspath__).
    pub fn path_from_py(path: &Bound<'_, PyAny>) -> PyResult<String> {
        path.extract::<String>().or_else(|_| {
            path.call_method0("__fspath__")
                .and_then(|m| m.extract::<String>())
        })
    }
}

#[cfg(target_os = "linux")]
pub use loader_impl::{
    PyQuantumLoader, config_from_args, parse_dtype, parse_null_handling, path_from_py,
};