Merge pull request #172 from PixelgenTechnologies/feature/exe-1873-up…

…date-polars-to-1.x-series

Feature/exe 1873 update polars to 1.x series

pyproject.toml

@@ -48,8 +48,7 @@ pyarrow = ">=14,<16"
 semver = "^3.0.0"
 ruamel-yaml = "^0.17.21"
 pydantic = ">=2.0,<2.5"
-# Pin to 0.20.19 to fix a bug in lazy collect projection_pushdown optimzation
-polars = "0.20.19"
+polars = ">=1.0.0,<2.0.0"
 importlib-resources = "^5.12.0"
 fsspec = "^2023.12.2"
 fastparquet = "^2023.8.0"

src/pixelator/analysis/analysis_engine.py

@@ -148,7 +148,7 @@ def edgelist_to_component_stream(
         .items()
     ):
         yield (
-            component_id,
+            str(component_id[0]),  # component id is a tuple here, hence the [0]
             Graph.from_edgelist(
                 edgelist=component_df.lazy(),
                 add_marker_counts=True,

src/pixelator/graph/community_detection.py

@@ -275,7 +275,7 @@ def recovered_component_info(
         merged.select(pl.col("component"), pl.col("component_new"))  # type: ignore
         .group_by(pl.col("component"))
         .agg(pl.col("component_new").unique().drop_nulls())
-        .collect(streaming=True, projection_pushdown=False)
+        .collect(streaming=True)
         .iter_rows()
     )
 

src/pixelator/graph/utils.py

@@ -439,7 +439,7 @@ def _update_edgelist_membership_lazy_frame(
     }
 
     def _map_edge_index_to_component_id():
-        return pl.col("edge_index").map_dict(edge_index_to_component_mapping)
+        return pl.col("edge_index").replace(edge_index_to_component_mapping)
 
     def _build_component_name_str():
         return pl.format(
@@ -457,7 +457,7 @@ def _build_component_name_str():
         .with_columns(_build_component_name_str().alias("component"))
     )
     edgelist_with_component_info = edgelist_with_component_info.drop(
-        columns=["edge_index", "component_index"]
+        ["edge_index", "component_index"]
     )
     return edgelist_with_component_info
 

src/pixelator/pixeldataset/__init__.py

@@ -166,7 +166,7 @@ def edgelist(self, value: pd.DataFrame) -> None:
     def edgelist_lazy(self) -> pl.LazyFrame:
         """Get the edge list as a lazy dataframe."""
         lz_edgelist = self._backend.edgelist_lazy
-        if "index" in lz_edgelist.columns:
+        if "index" in set(lz_edgelist.collect_schema().names()):
             warnings.warn(
                 "A column called `index` was identified in your edgelist. "
                 "This will be removed."
@@ -243,7 +243,7 @@ def graph(
             potential_component = self.edgelist_lazy.filter(
                 pl.col("component") == component_id
             )
-            if potential_component.fetch(1).is_empty():
+            if potential_component.head(1).collect().is_empty():
                 raise KeyError(f"{component_id} not found in edgelist")
             return Graph.from_edgelist(
                 potential_component,

src/pixelator/pixeldataset/precomputed_layouts.py

@@ -418,7 +418,7 @@ def data():
                         .sort_index(axis=1)
                     )
                     .with_columns(
-                        component=pl.lit(component),
+                        component=pl.lit(component[0]),
                         graph_projection=pl.lit("bipartite"),
                         layout=pl.lit(layout_algorithm),
                     )