Allow any index type when creating sparse matrix (#18)

Author: flying-sheep

.github/workflows/ci.yml

@@ -18,10 +18,7 @@ jobs:
       uses: Swatinem/rust-cache@v2
 
     - name: Test Rust package
-      run: |
-        cd ${GITHUB_WORKSPACE}/anndata-hdf5 && cargo test --no-fail-fast
-        cd ${GITHUB_WORKSPACE}/anndata && cargo test --no-fail-fast
-        cd ${GITHUB_WORKSPACE}/anndata-test-utils && cargo test --no-fail-fast
+      run: cargo test --no-fail-fast
 
     #- name: benchmark
     #  run: |

pyanndata/Cargo.toml

@@ -28,7 +28,10 @@ rayon = "1.11"
 
 [dependencies.pyo3]
 version = "0.25"
-features = ["extension-module", "multiple-pymethods", "anyhow"]
+features = ["multiple-pymethods", "anyhow"]
+
+[features]
+extension-module = ["pyo3/extension-module"]
 
 [lib]
 crate-type = ["lib"]

pyanndata/src/data/array.rs

@@ -86,37 +86,19 @@ pub(super) fn to_array(ob: &Bound<'_, PyAny>) -> PyResult<DynArray> {
     Ok(arr)
 }
 
-pub(super) fn to_csr(ob: &Bound<'_, PyAny>) -> PyResult<DynCsrMatrix> {
-    fn extract_csr_indicies(indicies: Bound<'_, PyAny>) -> PyResult<Vec<usize>> {
-        let res = match indicies
-            .getattr("dtype")?
-            .getattr("name")?
-            .extract::<&str>()?
-        {
-            "int32" => indicies
-                .extract::<PyReadonlyArrayDyn<i32>>()?
-                .as_array()
-                .iter()
-                .map(|x| (*x).try_into().unwrap())
-                .collect(),
-            "int64" => indicies
-                .extract::<PyReadonlyArrayDyn<i64>>()?
-                .as_array()
-                .iter()
-                .map(|x| (*x).try_into().unwrap())
-                .collect(),
-            other => panic!("CSR indicies type '{}' is not supported", other),
-        };
-        Ok(res)
-    }
+fn extract_array_as_usize(arr: Bound<'_, PyAny>) -> PyResult<Vec<usize>> {
+    arr.call_method1("astype", ("uintp",))?
+        .extract::<Vec<usize>>()
+}
 
+pub(super) fn to_csr(ob: &Bound<'_, PyAny>) -> PyResult<DynCsrMatrix> {
     if !isinstance_of_csr(ob)? {
         return Err(PyTypeError::new_err("not a csr matrix"));
     }
 
     let shape: Vec<usize> = ob.getattr("shape")?.extract()?;
-    let indices = extract_csr_indicies(ob.getattr("indices")?)?;
-    let indptr = extract_csr_indicies(ob.getattr("indptr")?)?;
+    let indices = extract_array_as_usize(ob.getattr("indices")?)?;
+    let indptr = extract_array_as_usize(ob.getattr("indptr")?)?;
     let ty = ob.getattr("data")?.getattr("dtype")?.getattr("name")?;
     let ty = ty.extract::<&str>()?;
 
@@ -139,36 +121,13 @@ pub(super) fn to_csr(ob: &Bound<'_, PyAny>) -> PyResult<DynCsrMatrix> {
 }
 
 pub(super) fn to_csr_noncanonical(ob: &Bound<'_, PyAny>) -> PyResult<DynCsrNonCanonical> {
-    fn extract_csr_indicies(indicies: Bound<'_, PyAny>) -> PyResult<Vec<usize>> {
-        let res = match indicies
-            .getattr("dtype")?
-            .getattr("name")?
-            .extract::<&str>()?
-        {
-            "int32" => indicies
-                .extract::<PyReadonlyArrayDyn<i32>>()?
-                .as_array()
-                .iter()
-                .map(|x| (*x).try_into().unwrap())
-                .collect(),
-            "int64" => indicies
-                .extract::<PyReadonlyArrayDyn<i64>>()?
-                .as_array()
-                .iter()
-                .map(|x| (*x).try_into().unwrap())
-                .collect(),
-            other => panic!("CSR indicies type '{}' is not supported", other),
-        };
-        Ok(res)
-    }
-
     if !isinstance_of_csr(ob)? {
         return Err(PyTypeError::new_err("not a csr matrix"));
     }
 
     let shape: Vec<usize> = ob.getattr("shape")?.extract()?;
-    let indices = extract_csr_indicies(ob.getattr("indices")?)?;
-    let indptr = extract_csr_indicies(ob.getattr("indptr")?)?;
+    let indices = extract_array_as_usize(ob.getattr("indices")?)?;
+    let indptr = extract_array_as_usize(ob.getattr("indptr")?)?;
     let ty = ob.getattr("data")?.getattr("dtype")?.getattr("name")?;
     let ty = ty.extract::<&str>()?;
 
@@ -190,36 +149,13 @@ pub(super) fn to_csr_noncanonical(ob: &Bound<'_, PyAny>) -> PyResult<DynCsrNonCa
 }
 
 pub(super) fn to_csc(ob: &Bound<'_, PyAny>) -> PyResult<DynCscMatrix> {
-    fn extract_csc_indicies(indicies: Bound<'_, PyAny>) -> PyResult<Vec<usize>> {
-        let res = match indicies
-            .getattr("dtype")?
-            .getattr("name")?
-            .extract::<&str>()?
-        {
-            "int32" => indicies
-                .extract::<PyReadonlyArrayDyn<i32>>()?
-                .as_array()
-                .iter()
-                .map(|x| (*x).try_into().unwrap())
-                .collect(),
-            "int64" => indicies
-                .extract::<PyReadonlyArrayDyn<i64>>()?
-                .as_array()
-                .iter()
-                .map(|x| (*x).try_into().unwrap())
-                .collect(),
-            other => panic!("CSC indicies type '{}' is not supported", other),
-        };
-        Ok(res)
-    }
-
     if !isinstance_of_csc(ob)? {
         return Err(PyTypeError::new_err("not a csc matrix"));
     }
 
     let shape: Vec<usize> = ob.getattr("shape")?.extract()?;
-    let indices = extract_csc_indicies(ob.getattr("indices")?)?;
-    let indptr = extract_csc_indicies(ob.getattr("indptr")?)?;
+    let indices = extract_array_as_usize(ob.getattr("indices")?)?;
+    let indptr = extract_array_as_usize(ob.getattr("indptr")?)?;
     let ty = ob.getattr("data")?.getattr("dtype")?.getattr("name")?;
     let ty = ty.extract::<&str>()?;
 

python/Cargo.toml

@@ -15,7 +15,10 @@ pyo3-log = "0.12"
 
 [dependencies.pyo3]
 version = "0.25"
-features = ["extension-module", "multiple-pymethods"]
+features = ["multiple-pymethods"]
+
+[features]
+extension-module = ["pyo3/extension-module", "pyanndata/extension-module"]
 
 [lib]
 name = "anndata_rs"

python/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["maturin>=1.4,<2.0"]
 build-backend = "maturin"
 
 [tool.maturin]
-features = ["pyo3/extension-module"]
+features = ["extension-module"]
 module-name = "anndata_rs"
 
 [project]