feat: add coverage

polars_bio/__init__.py

@@ -14,7 +14,7 @@
     sql,
 )
 from .polars_ext import PolarsRangesOperations as LazyFrame
-from .range_op import FilterOp, count_overlaps, coverage, merge, nearest, overlap
+from .range_op import FilterOp, cluster, count_overlaps, coverage, merge, nearest, overlap
 from .range_viz import visualize_intervals
 
 POLARS_BIO_MAX_THREADS = "datafusion.execution.target_partitions"
@@ -24,7 +24,9 @@
 __all__ = [
     "overlap",
     "nearest",
+    "coverage",
     "merge",
+    "cluster",
     "count_overlaps",
     "coverage",
     "ctx",

polars_bio/interval_op_helpers.py

@@ -16,14 +16,17 @@ def get_py_ctx() -> datafusion.context.SessionContext:
 
 def read_df_to_datafusion(
     py_ctx: datafusion.context.SessionContext,
-    df: Union[str, pl.DataFrame, pl.LazyFrame, pd.DataFrame],
-) -> datafusion.dataframe:
+    df: Union[str, pl.DataFrame, pl.LazyFrame, pd.DataFrame, datafusion.dataframe.DataFrame]
+) -> datafusion.dataframe.DataFrame:
+
     if isinstance(df, pl.DataFrame):
         return py_ctx.from_polars(df)
     elif isinstance(df, pd.DataFrame):
         return py_ctx.from_pandas(df)
     elif isinstance(df, pl.LazyFrame):
         return py_ctx.from_polars(df.collect())
+    elif isinstance(df, datafusion.dataframe.DataFrame):
+        return df
     elif isinstance(df, str):
         ext = Path(df).suffix
         if ext == ".csv":
@@ -46,8 +49,9 @@ def read_df_to_datafusion(
             return py_ctx.read_parquet(df)
     raise ValueError("Invalid `df` argument.")
 
-
-def df_to_lazyframe(df: datafusion.DataFrame) -> pl.LazyFrame:
+def df_to_lazyframe(
+    df: datafusion.dataframe.DataFrame
+) -> pl.LazyFrame:
     # TODO: make it actually lazy
     """
     def _get_lazy(
@@ -63,8 +67,10 @@ def _get_lazy(
 
 
 def convert_result(
-    df: datafusion.DataFrame, output_type: str, streaming: bool
-) -> Union[pl.LazyFrame, pl.DataFrame, pd.DataFrame]:
+    df: datafusion.dataframe.DataFrame,
+    output_type: str,
+    streaming: bool
+) -> Union[pl.LazyFrame, pl.DataFrame, pd.DataFrame, datafusion.dataframe.DataFrame]:
     # TODO: implement streaming
     if streaming:
         # raise NotImplementedError("streaming is not implemented")
@@ -75,4 +81,6 @@ def convert_result(
         return df.to_pandas()
     elif output_type == "polars.LazyFrame":
         return df_to_lazyframe(df)
+    elif output_type == "datafusion.DataFrame":
+        return df
     raise ValueError("Invalid `output_type` argument")

polars_bio/polars_ext.py

@@ -54,6 +54,59 @@ def nearest(
             cols1=cols1,
             cols2=cols2,
         )
+
+    def coverage(
+        self,
+        other_df: pl.LazyFrame,
+        suffixes: tuple[str, str] = ("", "_"),
+        cols1=["chrom", "start", "end"],
+        cols2=["chrom", "start", "end"],
+    ) -> pl.LazyFrame:
+        """
+        !!! note
+            Alias for [coverage](api.md#polars_bio.coverage)
+        """
+        return pb.coverage(
+            self._ldf,
+            other_df,
+            suffixes=suffixes,
+            cols1=cols1,
+            cols2=cols2,
+        )
+
+    def merge(
+        self,
+        overlap_filter: FilterOp = FilterOp.Strict,
+        min_dist: float = 0,
+        cols: Union[list[str], None] = None,
+    ) -> pl.LazyFrame:
+        """
+        !!! note
+            Alias for [merge](api.md#polars_bio.merge)
+        """
+        return pb.merge(
+            self._ldf,
+            overlap_filter=overlap_filter,
+            min_dist=min_dist,
+            cols=cols
+        )
+
+    def cluster(
+        self,
+        overlap_filter: FilterOp = FilterOp.Strict,
+        min_dist: float = 0,
+        cols: Union[list[str], None] = None,
+    ) -> pl.LazyFrame:
+        """
+        !!! note
+            Alias for [cluster](api.md#polars_bio.cluster)
+        """
+        return pb.cluster(
+            self._ldf,
+            overlap_filter=overlap_filter,
+            min_dist=min_dist,
+            cols=cols
+        )
 
     def count_overlaps(
         self,