Skip to content

Commit 8ff6341

Browse files
rustyconoverclaude
andcommitted
feat(scalar): deliver ConstParam as a typed Arrow scalar when annotated as one
A ConstParam was always converted to a Python value via as_py(), so workers that feed it back into pyarrow.compute (e.g. pc.multiply(value, factor)) make pyarrow re-infer an Arrow scalar on every batch. That inference is ~40-50x slower on Windows than passing a typed pa.scalar (and dominates per-batch cost for scalar UDFs — measured: 148us vs 3us for a 2048-row int64 multiply). Make delivery follow the annotation: when a ConstParam is annotated with a pyarrow scalar class (Annotated[pa.Int64Scalar, ConstParam(...)]), deliver the typed pa.Scalar straight through to compute(); Python-type annotations (int/str/float/...) still get as_py(), so existing workers are unaffected. - arguments.py: SCALAR_CLASS_TO_DATATYPE (scalar class -> Arrow type), exported. - scalar_function.py: _is_scalar_annotation(); _const_param_to_arg infers the Arrow type from the scalar class; per-class _const_scalar_params drives whether compute() receives the scalar or the Python value. - _test_fixtures/scalar/arithmetic.py: MultiplyFunction.factor is now Annotated[pa.Int64Scalar, ...] — the idiomatic fast pattern. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 4384825 commit 8ff6341

3 files changed

Lines changed: 96 additions & 5 deletions

File tree

‎vgi/_test_fixtures/scalar/arithmetic.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,17 @@ class Meta:
5151
def compute(
5252
cls,
5353
value: Annotated[pa.Int64Array, Param(doc="Integer value to multiply")],
54-
factor: Annotated[int, ConstParam("Multiplication factor")],
54+
factor: Annotated[pa.Int64Scalar, ConstParam("Multiplication factor")],
5555
) -> Annotated[pa.Int64Array, Returns()]:
56-
"""Multiply values by the constant factor."""
56+
"""Multiply values by the constant factor.
57+
58+
``factor`` is annotated as a ``pa.Int64Scalar`` (not a Python ``int``), so
59+
the SDK delivers the typed Arrow scalar straight through and it can be
60+
passed directly into ``pc.multiply``. Passing a bare Python int instead
61+
makes pyarrow infer an Arrow scalar on every call — ~40-50x slower on
62+
Windows than a typed scalar, and the dominant per-batch cost for a scalar
63+
UDF. This is the idiomatic fast pattern for compute-bound constants.
64+
"""
5765
return pc.multiply(value, factor)
5866

5967

‎vgi/arguments.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,46 @@
109109
}
110110

111111

112+
# Simple scalar classes that can be inferred automatically. Annotating a const
113+
# parameter with one of these — ``Annotated[pa.Int64Scalar, ConstParam(...)]`` —
114+
# makes the SDK deliver the typed ``pa.Scalar`` straight through to ``compute()``
115+
# instead of converting it to a Python object via ``as_py()``. That lets the
116+
# const be passed directly into ``pyarrow.compute`` with no per-call scalar
117+
# re-inference, which is ~40-50x slower on Windows than a typed scalar. Built
118+
# defensively so a pyarrow release missing a class name doesn't break import.
119+
def _build_scalar_class_to_datatype() -> dict[type, pa.DataType]:
120+
pairs: list[tuple[str, pa.DataType]] = [
121+
("Int8Scalar", pa.int8()),
122+
("Int16Scalar", pa.int16()),
123+
("Int32Scalar", pa.int32()),
124+
("Int64Scalar", pa.int64()),
125+
("UInt8Scalar", pa.uint8()),
126+
("UInt16Scalar", pa.uint16()),
127+
("UInt32Scalar", pa.uint32()),
128+
("UInt64Scalar", pa.uint64()),
129+
("HalfFloatScalar", pa.float16()),
130+
("FloatScalar", pa.float32()),
131+
("DoubleScalar", pa.float64()),
132+
("StringScalar", pa.string()),
133+
("LargeStringScalar", pa.large_string()),
134+
("BinaryScalar", pa.binary()),
135+
("LargeBinaryScalar", pa.large_binary()),
136+
("BooleanScalar", pa.bool_()),
137+
("Date32Scalar", pa.date32()),
138+
("Date64Scalar", pa.date64()),
139+
("NullScalar", pa.null()),
140+
]
141+
out: dict[type, pa.DataType] = {}
142+
for name, dt in pairs:
143+
cls = getattr(pa, name, None)
144+
if isinstance(cls, type):
145+
out[cls] = dt
146+
return out
147+
148+
149+
SCALAR_CLASS_TO_DATATYPE: dict[type, pa.DataType] = _build_scalar_class_to_datatype()
150+
151+
112152
# Sentinel for missing default value - proper type pattern
113153
class _MissingType:
114154
"""Sentinel type for missing default values.
@@ -158,6 +198,7 @@ def _accepts_none(annotated_inner_type: Any) -> bool:
158198
"Param",
159199
"PYTHON_TO_ARROW",
160200
"Returns",
201+
"SCALAR_CLASS_TO_DATATYPE",
161202
"TableInput",
162203
"TaggedUnion",
163204
"TypeBoundPredicate",

‎vgi/scalar_function.py‎

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
_PYTHON_TO_ARROW,
4949
ARRAY_CLASS_TO_DATATYPE,
5050
COMPLEX_ARRAY_CLASSES,
51+
SCALAR_CLASS_TO_DATATYPE,
5152
Arg,
5253
Arguments,
5354
ArgumentValidationError,
@@ -212,6 +213,25 @@ def _param_to_arg(param: Param, base_type: type, position: int) -> Arg[Any]:
212213
)
213214

214215

216+
def _is_scalar_annotation(base_type: object) -> bool:
217+
"""Whether a [`ConstParam`][] is annotated with a pyarrow scalar class.
218+
219+
``Annotated[pa.Int64Scalar, ConstParam(...)]`` (or the generic ``pa.Scalar``)
220+
signals that the typed Arrow scalar should be delivered to ``compute()``
221+
unchanged, instead of converting it to a Python value via ``as_py()``. The
222+
scalar can then flow straight into ``pyarrow.compute`` with no per-call scalar
223+
re-inference (the latter is ~40-50x slower on Windows than a typed scalar).
224+
225+
Args:
226+
base_type: The type from the ``Annotated`` first argument.
227+
228+
Returns:
229+
True iff ``base_type`` is ``pa.Scalar`` or a subclass.
230+
231+
"""
232+
return isinstance(base_type, type) and issubclass(base_type, pa.Scalar)
233+
234+
215235
def _const_param_to_arg(const_param: ConstParam, base_type: type, position: int) -> Arg[Any]:
216236
"""Convert [`ConstParam`][] dataclass to internal [`Arg`][] object.
217237
@@ -235,6 +255,15 @@ def _const_param_to_arg(const_param: ConstParam, base_type: type, position: int)
235255
elif base_type in _PYTHON_TO_ARROW:
236256
# Infer from Annotated first argument
237257
arrow_type = _PYTHON_TO_ARROW[base_type]
258+
elif base_type in SCALAR_CLASS_TO_DATATYPE:
259+
# Annotated[pa.Int64Scalar, ConstParam(...)] — infer from the scalar class
260+
# and deliver the typed scalar through (see _is_scalar_annotation).
261+
arrow_type = SCALAR_CLASS_TO_DATATYPE[base_type]
262+
elif _is_scalar_annotation(base_type):
263+
raise TypeError(
264+
f"Cannot infer Arrow type from {getattr(base_type, '__name__', base_type)}. "
265+
f"Use a concrete scalar class (e.g. pa.Int64Scalar) or set arrow_type=... on ConstParam()."
266+
)
238267
else:
239268
raise TypeError(
240269
f"Cannot infer Arrow type from {base_type}. "
@@ -791,6 +820,7 @@ def __getattr__(self, _name: str) -> Any:
791820

792821
_compute_params: dict[str, Arg[Any]] # Regular Param() arguments (arrays)
793822
_const_params: dict[str, Arg[Any]] # ConstParam() arguments (scalars)
823+
_const_scalar_params: set[str] # ConstParams annotated as pa.Scalar — delivered un-as_py()'d
794824
_setting_params: dict[str, str] # Setting params: param_name -> setting_key
795825
_secret_params: dict[str, Secret] # Secret params: param_name -> Secret instance
796826
_output_length_param: str | None # OutputLength param name (batch row count)
@@ -866,6 +896,7 @@ def __getattr__(self, name: str) -> Any:
866896

867897
compute_params: dict[str, Arg[Any]] = {}
868898
const_params: dict[str, Arg[Any]] = {}
899+
const_scalar_params: set[str] = set() # ConstParams annotated as pa.Scalar
869900
output_length_param: str | None = None # param that receives batch row count
870901
auth_param: str | None = None # param that receives AuthContext
871902
returns_output_type: pa.DataType | None = None
@@ -940,6 +971,8 @@ def __getattr__(self, name: str) -> Any:
940971
# _resolution_index points to Arguments.positional index
941972
arg._resolution_index = const_index
942973
const_params[name] = arg
974+
if _is_scalar_annotation(base_type):
975+
const_scalar_params.add(name)
943976
setattr(cls, name, _ConstArgDescriptor(arg, name))
944977
overall_position += 1
945978
const_index += 1
@@ -966,6 +999,7 @@ def __getattr__(self, name: str) -> Any:
966999

9671000
cls._compute_params = compute_params
9681001
cls._const_params = const_params
1002+
cls._const_scalar_params = const_scalar_params
9691003
cls._setting_params = setting_params
9701004
cls._secret_params = secret_params
9711005
cls._output_length_param = output_length_param
@@ -1051,14 +1085,22 @@ def _extract_compute_kwargs(
10511085
# Regular param: extract column by index
10521086
kwargs[name] = batch.column(col_idx)
10531087

1054-
# Const params: extract scalar values from arguments
1088+
# Const params: extract values from the constant arguments. Params annotated
1089+
# with a pyarrow scalar class (e.g. Annotated[pa.Int64Scalar, ConstParam(...)])
1090+
# receive the typed pa.Scalar unchanged, so it can be passed straight into
1091+
# pyarrow.compute with no per-call scalar re-inference; all others get the
1092+
# Python value via as_py() (backward compatible).
10551093
for name, arg in cls._const_params.items():
10561094
# Use _resolution_index for Arguments.positional lookup
10571095
arg_idx = cast(int, arg._resolution_index)
10581096
# Get the scalar value from arguments
10591097
scalar = bind_call.arguments.positional[arg_idx]
1060-
# Convert to Python value
1061-
kwargs[name] = scalar.as_py() if scalar is not None else None
1098+
if scalar is None:
1099+
kwargs[name] = None
1100+
elif name in cls._const_scalar_params:
1101+
kwargs[name] = scalar
1102+
else:
1103+
kwargs[name] = scalar.as_py()
10621104

10631105
# Setting params: extract pa.Scalar from settings RecordBatch
10641106
if bind_call.settings is not None and cls._setting_params:

0 commit comments

Comments
 (0)