chalk-remote-call-python 0.0.0__tar.gz

chalk_remote_call_python-0.0.0/PKG-INFO

@@ -0,0 +1,193 @@
+Metadata-Version: 2.4
+Name: chalk-remote-call-python
+Version: 0.0.0
+Summary: Chalk remote call Python runtime interface client
+Author: Chalk AI, Inc.
+Project-URL: Homepage, https://chalk.ai
+Project-URL: Documentation, https://docs.chalk.ai
+Project-URL: Changelog, https://docs.chalk.ai/docs/changelog
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python
+Classifier: Typing :: Typed
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
+Classifier: Topic :: Software Development :: Code Generators
+Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: <3.14,>=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: pyarrow>=14.0.0
+Provides-Extra: dev
+Requires-Dist: setuptools-rust>=1.7; extra == "dev"
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
+Requires-Dist: grpcio>=1.60.0; extra == "dev"
+Requires-Dist: grpcio-health-checking>=1.60.0; extra == "dev"
+Requires-Dist: grpcio-reflection>=1.60.0; extra == "dev"
+Requires-Dist: protobuf>=4.25.0; extra == "dev"
+Requires-Dist: grpcio-tools>=1.60.0; extra == "dev"
+
+# chalk-remote-call-python
+
+A Python runtime interface for Chalk's RemoteCallService. This package provides a gRPC server backed by a Rust implementation (tonic + PyO3) that lets you define a `handler(event, context)` function to process incoming Arrow-serialized requests.
+
+The server receives Arrow IPC record batches over a bidirectional gRPC stream, transforms them into a Python dict, invokes your handler, and streams the results back as Arrow IPC.
+
+## Requirements
+
+- Python >= 3.10
+- Rust toolchain (for building the native extension)
+- `pyarrow >= 14.0.0`
+
+## Usage
+
+### 1. Write a handler
+
+Create a Python module with a `handler` function:
+
+```python
+# my_handler.py
+import pyarrow as pa
+import pyarrow.compute as pc
+
+
+def on_startup():
+    """Optional -- runs once before the server starts accepting requests."""
+    print("Loading model weights...")
+
+
+def handler(event: dict[str, pa.Array], context: dict) -> pa.Array:
+    """Called for each incoming request.
+
+    Args:
+        event: dict mapping column names to pyarrow.Array values.
+        context: dict with request metadata:
+            - "peer": remote address string
+            - "metadata": dict of gRPC headers
+
+    Returns:
+        A pyarrow.Array, pyarrow.RecordBatch, pyarrow.Table, list, dict, or scalar.
+        The framework auto-converts the result to Arrow IPC for the response.
+    """
+    return pc.multiply(event["x"], event["y"])
+```
+
+### 2. Start the server
+
+```bash
+chalk-remote-call --handler my_handler.handler
+```
+
+Or via `python -m`:
+
+```bash
+python -m chalk_remote_call --handler my_handler.handler
+```
+
+### CLI options
+
+| Flag | Env var | Default | Description |
+|------|---------|---------|-------------|
+| `--handler` | | *(required)* | Dotted path to handler function |
+| `--port` | `CHALK_REMOTE_CALL_PORT` | `6666` | Port to listen on |
+| `--host` | `CHALK_REMOTE_CALL_HOST` | `[::]` | Host to bind to |
+| `--workers` | `CHALK_REMOTE_CALL_WORKERS` | `10` | Tokio runtime worker threads |
+| `--on-startup` | | | Dotted path to a startup function |
+| `--log-level` | | `INFO` | `DEBUG`, `INFO`, `WARNING`, or `ERROR` |
+
+### Environment variables
+
+| Variable | Description |
+|----------|-------------|
+| `CHALK_INPUT_ARGS` | Comma-separated list of column names (e.g. `x,y,z`). Renames incoming RecordBatch columns by index. If unset, the original column names are used. |
+
+### 3. Docker image
+
+```dockerfile
+FROM python:3.11-slim
+
+WORKDIR /app
+
+RUN pip install chalk-remote-call-python
+
+# Copy your handler code
+COPY my_handler.py .
+
+# The server listens on port 6666 by default
+EXPOSE 6666
+
+ENTRYPOINT ["chalk-remote-call", "--handler", "my_handler.handler"]
+```
+
+Build and run:
+
+```bash
+docker build -t my-chalk-handler .
+docker run -p 6666:6666 -e CHALK_INPUT_ARGS="x,y" my-chalk-handler
+```
+
+### Programmatic usage
+
+You can also start the server from Python:
+
+```python
+from chalk_remote_call import serve
+
+def handler(event, context):
+    return list(event.values())[0]
+
+serve(handler=handler, port=8080)
+```
+
+## Local Testing
+
+Use `grpcurl` to verify the server is running:
+
+```bash
+# Check health
+grpcurl -plaintext localhost:6666 grpc.health.v1.Health/Check
+
+# List services via reflection
+grpcurl -plaintext localhost:6666 list
+```
+
+## Architecture
+
+The server uses a Rust backend via PyO3:
+
+- **tonic** — gRPC server with health checking and reflection
+- **prost** — protobuf message handling
+- **arrow-rs** — Arrow IPC validation
+- **PyO3** — FFI bridge to call Python handler functions
+
+The Rust server runs on a tokio async runtime. When a request arrives, Arrow IPC bytes are passed to Python via PyO3 (`spawn_blocking` + GIL acquisition), where the handler is invoked. Results are passed back as IPC bytes.
+
+## Development
+
+### Setup
+
+```bash
+cd chalk-remote-call-python
+uv venv --python 3.11
+uv pip install -e ".[dev]"
+```
+
+### Running tests
+
+```bash
+pytest tests/ -v
+```
+
+### Rust code
+
+Rust source lives in `chalk-remote-call-rs/`:
+
+To check the Rust code independently:
+
+```bash
+PYO3_PYTHON=python3 cargo check --manifest-path chalk-remote-call-rs/chalk-remote-call-server/Cargo.toml
+```

chalk_remote_call_python-0.0.0/README.md

@@ -0,0 +1,160 @@
+# chalk-remote-call-python
+
+A Python runtime interface for Chalk's RemoteCallService. This package provides a gRPC server backed by a Rust implementation (tonic + PyO3) that lets you define a `handler(event, context)` function to process incoming Arrow-serialized requests.
+
+The server receives Arrow IPC record batches over a bidirectional gRPC stream, transforms them into a Python dict, invokes your handler, and streams the results back as Arrow IPC.
+
+## Requirements
+
+- Python >= 3.10
+- Rust toolchain (for building the native extension)
+- `pyarrow >= 14.0.0`
+
+## Usage
+
+### 1. Write a handler
+
+Create a Python module with a `handler` function:
+
+```python
+# my_handler.py
+import pyarrow as pa
+import pyarrow.compute as pc
+
+
+def on_startup():
+    """Optional -- runs once before the server starts accepting requests."""
+    print("Loading model weights...")
+
+
+def handler(event: dict[str, pa.Array], context: dict) -> pa.Array:
+    """Called for each incoming request.
+
+    Args:
+        event: dict mapping column names to pyarrow.Array values.
+        context: dict with request metadata:
+            - "peer": remote address string
+            - "metadata": dict of gRPC headers
+
+    Returns:
+        A pyarrow.Array, pyarrow.RecordBatch, pyarrow.Table, list, dict, or scalar.
+        The framework auto-converts the result to Arrow IPC for the response.
+    """
+    return pc.multiply(event["x"], event["y"])
+```
+
+### 2. Start the server
+
+```bash
+chalk-remote-call --handler my_handler.handler
+```
+
+Or via `python -m`:
+
+```bash
+python -m chalk_remote_call --handler my_handler.handler
+```
+
+### CLI options
+
+| Flag | Env var | Default | Description |
+|------|---------|---------|-------------|
+| `--handler` | | *(required)* | Dotted path to handler function |
+| `--port` | `CHALK_REMOTE_CALL_PORT` | `6666` | Port to listen on |
+| `--host` | `CHALK_REMOTE_CALL_HOST` | `[::]` | Host to bind to |
+| `--workers` | `CHALK_REMOTE_CALL_WORKERS` | `10` | Tokio runtime worker threads |
+| `--on-startup` | | | Dotted path to a startup function |
+| `--log-level` | | `INFO` | `DEBUG`, `INFO`, `WARNING`, or `ERROR` |
+
+### Environment variables
+
+| Variable | Description |
+|----------|-------------|
+| `CHALK_INPUT_ARGS` | Comma-separated list of column names (e.g. `x,y,z`). Renames incoming RecordBatch columns by index. If unset, the original column names are used. |
+
+### 3. Docker image
+
+```dockerfile
+FROM python:3.11-slim
+
+WORKDIR /app
+
+RUN pip install chalk-remote-call-python
+
+# Copy your handler code
+COPY my_handler.py .
+
+# The server listens on port 6666 by default
+EXPOSE 6666
+
+ENTRYPOINT ["chalk-remote-call", "--handler", "my_handler.handler"]
+```
+
+Build and run:
+
+```bash
+docker build -t my-chalk-handler .
+docker run -p 6666:6666 -e CHALK_INPUT_ARGS="x,y" my-chalk-handler
+```
+
+### Programmatic usage
+
+You can also start the server from Python:
+
+```python
+from chalk_remote_call import serve
+
+def handler(event, context):
+    return list(event.values())[0]
+
+serve(handler=handler, port=8080)
+```
+
+## Local Testing
+
+Use `grpcurl` to verify the server is running:
+
+```bash
+# Check health
+grpcurl -plaintext localhost:6666 grpc.health.v1.Health/Check
+
+# List services via reflection
+grpcurl -plaintext localhost:6666 list
+```
+
+## Architecture
+
+The server uses a Rust backend via PyO3:
+
+- **tonic** — gRPC server with health checking and reflection
+- **prost** — protobuf message handling
+- **arrow-rs** — Arrow IPC validation
+- **PyO3** — FFI bridge to call Python handler functions
+
+The Rust server runs on a tokio async runtime. When a request arrives, Arrow IPC bytes are passed to Python via PyO3 (`spawn_blocking` + GIL acquisition), where the handler is invoked. Results are passed back as IPC bytes.
+
+## Development
+
+### Setup
+
+```bash
+cd chalk-remote-call-python
+uv venv --python 3.11
+uv pip install -e ".[dev]"
+```
+
+### Running tests
+
+```bash
+pytest tests/ -v
+```
+
+### Rust code
+
+Rust source lives in `chalk-remote-call-rs/`:
+
+To check the Rust code independently:
+
+```bash
+PYO3_PYTHON=python3 cargo check --manifest-path chalk-remote-call-rs/chalk-remote-call-server/Cargo.toml
+```

chalk_remote_call_python-0.0.0/chalk_remote_call/__init__.py

@@ -0,0 +1,6 @@
+from importlib.metadata import version
+
+from chalk_remote_call.server import serve
+
+__version__ = version("chalk-remote-call-python")
+__all__ = ["serve"]

chalk_remote_call_python-0.0.0/chalk_remote_call/__main__.py

@@ -0,0 +1,3 @@
+from chalk_remote_call.cli import main
+
+main()

chalk_remote_call_python-0.0.0/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: chalk/runtime/v1/remote_python_call.proto
+# Protobuf Python Version: 4.25.3
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n)chalk/runtime/v1/remote_python_call.proto\x12\x10\x63halk.runtime.v1\"P\n\x13\x43\x61llFunctionRequest\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12%\n\x0e\x66\x65\x61ther_stream\x18\x02 \x01(\x0cR\rfeatherStream\"=\n\x14\x43\x61llFunctionResponse\x12%\n\x0e\x66\x65\x61ther_stream\x18\x01 \x01(\x0cR\rfeatherStream2x\n\x11RemoteCallService\x12\x63\n\x0c\x43\x61llFunction\x12%.chalk.runtime.v1.CallFunctionRequest\x1a&.chalk.runtime.v1.CallFunctionResponse\"\x00(\x01\x30\x01\x42\x8f\x01\n\x14\x63om.chalk.runtime.v1B\x15RemotePythonCallProtoP\x01\xa2\x02\x03\x43RX\xaa\x02\x10\x43halk.Runtime.V1\xca\x02\x10\x43halk\\Runtime\\V1\xe2\x02\x1c\x43halk\\Runtime\\V1\\GPBMetadata\xea\x02\x12\x43halk::Runtime::V1b\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'chalk.runtime.v1.remote_python_call_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  _globals['DESCRIPTOR']._options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n\024com.chalk.runtime.v1B\025RemotePythonCallProtoP\001\242\002\003CRX\252\002\020Chalk.Runtime.V1\312\002\020Chalk\\Runtime\\V1\342\002\034Chalk\\Runtime\\V1\\GPBMetadata\352\002\022Chalk::Runtime::V1'
+  _globals['_CALLFUNCTIONREQUEST']._serialized_start=63
+  _globals['_CALLFUNCTIONREQUEST']._serialized_end=143
+  _globals['_CALLFUNCTIONRESPONSE']._serialized_start=145
+  _globals['_CALLFUNCTIONRESPONSE']._serialized_end=206
+  _globals['_REMOTECALLSERVICE']._serialized_start=208
+  _globals['_REMOTECALLSERVICE']._serialized_end=328
+# @@protoc_insertion_point(module_scope)

chalk_remote_call_python-0.0.0/chalk_remote_call/_gen/chalk/runtime/v1/remote_python_call_pb2_grpc.py

@@ -0,0 +1,66 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+
+from chalk_remote_call._gen.chalk.runtime.v1 import remote_python_call_pb2 as chalk_dot_runtime_dot_v1_dot_remote__python__call__pb2
+
+
+class RemoteCallServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.CallFunction = channel.stream_stream(
+                '/chalk.runtime.v1.RemoteCallService/CallFunction',
+                request_serializer=chalk_dot_runtime_dot_v1_dot_remote__python__call__pb2.CallFunctionRequest.SerializeToString,
+                response_deserializer=chalk_dot_runtime_dot_v1_dot_remote__python__call__pb2.CallFunctionResponse.FromString,
+                )
+
+
+class RemoteCallServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def CallFunction(self, request_iterator, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_RemoteCallServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'CallFunction': grpc.stream_stream_rpc_method_handler(
+                    servicer.CallFunction,
+                    request_deserializer=chalk_dot_runtime_dot_v1_dot_remote__python__call__pb2.CallFunctionRequest.FromString,
+                    response_serializer=chalk_dot_runtime_dot_v1_dot_remote__python__call__pb2.CallFunctionResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'chalk.runtime.v1.RemoteCallService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+ # This class is part of an EXPERIMENTAL API.
+class RemoteCallService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def CallFunction(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(request_iterator, target, '/chalk.runtime.v1.RemoteCallService/CallFunction',
+            chalk_dot_runtime_dot_v1_dot_remote__python__call__pb2.CallFunctionRequest.SerializeToString,
+            chalk_dot_runtime_dot_v1_dot_remote__python__call__pb2.CallFunctionResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

chalk_remote_call_python-0.0.0/chalk_remote_call/_native.pyi

@@ -0,0 +1,11 @@
+from collections.abc import Callable
+from typing import Any
+
+def start_server(
+    handler: Callable[..., Any],
+    process_fn: Callable[[bytes, Callable[..., Any], list[str] | None, dict[str, Any]], list[bytes]],
+    host: str,
+    port: int,
+    workers: int,
+    arg_names: list[str] | None = None,
+) -> None: ...

chalk_remote_call_python-0.0.0/chalk_remote_call/arrow_utils.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import pyarrow as pa
+import pyarrow.ipc as ipc
+
+
+def decode_ipc_stream(data: bytes) -> list[pa.RecordBatch]:
+    """Decode Arrow IPC stream bytes into a list of RecordBatches."""
+    return list(ipc.open_stream(data))
+
+
+def encode_record_batch(batch: pa.RecordBatch) -> bytes:
+    """Encode a RecordBatch as Arrow IPC stream bytes."""
+    sink = pa.BufferOutputStream()
+    with ipc.new_stream(sink, batch.schema) as writer:
+        writer.write_batch(batch)
+    return sink.getvalue().to_pybytes()

chalk_remote_call_python-0.0.0/chalk_remote_call/cli.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+import sys
+
+from chalk_remote_call.handler_loader import load_handler, load_startup_function
+from chalk_remote_call.input_transform import parse_input_args
+from chalk_remote_call.server import serve
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="chalk-remote-call",
+        description="Start a gRPC server implementing the Chalk RemoteCallService.",
+    )
+    parser.add_argument(
+        "--handler",
+        required=True,
+        help="Dotted path to the handler function (e.g. my_module.handler)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=int(os.environ.get("CHALK_REMOTE_CALL_PORT", "6666")),
+        help="Port to listen on (default: 6666, env: CHALK_REMOTE_CALL_PORT)",
+    )
+    parser.add_argument(
+        "--host",
+        default=os.environ.get("CHALK_REMOTE_CALL_HOST", "[::]"),
+        help="Host to bind to (default: [::], env: CHALK_REMOTE_CALL_HOST)",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=int(os.environ.get("CHALK_REMOTE_CALL_WORKERS", "10")),
+        help="Number of worker threads (default: 10, env: CHALK_REMOTE_CALL_WORKERS)",
+    )
+    parser.add_argument(
+        "--on-startup",
+        default=None,
+        help="Dotted path to a startup function (e.g. my_module.setup)",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Log level (default: INFO)",
+    )
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=getattr(logging, args.log_level),
+        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    )
+
+    # Load handler
+    try:
+        handler, auto_startup = load_handler(args.handler)
+    except Exception as e:
+        print(f"Error loading handler: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    # Determine startup function
+    on_startup = None
+    if args.on_startup:
+        try:
+            on_startup = load_startup_function(args.on_startup)
+        except Exception as e:
+            print(f"Error loading startup function: {e}", file=sys.stderr)
+            sys.exit(1)
+    elif auto_startup is not None:
+        on_startup = auto_startup
+
+    # Parse input args
+    arg_names = parse_input_args()
+
+    serve(
+        handler=handler,
+        host=args.host,
+        port=args.port,
+        workers=args.workers,
+        on_startup=on_startup,
+        arg_names=arg_names,
+    )

chalk_remote_call_python-0.0.0/chalk_remote_call/handler_loader.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import importlib
+from collections.abc import Callable
+from typing import Any
+
+
+def load_handler(dotted_path: str) -> tuple[Callable[..., Any], Callable[[], None] | None]:
+    """Load a handler function from a dotted module path.
+
+    Returns (handler_fn, on_startup_fn_or_None).
+    The on_startup function is auto-discovered if the handler's module defines
+    a top-level `on_startup` callable.
+    """
+    if "." not in dotted_path:
+        raise ValueError(f"Handler path must be a dotted path like 'my_module.handler', got: {dotted_path!r}")
+
+    module_path, func_name = dotted_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    handler = getattr(module, func_name, None)
+    if handler is None:
+        raise AttributeError(f"Module {module_path!r} has no attribute {func_name!r}")
+    if not callable(handler):
+        raise TypeError(f"{dotted_path!r} is not callable")
+
+    on_startup = getattr(module, "on_startup", None)
+    if on_startup is not None and not callable(on_startup):
+        on_startup = None
+
+    return handler, on_startup
+
+
+def load_startup_function(dotted_path: str) -> Callable[[], None]:
+    """Load a startup function from a dotted module path."""
+    if "." not in dotted_path:
+        raise ValueError(f"Startup path must be a dotted path like 'my_module.setup', got: {dotted_path!r}")
+
+    module_path, func_name = dotted_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    func = getattr(module, func_name, None)
+    if func is None:
+        raise AttributeError(f"Module {module_path!r} has no attribute {func_name!r}")
+    if not callable(func):
+        raise TypeError(f"{dotted_path!r} is not callable")
+
+    return func

chalk_remote_call_python-0.0.0/chalk_remote_call/input_transform.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import os
+
+import pyarrow as pa
+
+
+def parse_input_args() -> list[str] | None:
+    """Parse CHALK_INPUT_ARGS env var into a list of column names.
+
+    Returns None if the env var is unset or empty.
+    """
+    raw = os.environ.get("CHALK_INPUT_ARGS", "").strip()
+    if not raw:
+        return None
+    return [name.strip() for name in raw.split(",") if name.strip()]
+
+
+def transform(batch: pa.RecordBatch, arg_names: list[str] | None) -> dict[str, pa.Array]:
+    """Transform a RecordBatch into a dict of named arrays.
+
+    If arg_names is provided, columns are renamed by index to the given names
+    (ignoring the original column names). If arg_names is None, the original
+    column names are used.
+    """
+    if arg_names is not None:
+        if len(arg_names) != batch.num_columns:
+            raise ValueError(
+                f"CHALK_INPUT_ARGS specifies {len(arg_names)} names but RecordBatch has {batch.num_columns} columns"
+            )
+        return {name: batch.column(i) for i, name in enumerate(arg_names)}
+    return {batch.schema.field(i).name: batch.column(i) for i in range(batch.num_columns)}