fix tests

Author: kaizhang

anndata/src/container/base.rs

@@ -178,8 +178,11 @@ impl<B: Backend> InnerDataFrameElem<B> {
     /// Set a column with a Series.
     //TODO: this is not efficient. We should be able to replace a column without reading the whole dataframe.
     pub fn set_column<C: IntoColumn>(&mut self, name: &str, new_col: C) -> Result<()> {
+        let mut column = new_col.into_column();
+        column.rename(name.into());
+
         let mut df = self.data()?.clone();
-        df.replace(name.into(), new_col.into_column())?;
+        df.with_column(column)?;
         self.save(df)
     }
 

anndata/src/data/array/dataframe.rs

@@ -93,8 +93,10 @@ impl Writable for DataFrame {
 
 impl Readable for DataFrame {
     fn read<B: Backend>(container: &DataContainer<B>) -> Result<Self> {
+        let shape = DataFrame::get_shape(container)?;
         let columns: Vec<String> = container.get_attr("column-order")?;
-        Ok(DataFrame::new_infer_height(
+        Ok(DataFrame::new(
+            shape[0],
             columns
                 .into_iter()
                 .map(|name| {

pyanndata/src/anndata/memory.rs

@@ -253,7 +253,7 @@ impl<'py> AnnDataOp for PyAnnData<'py> {
 
     fn set_obs(&self, obs: DataFrame) -> Result<()> {
         let py = self.py();
-        let is_empty = obs.height() == 0;
+        let is_empty = obs.height() == 0 || obs.width() == 0;
         let index = self.getattr("obs")?.getattr("index")?;
         let mut df = PyDataFrame(obs).into_pyobject(py)?
             .call_method0("to_pandas")?;
@@ -270,7 +270,7 @@ impl<'py> AnnDataOp for PyAnnData<'py> {
 
     fn set_var(&self, var: DataFrame) -> Result<()> {
         let py = self.py();
-        let is_empty = var.height() == 0;
+        let is_empty = var.height() == 0 || var.width() == 0;
         let index = self.getattr("var")?.getattr("index")?;
         let mut df = PyDataFrame(var).into_pyobject(py)?
             .call_method0("to_pandas")?;

python/tests/test_base.py

@@ -130,11 +130,11 @@ def test_to_memory(tmp_path, backend):
     adata.obs_names = ['a', 'b']
     adata.var_names = ['a', 'b', 'c']
 
-    adata.to_memory()
+    adata = adata.to_memory()
     assert adata.n_obs == 2
     assert adata.n_vars == 3
-    assert adata.obs_names == ['a', 'b']
-    assert adata.var_names == ['a', 'b', 'c']
+    assert adata.obs_names.tolist() == ['a', 'b']
+    assert adata.var_names.tolist() == ['a', 'b', 'c']
 
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_resize(tmp_path, backend):