Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions docs/wiki/API-Reference.md
Original file line number Diff line number Diff line change
@@ -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:

Expand Down Expand Up @@ -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
```

Expand Down
2 changes: 1 addition & 1 deletion docs/wiki/Concepts.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
6 changes: 3 additions & 3 deletions docs/wiki/Query-Optimization.md
Original file line number Diff line number Diff line change
Expand Up @@ -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, {
Expand Down
4 changes: 2 additions & 2 deletions polars_io_tools/io_sources/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 *
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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(),
Expand Down Expand Up @@ -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__)

Expand Down Expand Up @@ -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],
*,
Expand Down Expand Up @@ -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)),
Expand Down Expand Up @@ -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))}),
Expand All @@ -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
Expand Down Expand Up @@ -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()}
Expand Down Expand Up @@ -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)
4 changes: 2 additions & 2 deletions polars_io_tools/io_sources/ts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",)

Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions polars_io_tools/testing/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
)
Expand Down
6 changes: 3 additions & 3 deletions polars_io_tools/testing/predicate_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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"],
)
Expand Down
10 changes: 5 additions & 5 deletions polars_io_tools/tests/io_sources/test_pickle.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)],
Expand All @@ -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()}),
Expand Down
Loading
Loading