From eb906c1c8d8d34f4c732c9d90377d550d1105228 Mon Sep 17 00:00:00 2001 From: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:18:19 -0400 Subject: [PATCH] Rename multi_source to pushdown_combine The name `multi_source` described the inputs (multiple sources) rather than what the function does. Its defining feature is coordinated predicate pushdown into a user-supplied `combine` function; the multi-source aspect is incidental. Rename to `pushdown_combine` to name the behavior, consistent with polars/polars-io-tools naming. - Move io_sources/multi_source.py -> pushdown_combine.py - Move tests/io_sources/test_multi_source.py -> test_pushdown_combine.py - Update the function, __all__, imports, test class names, and docs - FilterSpec is unchanged Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com> --- README.md | 2 +- docs/wiki/API-Reference.md | 6 +- docs/wiki/Concepts.md | 2 +- docs/wiki/Query-Optimization.md | 6 +- polars_io_tools/io_sources/__init__.py | 4 +- .../{multi_source.py => pushdown_combine.py} | 22 +- polars_io_tools/io_sources/ts.py | 4 +- polars_io_tools/testing/__init__.py | 4 +- polars_io_tools/testing/predicate_tracker.py | 6 +- .../tests/io_sources/test_pickle.py | 10 +- ...lti_source.py => test_pushdown_combine.py} | 194 +++++++++--------- .../tests/pushdown/test_join_lookup.py | 2 +- polars_io_tools/tests/test_testing_utils.py | 22 +- 13 files changed, 142 insertions(+), 142 deletions(-) rename polars_io_tools/io_sources/{multi_source.py => pushdown_combine.py} (97%) rename polars_io_tools/tests/io_sources/{test_multi_source.py => test_pushdown_combine.py} (96%) diff --git a/README.md b/README.md index 9edb576..acce552 100644 --- a/README.md +++ b/README.md @@ -79,7 +79,7 @@ For a guided walkthrough, start with the Delta Lake or ClickHouse, including streaming/chunked writes and transparent handling of types the target store cannot represent natively. - **Pushdown-preserving query building** — `filtered_join`, `filtered_join_asof`, - `join_between`, `multi_source`, `concat_named`, and `ts_with_columns` express joins, + `join_between`, `pushdown_combine`, `concat_named`, and `ts_with_columns` express joins, multi-source composition, and rolling/lookback time-series logic without blocking the filter pushdown that those operations normally defeat. - **Caching** — `cache` keeps an in-memory, column- and partition-level cache for diff --git a/docs/wiki/API-Reference.md b/docs/wiki/API-Reference.md index 7871c49..94bc95e 100644 --- a/docs/wiki/API-Reference.md +++ b/docs/wiki/API-Reference.md @@ -1,7 +1,7 @@ # API reference The public surface of `polars-io-tools`. Importing the package (`import polars_io_tools`) registers the `piot` namespace and re-exports the functions below at -the top level (`from polars_io_tools import scan_db, multi_source, ...`). +the top level (`from polars_io_tools import scan_db, pushdown_combine, ...`). Most operations are available two ways: @@ -253,10 +253,10 @@ Top-level form of `lf.piot.sink_clickhouse`; see the namespace entry above. ## Composing frames -### `multi_source` +### `pushdown_combine` ```python -multi_source(sources, combine, *, combine_kwargs=None, sources_as_kwargs=False, +pushdown_combine(sources, combine, *, combine_kwargs=None, sources_as_kwargs=False, log_explain=False) -> pl.LazyFrame ``` diff --git a/docs/wiki/Concepts.md b/docs/wiki/Concepts.md index 73dc9ff..0a75b9a 100644 --- a/docs/wiki/Concepts.md +++ b/docs/wiki/Concepts.md @@ -76,7 +76,7 @@ it cannot see in advance. Three common cases motivate most of this library: In each case the trick is the same: intercept the predicate at a custom source, transform it into something that is safe to apply earlier, and re-apply the exact original filter at -the end so the result is identical to the naive version. [`multi_source`](Query-Optimization#combine-sources-with-coordinated-filter-pushdown) +the end so the result is identical to the naive version. [`pushdown_combine`](Query-Optimization#combine-sources-with-coordinated-filter-pushdown) generalises this to arbitrary compositions, with a per-source `FilterSpec` describing how each output filter maps onto each source. diff --git a/docs/wiki/Query-Optimization.md b/docs/wiki/Query-Optimization.md index b610814..ca3c73c 100644 --- a/docs/wiki/Query-Optimization.md +++ b/docs/wiki/Query-Optimization.md @@ -79,16 +79,16 @@ For overlapping intervals where each match should produce a row, use Polars' ## Combine sources with coordinated filter pushdown -`multi_source` builds a single `LazyFrame` from several named sources plus a `combine` +`pushdown_combine` builds a single `LazyFrame` from several named sources plus a `combine` function. When the result is filtered, each source receives a *transformed* version of the filter described by its `FilterSpec` — a renamed column, an expanded date range, or a remapped value — and the original filter is re-applied after `combine` runs. ```python from datetime import timedelta -from polars_io_tools import multi_source, FilterSpec +from polars_io_tools import pushdown_combine, FilterSpec -lf = multi_source( +lf = pushdown_combine( sources={ "prices": (prices_lf, {"date": FilterSpec(), "id": FilterSpec()}), "signals": (signals_lf, { diff --git a/polars_io_tools/io_sources/__init__.py b/polars_io_tools/io_sources/__init__.py index 694a94f..865b43b 100644 --- a/polars_io_tools/io_sources/__init__.py +++ b/polars_io_tools/io_sources/__init__.py @@ -18,7 +18,7 @@ from .lazy_iter_rows import * from .lazy_narwhals_reader import * from .lazy_sql_reader import * -from .multi_source import * +from .pushdown_combine import * from .pushdown_pivot import * from .pushdown_unpivot import * from .sql_dialects import * @@ -37,7 +37,7 @@ from .lazy_clickhouse_writer import sink_clickhouse # noqa: TC004 from .lazy_data_generator import scan_synthetic_panel, scan_synthetic_regression from .lazy_iter_rows import iter_rows # noqa: TC004 - from .multi_source import FilterSpec, multi_source + from .pushdown_combine import FilterSpec, pushdown_combine from .pushdown_pivot import pushdown_pivot from .pushdown_unpivot import pushdown_unpivot from .ts import ts_with_columns # noqa: TC004 diff --git a/polars_io_tools/io_sources/multi_source.py b/polars_io_tools/io_sources/pushdown_combine.py similarity index 97% rename from polars_io_tools/io_sources/multi_source.py rename to polars_io_tools/io_sources/pushdown_combine.py index 8cd869b..d591353 100644 --- a/polars_io_tools/io_sources/multi_source.py +++ b/polars_io_tools/io_sources/pushdown_combine.py @@ -1,16 +1,16 @@ """ -Multi-source LazyFrame composition with coordinated filter pushdown. +Coordinated filter pushdown across a user-defined combine of multiple sources. -This module provides the ``multi_source`` function for creating LazyFrames that +This module provides the ``pushdown_combine`` function for creating LazyFrames that combine multiple data sources while automatically propagating and transforming filters to each source appropriately. Example usage:: import polars_io_tools as cpl - from polars_io_tools import multi_source, FilterSpec + from polars_io_tools import pushdown_combine, FilterSpec - lf = multi_source( + lf = pushdown_combine( sources={ "left": (left_lf, { "date": FilterSpec(), @@ -52,7 +52,7 @@ from .set_visitor import convert_expr_to_valid_values from .util import collect_lf_in_io_source, register_io_source_with_is_pure -__all__ = ("FilterSpec", "multi_source") +__all__ = ("FilterSpec", "pushdown_combine") log = logging.getLogger(__name__) @@ -204,7 +204,7 @@ def _is_temporal_dtype(dtype: pl.DataType) -> bool: return dtype in (pl.Date, pl.Datetime) or isinstance(dtype, pl.Datetime) -def multi_source( +def pushdown_combine( sources: dict[str, tuple[pl.LazyFrame, dict[str, FilterSpec]]], combine: Callable[..., pl.LazyFrame], *, @@ -303,7 +303,7 @@ def combine(prices, rates, multiplier): Examples: Basic usage with lookback:: - lf = multi_source( + lf = pushdown_combine( sources={ "prices": (prices_lf, { "date": FilterSpec(lookback=timedelta(days=5)), @@ -332,7 +332,7 @@ def combine_with_region_mapping(sources, region_to_code): REGION_TO_CODE = {"NORTH_AMERICA": "NA", "EUROPE": "EU"} - lf = multi_source( + lf = pushdown_combine( sources={ "primary": (primary_lf, {"date": FilterSpec(), "region": FilterSpec()}), "reference": (reference_lf, {"date": FilterSpec(lookback=timedelta(days=5))}), @@ -344,7 +344,7 @@ def combine_with_region_mapping(sources, region_to_code): # Compute output schema for the IO source. # Wrap in a lambda so that the schema (and the per-source `lf.collect_schema()` # calls it requires) is only resolved when Polars actually needs it (i.e. at - # collect time), rather than eagerly when `multi_source` is constructed. + # collect time), rather than eagerly when `pushdown_combine` is constructed. output_schema = lambda: _compute_output_schema(sources, combine, combine_kwargs, sources_as_kwargs) # Collect all output columns that have FilterSpecs across all sources @@ -389,7 +389,7 @@ def source_generator( log.debug(f"Extracted values: {extracted_values}") # Get source schemas for dtype checking. Resolved lazily here (at - # collect time) rather than eagerly at multi_source construction so + # collect time) rather than eagerly at pushdown_combine construction so # we don't force schema resolution on each source until we actually # need it. source_schemas: dict[str, dict[str, pl.DataType]] = {name: lf.collect_schema() for name, (lf, _) in sources.items()} @@ -529,7 +529,7 @@ def source_generator( try: yield from collect_lf_in_io_source(result_lf, batch_size) except Exception as e: - err_msg = f"Failed during collection in multi_source.\nPolars plan:\n{result_lf.explain()}\nError: {e.__class__.__name__}: {e}" + err_msg = f"Failed during collection in pushdown_combine.\nPolars plan:\n{result_lf.explain()}\nError: {e.__class__.__name__}: {e}" raise RuntimeError(err_msg) from e return register_io_source_with_is_pure(source_generator, schema=output_schema) diff --git a/polars_io_tools/io_sources/ts.py b/polars_io_tools/io_sources/ts.py index 6fc9b59..4590eb1 100644 --- a/polars_io_tools/io_sources/ts.py +++ b/polars_io_tools/io_sources/ts.py @@ -4,7 +4,7 @@ import polars as pl -from .multi_source import FilterSpec, multi_source +from .pushdown_combine import FilterSpec, pushdown_combine __all__ = ("ts_with_columns",) @@ -141,7 +141,7 @@ def combine(sources: dict[str, pl.LazyFrame]) -> pl.LazyFrame: lf = lf.with_columns(expressions) return lf - return multi_source( + return pushdown_combine( sources={"main": (self, filter_specs)}, combine=combine, log_explain=log_explain, diff --git a/polars_io_tools/testing/__init__.py b/polars_io_tools/testing/__init__.py index d743a04..8fec18e 100644 --- a/polars_io_tools/testing/__init__.py +++ b/polars_io_tools/testing/__init__.py @@ -13,8 +13,8 @@ df = pl.DataFrame({"date": dates, "val": values}) tracker = PredicateTracker(df) - # Use the LazyFrame in your multi_source or IO source - lf = multi_source( + # Use the LazyFrame in your pushdown_combine or IO source + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], ) diff --git a/polars_io_tools/testing/predicate_tracker.py b/polars_io_tools/testing/predicate_tracker.py index 6b59368..cd3a0ef 100644 --- a/polars_io_tools/testing/predicate_tracker.py +++ b/polars_io_tools/testing/predicate_tracker.py @@ -46,13 +46,13 @@ class PredicateTracker: A utility class that creates an IO source which tracks pushed-down predicates. This is useful for testing that filters are correctly pushed down to sources - in multi_source or other IO source implementations. + in pushdown_combine or other IO source implementations. Args: df (pl.DataFrame): The DataFrame to use as the underlying data source. Attributes: - lazy_frame (pl.LazyFrame): The LazyFrame that can be used in multi_source or other operations. + lazy_frame (pl.LazyFrame): The LazyFrame that can be used in pushdown_combine or other operations. last_predicate (pl.Expr | None): The last predicate that was pushed down during collection. last_with_columns (list[str] | None): The last column projection that was pushed down. call_count (int): Number of times the source has been called. @@ -61,7 +61,7 @@ class PredicateTracker: :: tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], ) diff --git a/polars_io_tools/tests/io_sources/test_pickle.py b/polars_io_tools/tests/io_sources/test_pickle.py index 1776c0d..c206684 100644 --- a/polars_io_tools/tests/io_sources/test_pickle.py +++ b/polars_io_tools/tests/io_sources/test_pickle.py @@ -192,11 +192,11 @@ def test_concat_named_pickle_basic(self): assert_frame_equal(result, expected) -class TestMultiSourcePickle: - """Tests for multi_source pickle support.""" +class TestPushdownCombinePickle: + """Tests for pushdown_combine pickle support.""" - def test_multi_source_pickle_basic(self): - """multi_source LazyFrames can be pickled and unpickled.""" + def test_pushdown_combine_pickle_basic(self): + """pushdown_combine LazyFrames can be pickled and unpickled.""" left_df = pl.DataFrame( { "date": [date(2025, 1, 1), date(2025, 1, 2)], @@ -212,7 +212,7 @@ def test_multi_source_pickle_basic(self): } ) - lf = cpl.multi_source( + lf = cpl.pushdown_combine( sources={ "left": (left_df.lazy(), {"date": cpl.FilterSpec(), "id": cpl.FilterSpec()}), "right": (right_df.lazy(), {"date": cpl.FilterSpec(), "id": cpl.FilterSpec()}), diff --git a/polars_io_tools/tests/io_sources/test_multi_source.py b/polars_io_tools/tests/io_sources/test_pushdown_combine.py similarity index 96% rename from polars_io_tools/tests/io_sources/test_multi_source.py rename to polars_io_tools/tests/io_sources/test_pushdown_combine.py index 3ec0a8c..a539cf4 100644 --- a/polars_io_tools/tests/io_sources/test_multi_source.py +++ b/polars_io_tools/tests/io_sources/test_pushdown_combine.py @@ -1,7 +1,7 @@ """ -Tests for the multi_source function and FilterSpec class. +Tests for the pushdown_combine function and FilterSpec class. -This module tests the coordinated filter pushdown capabilities of multi_source, +This module tests the coordinated filter pushdown capabilities of pushdown_combine, including: - Basic filter propagation - Lookback/lookahead temporal expansion @@ -19,12 +19,12 @@ from polars_io_tools.io_sources.base import BinaryExprNode, FunctionNode from polars_io_tools.io_sources.enum import BooleanFunctionType, OperatorType -from polars_io_tools.io_sources.multi_source import ( +from polars_io_tools.io_sources.pushdown_combine import ( FilterSpec, _apply_value_mapping, _compute_output_schema, _get_source_col, - multi_source, + pushdown_combine, ) from polars_io_tools.testing import PredicateAnalyzer, PredicateTracker @@ -188,15 +188,15 @@ def combine(s): assert set(schema.keys()) == {"a", "b"} -class TestMultiSourceBasic: - """Test basic multi_source functionality.""" +class TestPushdownCombineBasic: + """Test basic pushdown_combine functionality.""" def test_no_filters(self): - """multi_source works without any filters applied.""" + """pushdown_combine works without any filters applied.""" left = pl.LazyFrame({"id": [1, 2], "val": [10, 20]}) right = pl.LazyFrame({"id": [1, 2], "other": [100, 200]}) - lf = multi_source( + lf = pushdown_combine( sources={ "left": (left, {}), "right": (right, {}), @@ -214,7 +214,7 @@ def test_simple_date_filter(self): left = pl.LazyFrame({"date": dates, "val": list(range(10))}) right = pl.LazyFrame({"date": dates, "other": list(range(10, 20))}) - lf = multi_source( + lf = pushdown_combine( sources={ "left": (left, {"date": FilterSpec()}), "right": (right, {"date": FilterSpec()}), @@ -243,7 +243,7 @@ def test_dt_date_filter_prunes_datetime_source(self): df = pl.DataFrame({"ts": ts, "val": list(range(90))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"ts": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -261,7 +261,7 @@ def test_simple_equality_filter(self): left = pl.LazyFrame({"id": ["A", "B", "C"], "val": [1, 2, 3]}) right = pl.LazyFrame({"id": ["A", "B", "C"], "other": [10, 20, 30]}) - lf = multi_source( + lf = pushdown_combine( sources={ "left": (left, {"id": FilterSpec()}), "right": (right, {"id": FilterSpec()}), @@ -278,7 +278,7 @@ def test_is_in_filter(self): left = pl.LazyFrame({"id": ["A", "B", "C", "D"], "val": [1, 2, 3, 4]}) right = pl.LazyFrame({"id": ["A", "B", "C", "D"], "other": [10, 20, 30, 40]}) - lf = multi_source( + lf = pushdown_combine( sources={ "left": (left, {"id": FilterSpec()}), "right": (right, {"id": FilterSpec()}), @@ -291,8 +291,8 @@ def test_is_in_filter(self): assert_frame_equal(result, expected, check_row_order=False) -class TestMultiSourceLookback: - """Test lookback functionality in multi_source.""" +class TestPushdownCombineLookback: + """Test lookback functionality in pushdown_combine.""" def test_lookback_with_date_equality(self): """Lookback works correctly with date equality filter (date == specific_date).""" @@ -306,7 +306,7 @@ def test_lookback_with_date_equality(self): def combine_with_lag(s): return s["data"].with_columns(pl.col("val").shift(3).alias("val_lag3")) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=4))})}, combine=combine_with_lag, ) @@ -330,7 +330,7 @@ def test_lookahead_with_date_equality(self): def combine_with_lead(s): return s["data"].with_columns(pl.col("val").shift(-3).alias("val_lead3")) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookahead=timedelta(days=4))})}, combine=combine_with_lead, ) @@ -357,7 +357,7 @@ def combine_with_lag_and_lead(s): pl.col("val").shift(-2).alias("val_lead2"), ) - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -384,7 +384,7 @@ def test_lookback_expands_date_range(self): right = pl.LazyFrame({"date": dates, "other": list(range(10))}) - lf = multi_source( + lf = pushdown_combine( sources={ "left": (left_tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=3))}), "right": (right, {"date": FilterSpec()}), @@ -420,7 +420,7 @@ def test_lookback_for_rolling_computation(self): def combine_with_rolling(s: dict[str, pl.LazyFrame]) -> pl.LazyFrame: return s["data"].sort("date").with_columns(pl.col("val").rolling_sum(window_size=3, min_samples=1).alias("rolling_sum")) - lf = multi_source( + lf = pushdown_combine( sources={ "data": (data, {"date": FilterSpec(lookback=timedelta(days=3))}), }, @@ -439,8 +439,8 @@ def combine_with_rolling(s: dict[str, pl.LazyFrame]) -> pl.LazyFrame: assert_frame_equal(result, expected) -class TestMultiSourceLookahead: - """Test lookahead functionality in multi_source.""" +class TestPushdownCombineLookahead: + """Test lookahead functionality in pushdown_combine.""" def test_lookahead_expands_date_range(self): """Lookahead expands the upper date range for the source.""" @@ -450,7 +450,7 @@ def test_lookahead_expands_date_range(self): df = pl.DataFrame({"date": dates, "val": values}) left_tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={ "left": (left_tracker.lazy_frame, {"date": FilterSpec(lookahead=timedelta(days=2))}), }, @@ -477,7 +477,7 @@ def test_combined_lookback_lookahead(self): df = pl.DataFrame({"date": dates, "val": values}) left_tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={ "left": ( left_tracker.lazy_frame, @@ -501,8 +501,8 @@ def test_combined_lookback_lookahead(self): assert_frame_equal(result, expected) -class TestMultiSourceValueMapping: - """Test value mapping functionality in multi_source.""" +class TestPushdownCombineValueMapping: + """Test value mapping functionality in pushdown_combine.""" def test_dict_value_mapping(self): """Dict value mapping transforms filter values.""" @@ -521,7 +521,7 @@ def test_dict_value_mapping(self): REGION_TO_CODE = {"NORTH_AMERICA": "NA", "EUROPE": "EU"} CODE_TO_REGION = {"NA": "NORTH_AMERICA", "EU": "EUROPE"} # Reverse mapping for combine - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -549,7 +549,7 @@ def test_callable_value_mapping(self): tracker = PredicateTracker(source_df) # Output uses uppercase names, source uses lowercase - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -577,7 +577,7 @@ def test_is_in_with_value_mapping(self): REGION_TO_CODE = {"NORTH_AMERICA": "NA", "EUROPE": "EU", "ASIA_PACIFIC": "APAC"} CODE_TO_REGION = {"NA": "NORTH_AMERICA", "EU": "EUROPE", "APAC": "ASIA_PACIFIC"} - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -612,7 +612,7 @@ def test_unmapped_value_not_pushed_down(self): REGION_TO_CODE = {"NORTH_AMERICA": "NA", "EUROPE": "EU"} CODE_TO_REGION = {"NA": "NORTH_AMERICA", "EU": "EUROPE", "APAC": "ASIA_PACIFIC", "LATAM": "LATIN_AMERICA"} - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -646,7 +646,7 @@ def test_mixed_mapped_and_unmapped_values_in_is_in(self): REGION_TO_CODE = {"NORTH_AMERICA": "NA", "EUROPE": "EU"} CODE_TO_REGION = {"NA": "NORTH_AMERICA", "EU": "EUROPE", "APAC": "ASIA_PACIFIC", "LATAM": "LATIN_AMERICA"} - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -666,7 +666,7 @@ def test_mixed_mapped_and_unmapped_values_in_is_in(self): assert_frame_equal(result, expected, check_row_order=False) -class TestMultiSourceColumnRemapping: +class TestPushdownCombineColumnRemapping: """Test source column name remapping.""" def test_different_source_column_name(self): @@ -679,7 +679,7 @@ def test_different_source_column_name(self): ) tracker = PredicateTracker(source_df) - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -694,8 +694,8 @@ def test_different_source_column_name(self): assert_frame_equal(result, expected) -class TestMultiSourceMultipleSources: - """Test multi_source with multiple sources having different specs.""" +class TestPushdownCombineMultipleSources: + """Test pushdown_combine with multiple sources having different specs.""" def test_different_lookback_per_source(self): """Different sources can have different lookback values.""" @@ -707,7 +707,7 @@ def test_different_lookback_per_source(self): source1_tracker = PredicateTracker(df1) source2_tracker = PredicateTracker(df2) - lf = multi_source( + lf = pushdown_combine( sources={ "source1": (source1_tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=5))}), "source2": (source2_tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=2))}), @@ -750,7 +750,7 @@ def test_mixed_filter_types_per_source(self): source1_tracker = PredicateTracker(df1) source2_tracker = PredicateTracker(df2) - lf = multi_source( + lf = pushdown_combine( sources={ "source1": ( source1_tracker.lazy_frame, @@ -788,14 +788,14 @@ def test_mixed_filter_types_per_source(self): assert_frame_equal(result, expected) -class TestMultiSourceEdgeCases: +class TestPushdownCombineEdgeCases: """Test edge cases and error handling.""" def test_empty_result(self): - """multi_source handles filters that result in empty output.""" + """pushdown_combine handles filters that result in empty output.""" left = pl.LazyFrame({"date": [date(2024, 1, 1)], "val": [1]}) - lf = multi_source( + lf = pushdown_combine( sources={"left": (left, {"date": FilterSpec()})}, combine=lambda s: s["left"], ) @@ -805,10 +805,10 @@ def test_empty_result(self): assert len(result) == 0 def test_no_filter_specs(self): - """multi_source works when no FilterSpecs are provided.""" + """pushdown_combine works when no FilterSpecs are provided.""" left = pl.LazyFrame({"id": [1, 2], "val": [10, 20]}) - lf = multi_source( + lf = pushdown_combine( sources={"left": (left, {})}, combine=lambda s: s["left"], ) @@ -820,7 +820,7 @@ def test_source_col_not_in_schema(self): """Gracefully handle when source_col doesn't exist in source schema.""" left = pl.LazyFrame({"id": [1, 2], "val": [10, 20]}) - lf = multi_source( + lf = pushdown_combine( sources={ "left": ( left, @@ -845,7 +845,7 @@ def test_multiple_filters_combined(self): } ) - lf = multi_source( + lf = pushdown_combine( sources={ "left": ( left, @@ -869,7 +869,7 @@ def test_datetime_column(self): datetimes = [datetime(2024, 1, 1, i) for i in range(24)] left = pl.LazyFrame({"ts": datetimes, "val": list(range(24))}) - lf = multi_source( + lf = pushdown_combine( sources={"left": (left, {"ts": FilterSpec(lookback=timedelta(hours=3))})}, combine=lambda s: s["left"], ) @@ -889,7 +889,7 @@ def test_column_selection_with_columns(self): } ) - lf = multi_source( + lf = pushdown_combine( sources={"left": (left, {"date": FilterSpec()})}, combine=lambda s: s["left"], ) @@ -903,7 +903,7 @@ def test_row_limit(self): dates = [date(2024, 1, i) for i in range(1, 11)] left = pl.LazyFrame({"date": dates, "val": list(range(10))}) - lf = multi_source( + lf = pushdown_combine( sources={"left": (left, {"date": FilterSpec()})}, combine=lambda s: s["left"], ) @@ -912,10 +912,10 @@ def test_row_limit(self): assert len(result) == 3 -class TestMultiSourceComplexJoinUseCase: - """Test multi_source with a complex multi-source join scenario.""" +class TestPushdownCombineComplexJoinUseCase: + """Test pushdown_combine with a complex multi-source join scenario.""" - def test_multi_source_join_with_mapping_and_lookback(self): + def test_pushdown_combine_join_with_mapping_and_lookback(self): """ Test a complex join pattern where: - Multiple sources are joined @@ -963,7 +963,7 @@ def combine(s: dict[str, pl.LazyFrame]) -> pl.LazyFrame: return df - lf = multi_source( + lf = pushdown_combine( sources={ "primary": ( primary_data, @@ -1010,7 +1010,7 @@ def test_date_gte_filter_structure(self): df = pl.DataFrame({"date": dates, "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -1038,7 +1038,7 @@ def test_date_between_filter_structure(self): df = pl.DataFrame({"date": dates, "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -1068,7 +1068,7 @@ def test_lookback_expands_date_range_correctly(self): df = pl.DataFrame({"date": dates, "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=3))})}, combine=lambda s: s["data"], ) @@ -1105,7 +1105,7 @@ def combine_with_lag(s): # Add a column that shows the value from 3 days ago return s["data"].with_columns(pl.col("val").shift(3).alias("val_lag3")) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=4))})}, combine=combine_with_lag, ) @@ -1149,7 +1149,7 @@ def test_lookback_with_is_between_and_complex_combine(self): def combine_with_lag(s): return s["data"].with_columns(pl.col("val").shift(3).alias("val_lag3")) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=4))})}, combine=combine_with_lag, ) @@ -1195,7 +1195,7 @@ def test_lookback_optimization_with_identity_combine(self): tracker = PredicateTracker(df) # Identity combine - Polars can optimize through this - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=4))})}, combine=lambda s: s["data"], ) @@ -1227,7 +1227,7 @@ def test_lookahead_expands_upper_bound(self): df = pl.DataFrame({"date": dates, "val": list(range(14))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookahead=timedelta(days=3))})}, combine=lambda s: s["data"], ) @@ -1256,7 +1256,7 @@ def test_discrete_equality_filter_structure(self): df = pl.DataFrame({"group": ["A", "B", "C"], "val": [1, 2, 3]}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"group": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -1284,7 +1284,7 @@ def test_discrete_is_in_filter_structure(self): df = pl.DataFrame({"group": ["A", "B", "C", "D"], "val": [1, 2, 3, 4]}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"group": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -1319,7 +1319,7 @@ def test_value_mapping_transforms_pushed_value(self): CODE_TO_REGION = {"NA": "NORTH_AMERICA", "EU": "EUROPE"} - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -1366,7 +1366,7 @@ def test_value_mapping_with_is_in(self): CODE_TO_REGION = {"NA": "NORTH_AMERICA", "EU": "EUROPE", "APAC": "ASIA_PACIFIC"} REGION_TO_CODE = {"NORTH_AMERICA": "NA", "EUROPE": "EU", "ASIA_PACIFIC": "APAC"} - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -1410,7 +1410,7 @@ def test_column_remapping_in_pushed_filter(self): ) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -1448,7 +1448,7 @@ def test_multiple_sources_get_different_predicates(self): source1_tracker = PredicateTracker(df1) source2_tracker = PredicateTracker(df2) - lf = multi_source( + lf = pushdown_combine( sources={ "source1": (source1_tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=5))}), "source2": (source2_tracker.lazy_frame, {"date": FilterSpec()}), # No lookback @@ -1481,7 +1481,7 @@ def test_multiple_sources_get_different_predicates(self): assert lower2 == date(2024, 1, 7), f"Source2 should have no lookback (Jan 7), got {lower2}" -class TestMultiSourceRobustness: +class TestPushdownCombineRobustness: """Additional tests for edge cases and robustness.""" def test_filter_on_column_without_spec(self): @@ -1495,7 +1495,7 @@ def test_filter_on_column_without_spec(self): ) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, # Only date has FilterSpec combine=lambda s: s["data"], ) @@ -1513,7 +1513,7 @@ def test_complex_predicate_with_or(self): df = pl.DataFrame({"date": dates, "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -1531,7 +1531,7 @@ def test_nested_filters_with_and(self): df = pl.DataFrame({"date": dates, "group": groups, "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -1561,7 +1561,7 @@ def test_filter_with_null_values(self): ) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -1579,7 +1579,7 @@ def test_multiple_collects(self): df = pl.DataFrame({"date": dates, "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -1601,7 +1601,7 @@ def test_empty_source(self): df = pl.DataFrame({"date": [], "val": []}, schema={"date": pl.Date, "val": pl.Int64}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], ) @@ -1615,7 +1615,7 @@ def test_very_large_lookback(self): df = pl.DataFrame({"date": dates, "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={ "data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=365))}) # 1 year }, @@ -1640,7 +1640,7 @@ def combine_with_computed(s): pl.lit("constant").alias("const_col"), ) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=combine_with_computed, ) @@ -1663,7 +1663,7 @@ def test_combine_drops_columns(self): ) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"].drop("drop_col"), ) @@ -1687,7 +1687,7 @@ def test_combine_kwargs_basic(self): def combine_with_multiplier(sources, multiplier): return sources["data"].with_columns((pl.col("val") * multiplier).alias("scaled_val")) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=combine_with_multiplier, combine_kwargs={"multiplier": 10}, @@ -1716,7 +1716,7 @@ def test_combine_kwargs_with_mapping_dict(self): def combine_with_mapping(sources, code_to_region): return sources["data"].with_columns(pl.col("code").replace(code_to_region).alias("region")) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=combine_with_mapping, combine_kwargs={"code_to_region": CODE_TO_REGION}, @@ -1741,7 +1741,7 @@ def combine_with_multiple_args(sources, multiplier, suffix, constant): pl.lit(constant).alias("const_col"), ) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=combine_with_multiple_args, combine_kwargs={"multiplier": 2, "suffix": "doubled", "constant": "hello"}, @@ -1761,7 +1761,7 @@ def test_combine_kwargs_with_none(self): tracker = PredicateTracker(df) # Lambda that takes only sources dict (no kwargs) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], combine_kwargs=None, # Explicit None @@ -1776,7 +1776,7 @@ def test_combine_kwargs_empty_dict(self): df = pl.DataFrame({"date": dates, "val": [1, 2, 3, 4, 5]}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=lambda s: s["data"], combine_kwargs={}, # Empty dict @@ -1799,7 +1799,7 @@ def combine_with_join_type(sources, join_how): return sources["left"].join(sources["right"], on="key", how=join_how) # Test with inner join - lf_inner = multi_source( + lf_inner = pushdown_combine( sources={ "left": (left_tracker.lazy_frame, {"date": FilterSpec()}), "right": (right_tracker.lazy_frame, {}), @@ -1812,7 +1812,7 @@ def combine_with_join_type(sources, join_how): assert len(result_inner) == 3 # Only A, C, E match # Test with left join - lf_left = multi_source( + lf_left = pushdown_combine( sources={ "left": (left_tracker.lazy_frame, {"date": FilterSpec()}), "right": (right_tracker.lazy_frame, {}), @@ -1859,7 +1859,7 @@ def combine_primary_and_reference(sources, category_to_code): # Join with reference on date and category_code return primary_with_code.join(sources["reference"], on=["date", "category_code"], how="left") - lf = multi_source( + lf = pushdown_combine( sources={ "primary": ( primary_tracker.lazy_frame, @@ -1905,7 +1905,7 @@ def test_sources_as_kwargs_basic(self): def combine(left, right): return left.join(right, on="date") - lf = multi_source( + lf = pushdown_combine( sources={ "left": (left_tracker.lazy_frame, {"date": FilterSpec()}), "right": (right_tracker.lazy_frame, {"date": FilterSpec()}), @@ -1930,7 +1930,7 @@ def test_sources_as_kwargs_with_combine_kwargs(self): def combine(data, multiplier, suffix): return data.with_columns((pl.col("val") * multiplier).alias(f"scaled_{suffix}")) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=combine, combine_kwargs={"multiplier": 10, "suffix": "x10"}, @@ -1953,7 +1953,7 @@ def test_sources_as_kwargs_with_lookback(self): def combine(data): return data.with_columns(pl.col("val").shift(3).alias("val_lag3")) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=4))})}, combine=combine, sources_as_kwargs=True, @@ -1979,7 +1979,7 @@ def test_sources_as_kwargs_multiple_sources(self): def combine(prices, volumes, metadata): return prices.join(volumes, on="date").join(metadata, on="date") - lf = multi_source( + lf = pushdown_combine( sources={ "prices": (prices_tracker.lazy_frame, {"date": FilterSpec()}), "volumes": (volumes_tracker.lazy_frame, {"date": FilterSpec()}), @@ -2004,7 +2004,7 @@ def test_sources_as_kwargs_false_is_default(self): def combine(sources): return sources["data"] - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec()})}, combine=combine, # sources_as_kwargs not specified, should default to False @@ -2030,7 +2030,7 @@ def test_sources_as_kwargs_with_value_mapping(self): def combine(data): return data.with_columns(pl.col("region_code").replace(CODE_TO_REGION).alias("region")).drop("region_code") - lf = multi_source( + lf = pushdown_combine( sources={ "data": ( tracker.lazy_frame, @@ -2063,7 +2063,7 @@ def combine(primary, reference, category_mapping): primary_with_code = primary.with_columns(pl.col("category").replace(category_mapping).alias("category_code")) return primary_with_code.join(reference, on=["date", "category_code"], how="left") - lf = multi_source( + lf = pushdown_combine( sources={ "primary": (primary_tracker.lazy_frame, {"date": FilterSpec(), "category": FilterSpec()}), "reference": ( @@ -2120,7 +2120,7 @@ def test_eq_covers_full_day(self): df = self._make_intraday_df() tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"main": (tracker.lazy_frame, {"data_date": FilterSpec(source_col="timestamp")})}, combine=self._combine, ) @@ -2133,7 +2133,7 @@ def test_eq_with_lookback_covers_full_day(self): df = self._make_intraday_df() tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={ "main": ( tracker.lazy_frame, @@ -2168,7 +2168,7 @@ def test_eq_with_subday_lookback_fetches_prior_evening(self): ) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={ "main": ( tracker.lazy_frame, @@ -2196,7 +2196,7 @@ def test_le_covers_full_day(self): df = self._make_intraday_df() tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"main": (tracker.lazy_frame, {"data_date": FilterSpec(source_col="timestamp")})}, combine=self._combine, ) @@ -2209,7 +2209,7 @@ def test_ge_starts_at_midnight(self): df = self._make_intraday_df() tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"main": (tracker.lazy_frame, {"data_date": FilterSpec(source_col="timestamp")})}, combine=self._combine, ) @@ -2222,7 +2222,7 @@ def test_lt_excludes_full_day(self): df = self._make_intraday_df() tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"main": (tracker.lazy_frame, {"data_date": FilterSpec(source_col="timestamp")})}, combine=self._combine, ) @@ -2235,7 +2235,7 @@ def test_gt_starts_next_day(self): df = self._make_intraday_df() tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"main": (tracker.lazy_frame, {"data_date": FilterSpec(source_col="timestamp")})}, combine=self._combine, ) @@ -2248,7 +2248,7 @@ def test_is_in_covers_full_days(self): df = self._make_intraday_df() tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"main": (tracker.lazy_frame, {"data_date": FilterSpec(source_col="timestamp")})}, combine=self._combine, ) @@ -2261,7 +2261,7 @@ def test_datetime_filter_on_datetime_source_unchanged(self): df = self._make_intraday_df() tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"main": (tracker.lazy_frame, {"timestamp": FilterSpec()})}, combine=lambda s: s["main"], ) diff --git a/polars_io_tools/tests/pushdown/test_join_lookup.py b/polars_io_tools/tests/pushdown/test_join_lookup.py index 325a796..26d3c28 100644 --- a/polars_io_tools/tests/pushdown/test_join_lookup.py +++ b/polars_io_tools/tests/pushdown/test_join_lookup.py @@ -14,7 +14,7 @@ is achievable. This is the natural alternative to user-side workarounds like cubist's -``multi_source(value_mapping=...)`` for many-to-one mappings expressed as +``pushdown_combine(value_mapping=...)`` for many-to-one mappings expressed as a small dataframe. """ diff --git a/polars_io_tools/tests/test_testing_utils.py b/polars_io_tools/tests/test_testing_utils.py index c306081..02da616 100644 --- a/polars_io_tools/tests/test_testing_utils.py +++ b/polars_io_tools/tests/test_testing_utils.py @@ -9,7 +9,7 @@ import polars as pl import pytest -from polars_io_tools.io_sources.multi_source import FilterSpec, multi_source +from polars_io_tools.io_sources.pushdown_combine import FilterSpec, pushdown_combine from polars_io_tools.testing import PredicateAnalyzer, PredicateTracker, io_source_assert @@ -204,7 +204,7 @@ def test_find_temporal_filters_with_lookback(self): df = pl.DataFrame({"date": [date(2024, 1, i) for i in range(1, 11)], "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=3))})}, combine=lambda s: s["data"], ) @@ -252,7 +252,7 @@ def test_count_filters_on_column(self): df = pl.DataFrame({"date": [date(2024, 1, i) for i in range(1, 11)], "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=3))})}, combine=lambda s: s["data"], ) @@ -285,7 +285,7 @@ def test_singular_methods_return_first(self): df = pl.DataFrame({"date": [date(2024, 1, i) for i in range(1, 11)], "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=3))})}, combine=lambda s: s["data"], ) @@ -303,15 +303,15 @@ def test_singular_methods_return_first(self): assert single is plural[0] -class TestIntegrationWithMultiSource: - """Integration tests with multi_source.""" +class TestIntegrationWithPushdownCombine: + """Integration tests with pushdown_combine.""" - def test_tracker_with_multi_source(self): - """PredicateTracker works with multi_source.""" + def test_tracker_with_pushdown_combine(self): + """PredicateTracker works with pushdown_combine.""" df = pl.DataFrame({"date": [date(2024, 1, i) for i in range(1, 11)], "val": list(range(10))}) tracker = PredicateTracker(df) - lf = multi_source( + lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"date": FilterSpec(lookback=timedelta(days=3))})}, combine=lambda s: s["data"], ) @@ -329,11 +329,11 @@ def test_tracker_with_multi_source(self): assert date(2024, 1, 5) in lower_bounds # Original filter def test_tracker_with_discrete_filter(self): - """PredicateTracker works with discrete filters in multi_source.""" + """PredicateTracker works with discrete filters in pushdown_combine.""" df = pl.DataFrame({"category": ["A", "B", "C", "A", "B"], "val": [1, 2, 3, 4, 5]}) tracker = PredicateTracker(df) - result_lf = multi_source( + result_lf = pushdown_combine( sources={"data": (tracker.lazy_frame, {"category": FilterSpec()})}, combine=lambda s: s["data"], )