Skip to content

Commit 8cf1d64

Browse files
rustyconoverclaude
andcommitted
test(fixtures): a partition column that is NOT first, exposed as a table too
Every partitioned fixture so far declares its partition column at index 0 and is reachable only as a function call. Both choices hid a client bug: * With the partition column first, the sink's SCAN-LOCAL column index and the worker-schema index coincide, so a client comparing the two without mapping looks correct. Anywhere else it resolves the wrong column — or, in the C++ extension, threw a FATAL InternalException that invalidated the database. * A table's scan function is built through a different path than a direct function call, so a client can install get_partition_info for one and not the other, and a catalog table then silently never plans PARTITIONED_AGGREGATE however it is declared. TrailingPartitionSalesFunction closes both: same contract and the same deterministic sales values as country_partitioned_sales, with `country` declared LAST (index 3), registered as a function AND as the catalog table example.data.trailing_partition_sales. GROUP BY country projects only country+sales, so the sink asks about scan-local 0 while the declared index is 3 — the mismatch, made reproducible. It declares projection_pushdown = True deliberately. Without pushdown the scan emits every base column and the projection above it carries base-column indices, which DuckDB's own CanUsePartitionedAggregate then maps a second time — duckdb/duckdb#24327, fixed on main and v2.0 by #24557 but NOT backported to v1.5, so that shape still crashes the planner on 1.5-based builds. Worth a no-pushdown variant once a build carries the backport. Drives the new cases in the C++ suite's table/partition_columns.test (Query-farm/vgi 48d4196); full integration suite green there, 289 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 27c1b17 commit 8cf1d64

3 files changed

Lines changed: 144 additions & 1 deletion

File tree

‎vgi/_test_fixtures/table/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@
113113
OverlappingRangePartitionedFunction,
114114
PartitionedWithExplicitOverrideFunction,
115115
RegionYearPartitionedFunction,
116+
TrailingPartitionSalesFunction,
116117
)
117118
from vgi._test_fixtures.table.partition_columns_broken import (
118119
BrokenMissingPartitionValuesFunction,
@@ -264,6 +265,7 @@
264265
"RFF_SIMPLE_COLUMNS",
265266
"RFF_STRUCT_COLUMNS",
266267
"RegionYearPartitionedFunction",
268+
"TrailingPartitionSalesFunction",
267269
"RepeatValueIntFunction",
268270
"RepeatValueStrFunction",
269271
"RffMultiScanFunction",

‎vgi/_test_fixtures/table/partition_columns.py‎

Lines changed: 126 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,11 @@
2323
* :class:`RegionYearPartitionedFunction` — multi-column SINGLE_VALUE.
2424
Each chunk has a single ``(region, year)`` tuple.
2525
26+
* :class:`TrailingPartitionSalesFunction` — same contract as the
27+
country fixture, but the partition column is declared LAST rather
28+
than first, so the worker-schema and scan-local index spaces do not
29+
coincide. Also exposed as a catalog table.
30+
2631
* :class:`PartitionedWithProjectedOutColumnFunction` — declares
2732
partition on ``category`` but DOES NOT include ``category`` in the
2833
emitted batch. Uses the explicit ``partition_values=`` override on
@@ -49,7 +54,7 @@
4954

5055
import struct
5156
from dataclasses import dataclass
52-
from typing import Annotated, ClassVar, cast
57+
from typing import Annotated, Any, ClassVar, cast
5358

5459
import pyarrow as pa
5560
from vgi_rpc import ArrowSerializableDataclass
@@ -285,6 +290,126 @@ def process(
285290
state.current_idx = rpp
286291

287292

293+
# =============================================================================
294+
# SINGLE_VALUE_PARTITIONS with the partition column NOT first
295+
# =============================================================================
296+
297+
298+
@dataclass(slots=True, frozen=True)
299+
class _TrailingPartitionArgs:
300+
"""Arguments for ``trailing_partition_sales``."""
301+
302+
rows_per_country: Annotated[int, Arg(0, doc="Rows to emit per country partition", ge=1)]
303+
304+
305+
@dataclass(kw_only=True)
306+
class _TrailingPartitionState(ArrowSerializableDataclass):
307+
"""Per-worker cursor over countries (same shape as the country fixture)."""
308+
309+
current_country: str | None = None
310+
current_country_idx: int = -1
311+
current_idx: int = 0
312+
313+
314+
@bind_fixed_schema
315+
@_cardinality_from_count
316+
class TrailingPartitionSalesFunction(TableFunctionGenerator[_TrailingPartitionArgs, _TrailingPartitionState]):
317+
"""``country`` is single-valued per chunk, but sits LAST in the schema.
318+
319+
Identical contract to :class:`CountryPartitionedSalesFunction`; the only
320+
difference is where the partition column lives. That difference is the
321+
point: every other partitioned fixture declares its partition column at
322+
index 0, which makes two distinct index spaces accidentally agree.
323+
324+
``CanUsePartitionedAggregate`` asks ``get_partition_info`` about
325+
WORKER-SCHEMA indices, but the sink later asks ``get_partition_data``
326+
about SCAN-LOCAL ones (positions in the scan's own ``column_ids``, after
327+
projection pushdown). ``GROUP BY country`` projects just ``country`` and
328+
``sales``, so here the sink asks about scan-local 0 while the declared
329+
index is 3 — and a client that compares them without mapping raises
330+
"sink requested partition column 0 that is not in the declared partition
331+
set", which is fatal. With the partition column at index 0 both spaces
332+
say 0 and the bug is invisible.
333+
334+
Also registered as a catalog TABLE (``example.data.trailing_partition_sales``)
335+
because the two paths install their scan functions separately: a client can
336+
wire ``get_partition_info`` for direct function calls and miss it for
337+
catalog tables, in which case the planner silently never picks
338+
``PARTITIONED_AGGREGATE`` for a table however it is declared.
339+
"""
340+
341+
FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
342+
cast(
343+
"list[pa.Field[Any]]",
344+
[
345+
pa.field("seq", pa.int64()),
346+
pa.field("label", pa.string()),
347+
pa.field("sales", pa.int64()),
348+
partition_field("country", pa.string()),
349+
],
350+
)
351+
)
352+
353+
class Meta:
354+
name = "trailing_partition_sales"
355+
projection_pushdown = True
356+
description = (
357+
"Per-country sales rows, one Arrow batch per country, with the "
358+
"SINGLE_VALUE partition column declared LAST in the schema "
359+
"instead of first."
360+
)
361+
categories = ["generator", "partitioning"]
362+
partition_kind = PartitionKind.SINGLE_VALUE_PARTITIONS
363+
examples = [
364+
FunctionExample(
365+
sql="SELECT country, SUM(sales) FROM trailing_partition_sales(100) GROUP BY country",
366+
description="Partitioned aggregate over a non-leading partition column",
367+
),
368+
]
369+
370+
@classmethod
371+
def on_init(cls, params: InitParams[_TrailingPartitionArgs]) -> GlobalInitResponse:
372+
items = [struct.pack(_QUEUE_ITEM_FMT, i) for i in range(len(_COUNTRIES))]
373+
params.storage.queue_push(items)
374+
return GlobalInitResponse()
375+
376+
@classmethod
377+
def initial_state(cls, params: ProcessParams[_TrailingPartitionArgs]) -> _TrailingPartitionState:
378+
return _TrailingPartitionState()
379+
380+
@classmethod
381+
def process(
382+
cls,
383+
params: ProcessParams[_TrailingPartitionArgs],
384+
state: _TrailingPartitionState,
385+
out: OutputCollector,
386+
) -> None:
387+
if state.current_country is None or state.current_idx >= params.args.rows_per_country:
388+
item = params.storage.queue_pop()
389+
if item is None:
390+
out.finish()
391+
return
392+
(state.current_country_idx,) = struct.unpack(_QUEUE_ITEM_FMT, item)
393+
state.current_country = _COUNTRIES[state.current_country_idx]
394+
state.current_idx = 0
395+
396+
rpc = params.args.rows_per_country
397+
# Same deterministic sales values as country_partitioned_sales, so a
398+
# test can assert the two agree column-for-column despite the layout.
399+
base = state.current_country_idx * 1_000_000
400+
batch = pa.RecordBatch.from_pydict(
401+
{
402+
"seq": list(range(rpc)),
403+
"label": [f"{state.current_country}-{i}" for i in range(rpc)],
404+
"sales": [base + i for i in range(rpc)],
405+
"country": [state.current_country] * rpc,
406+
},
407+
schema=cls.FIXED_SCHEMA,
408+
)
409+
out.emit(batch)
410+
state.current_idx = rpc
411+
412+
288413
# =============================================================================
289414
# Projected-out partition column — exercises explicit override path
290415
# =============================================================================

‎vgi/_test_fixtures/worker.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,7 @@
243243
SplitZeroFunction,
244244
StructSettingsFunction,
245245
TenThousandFunction,
246+
TrailingPartitionSalesFunction,
246247
TxCachedValueFunction,
247248
TypedProbeFunction,
248249
UnionVarargsFunction,
@@ -528,6 +529,7 @@ def _build_fixture_stats() -> tuple[
528529
OverlappingRangePartitionedFunction,
529530
PartitionedWithExplicitOverrideFunction,
530531
RegionYearPartitionedFunction,
532+
TrailingPartitionSalesFunction,
531533
# Deliberately-broken batch_index fixtures (see
532534
# vgi/_test_fixtures/table/batch_index_broken.py). Registered
533535
# so SQL integration tests in batch_index_contract.test can
@@ -883,6 +885,20 @@ def _build_fixture_stats() -> tuple[
883885
columns=schema(n=pa.int64()),
884886
comment="123456 integers; stats served by the sequence function, not the table",
885887
),
888+
# PartitionColumns as a CATALOG TABLE. A table's scan function is
889+
# built through a different path than a direct function call, so a
890+
# client can support partitioned aggregates for one and silently
891+
# not the other; only a table exercises the catalog path. Paired
892+
# with a partition column declared last — see the fixture docstring.
893+
Table(
894+
name="trailing_partition_sales",
895+
function=TrailingPartitionSalesFunction,
896+
arguments=Arguments(positional=(pa.scalar(100),)),
897+
comment=(
898+
"Per-country sales, SINGLE_VALUE partition column declared last; "
899+
"GROUP BY country must plan as PARTITIONED_AGGREGATE"
900+
),
901+
),
886902
# Result-cache fixtures, exposed as function-backed tables so the
887903
# catalog-attached path (SELECT ... FROM ex.data.<name>) exercises
888904
# the C++ result cache. See vgi/_test_fixtures/table/cache.py.

0 commit comments

Comments
 (0)