Skip to content

Commit a8e8192

Browse files
rustyconoverclaude
andcommitted
feat: per-partition result-cache opt-in + fixtures
Add `CacheControl.partition_scope` (renders `vgi.cache.partition_scope`), the worker-side opt-in for the extension's new per-partition result cache (cache a SINGLE_VALUE_PARTITIONS result split by partition value so a later `=`/`IN` scan reuses per-partition entries). Fixtures backing the extension tests: - cache_partition_scope: single-column, filter_pushdown + auto_apply_filters. - cache_partition_parallel: work-queue fan-out (parallel capture) + a NULL partition. - cache_partition_multicol: (region, year) multi-column SINGLE_VALUE. - cache_partition_proj: projection pushdown + explicit partition_values. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent f807d81 commit a8e8192

4 files changed

Lines changed: 366 additions & 0 deletions

File tree

‎vgi/_test_fixtures/table/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,10 @@
4444
CacheOrderedFunction,
4545
CacheParallelFunction,
4646
CachePartitionedFunction,
47+
CachePartitionScopeFunction,
48+
CachePartitionParallelFunction,
49+
CachePartitionMultiColFunction,
50+
CachePartitionProjFunction,
4751
CachePoisonFunction,
4852
CacheProjectionFunction,
4953
CacheRevalidatableFunction,
@@ -182,6 +186,10 @@
182186
"CacheOrderedFunction",
183187
"CacheParallelFunction",
184188
"CachePartitionedFunction",
189+
"CachePartitionScopeFunction",
190+
"CachePartitionParallelFunction",
191+
"CachePartitionMultiColFunction",
192+
"CachePartitionProjFunction",
185193
"CachePoisonFunction",
186194
"CacheProjectionFunction",
187195
"CacheTypesFunction",

‎vgi/_test_fixtures/table/cache.py‎

Lines changed: 340 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1523,3 +1523,343 @@ def process(
15231523
cast(VgiOutputCollector, out).emit(batch, cache_control=cache_control)
15241524
state.advertised = True
15251525
state.country_idx += 1
1526+
1527+
1528+
# ---------------------------------------------------------------------------
1529+
# cache_partition_scope — per-PARTITION result caching opt-in
1530+
# ---------------------------------------------------------------------------
1531+
# SINGLE_VALUE_PARTITIONS + vgi.cache.partition_scope: the client ALSO caches the
1532+
# result split by partition value (one entry per country), so a later =/IN scan on
1533+
# `country` serves the requested partitions from cache without the worker.
1534+
# filter_pushdown + auto_apply_filters means a WHERE country=... predicate reaches
1535+
# the worker as a real filter (so input.filters is populated → the client can
1536+
# enumerate the requested set) and the framework prunes emitted batches to it (so a
1537+
# fall-through worker scan is row-correct, since DuckDB does NOT re-filter a pushed
1538+
# predicate). partition_scope is advertised on EVERY batch so the client latches it
1539+
# even if the first country is filtered away on a fall-through scan.
1540+
1541+
1542+
@dataclass(slots=True, frozen=True)
1543+
class CachePartitionScopeArgs:
1544+
"""Arguments for CachePartitionScopeFunction."""
1545+
1546+
rows_per_country: Annotated[int, Arg(0, doc="Rows per country partition", ge=1)]
1547+
1548+
1549+
@init_single_worker
1550+
@bind_fixed_schema
1551+
class CachePartitionScopeFunction(TableFunctionGenerator[CachePartitionScopeArgs, _CachePartitionedState]):
1552+
"""Per-partition cacheable single-value-partitioned result (country + sales)."""
1553+
1554+
FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
1555+
[_partition_field("country", pa.string()), pa.field("sales", pa.int64())]
1556+
)
1557+
1558+
class Meta:
1559+
"""Metadata for CachePartitionScopeFunction."""
1560+
1561+
name = "cache_partition_scope"
1562+
description = "Per-partition cacheable single-value-partitioned result (vgi.cache.partition_scope)"
1563+
categories = ["generator", "cache", "testing", "partitioning"]
1564+
tags = {"category": "cache", "type": "partitioned"}
1565+
partition_kind = _PartitionKind.SINGLE_VALUE_PARTITIONS
1566+
filter_pushdown = True
1567+
auto_apply_filters = True
1568+
examples = [
1569+
FunctionExample(
1570+
sql="SELECT * FROM cache_partition_scope(10) WHERE country = 'US'",
1571+
description="Per-partition cache serve for one country",
1572+
),
1573+
]
1574+
1575+
FunctionArguments = CachePartitionScopeArgs
1576+
1577+
@classmethod
1578+
def initial_state(cls, params: ProcessParams[CachePartitionScopeArgs]) -> _CachePartitionedState:
1579+
"""Create initial state."""
1580+
return _CachePartitionedState()
1581+
1582+
@classmethod
1583+
def process(
1584+
cls,
1585+
params: ProcessParams[CachePartitionScopeArgs],
1586+
state: _CachePartitionedState,
1587+
out: OutputCollector,
1588+
) -> None:
1589+
"""Emit one single-country batch per tick; advertise per-partition cacheability."""
1590+
if state.country_idx >= len(_CACHE_COUNTRIES):
1591+
out.finish()
1592+
return
1593+
country = _CACHE_COUNTRIES[state.country_idx]
1594+
rpc = params.args.rows_per_country
1595+
base = state.country_idx * 1_000_000
1596+
batch = pa.RecordBatch.from_pydict(
1597+
{"country": [country] * rpc, "sales": [base + i for i in range(rpc)]},
1598+
schema=params.output_schema,
1599+
)
1600+
# Advertise on every batch (extension latches the first it sees) so the opt-in
1601+
# survives a fall-through scan whose leading country was filtered to 0 rows.
1602+
cast(VgiOutputCollector, out).emit(
1603+
batch, cache_control=CacheControl(ttl=_DEFAULT_TTL_SECONDS, partition_scope=True)
1604+
)
1605+
state.country_idx += 1
1606+
1607+
1608+
# ---------------------------------------------------------------------------
1609+
# cache_partition_parallel — work-queue fan-out (PARALLEL capture) + a NULL partition
1610+
# ---------------------------------------------------------------------------
1611+
# Unlike cache_partition_scope (single-worker), this uses the work-queue pattern so a
1612+
# `threads=N` + `pool false` scan fans partitions across N workers → the per-partition
1613+
# split at commit must bucket batches drawn from MULTIPLE capture substreams. Also emits
1614+
# one NULL-valued partition (SINGLE_VALUE permits NULL) so capture/serve of a NULL tuple
1615+
# and the correct non-enumerability of `IS NULL` are exercised.
1616+
_PSCOPE_COUNTRIES: list[str | None] = ["AU", "CA", "US", None]
1617+
_PSCOPE_QUEUE_FMT = ">i" # int32 country index
1618+
1619+
1620+
@dataclass(slots=True, frozen=True)
1621+
class CachePartitionParallelArgs:
1622+
"""Arguments for CachePartitionParallelFunction."""
1623+
1624+
rows_per_country: Annotated[int, Arg(0, doc="Rows per country partition", ge=1)]
1625+
1626+
1627+
@dataclass(kw_only=True)
1628+
class _CachePartitionParallelState(ArrowSerializableDataclass):
1629+
"""Per-worker cursor over the work queue."""
1630+
1631+
country_idx: int = -1
1632+
current_idx: int = 0
1633+
started: bool = False
1634+
1635+
1636+
@bind_fixed_schema
1637+
class CachePartitionParallelFunction(
1638+
TableFunctionGenerator[CachePartitionParallelArgs, _CachePartitionParallelState]
1639+
):
1640+
"""Per-partition cacheable; work-queue fan-out across workers; includes a NULL partition."""
1641+
1642+
FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
1643+
[_partition_field("country", pa.string()), pa.field("sales", pa.int64())]
1644+
)
1645+
1646+
class Meta:
1647+
"""Metadata for CachePartitionParallelFunction."""
1648+
1649+
name = "cache_partition_parallel"
1650+
description = "Per-partition cacheable; work-queue fan-out (parallel capture); one NULL partition"
1651+
categories = ["generator", "cache", "testing", "partitioning"]
1652+
tags = {"category": "cache", "type": "partitioned"}
1653+
partition_kind = _PartitionKind.SINGLE_VALUE_PARTITIONS
1654+
filter_pushdown = True
1655+
auto_apply_filters = True
1656+
1657+
FunctionArguments = CachePartitionParallelArgs
1658+
1659+
@classmethod
1660+
def on_init(cls, params: InitParams[CachePartitionParallelArgs]) -> GlobalInitResponse:
1661+
"""Primary worker enqueues one (country_idx) item per partition."""
1662+
items = [struct.pack(_PSCOPE_QUEUE_FMT, i) for i in range(len(_PSCOPE_COUNTRIES))]
1663+
params.storage.queue_push(items)
1664+
return GlobalInitResponse()
1665+
1666+
@classmethod
1667+
def initial_state(cls, params: ProcessParams[CachePartitionParallelArgs]) -> _CachePartitionParallelState:
1668+
"""Create initial state."""
1669+
return _CachePartitionParallelState()
1670+
1671+
@classmethod
1672+
def process(
1673+
cls,
1674+
params: ProcessParams[CachePartitionParallelArgs],
1675+
state: _CachePartitionParallelState,
1676+
out: OutputCollector,
1677+
) -> None:
1678+
"""Pop one country from the queue; emit its single-valued batch (auto pv)."""
1679+
if not state.started or state.current_idx >= params.args.rows_per_country:
1680+
item = params.storage.queue_pop()
1681+
if item is None:
1682+
out.finish()
1683+
return
1684+
(state.country_idx,) = struct.unpack(_PSCOPE_QUEUE_FMT, item)
1685+
state.current_idx = 0
1686+
state.started = True
1687+
country = _PSCOPE_COUNTRIES[state.country_idx]
1688+
rpc = params.args.rows_per_country
1689+
base = state.country_idx * 1_000_000
1690+
batch = pa.RecordBatch.from_pydict(
1691+
{"country": [country] * rpc, "sales": [base + i for i in range(rpc)]},
1692+
schema=params.output_schema,
1693+
)
1694+
# Explicit pv so the NULL partition is unambiguous (auto-extract of an all-NULL
1695+
# column is fine too, but explicit keeps the NULL scalar's type pinned).
1696+
cast(VgiOutputCollector, out).emit(
1697+
batch,
1698+
cache_control=CacheControl(ttl=_DEFAULT_TTL_SECONDS, partition_scope=True),
1699+
partition_values={"country": (pa.scalar(country, pa.string()), pa.scalar(country, pa.string()))},
1700+
)
1701+
state.current_idx = rpc
1702+
1703+
1704+
# ---------------------------------------------------------------------------
1705+
# cache_partition_multicol — MULTI-COLUMN (region, year) SINGLE_VALUE partitions
1706+
# ---------------------------------------------------------------------------
1707+
# Exercises the cross-product enumeration (region IN × year IN), 2-column tuple
1708+
# canonicalization, and the partial-constraint case (region constrained, year free →
1709+
# NOT enumerable → falls through).
1710+
# Years are NON-contiguous on purpose: DuckDB rewrites `year IN (2020, 2021)` (contiguous
1711+
# ints) into a BETWEEN range (non-enumerable), so a gap keeps the pushed filter an IN_FILTER
1712+
# and the cross-product enumeration path is actually exercised.
1713+
_PSCOPE_REGIONS: list[str] = ["EU", "US"]
1714+
_PSCOPE_YEARS: list[int] = [2020, 2022]
1715+
_PSCOPE_RY: list[tuple[str, int]] = [(r, y) for r in _PSCOPE_REGIONS for y in _PSCOPE_YEARS]
1716+
1717+
1718+
@dataclass(slots=True, frozen=True)
1719+
class CachePartitionMultiColArgs:
1720+
"""Arguments for CachePartitionMultiColFunction."""
1721+
1722+
rows_per_partition: Annotated[int, Arg(0, doc="Rows per (region, year) partition", ge=1)]
1723+
1724+
1725+
@dataclass(kw_only=True)
1726+
class _CachePartitionMultiColState(ArrowSerializableDataclass):
1727+
"""Cursor over the fixed (region, year) list."""
1728+
1729+
idx: int = 0
1730+
1731+
1732+
@init_single_worker
1733+
@bind_fixed_schema
1734+
class CachePartitionMultiColFunction(
1735+
TableFunctionGenerator[CachePartitionMultiColArgs, _CachePartitionMultiColState]
1736+
):
1737+
"""Per-partition cacheable over TWO SINGLE_VALUE partition columns (region, year)."""
1738+
1739+
FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
1740+
[
1741+
_partition_field("region", pa.string()),
1742+
_partition_field("year", pa.int64()),
1743+
pa.field("amount", pa.int64()),
1744+
]
1745+
)
1746+
1747+
class Meta:
1748+
"""Metadata for CachePartitionMultiColFunction."""
1749+
1750+
name = "cache_partition_multicol"
1751+
description = "Per-partition cacheable over (region, year) SINGLE_VALUE partition columns"
1752+
categories = ["generator", "cache", "testing", "partitioning"]
1753+
tags = {"category": "cache", "type": "partitioned"}
1754+
partition_kind = _PartitionKind.SINGLE_VALUE_PARTITIONS
1755+
filter_pushdown = True
1756+
auto_apply_filters = True
1757+
1758+
FunctionArguments = CachePartitionMultiColArgs
1759+
1760+
@classmethod
1761+
def initial_state(cls, params: ProcessParams[CachePartitionMultiColArgs]) -> _CachePartitionMultiColState:
1762+
"""Create initial state."""
1763+
return _CachePartitionMultiColState()
1764+
1765+
@classmethod
1766+
def process(
1767+
cls,
1768+
params: ProcessParams[CachePartitionMultiColArgs],
1769+
state: _CachePartitionMultiColState,
1770+
out: OutputCollector,
1771+
) -> None:
1772+
"""Emit one single-valued (region, year) batch per tick (pv auto-extracted)."""
1773+
if state.idx >= len(_PSCOPE_RY):
1774+
out.finish()
1775+
return
1776+
region, year = _PSCOPE_RY[state.idx]
1777+
rpp = params.args.rows_per_partition
1778+
base = state.idx * 1000
1779+
batch = pa.RecordBatch.from_pydict(
1780+
{"region": [region] * rpp, "year": [year] * rpp, "amount": [base + i for i in range(rpp)]},
1781+
schema=params.output_schema,
1782+
)
1783+
cast(VgiOutputCollector, out).emit(
1784+
batch, cache_control=CacheControl(ttl=_DEFAULT_TTL_SECONDS, partition_scope=True)
1785+
)
1786+
state.idx += 1
1787+
1788+
1789+
# ---------------------------------------------------------------------------
1790+
# cache_partition_proj — projection pushdown + per-partition cache
1791+
# ---------------------------------------------------------------------------
1792+
# projection_pushdown=True makes projection part of the cache key; explicit pv means the
1793+
# partition value survives even when the partition column is projected OUT of the emitted
1794+
# batch (the auto-extract-impossible case). `extra` is a non-partition column to project
1795+
# away while keeping `country` (so a WHERE country=X can still push).
1796+
_PSCOPE_PROJ_COUNTRIES: list[str] = ["CA", "US"]
1797+
1798+
1799+
@dataclass(slots=True, frozen=True)
1800+
class CachePartitionProjArgs:
1801+
"""Arguments for CachePartitionProjFunction."""
1802+
1803+
rows_per_country: Annotated[int, Arg(0, doc="Rows per country partition", ge=1)]
1804+
1805+
1806+
@init_single_worker
1807+
@bind_fixed_schema
1808+
class CachePartitionProjFunction(TableFunctionGenerator[CachePartitionProjArgs, _CachePartitionedState]):
1809+
"""Per-partition cacheable with projection pushdown (explicit pv survives a dropped column)."""
1810+
1811+
FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
1812+
[
1813+
_partition_field("country", pa.string()),
1814+
pa.field("sales", pa.int64()),
1815+
pa.field("extra", pa.int64()),
1816+
]
1817+
)
1818+
1819+
class Meta:
1820+
"""Metadata for CachePartitionProjFunction."""
1821+
1822+
name = "cache_partition_proj"
1823+
description = "Per-partition cacheable with projection pushdown + explicit partition_values"
1824+
categories = ["generator", "cache", "testing", "partitioning"]
1825+
tags = {"category": "cache", "type": "partitioned"}
1826+
partition_kind = _PartitionKind.SINGLE_VALUE_PARTITIONS
1827+
projection_pushdown = True
1828+
filter_pushdown = True
1829+
auto_apply_filters = True
1830+
1831+
FunctionArguments = CachePartitionProjArgs
1832+
1833+
@classmethod
1834+
def initial_state(cls, params: ProcessParams[CachePartitionProjArgs]) -> _CachePartitionedState:
1835+
"""Create initial state."""
1836+
return _CachePartitionedState()
1837+
1838+
@classmethod
1839+
def process(
1840+
cls,
1841+
params: ProcessParams[CachePartitionProjArgs],
1842+
state: _CachePartitionedState,
1843+
out: OutputCollector,
1844+
) -> None:
1845+
"""Emit only the projected columns per partition; supply pv explicitly."""
1846+
if state.country_idx >= len(_PSCOPE_PROJ_COUNTRIES):
1847+
out.finish()
1848+
return
1849+
country = _PSCOPE_PROJ_COUNTRIES[state.country_idx]
1850+
rpc = params.args.rows_per_country
1851+
base = state.country_idx * 1_000_000
1852+
full = {
1853+
"country": [country] * rpc,
1854+
"sales": [base + i for i in range(rpc)],
1855+
"extra": [base + 500 + i for i in range(rpc)],
1856+
}
1857+
# Emit only the projected columns (params.output_schema reflects the pushdown).
1858+
cols = {f.name: full[f.name] for f in params.output_schema}
1859+
batch = pa.RecordBatch.from_pydict(cols, schema=params.output_schema)
1860+
cast(VgiOutputCollector, out).emit(
1861+
batch,
1862+
cache_control=CacheControl(ttl=_DEFAULT_TTL_SECONDS, partition_scope=True),
1863+
partition_values={"country": (pa.scalar(country, pa.string()), pa.scalar(country, pa.string()))},
1864+
)
1865+
state.country_idx += 1

‎vgi/_test_fixtures/worker.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,10 @@
141141
CacheOrderedFunction,
142142
CacheParallelFunction,
143143
CachePartitionedFunction,
144+
CachePartitionScopeFunction,
145+
CachePartitionParallelFunction,
146+
CachePartitionMultiColFunction,
147+
CachePartitionProjFunction,
144148
CachePoisonFunction,
145149
CacheProjectionFunction,
146150
CacheRevalidatableFunction,
@@ -487,6 +491,10 @@ def _build_fixture_stats() -> tuple[
487491
CacheTypesFunction,
488492
CacheFilteredFunction,
489493
CachePartitionedFunction,
494+
CachePartitionScopeFunction,
495+
CachePartitionParallelFunction,
496+
CachePartitionMultiColFunction,
497+
CachePartitionProjFunction,
490498
SecretDemoFunction,
491499
MultiSecretDemoFunction,
492500
ScopedSecretDemoFunction,

0 commit comments

Comments
 (0)