From bdf400eb6e3f99056e3fefe87d1434ec61ea1af0 Mon Sep 17 00:00:00 2001 From: Wojciech Sitek Date: Thu, 12 Feb 2026 21:28:07 +0100 Subject: [PATCH 1/2] feat: Enhance on_cols support in range operations - Updated range operation functions to accept and validate on_cols parameter, ensuring it is a list of strings or None. - Modified query generation to include additional WHERE conditions based on on_cols. - Updated related structures and functions to accommodate on_cols, improving flexibility in overlap queries. Bump polars_bio version to 0.22.0 in Cargo.lock. --- polars_bio/range_op.py | 16 +-- polars_bio/range_op_helpers.py | 11 +- src/operation.rs | 16 +++ src/option.rs | 2 +- src/query.rs | 0 tests/test_on_cols.py | 228 +++++++++++++++++++++++++++++++++ 6 files changed, 259 insertions(+), 14 deletions(-) delete mode 100644 src/query.rs create mode 100644 tests/test_on_cols.py diff --git a/polars_bio/range_op.py b/polars_bio/range_op.py index c15f9d54..3b2f23d4 100644 --- a/polars_bio/range_op.py +++ b/polars_bio/range_op.py @@ -188,9 +188,6 @@ def overlap( 1 chr1 3 8 chr1 4 8 ``` - - Todo: - Support for on_cols. """ _validate_overlap_input(cols1, cols2, on_cols, suffixes, output_type) @@ -206,6 +203,7 @@ def overlap( suffixes=suffixes, columns_1=cols1, columns_2=cols2, + on_cols=on_cols, overlap_alg=algorithm, overlap_low_memory=low_memory, ) @@ -273,9 +271,6 @@ def nearest( This enables efficient processing of large datasets without loading the entire output dataset into memory. Example: - - Todo: - Support for on_cols. """ _validate_overlap_input(cols1, cols2, on_cols, suffixes, output_type) @@ -294,6 +289,7 @@ def nearest( nearest_k=k, include_overlaps=overlap, compute_distance=distance, + on_cols=on_cols, ) return range_operation( df1, @@ -351,9 +347,6 @@ def coverage( This enables efficient processing of large datasets without loading the entire output dataset into memory. Example: - - Todo: - Support for on_cols. """ _validate_overlap_input(cols1, cols2, on_cols, suffixes, output_type) @@ -369,6 +362,7 @@ def coverage( suffixes=suffixes, columns_1=cols1, columns_2=cols2, + on_cols=on_cols, ) return range_operation( df2, @@ -473,6 +467,7 @@ def count_overlaps( suffixes=suffixes, columns_1=cols1, columns_2=cols2, + on_cols=on_cols, ) return range_operation(df2, df1, range_options, output_type, ctx) df1 = read_df_to_datafusion(my_ctx, df1) @@ -594,9 +589,6 @@ def merge( and `datafusion.bio.coordinate_system_check` is "true" (default). Example: - - Todo: - Support for on_cols. """ suffixes = ("_1", "_2") _validate_overlap_input(cols, cols, on_cols, suffixes, output_type) diff --git a/polars_bio/range_op_helpers.py b/polars_bio/range_op_helpers.py index b54739e3..8891fc0e 100644 --- a/polars_bio/range_op_helpers.py +++ b/polars_bio/range_op_helpers.py @@ -391,7 +391,16 @@ def _validate_overlap_input( Note: Coordinate system is now determined from DataFrame metadata, not from an explicit parameter. """ - assert on_cols is None, "on_cols is not supported yet" + if on_cols is not None: + if not isinstance(on_cols, list): + raise TypeError( + f"on_cols must be a list of column names or None, got {type(on_cols)}" + ) + if not all(isinstance(col, str) for col in on_cols): + raise TypeError("All elements in on_cols must be strings") + if len(on_cols) == 0: + raise ValueError("on_cols cannot be an empty list, use None instead") + assert output_type in [ "polars.LazyFrame", "polars.DataFrame", diff --git a/src/operation.rs b/src/operation.rs index 44ca1cf2..085e3210 100644 --- a/src/operation.rs +++ b/src/operation.rs @@ -287,6 +287,22 @@ async fn do_count_overlaps_coverage_naive( right_table: String, coverage: bool, ) -> datafusion::dataframe::DataFrame { + // Validate that on_cols is not provided for naive operations + if let Some(ref cols) = range_opts.on_cols { + if !cols.is_empty() { + let op_name = if coverage { "coverage" } else { "count_overlaps" }; + let suggestion = if coverage { + "The coverage operation does not support on_cols filtering. \ + Consider filtering the input DataFrames before calling coverage." + } else { + "Set naive_query=False to use the SQL-based path which supports on_cols." + }; + panic!( + "on_cols parameter {:?} is not supported in naive {} operation. {}", + cols, op_name, suggestion + ); + } + } let columns_1 = range_opts.columns_1.unwrap(); let columns_2 = range_opts.columns_2.unwrap(); let session = ctx.clone(); diff --git a/src/option.rs b/src/option.rs index e1b2dc77..ca73f76d 100644 --- a/src/option.rs +++ b/src/option.rs @@ -17,7 +17,7 @@ pub struct RangeOptions { #[pyo3(get, set)] pub columns_2: Option>, #[pyo3(get, set)] - on_cols: Option>, + pub on_cols: Option>, #[pyo3(get, set)] pub overlap_alg: Option, #[pyo3(get, set)] diff --git a/src/query.rs b/src/query.rs deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/test_on_cols.py b/tests/test_on_cols.py new file mode 100644 index 00000000..c6f6c8a5 --- /dev/null +++ b/tests/test_on_cols.py @@ -0,0 +1,228 @@ +"""Tests for on_cols parameter in genomic range operations.""" +import pandas as pd +import polars as pl +import pytest + +import polars_bio as pb + +# Set up test context +pb.ctx.set_option("datafusion.execution.parquet.schema_force_view_types", "true", False) + + +@pytest.fixture +def df1_with_strand(): + """Create a test DataFrame with strand information.""" + data = { + "chrom": ["chr1", "chr1", "chr1", "chr2"], + "start": [10, 20, 30, 10], + "end": [20, 30, 40, 20], + "strand": ["+", "+", "-", "+"], + "name": ["a", "b", "c", "d"], + } + df = pd.DataFrame(data) + df.attrs["coordinate_system_zero_based"] = True + return df + + +@pytest.fixture +def df2_with_strand(): + """Create a second test DataFrame with strand information.""" + data = { + "chrom": ["chr1", "chr1", "chr1", "chr2"], + "start": [15, 25, 25, 15], + "end": [25, 35, 35, 25], + "strand": ["+", "-", "+", "+"], + "name": ["x", "y", "z", "w"], + } + df = pd.DataFrame(data) + df.attrs["coordinate_system_zero_based"] = True + return df + + +class TestOnColsOverlap: + """Test on_cols parameter with overlap operation.""" + + def test_overlap_without_on_cols(self, df1_with_strand, df2_with_strand): + """Test overlap without on_cols returns all overlaps.""" + result = pb.overlap( + df1_with_strand, + df2_with_strand, + output_type="pandas.DataFrame", + ) + # Should return overlaps regardless of strand + # chr1: (10-20) overlaps (15-25), (20-30) overlaps (15-25) and (25-35), (30-40) overlaps (25-35) + # chr2: (10-20) overlaps (15-25) + assert len(result) >= 4, f"Expected at least 4 overlaps, got {len(result)}" + + def test_overlap_with_strand_on_cols(self, df1_with_strand, df2_with_strand): + """Test overlap with on_cols=['strand'] filters by strand.""" + result = pb.overlap( + df1_with_strand, + df2_with_strand, + on_cols=["strand"], + output_type="pandas.DataFrame", + ) + # Should only return overlaps where strand matches + # chr1 + strand: (10-20) with (15-25), (20-30) with (25-35) + # chr1 - strand: none (c at 30-40 with - strand, but x,z,w are +) + # chr2 + strand: (10-20) with (15-25) + assert len(result) >= 2, f"Expected at least 2 strand-matched overlaps, got {len(result)}" + + # Verify all results have matching strands + if len(result) > 0: + assert (result["strand_1"] == result["strand_2"]).all(), \ + "All overlaps should have matching strands" + + def test_overlap_with_on_cols_polars_lazyframe(self, df1_with_strand, df2_with_strand): + """Test overlap with on_cols returns LazyFrame correctly.""" + result = pb.overlap( + df1_with_strand, + df2_with_strand, + on_cols=["strand"], + output_type="polars.LazyFrame", + ) + assert isinstance(result, pl.LazyFrame), "Should return LazyFrame" + + result_df = result.collect().to_pandas() + if len(result_df) > 0: + assert (result_df["strand_1"] == result_df["strand_2"]).all(), \ + "All overlaps should have matching strands" + + +class TestOnColsNearest: + """Test on_cols parameter with nearest operation.""" + + def test_nearest_with_strand_on_cols(self, df1_with_strand, df2_with_strand): + """Test nearest with on_cols=['strand'] filters by strand.""" + result = pb.nearest( + df1_with_strand, + df2_with_strand, + on_cols=["strand"], + output_type="pandas.DataFrame", + ) + # Should only find nearest intervals with matching strand + assert len(result) > 0, "Should find nearest intervals" + + # Filter out null results (no match found) + non_null = result[result["strand_2"].notna()] + if len(non_null) > 0: + assert (non_null["strand_1"] == non_null["strand_2"]).all(), \ + "All nearest matches should have matching strands" + + +class TestOnColsCoverage: + """Test on_cols parameter with coverage operation.""" + + def test_coverage_with_on_cols(self, df1_with_strand, df2_with_strand): + """Test coverage with on_cols=['strand'].""" + result = pb.coverage( + df1_with_strand, + df2_with_strand, + on_cols=["strand"], + output_type="pandas.DataFrame", + ) + # Should calculate coverage only for matching strands + assert len(result) == len(df1_with_strand), \ + "Coverage should return one row per input interval" + assert "coverage" in result.columns, "Should have coverage column" + + +class TestOnColsCountOverlaps: + """Test on_cols parameter with count_overlaps operation.""" + + def test_count_overlaps_with_on_cols(self, df1_with_strand, df2_with_strand): + """Test count_overlaps with on_cols=['strand'].""" + result = pb.count_overlaps( + df1_with_strand, + df2_with_strand, + on_cols=["strand"], + output_type="pandas.DataFrame", + ) + # Should count only overlaps with matching strands + assert len(result) == len(df1_with_strand), \ + "count_overlaps should return one row per input interval" + assert "count" in result.columns, "Should have count column" + + +class TestOnColsValidation: + """Test validation of on_cols parameter.""" + + def test_on_cols_must_be_list(self, df1_with_strand, df2_with_strand): + """Test that on_cols must be a list.""" + with pytest.raises(TypeError, match="on_cols must be a list"): + pb.overlap( + df1_with_strand, + df2_with_strand, + on_cols="strand", # Should be a list, not a string + output_type="pandas.DataFrame", + ) + + def test_on_cols_elements_must_be_strings(self, df1_with_strand, df2_with_strand): + """Test that all on_cols elements must be strings.""" + with pytest.raises(TypeError, match="All elements in on_cols must be strings"): + pb.overlap( + df1_with_strand, + df2_with_strand, + on_cols=[123], # Should be strings, not integers + output_type="pandas.DataFrame", + ) + + def test_on_cols_cannot_be_empty_list(self, df1_with_strand, df2_with_strand): + """Test that on_cols cannot be an empty list.""" + with pytest.raises(ValueError, match="on_cols cannot be an empty list"): + pb.overlap( + df1_with_strand, + df2_with_strand, + on_cols=[], + output_type="pandas.DataFrame", + ) + + def test_on_cols_none_is_valid(self, df1_with_strand, df2_with_strand): + """Test that on_cols=None is valid (default behavior).""" + result = pb.overlap( + df1_with_strand, + df2_with_strand, + on_cols=None, + output_type="pandas.DataFrame", + ) + assert len(result) > 0, "Should work with on_cols=None" + + +class TestOnColsMultipleColumns: + """Test on_cols with multiple columns.""" + + def test_overlap_with_multiple_on_cols(self): + """Test overlap with multiple columns in on_cols.""" + data1 = { + "chrom": ["chr1"] * 4, + "start": [10, 20, 30, 40], + "end": [20, 30, 40, 50], + "strand": ["+", "+", "-", "+"], + "sample": ["A", "A", "B", "B"], + } + df1 = pd.DataFrame(data1) + df1.attrs["coordinate_system_zero_based"] = True + + data2 = { + "chrom": ["chr1"] * 4, + "start": [15, 25, 25, 45], + "end": [25, 35, 35, 55], + "strand": ["+", "+", "-", "+"], + "sample": ["A", "B", "B", "A"], + } + df2 = pd.DataFrame(data2) + df2.attrs["coordinate_system_zero_based"] = True + + result = pb.overlap( + df1, + df2, + on_cols=["strand", "sample"], + output_type="pandas.DataFrame", + ) + + # Should only match when both strand AND sample match + if len(result) > 0: + assert (result["strand_1"] == result["strand_2"]).all(), \ + "All overlaps should have matching strands" + assert (result["sample_1"] == result["sample_2"]).all(), \ + "All overlaps should have matching samples" From 74911714fbfb935b6170c7e2993d1be832ac3e49 Mon Sep 17 00:00:00 2001 From: Wojciech Sitek Date: Thu, 12 Feb 2026 22:42:18 +0100 Subject: [PATCH 2/2] feat: Implement on_cols validation for naive operations - Added validation to ensure that the on_cols parameter is not provided for naive operations (coverage and count_overlaps). - Updated error handling to return informative messages when on_cols is used incorrectly. - Enhanced tests to verify that appropriate errors are raised when on_cols is used with naive_query=True, while allowing it in non-naive operations. --- tests/test_on_cols.py | 71 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 58 insertions(+), 13 deletions(-) diff --git a/tests/test_on_cols.py b/tests/test_on_cols.py index c6f6c8a5..4b48074f 100644 --- a/tests/test_on_cols.py +++ b/tests/test_on_cols.py @@ -114,34 +114,79 @@ class TestOnColsCoverage: """Test on_cols parameter with coverage operation.""" def test_coverage_with_on_cols(self, df1_with_strand, df2_with_strand): - """Test coverage with on_cols=['strand'].""" - result = pb.coverage( - df1_with_strand, - df2_with_strand, - on_cols=["strand"], - output_type="pandas.DataFrame", - ) - # Should calculate coverage only for matching strands - assert len(result) == len(df1_with_strand), \ - "Coverage should return one row per input interval" - assert "coverage" in result.columns, "Should have coverage column" + """Test coverage with on_cols=['strand'] raises error.""" + # Coverage operation does not support on_cols filtering + with pytest.raises(Exception) as exc_info: + pb.coverage( + df1_with_strand, + df2_with_strand, + on_cols=["strand"], + output_type="pandas.DataFrame", + ) + # Check that error message mentions on_cols and coverage + error_msg = str(exc_info.value) + assert "on_cols" in error_msg.lower(), "Error should mention on_cols" + assert "coverage" in error_msg.lower(), "Error should mention coverage operation" class TestOnColsCountOverlaps: """Test on_cols parameter with count_overlaps operation.""" - def test_count_overlaps_with_on_cols(self, df1_with_strand, df2_with_strand): - """Test count_overlaps with on_cols=['strand'].""" + def test_count_overlaps_with_on_cols_naive(self, df1_with_strand, df2_with_strand): + """Test count_overlaps with on_cols=['strand'] raises error with naive_query=True.""" + # Naive count_overlaps does not support on_cols filtering + with pytest.raises(Exception) as exc_info: + pb.count_overlaps( + df1_with_strand, + df2_with_strand, + on_cols=["strand"], + naive_query=True, # Explicit naive mode + output_type="pandas.DataFrame", + ) + # Check that error message mentions on_cols and count_overlaps + error_msg = str(exc_info.value) + assert "on_cols" in error_msg.lower(), "Error should mention on_cols" + assert "count_overlaps" in error_msg.lower(), "Error should mention count_overlaps operation" + + def test_count_overlaps_with_on_cols_default(self, df1_with_strand, df2_with_strand): + """Test count_overlaps with on_cols=['strand'] raises error by default (naive_query=True).""" + # Default naive_query=True, so should also raise error + with pytest.raises(Exception) as exc_info: + pb.count_overlaps( + df1_with_strand, + df2_with_strand, + on_cols=["strand"], + output_type="pandas.DataFrame", + ) + # Check that error message mentions on_cols + error_msg = str(exc_info.value) + assert "on_cols" in error_msg.lower(), "Error should mention on_cols" + + def test_count_overlaps_with_on_cols_non_naive(self, df1_with_strand, df2_with_strand): + """Test count_overlaps with on_cols=['strand'] works with naive_query=False.""" + # Non-naive path should support on_cols filtering result = pb.count_overlaps( df1_with_strand, df2_with_strand, on_cols=["strand"], + naive_query=False, # Use SQL-based path which supports on_cols output_type="pandas.DataFrame", ) # Should count only overlaps with matching strands assert len(result) == len(df1_with_strand), \ "count_overlaps should return one row per input interval" assert "count" in result.columns, "Should have count column" + + # Verify the counts are different from non-filtered version + result_no_filter = pb.count_overlaps( + df1_with_strand, + df2_with_strand, + naive_query=False, + output_type="pandas.DataFrame", + ) + # At least some counts should differ when filtering by strand + assert not (result["count"] == result_no_filter["count"]).all(), \ + "Filtered counts should differ from unfiltered counts" class TestOnColsValidation: