@@ -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
0 commit comments