fix: count_overlaps for non-parquet sources (#110)

Author: mwiewior

polars_bio/range_op_helpers.py

@@ -58,7 +58,6 @@ def range_operation(
             merged_schema = pl.Schema(
                 {**_get_schema(df1, ctx, None, read_options1), **{"count": pl.Int32}}
             )
-            # print(merged_schema)
         else:
             df_schema1 = _get_schema(df1, ctx, range_options.suffixes[0], read_options1)
             df_schema2 = _get_schema(df2, ctx, range_options.suffixes[1], read_options2)
@@ -121,9 +120,10 @@ def range_operation(
                 raise ValueError(
                     "Input and output dataframes must be of the same type: either polars or pandas"
                 )
-            return range_operation_frame_wrapper(
-                ctx, df1, df2, range_options
-            ).to_pandas()
+            df = range_operation_frame_wrapper(ctx, df1, df2, range_options)
+            print(range_options.range_op)
+            print(df.schema())
+            return df.to_pandas()
     else:
         raise ValueError(
             "Both dataframes must be of the same type: either polars or pandas or a path to a file"

polars_bio/range_wrappers.py

@@ -21,7 +21,7 @@ def range_operation_frame_wrapper(
     df2,
     range_options: RangeOptions,
     limit: Union[int, None] = None,
-) -> datafusion.dataframe:
+) -> datafusion.DataFrame:
     if range_options.range_op != RangeOp.CountOverlaps:
         return range_operation_frame(ctx, df1, df2, range_options)
     py_ctx = datafusion.SessionContext()

src/context.rs

@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::sync::Arc;
 
 use datafusion::config::ConfigOptions;
 use datafusion::prelude::SessionConfig;
@@ -9,8 +8,6 @@ use log::debug;
 use pyo3::{pyclass, pymethods, PyResult};
 use sequila_core::session_context::SequilaConfig;
 
-use crate::udtf::CountOverlapsFunction;
-
 #[pyclass(name = "BioSessionContext")]
 // #[derive(Clone)]
 pub struct PyBioSessionContext {
@@ -28,10 +25,6 @@ impl PyBioSessionContext {
     pub fn new(seed: String, catalog_dir: String) -> PyResult<Self> {
         let ctx = create_context().unwrap();
         let session_config: HashMap<String, String> = HashMap::new();
-        ctx.session.register_udtf(
-            "count_overlaps",
-            Arc::new(CountOverlapsFunction::new(ctx.session.clone())),
-        );
 
         Ok(PyBioSessionContext {
             ctx,

src/lib.rs

@@ -60,13 +60,17 @@ fn range_operation_frame(
             )
             .limit(0, Some(l))?,
         )),
-        _ => Ok(PyDataFrame::new(do_range_operation(
-            ctx,
-            &rt,
-            range_options,
-            LEFT_TABLE.to_string(),
-            RIGHT_TABLE.to_string(),
-        ))),
+        _ => {
+            let df = do_range_operation(
+                ctx,
+                &rt,
+                range_options,
+                LEFT_TABLE.to_string(),
+                RIGHT_TABLE.to_string(),
+            );
+            let py_df = PyDataFrame::new(df);
+            Ok(py_df)
+        },
     }
 }
 

src/operation.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use datafusion::catalog_common::TableReference;
 use exon::ExonSession;
 use log::{debug, info};
@@ -7,6 +9,7 @@ use tokio::runtime::Runtime;
 use crate::context::set_option_internal;
 use crate::option::{FilterOp, RangeOp, RangeOptions};
 use crate::query::{count_overlaps_query, nearest_query, overlap_query};
+use crate::udtf::CountOverlapsProvider;
 use crate::utils::default_cols_to_string;
 use crate::DEFAULT_COLUMN_NAMES;
 
@@ -150,17 +153,30 @@ async fn do_count_overlaps_naive(
 ) -> datafusion::dataframe::DataFrame {
     let columns_1 = range_opts.columns_1.unwrap();
     let columns_2 = range_opts.columns_2.unwrap();
-    let query = format!(
-        "SELECT * FROM count_overlaps('{}', '{}', '{}', '{}', '{}', '{}', '{}', '{}' , false )",
+    let session = &ctx.session;
+    let right_table_ref = TableReference::from(right_table.clone());
+    let right_schema = session
+        .table(right_table_ref.clone())
+        .await
+        .unwrap()
+        .schema()
+        .as_arrow()
+        .clone();
+    let count_overlaps_provider = CountOverlapsProvider::new(
+        Arc::new(session.clone()),
         left_table,
         right_table,
-        columns_1[0],
-        columns_1[1],
-        columns_1[2],
-        columns_2[0],
-        columns_2[1],
-        columns_2[2]
+        right_schema,
+        columns_1,
+        columns_2,
+        range_opts.filter_op.unwrap(),
+        false,
     );
+    session.deregister_table("count_overlaps").unwrap();
+    session
+        .register_table("count_overlaps", Arc::new(count_overlaps_provider))
+        .unwrap();
+    let query = "SELECT * FROM count_overlaps";
     debug!("Query: {}", query);
     ctx.sql(&query).await.unwrap()
 }