[python] Fix output type when exporting PointCloudDataFrame to Spat…

…ialData points (#3408) * Update export `PointCloudDataFrame` validation tests * Fix point cloud export type error
single-cell-data · Dec 9, 2024 · 504df5f · 504df5f
1 parent c4fe227
commit 504df5f
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/apis/python/src/tiledbsoma/io/spatial/outgest.py b/apis/python/src/tiledbsoma/io/spatial/outgest.py
@@ -13,6 +13,13 @@
 except ImportError as err:
     warnings.warn("Experimental spatial outgestor requires the spatialdata package.")
     raise err
+
+try:
+    import dask.dataframe as dd
+except ImportError as err:
+    warnings.warn("Experimental spatial outgestor requires the dask package.")
+    raise err
+
 try:
     import geopandas as gpd
 except ImportError as err:
@@ -101,7 +108,7 @@ def to_spatial_data_points(
     scene_dim_map: Dict[str, str],
     transform: Optional[somacore.CoordinateTransform],
     soma_joinid_name: str,
-) -> pd.DataFrame:
+) -> dd.DataFrame:
     """Export a :class:`PointCloudDataFrame` to a :class:`spatialdata.ShapesModel.
 
     Args:
@@ -134,8 +141,7 @@ def to_spatial_data_points(
     # Read the pandas dataframe, rename SOMA_JOINID, add metadata, and return.
     df: pd.DataFrame = points.read().concat().to_pandas()
     df.rename(columns={SOMA_JOINID: soma_joinid_name}, inplace=True)
-    df.attrs["transform"] = transforms
-    return df
+    return sd.models.PointsModel.parse(df, transformations=transforms)
 
 
 def to_spatial_data_shapes(

diff --git a/apis/python/tests/test_export_point_cloud_dataframe.py b/apis/python/tests/test_export_point_cloud_dataframe.py
@@ -11,7 +11,7 @@
 
 gpd = pytest.importorskip("geopandas")
 soma_outgest = pytest.importorskip("tiledbsoma.io.spatial.outgest")
-spatialdata = pytest.importorskip("spatialdata")
+sd = pytest.importorskip("spatialdata")
 
 
 @pytest.fixture(scope="module")
@@ -38,12 +38,12 @@ def sample_point_cloud_dataframe_2d(tmp_path_factory):
 @pytest.mark.parametrize(
     "transform,expected_transformation",
     [
-        (None, {"point_cloud": spatialdata.transformations.Identity()}),
+        (None, {"point_cloud": sd.transformations.Identity()}),
         (
             somacore.IdentityTransform(
                 ("x_scene", "y_scene"), ("x_points", "y_points")
             ),
-            {"scene0": spatialdata.transformations.Identity()},
+            {"scene0": sd.transformations.Identity()},
         ),
     ],
 )
@@ -62,7 +62,8 @@ def test_export_to_shapes_2d(
     )
 
     # Check this is valid storage for the SpatialData "Shapes" model.
-    spatialdata.models.ShapesModel.validate(shape)
+    model_schema = sd.models.get_model(shape)
+    assert model_schema == sd.models.ShapesModel
 
     # Check the dataframe.
     expected = gpd.GeoDataFrame.from_dict(
@@ -85,12 +86,12 @@ def test_export_to_shapes_2d(
 @pytest.mark.parametrize(
     "transform,expected_transformation",
     [
-        (None, {"point_cloud": spatialdata.transformations.Identity()}),
+        (None, {"point_cloud": sd.transformations.Identity()}),
         (
             somacore.IdentityTransform(
                 ("x_scene", "y_scene"), ("x_points", "y_points")
             ),
-            {"scene0": spatialdata.transformations.Identity()},
+            {"scene0": sd.transformations.Identity()},
         ),
     ],
 )
@@ -109,7 +110,8 @@ def test_export_to_points_2d(
     )
 
     # Check this is valid storage for the SpatialData "Points" model.
-    spatialdata.models.PointsModel.validate(points)
+    model_schema = sd.models.get_model(points)
+    assert model_schema == sd.models.PointsModel
 
     # Check the dataframe.
     expected = pd.DataFrame.from_dict(