diff --git a/.agents/server-map.md b/.agents/server-map.md index 9c9b22d7..8c6c97b6 100644 --- a/.agents/server-map.md +++ b/.agents/server-map.md @@ -52,6 +52,50 @@ Cite file and class or function names in your answers. Do not cite line numbers, `src/DataTypes/DataTypeDecimalBase.{cpp,h}`, `src/DataTypes/DataTypesDecimal.h`, `src/DataTypes/Serializations/SerializationDecimal.{cpp,h}`. Stored as fixed-width two's-complement integer (32/64/128/256-bit) with precision and scale carried in the type string. +### QBit + +- **Confirmed at v26.3.9.8-lts:** The type implementation is + `src/DataTypes/DataTypeQBit.{cpp,h}` in `DataTypeQBit`. + `DataTypeQBit::createColumn` creates a `ColumnQBit` backed by a tuple containing + 16, 32, or 64 `FixedString(ceil(dimension / 8))` columns. + `DataTypeQBit::doGetSerialization` returns `SerializationQBit`. The column + wrapper and bit-plane storage accessors are in `src/Columns/ColumnQBit.{cpp,h}`. +- **Confirmed at v26.3.9.8-lts:** Binary and Native serialization are implemented + by `SerializationQBit` in + `src/DataTypes/Serializations/SerializationQBit.{cpp,h}`. Row-wise binary uses + `serializeBinary`, `deserializeBinary`, `serializeFloatsFromQBit`, + `deserializeFloatsToQBit`, `transposeBits`, and `untransposeBitPlane`. Its wire + value is a VarUInt dimension followed by that many BFloat16, Float32, or Float64 + values. +- **Confirmed at v26.3.9.8-lts:** Native uses + `SerializationQBit::serializeBinaryBulkWithMultipleStreams` and + `deserializeBinaryBulkWithMultipleStreams`, reached through + `NativeWriter::writeData` and `NativeReader::readData` in + `src/Formats/NativeWriter.cpp` and `src/Formats/NativeReader.cpp`. These methods + delegate to the nested `SerializationTuple` implementation in + `src/DataTypes/Serializations/SerializationTuple.cpp`; each bit plane is emitted + as raw fixed-width bytes by `SerializationFixedString::serializeBinaryBulk` in + `src/DataTypes/Serializations/SerializationFixedString.cpp`. Native therefore + carries bit-transposed planes, not row-wise float arrays. +- **Confirmed at v26.3.9.8-lts:** The binary type-schema encoding used by Dynamic + and related self-describing serializations is in + `src/DataTypes/DataTypesBinaryEncoding.{cpp,h}`. QBit uses type tag `0x36`, + followed by the encoded element type and a VarUInt dimension. +- **Confirmed at v26.3.9.8-lts:** Focused tests are + `src/DataTypes/Serializations/tests/gtest_qbit_serialization.cpp` + (`QBitSerialization.FieldBinarySerializationFloat32` and + `QBitSerialization.RejectInvalidElementType`) and + `tests/queries/0_stateless/03371_qbit_read_write.{sh,reference}`, which exercises + text, RowBinary, and Native round trips. +- **Confirmed at v26.3.9.8-lts:** Additional coverage is in + `tests/queries/0_stateless/03363_qbit_create_insert_select.{sql,reference}` for all + three element widths and byte-aligned/non-byte-aligned dimensions, + `03368_qbit_subcolumns.{sql,reference}` for bit-plane subcolumns, + `03372_qbit_mergetree_1.{sql,reference}` and + `03372_qbit_mergetree_2.{sql,reference}` for stored data, + `03373_qbit_dynamic.{sql,reference}` for Dynamic, and + `03374_qbit_nullable.{sql,reference}` for Nullable. + ### Enum8 / Enum16 `src/DataTypes/DataTypeEnum.{cpp,h}`, `src/DataTypes/EnumValues.{cpp,h}`, `src/DataTypes/Serializations/SerializationEnum.{cpp,h}`. Name-to-value map lives in the type string. Wire format is the underlying Int8/Int16. diff --git a/.github/workflows/binding_ci.yml b/.github/workflows/binding_ci.yml new file mode 100644 index 00000000..cdb8ef3d --- /dev/null +++ b/.github/workflows/binding_ci.yml @@ -0,0 +1,91 @@ +name: 'Rust Binding CI' + +on: + pull_request: + branches: + - main + paths: + - 'rust/**' + - 'clickhouse_connect/driver/rustcodec.py' + - 'clickhouse_connect/driver/rustnumpy.py' + - 'tests/integration_tests/test_rust_codec.py' + - '.github/workflows/binding_ci.yml' + push: + branches: + - main + paths: + - 'rust/**' + - 'clickhouse_connect/driver/rustcodec.py' + - 'clickhouse_connect/driver/rustnumpy.py' + - 'tests/integration_tests/test_rust_codec.py' + - '.github/workflows/binding_ci.yml' + workflow_dispatch: + +jobs: + rust-checks: + runs-on: ubuntu-latest + name: Cargo fmt + clippy + steps: + - uses: actions/checkout@v6 + # pinned so new stable clippy lints don't break unrelated PRs, bump deliberately + - uses: dtolnay/rust-toolchain@1.97.0 + with: + components: rustfmt, clippy + - uses: Swatinem/rust-cache@v2 + with: + workspaces: rust + - name: Format check + working-directory: rust + run: cargo fmt --all -- --check + - name: Clippy (deny warnings) + working-directory: rust + run: cargo clippy --all-targets -- -D warnings + + binding-tests: + runs-on: ubuntu-latest + name: Build binding and run codec tests + needs: rust-checks + steps: + - name: Checkout + uses: actions/checkout@v6 + - name: Start ClickHouse (latest) in Docker + env: + CLICKHOUSE_CONNECT_TEST_CH_VERSION: latest + COMPOSE_PROJECT_NAME: clickhouse-connect-binding-ci + run: docker compose -f docker-compose.yml up -d --wait clickhouse + - name: Set up Python 3.12 + uses: actions/setup-python@v6 + with: + python-version: '3.12' + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + with: + workspaces: rust + - name: Install Test Dependencies + run: | + python -m pip install --upgrade pip + pip install -r tests/test_requirements.txt + pip install maturin + - name: Build cython extensions + run: python setup.py build_ext --inplace + - name: "Add distribution info" # This lets SQLAlchemy find entry points + run: python setup.py develop + - name: Build and install the clickhouse-connect-core wheel + run: | + maturin build --release -m rust/ch-core-py/Cargo.toml -o core-wheelhouse + pip install core-wheelhouse/*.whl + - name: Extension smoke check + run: python -c "import _ch_core; print(_ch_core.__version__, _ch_core.BINDING_API_VERSION)" + - name: Run binding tests + run: pytest rust/ch-core-py/tests -n 4 + - name: Run unit tests + run: pytest tests/unit_tests -n 4 + - name: Run rust codec integration tests + env: + CLICKHOUSE_CONNECT_TEST_DOCKER: 'False' + run: pytest tests/integration_tests/test_rust_codec.py -n 4 + - name: Stop ClickHouse + if: ${{ always() && hashFiles('docker-compose.yml') != '' }} + env: + COMPOSE_PROJECT_NAME: clickhouse-connect-binding-ci + run: docker compose -f docker-compose.yml down --volumes --remove-orphans diff --git a/.github/workflows/on_push.yml b/.github/workflows/on_push.yml index 3a5c0966..47913eb0 100644 --- a/.github/workflows/on_push.yml +++ b/.github/workflows/on_push.yml @@ -87,8 +87,8 @@ jobs: "$RUNNER_TEMP/consumer/bin/mypy" --strict --follow-imports=silent consumer_smoke.py - name: Public type completeness ratchet # SQLAlchemy must be installed while verifytypes scans the public cc_sqlalchemy annotations. - # The pre-feature baseline is 1120 with that optional surface resolved. - run: python scripts/check_public_types.py clickhouse_connect --max-untyped 1120 --python "$RUNNER_TEMP/consumer/bin/python" + # The merged baseline is 1118 with that optional surface resolved. + run: python scripts/check_public_types.py clickhouse_connect --max-untyped 1118 --python "$RUNNER_TEMP/consumer/bin/python" - name: SQLAlchemy Select typing smoke test run: | cp tests/type_check/sqlalchemy_select_smoke.py "$RUNNER_TEMP" diff --git a/.github/workflows/publish_core.yml b/.github/workflows/publish_core.yml index c1944a72..afbbdca6 100644 --- a/.github/workflows/publish_core.yml +++ b/.github/workflows/publish_core.yml @@ -28,9 +28,9 @@ jobs: fail-fast: false matrix: include: - - { runner: ubuntu-latest, arch: x86_64, libc: manylinux, manylinux: auto } + - { runner: ubuntu-latest, arch: x86_64, libc: manylinux, manylinux: 2_28 } - { runner: ubuntu-latest, arch: x86_64, libc: musllinux, manylinux: musllinux_1_2 } - - { runner: ubuntu-24.04-arm, arch: aarch64, libc: manylinux, manylinux: auto } + - { runner: ubuntu-24.04-arm, arch: aarch64, libc: manylinux, manylinux: 2_28 } - { runner: ubuntu-24.04-arm, arch: aarch64, libc: musllinux, manylinux: musllinux_1_2 } steps: - uses: actions/checkout@v6 @@ -38,7 +38,7 @@ jobs: uses: PyO3/maturin-action@v1 with: manylinux: ${{ matrix.manylinux }} - args: --release -m ${{ env.MATURIN_MANIFEST }} -o wheelhouse --find-interpreter + args: --release -m ${{ env.MATURIN_MANIFEST }} -o wheelhouse -i 3.10 3.11 3.12 3.13 3.14 - uses: actions/upload-artifact@v7 with: name: core-${{ matrix.libc }}-${{ matrix.arch }} @@ -127,7 +127,7 @@ jobs: uses: PyO3/maturin-action@v1 with: target: aarch64-pc-windows-msvc - args: --release -m ${{ env.MATURIN_MANIFEST }} -o wheelhouse + args: --release -m ${{ env.MATURIN_MANIFEST }} -o wheelhouse -i python3.10 - uses: actions/upload-artifact@v7 with: name: core-windows-arm64-3.10 diff --git a/.gitignore b/.gitignore index 811c5635..37a50307 100644 --- a/.gitignore +++ b/.gitignore @@ -32,3 +32,10 @@ test.env # Local ClickHouse server source checkout for agent reference (see AGENTS.md) /.server-src/ .server-ref + +# Rust build artifacts +rust/target/ +*.so +*.dylib +/.cargo/config.toml +.claude diff --git a/CHANGELOG.md b/CHANGELOG.md index d0924ce9..c7a67c2c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,16 @@ ## UNRELEASED +## 1.8.0rc2, 2026-08-20 + +Follow-up release candidate to 1.8.0rc1, rebased on 1.7.2 so all bug fixes from that stable release are included. The optional Rust codec itself is unchanged. + +## 1.8.0rc1, 2026-08-12 + +### Improvements + +- Added an experimental `native_codec` client option that selects the codec for FORMAT Native query decode and insert encode. `python` is the default and uses the existing codec. `rust` prefers the compiled Rust codec and falls back to the Python codec for unsupported options and types, while `rust_strict` raises instead of falling back. Results and dtypes match the Python codec, and the Arrow methods are unaffected. The compiled codec ships as the separate clickhouse-connect-core wheel, installed with `pip install clickhouse-connect[rust]`. See the rust-codec documentation page for details. This is early access for benchmarking and is not yet a supported path. + ## 1.7.2, 2026-08-19 ### Bug Fixes @@ -63,6 +73,8 @@ - `AsyncClient` initialization no longer overwrites user-supplied session settings with generated defaults. A client created with `settings={'date_time_input_format': 'basic'}` previously had that value replaced by the generated `best_effort` default. User settings now always win, matching the sync client. - An `AsyncClient` created with both client certificates and an access token now sends the mutual TLS authentication headers and the `Authorization: Bearer` header together, matching the sync client. The certificates previously suppressed the token at construction, while the `token_provider` option re-added its token right after initialization, so the two async token paths disagreed with each other. The server resolves the credential precedence. - Dict-valued settings such as `additional_table_filters` no longer crash with `DB::Exception: Cannot parse quoted string` when passed through `query()`'s `settings` parameter. The value was rendered with Python's own `str()`/`repr()` of the dict, which mixes single and double quotes and is not valid ClickHouse map-literal syntax; it is now rendered as a properly single-quoted, escaped ClickHouse map literal. Closes [#501](https://github.com/ClickHouse/clickhouse-connect/issues/501). +- Explicit NaN and infinity values in nullable `BFloat16` row inserts are now stored instead of being written as 0. +- Inserting an `Array(Dynamic)` column no longer raises `ZeroDivisionError` when a sampled row holds an empty array. The insert block size estimate now treats an empty sample as minimal instead of dividing by its length. ### Improvements - Async clients now emit URL query parameters in the same order as the sync client on every request. The parameter names and values are unchanged, so this is only visible to systems that match or sign the exact request URL. diff --git a/DISTRIBUTION_PLAN.md b/DISTRIBUTION_PLAN.md new file mode 100644 index 00000000..ec22d994 --- /dev/null +++ b/DISTRIBUTION_PLAN.md @@ -0,0 +1,71 @@ +# Rust Codec Distribution Plan + +How the Rust native codec is packaged, versioned, released, and supported across ch-core-rs, the Python binding, and clickhouse-connect. + +## Decisions + +- Users opt in with a single switch on the client. The Rust path uses the new, correct semantics. Known behavior differences from the Python codec are enumerated and documented. Features the Rust path does not implement fall back to the Python codec. +- ch-core-rs is open source on GitHub. It is not published to crates.io. Consumers depend on it as a git dependency pinned to a release tag. The repo is currently internal. Making it public is in progress and gates the first PyPI release. +- The core repo stays pure Rust with no language-specific dependencies. Each language client owns its binding in its own repo. +- The Python binding ships as a separate PyPI wheel, installed via the `rust` extra. It is not compiled into the clickhouse-connect wheel. +- Long term: deprecation warnings on the Python codec late in 1.x, Rust becomes the only codec in 2.0. + +## Artifacts and repos + +| Piece | Lives in | Ships as | Releases when | +|---|---|---|---| +| ch-core-rs | own public repo, pure Rust | git tags only | core logic changes | +| ch-core-py binding and the rustcodec.py seam | clickhouse-connect repo | `clickhouse-connect-core` wheel on PyPI, module name `_ch_core` | core repin or binding change | +| clickhouse-connect | clickhouse-connect repo | pure/cython package with a `rust` extra | driver features, seam or floor changes | + +One repo can publish two PyPI artifacts. Living in the connect repo does not mean shipping in the connect wheel. + +## The two seams + +**Rust seam (ch-core-rs into the binding).** Compile time only. The core is statically linked into the extension when the wheel builds. Users never resolve it and no runtime skew is possible. Pin by git tag in ch-core-py's Cargo.toml and commit Cargo.lock for reproducible builds. + +**Python seam (`clickhouse-connect-core` wheel and clickhouse-connect).** The only seam users see. clickhouse-connect pins a compatible range, for example `clickhouse-connect-core>=1.2,<1.3`. The pin floor encodes the oldest core wheel this driver knows how to drive. Bump the floor only when connect starts using a new binding capability. + +## Packaging + +- PyPI name `clickhouse-connect-core`, module name `_ch_core`. The name should be recognizable in a user's pip list. +- Install: `pip install clickhouse-connect[rust]`. +- Build with maturin, one wheel per platform per Python version (cp310 through cp314), matching the main package's matrix. +- Not abi3: the binding's hot paths use non-limited C API (`PyTuple_SET_ITEM`, `PyList_SET_ITEM`, presized dict construction) on purpose. Moving to the limited API would tax exactly the paths the codec exists to accelerate. Revisit only if the wheel matrix becomes a real maintenance cost. +- If the switch is enabled without the wheel installed, raise a clear error naming the install command. +- While ch-core-rs remains non-public, wheel builds need repo access from CI and the sdist cannot build for outside users. Making the repo public precedes the first PyPI release. + +## Runtime handshake + +- `_ch_core` exports a binding API version constant. +- rustcodec.py checks it at import. Too old raises a legible message naming the required `clickhouse-connect-core` version. Never a crash or silent misbehavior. +- The client's diagnostic output includes the `_ch_core` version alongside the driver version so bug reports arrive with both. + +## Release workflows + +**Core bugfix.** Fix in ch-core-rs, tag a patch release. In the connect repo bump the git tag in ch-core-py, publish a `clickhouse-connect-core` patch wheel. No clickhouse-connect release. Users run `pip install -U clickhouse-connect[rust]`. + +**Transparent core improvement** (faster decode, internal wins). Same as a bugfix but a minor bump. Users get it for free with a wheel upgrade. + +**User-facing core feature** (new setting, new type, new capability). Core minor bump, binding exposes it, `clickhouse-connect-core` minor release. Then a clickhouse-connect release that uses it and raises the pin floor. The connect release is the feature's public API. + +**Driver-only change.** Normal clickhouse-connect release. The wheel is untouched. + +## Issue handling + +- Users file everything on clickhouse-connect. The core repo tracker is for maintainers and binding authors. +- If a root cause lands in ch-core-rs, fix it there but keep and close the loop in the original clickhouse-connect issue, noting the `clickhouse-connect-core` version that carries the fix. + +## Core repo contract + +- Semver on tags. No breaking changes on patch or minor. +- A changelog maintained per release. Downstream bindings in multiple languages will depend on reading it. +- No PyO3, napi, or other language-specific dependencies in the core crate. Language artifacts are built by the binding repos. +- CI tests the crate on supported platforms. It builds no wheels or language artifacts. + +## Rollout + +1. Next minor: ship the opt-in switch, the `rust` extra, the known-differences documentation, and the fallback behavior. +2. During 1.x: promote the Rust path as it proves out. Wheel-only releases carry core fixes and wins to opted-in users. +3. Late 1.x: DeprecationWarning on the Python codec path. +4. 2.0: Rust codec becomes the only codec. The differences list becomes the documented behavior. diff --git a/clickhouse_connect/_version.py b/clickhouse_connect/_version.py index 2196826f..76a9c93f 100644 --- a/clickhouse_connect/_version.py +++ b/clickhouse_connect/_version.py @@ -1 +1 @@ -version = "1.7.2" +version = "1.8.0rc2" diff --git a/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py b/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py index 33c3a580..5460b6b7 100644 --- a/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py +++ b/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py @@ -259,6 +259,10 @@ class MultiLineString(ChSqlaType, UserDefinedType): # type: ignore[misc] python_type = list +class Geometry(ChSqlaType, UserDefinedType): # type: ignore[misc] + python_type = object + + class Date(ChSqlaType, SqlaDate): # type: ignore[misc] pass diff --git a/clickhouse_connect/common.py b/clickhouse_connect/common.py index a6361bc3..15560829 100644 --- a/clickhouse_connect/common.py +++ b/clickhouse_connect/common.py @@ -1,4 +1,6 @@ import getpass +import logging +import os import sys from collections.abc import Sequence from dataclasses import dataclass @@ -7,6 +9,10 @@ from clickhouse_connect._version import version as _version_string from clickhouse_connect.driver.exceptions import ProgrammingError +logger: logging.Logger = logging.getLogger(__name__) + +_NATIVE_CODEC_OPTIONS = ("python", "rust", "rust_strict") + def version() -> str: return _version_string @@ -68,6 +74,17 @@ def _init_common(name: str, options: Sequence[Any], default: Any) -> None: _common_settings[name] = CommonSetting(name, options, default) +def _native_codec_env_default() -> str: + raw = os.environ.get("CLICKHOUSE_CONNECT_NATIVE_CODEC") + if raw is None: + return "python" + value = raw.strip().lower() + if value in _NATIVE_CODEC_OPTIONS: + return value + logger.warning("Ignoring invalid CLICKHOUSE_CONNECT_NATIVE_CODEC=%r; using 'python'", raw) + return "python" + + _init_common("autogenerate_session_id", (True, False), True) _init_common("autogenerate_query_id", (True, False), True) _init_common("dict_parameter_format", ("json", "map"), "json") @@ -79,6 +96,10 @@ def _init_common(name: str, options: Sequence[Any], default: Any) -> None: _init_common("readonly", (0, 1), 0) # Deprecated no-op retained for 1.x compatibility _init_common("send_os_user", (True, False), True) +# Selects the codec for client-managed FORMAT Native query and insert paths. Seeded by +# CLICKHOUSE_CONNECT_NATIVE_CODEC; overridable per client via the native_codec kwarg. +_init_common("native_codec", _NATIVE_CODEC_OPTIONS, _native_codec_env_default()) + # Include integration tags (library name/version) in the User-Agent, e.g.: # pandas/2.2.5; polars/0.20.x; sqlalchemy/2.0.x. These tags are only included # when using relevant API methods. diff --git a/clickhouse_connect/datatypes/base.py b/clickhouse_connect/datatypes/base.py index ab5e1774..2be9aa83 100644 --- a/clickhouse_connect/datatypes/base.py +++ b/clickhouse_connect/datatypes/base.py @@ -123,6 +123,8 @@ def data_size(self, sample: Collection) -> int: def _data_size(self, sample: Collection) -> int: if self.byte_size: return self.byte_size + if len(sample) == 0: + return 1 total = 0 for x in sample: total += len(str(x)) diff --git a/clickhouse_connect/datatypes/container.py b/clickhouse_connect/datatypes/container.py index 951e618d..424b08fd 100644 --- a/clickhouse_connect/datatypes/container.py +++ b/clickhouse_connect/datatypes/container.py @@ -8,6 +8,7 @@ from clickhouse_connect.driver.binding import quote_identifier from clickhouse_connect.driver.common import first_value, must_swap from clickhouse_connect.driver.ctypes import data_conv +from clickhouse_connect.driver.exceptions import DataError from clickhouse_connect.driver.insert import InsertContext from clickhouse_connect.driver.query import QueryContext from clickhouse_connect.driver.types import ByteSource @@ -116,17 +117,18 @@ def __init__(self, type_def: TypeDef): self._insert_name = f"Tuple({', '.join(v.insert_name for v in self.element_types)})" def _data_size(self, sample: Collection) -> int: - if len(sample) == 0: + rows = [x for x in sample if x is not None] if self.nullable else list(sample) + if len(rows) == 0: return 0 elem_size = 0 - is_dict = self.element_names and isinstance(first_value(list(sample), self.nullable), dict) + is_dict = self.element_names and isinstance(first_value(rows, self.nullable), dict) for ix, e_type in enumerate(self.element_types): if e_type.byte_size > 0: elem_size += e_type.byte_size elif is_dict: - elem_size += e_type.data_size([x.get(self.element_names[ix], None) for x in sample]) + elem_size += e_type.data_size([x.get(self.element_names[ix], None) for x in rows]) else: - elem_size += e_type.data_size([x[ix] for x in sample]) + elem_size += e_type.data_size([x[ix] for x in rows]) return elem_size def read_column_prefix(self, source: ByteSource, ctx: QueryContext): @@ -153,11 +155,27 @@ def write_column_prefix(self, dest: bytearray): for e_type in self.element_types: e_type.write_column_prefix(dest) + def _row_length_error(self, column: Sequence) -> DataError: + expected = len(self.element_types) + for row_number, row in enumerate(column): + try: + actual = len(row) + except TypeError: + return DataError(f"{self.name} expects a sequence with {expected} elements at row {row_number}") + if actual != expected: + return DataError(f"{self.name} expects {expected} elements, got {actual} at row {row_number}") + return DataError(f"{self.name} rows have inconsistent lengths") + def write_column_data(self, column: Sequence, dest: bytearray, ctx: InsertContext): if self.element_names and isinstance(first_value(column, self.nullable), dict): columns = self.convert_dict_insert(column) else: - columns = list(zip(*column)) + try: + columns = list(zip(*column, strict=True)) + except ValueError: + raise self._row_length_error(column) from None + if len(column) and len(columns) != len(self.element_types): + raise self._row_length_error(column) for e_type, elem_column in zip(self.element_types, columns): e_type.write_column_data(elem_column, dest, ctx) diff --git a/clickhouse_connect/datatypes/geometric.py b/clickhouse_connect/datatypes/geometric.py index f7773dd3..e0c7a21d 100644 --- a/clickhouse_connect/datatypes/geometric.py +++ b/clickhouse_connect/datatypes/geometric.py @@ -1,4 +1,4 @@ -from collections.abc import Sequence +from collections.abc import Collection, Sequence from typing import Any from clickhouse_connect.datatypes.base import ClickHouseType @@ -10,14 +10,21 @@ RING_DATA_TYPE: ClickHouseType POLYGON_DATA_TYPE: ClickHouseType MULTI_POLYGON_DATA_TYPE: ClickHouseType +GEOMETRY_DATA_TYPE: ClickHouseType # ruff: noqa: F821 (Undefine name) class Point(ClickHouseType): + def _data_size(self, sample: Collection) -> int: + return POINT_DATA_TYPE._data_size(sample) + def write_column(self, column: Sequence, dest: bytearray, ctx: InsertContext): return POINT_DATA_TYPE.write_column(column, dest, ctx) + def write_column_data(self, column: Sequence, dest: bytearray, ctx: InsertContext): + return POINT_DATA_TYPE.write_column_data(column, dest, ctx) + def read_column_prefix(self, source: ByteSource, ctx: QueryContext): return POINT_DATA_TYPE.read_column_prefix(source, ctx) @@ -26,9 +33,15 @@ def read_column_data(self, source: ByteSource, num_rows: int, ctx: QueryContext, class Ring(ClickHouseType): + def _data_size(self, sample: Collection) -> int: + return RING_DATA_TYPE._data_size(sample) + def write_column(self, column: Sequence, dest: bytearray, ctx: InsertContext): return RING_DATA_TYPE.write_column(column, dest, ctx) + def write_column_data(self, column: Sequence, dest: bytearray, ctx: InsertContext): + return RING_DATA_TYPE.write_column_data(column, dest, ctx) + def read_column_prefix(self, source: ByteSource, ctx: QueryContext): return RING_DATA_TYPE.read_column_prefix(source, ctx) @@ -37,9 +50,15 @@ def read_column_data(self, source: ByteSource, num_rows: int, ctx: QueryContext, class Polygon(ClickHouseType): + def _data_size(self, sample: Collection) -> int: + return POLYGON_DATA_TYPE._data_size(sample) + def write_column(self, column: Sequence, dest: bytearray, ctx: InsertContext): return POLYGON_DATA_TYPE.write_column(column, dest, ctx) + def write_column_data(self, column: Sequence, dest: bytearray, ctx: InsertContext): + return POLYGON_DATA_TYPE.write_column_data(column, dest, ctx) + def read_column_prefix(self, source: ByteSource, ctx: QueryContext): return POLYGON_DATA_TYPE.read_column_prefix(source, ctx) @@ -48,9 +67,15 @@ def read_column_data(self, source: ByteSource, num_rows: int, ctx: QueryContext, class MultiPolygon(ClickHouseType): + def _data_size(self, sample: Collection) -> int: + return MULTI_POLYGON_DATA_TYPE._data_size(sample) + def write_column(self, column: Sequence, dest: bytearray, ctx: InsertContext): return MULTI_POLYGON_DATA_TYPE.write_column(column, dest, ctx) + def write_column_data(self, column: Sequence, dest: bytearray, ctx: InsertContext): + return MULTI_POLYGON_DATA_TYPE.write_column_data(column, dest, ctx) + def read_column_prefix(self, source: ByteSource, ctx: QueryContext): return MULTI_POLYGON_DATA_TYPE.read_column_prefix(source, ctx) @@ -64,3 +89,22 @@ class LineString(Ring): class MultiLineString(Polygon): pass + + +class Geometry(ClickHouseType): + """ClickHouse Geometry, a fixed Variant over the six geo types.""" + + def data_size(self, sample: Collection[Any]) -> int: + return GEOMETRY_DATA_TYPE.data_size(sample) + + def write_column_prefix(self, dest: bytearray) -> None: + return GEOMETRY_DATA_TYPE.write_column_prefix(dest) + + def write_column_data(self, column: Sequence, dest: bytearray, ctx: InsertContext): + return GEOMETRY_DATA_TYPE.write_column_data(column, dest, ctx) + + def read_column_prefix(self, source: ByteSource, ctx: QueryContext): + return GEOMETRY_DATA_TYPE.read_column_prefix(source, ctx) + + def read_column_data(self, source: ByteSource, num_rows: int, ctx: QueryContext, read_state: Any) -> Sequence: + return GEOMETRY_DATA_TYPE.read_column_data(source, num_rows, ctx, read_state) diff --git a/clickhouse_connect/datatypes/numeric.py b/clickhouse_connect/datatypes/numeric.py index 90d2abea..a414103a 100644 --- a/clickhouse_connect/datatypes/numeric.py +++ b/clickhouse_connect/datatypes/numeric.py @@ -236,7 +236,7 @@ def _write_column_binary( if self.nullable: first = next((x for x in column if x is not None), None) if isinstance(first, float): - column = [0 if (x is None or isnan(x) or isinf(x)) else x for x in column] + column = [0 if x is None else x for x in column] else: column = [0 if x is None else float(x) for x in column] elif not isinstance(column[0], float): @@ -457,45 +457,57 @@ class Decimal256(BigDecimal): dec_size = 256 -class IntervalNanosecond(Int32): +class Interval(Int64, registered=False): + def _finalize_column(self, column: Sequence, ctx: QueryContext) -> Sequence: + if self.read_format(ctx) == "string": + return [str(x) for x in column] + if ctx.use_extended_dtypes and self.nullable: + # base_type is the interval type name, which is not a pandas dtype + return options.pd.array(column, dtype="Int64") + if ctx.use_numpy and self.nullable and (not ctx.use_none): + return options.np.array(column, dtype=self.np_type) + return column + + +class IntervalNanosecond(Interval): pass -class IntervalMicrosecond(Int32): +class IntervalMicrosecond(Interval): pass -class IntervalMillisecond(Int32): +class IntervalMillisecond(Interval): pass -class IntervalSecond(Int32): +class IntervalSecond(Interval): pass -class IntervalMinute(Int32): +class IntervalMinute(Interval): pass -class IntervalHour(Int32): +class IntervalHour(Interval): pass -class IntervalDay(Int32): +class IntervalDay(Interval): pass -class IntervalWeek(Int32): +class IntervalWeek(Interval): pass -class IntervalMonth(Int32): +class IntervalMonth(Interval): pass -class IntervalQuarter(Int32): +class IntervalQuarter(Interval): pass -class IntervalYear(Int32): +class IntervalYear(Interval): pass diff --git a/clickhouse_connect/datatypes/postinit.py b/clickhouse_connect/datatypes/postinit.py index 0f41d678..0ea22215 100644 --- a/clickhouse_connect/datatypes/postinit.py +++ b/clickhouse_connect/datatypes/postinit.py @@ -18,8 +18,10 @@ ring = f"Array({point})" polygon = f"Array({ring})" multi_polygon = f"Array({polygon})" +geometry = "Variant(LineString, MultiLineString, MultiPolygon, Point, Polygon, Ring)" geometric.POINT_DATA_TYPE = registry.get_from_name(point) geometric.RING_DATA_TYPE = registry.get_from_name(ring) geometric.POLYGON_DATA_TYPE = registry.get_from_name(polygon) geometric.MULTI_POLYGON_DATA_TYPE = registry.get_from_name(multi_polygon) +geometric.GEOMETRY_DATA_TYPE = registry.get_from_name(geometry) diff --git a/clickhouse_connect/datatypes/registry.py b/clickhouse_connect/datatypes/registry.py index cb035782..5d398363 100644 --- a/clickhouse_connect/datatypes/registry.py +++ b/clickhouse_connect/datatypes/registry.py @@ -42,6 +42,8 @@ def parse_name(name: str) -> tuple[str, str, TypeDef]: elif base.startswith("JSON") and len(base) > 4 and base[4] == "(": keys, values = parse_columns(base[4:]) base = "JSON" + elif base == "GEOMETRY": + base = "Geometry" elif base == "Point": values = ("Float64", "Float64") else: diff --git a/clickhouse_connect/datatypes/special.py b/clickhouse_connect/datatypes/special.py index f5f5df9f..bbcb2ade 100644 --- a/clickhouse_connect/datatypes/special.py +++ b/clickhouse_connect/datatypes/special.py @@ -4,6 +4,7 @@ from clickhouse_connect.datatypes.base import ArrayType, ClickHouseType, TypeDef, UnsupportedType from clickhouse_connect.datatypes.registry import get_from_name +from clickhouse_connect.driver import options from clickhouse_connect.driver.common import first_value from clickhouse_connect.driver.ctypes import data_conv from clickhouse_connect.driver.insert import InsertContext @@ -78,6 +79,12 @@ def __init__(self, type_def: TypeDef): super().__init__(type_def) self.nullable = True + def _finalize_column(self, column: Sequence, ctx: QueryContext) -> Sequence: + # base_type is not a pandas dtype; every value is null regardless of token + if ctx.use_extended_dtypes: + return options.np.full(len(column), None, dtype="object") + return super()._finalize_column(column, ctx) + def _write_column_binary(self, column: Sequence | MutableSequence, dest: bytearray, _ctx): dest += bytes(0x30 for _ in range(len(column))) diff --git a/clickhouse_connect/datatypes/temporal.py b/clickhouse_connect/datatypes/temporal.py index 643d4676..b942f911 100644 --- a/clickhouse_connect/datatypes/temporal.py +++ b/clickhouse_connect/datatypes/temporal.py @@ -274,13 +274,17 @@ def _read_binary_tz(self, column: Sequence, tz_info: tzinfo): def _read_binary_naive(self, column: Sequence): return data_conv.read_datetime64_naive_col(column, self.prec) + def _datetime64_ticks(self, value: datetime, active_tz: tzinfo | None = None) -> int: + timestamp = _localized_timestamp(value, active_tz) if active_tz is not None else value.timestamp() + seconds = floor(timestamp) + return ((seconds * 1000000 + value.microsecond) * self.prec) // 1000000 + def _write_column_binary(self, column: Sequence | MutableSequence, dest: bytearray, ctx: InsertContext): first = first_value(column, self.nullable) if isinstance(first, int) or self.write_format(ctx) == "int": if self.nullable: column = [x if x else 0 for x in column] else: - prec = self.prec server_mode = common.get_setting("naive_datetime_insert") == "server" active_tz = (self.tzinfo or ctx.server_tz) if server_mode else None if isinstance(first, str): @@ -292,22 +296,13 @@ def _write_column_binary(self, column: Sequence | MutableSequence, dest: bytearr v = 0 else: dt = datetime.fromisoformat(x) - timestamp = _localized_timestamp(dt, active_tz) if active_tz is not None else dt.timestamp() - v = ((floor(timestamp) * 1000000 + dt.microsecond) * prec) // 1000000 + v = self._datetime64_ticks(dt, active_tz) column.append(v) - elif active_tz is not None: - if self.nullable: - column = [ - ((floor(_localized_timestamp(x, active_tz)) * 1000000 + x.microsecond) * prec) // 1000000 if x else 0 - for x in column - ] - else: - column = [((floor(_localized_timestamp(x, active_tz)) * 1000000 + x.microsecond) * prec) // 1000000 for x in column] elif self.nullable: - column = [((floor(x.timestamp()) * 1000000 + x.microsecond) * prec) // 1000000 if x else 0 for x in column] + column = [self._datetime64_ticks(x, active_tz) if x else 0 for x in column] else: - column = [((floor(x.timestamp()) * 1000000 + x.microsecond) * prec) // 1000000 for x in column] + column = [self._datetime64_ticks(x, active_tz) for x in column] write_array("q", column, dest, ctx.column_name) @@ -354,6 +349,19 @@ class TimeBase(ClickHouseType, registered=False): _MICROS_PER_SECOND = 1_000_000 _NANOS_PER_SECOND = 1_000_000_000 _SECONDS_PER_DAY = 86_400 + _NUMPY_TIME_UNITS = { + "W": (604_800, 1), + "D": (86_400, 1), + "h": (3_600, 1), + "m": (60, 1), + "s": (1, 1), + "ms": (1, 1_000), + "us": (1, 1_000_000), + "ns": (1, 1_000_000_000), + "ps": (1, 1_000_000_000_000), + "fs": (1, 1_000_000_000_000_000), + "as": (1, 1_000_000_000_000_000_000), + } _array_type: str byte_size: int @@ -469,6 +477,45 @@ def _numerical_to_ticks(self, value: int | float | numpy.int64) -> int: self._validate_standard_range(value, value) return value + @classmethod + def _numpy_timedelta_to_ticks(cls, td: numpy.timedelta64, precision: int) -> int: + """Rescale a NumPy duration with exact integer math and no intermediate overflow.""" + unit, multiplier = options.np.datetime_data(td.dtype) + try: + numerator, denominator = cls._NUMPY_TIME_UNITS[unit] + except KeyError as ex: + raise ValueError(f"Unsupported NumPy timedelta64 unit {unit!r}; only fixed-duration time units are supported") from ex + raw = int(td.astype("int64")) + scaled = raw * int(multiplier) * numerator * precision + ticks = abs(scaled) // denominator + return -ticks if scaled < 0 else ticks + + @staticmethod + def _normalize_timedelta_value(value: Any) -> Any: + if options.np is not None and isinstance(value, options.np.timedelta64): + return None if options.np.isnat(value) else value + if options.pd is not None: + if value is options.pd.NaT: + return None + if isinstance(value, options.pd.Timedelta): + return value.to_timedelta64() + return value + + def write_column_data(self, column: Sequence, dest: bytearray, ctx: InsertContext): + if self.nullable: + # Treat NumPy/pandas NaT like None before the base implementation + # writes the Native null map. This also composes inside Array and + # LowCardinality because their writers delegate back here. A python + # timedelta sample can still precede NaT values, so it triggers the + # same pass. + sample = first_value(column, True) + is_timedelta = isinstance(sample, timedelta) + is_numpy_time = options.np is not None and isinstance(sample, options.np.timedelta64) + is_pandas_nat = options.pd is not None and sample is options.pd.NaT + if is_timedelta or is_numpy_time or is_pandas_nat: + column = [self._normalize_timedelta_value(value) for value in column] + super().write_column_data(column, dest, ctx) + def _active_null(self, ctx: QueryContext): """Return appropriate null value based on context.""" fmt = self.read_format(ctx) @@ -609,11 +656,11 @@ def _ticks_to_string(self, ticks: int) -> str: return f"{sign}{h:03d}:{m:02d}:{s:02d}" def _timedelta_to_ticks(self, td: timedelta | numpy.timedelta64) -> int: - """Convert timedelta to ticks (seconds), flooring fractional seconds.""" + """Convert timedelta to ticks (seconds), truncating fractional seconds toward zero.""" if isinstance(td, timedelta): total = int(td.total_seconds()) else: - total = td.astype("timedelta64[s]").astype(int) + total = self._numpy_timedelta_to_ticks(td, 1) self._validate_standard_range(total, td) return total @@ -701,10 +748,13 @@ def _ticks_to_string(self, ticks: int) -> str: def _timedelta_to_ticks(self, td: timedelta | numpy.timedelta64) -> int: """Convert timedelta to ticks with sub-second precision.""" if isinstance(td, timedelta): - total_us = int(td.total_seconds()) * self._MICROS_PER_SECOND + td.microseconds - ticks = (total_us * self.precision) // self._MICROS_PER_SECOND + total_us = (td.days * self._SECONDS_PER_DAY + td.seconds) * self._MICROS_PER_SECOND + td.microseconds + scaled = total_us * self.precision + ticks = abs(scaled) // self._MICROS_PER_SECOND + if scaled < 0: + ticks = -ticks else: - ticks = td.astype("timedelta64[s]").astype(int) + ticks = self._numpy_timedelta_to_ticks(td, self.precision) self._validate_standard_range(ticks, td) return ticks diff --git a/clickhouse_connect/driver/__init__.py b/clickhouse_connect/driver/__init__.py index 880ce2fa..eab2f2a7 100644 --- a/clickhouse_connect/driver/__init__.py +++ b/clickhouse_connect/driver/__init__.py @@ -265,6 +265,14 @@ def create_client( instead of as URL parameters. When False, large parameter payloads are still automatically sent as form data to avoid exceeding URL length limits, except for queries using binary parameter binds, which are only form-encoded when this is True. Only available for query operations (not inserts). Default: False + :param native_codec: selects the codec for client-managed FORMAT Native query and insert paths. 'python' (default) + uses the Python/Cython codec. 'rust' prefers the compiled _ch_core codec and routes unsupported options or types + to Python. 'rust_strict' raises instead of routing. It covers query, query_np, query_df, their block and row + stream variants, and inserts including insert_df. Falls back to the common setting 'native_codec', which can be + seeded by the CLICKHOUSE_CONNECT_NATIVE_CODEC environment variable. Does not affect the Arrow methods + (query_arrow, query_df_arrow, insert_arrow), raw, JSON, or caller-provided byte streams. The rust codecs + require the separate clickhouse-connect-core wheel, installed via the rust extra: + pip install clickhouse-connect[rust]. Ignored for interface="chdb". :return: ClickHouse Connect Client instance """ if _is_chdb_target(interface, dsn): @@ -418,6 +426,14 @@ async def create_async_client( instead of as URL parameters. When False, large parameter payloads are still automatically sent as form data to avoid exceeding URL length limits, except for queries using binary parameter binds, which are only form-encoded when this is True. Only available for query operations (not inserts). Default: False + :param native_codec: selects the codec for client-managed FORMAT Native query and insert paths. 'python' (default) + uses the Python/Cython codec. 'rust' prefers the compiled _ch_core codec and routes unsupported options or types + to Python. 'rust_strict' raises instead of routing. It covers query, query_np, query_df, their block and row + stream variants, and inserts including insert_df. Falls back to the common setting 'native_codec', which can be + seeded by the CLICKHOUSE_CONNECT_NATIVE_CODEC environment variable. Does not affect the Arrow methods + (query_arrow, query_df_arrow, insert_arrow), raw, JSON, or caller-provided byte streams. The rust codecs + require the separate clickhouse-connect-core wheel, installed via the rust extra: + pip install clickhouse-connect[rust]. :return: ClickHouse Connect AsyncClient instance """ if _is_chdb_target(interface, dsn): diff --git a/clickhouse_connect/driver/_backendclient.py b/clickhouse_connect/driver/_backendclient.py index c07b89a9..fe006690 100644 --- a/clickhouse_connect/driver/_backendclient.py +++ b/clickhouse_connect/driver/_backendclient.py @@ -18,21 +18,23 @@ from clickhouse_connect.driver.binding import bind_query from clickhouse_connect.driver.client import Client from clickhouse_connect.driver.ctypes import RespBuffCls +from clickhouse_connect.driver.exceptions import Error from clickhouse_connect.driver.external import ExternalData from clickhouse_connect.driver.insert import InsertContext from clickhouse_connect.driver.query import QueryContext, QueryResult +from clickhouse_connect.driver.streaming import _SyncStreamingInsertSource from clickhouse_connect.driver.summary import QuerySummary if TYPE_CHECKING: from clickhouse_connect.driver._backend.contracts import SyncBackend - from clickhouse_connect.driver.transform import NativeTransform + from clickhouse_connect.driver.transform import Transform logger = logging.getLogger(__name__) class SyncBackendClient(Client): _backend: SyncBackend - _transform: NativeTransform + _transform: Transform _write_format = "Native" _rename_response_column: str | None = None @@ -67,17 +69,40 @@ def data_insert(self, context: InsertContext) -> QuerySummary: if context.compression is None: context.compression = self.write_compression - block_gen = self._transform.build_insert(context) + threaded = self._transform.threaded_insert + active_source = None + if threaded: + active_source = _SyncStreamingInsertSource(transform=self._transform, context=context, maxsize=10) + active_source.start_producer() + block_gen = active_source.gen + else: + block_gen = self._transform.build_insert(context) def rebuild_block_gen(): + nonlocal active_source + recorded = context.insert_exception + if isinstance(recorded, Error): + # Deterministic client-side refusal; a rebuilt insert would fail identically. + context.insert_exception = None + raise recorded + # Reset so a failure on the rebuilt attempt is not masked by the first attempt's error. + context.insert_exception = None + if active_source is not None: + active_source.close(timeout=None) context.current_row = 0 context.current_block = 0 + if threaded: + active_source = _SyncStreamingInsertSource(transform=self._transform, context=context, maxsize=10) + active_source.start_producer() + return active_source.gen return self._transform.build_insert(context) runtime = QueryRuntime(database=self.database, settings=self._validate_settings(context.settings)) try: return QuerySummary(self._backend.execute_data_insert(context, runtime, block_gen, rebuild_block_gen)) finally: + if active_source is not None: + active_source.close() context.data = None def raw_insert( diff --git a/clickhouse_connect/driver/asyncclient.py b/clickhouse_connect/driver/asyncclient.py index bc4c715e..97a7cda0 100644 --- a/clickhouse_connect/driver/asyncclient.py +++ b/clickhouse_connect/driver/asyncclient.py @@ -56,7 +56,7 @@ dict_copy, # noqa: F401 (compatibility re-export) ) from clickhouse_connect.driver.ctypes import RespBuffCls -from clickhouse_connect.driver.exceptions import DataError, ProgrammingError +from clickhouse_connect.driver.exceptions import DataError, Error, ProgrammingError from clickhouse_connect.driver.external import ExternalData from clickhouse_connect.driver.insert import InsertContext from clickhouse_connect.driver.options import check_arrow, check_numpy, check_pandas, check_polars @@ -67,15 +67,17 @@ TzSource, arrow_buffer, ) +from clickhouse_connect.driver.rustcodec import NativeCodec, _make_native_transform from clickhouse_connect.driver.streaming import ( QueuedStreamSource, + ReadAheadSource, StreamingFileAdapter, StreamingInsertSource, StreamingResponseSource, start_streaming_response, ) from clickhouse_connect.driver.summary import QuerySummary -from clickhouse_connect.driver.transform import NativeTransform +from clickhouse_connect.driver.transform import NativeTransform, Transform from clickhouse_connect.driver.types import Closable logger = logging.getLogger(__name__) @@ -156,6 +158,7 @@ def __init__( form_encode_query_params: bool = False, rename_response_column: str | None = None, headers: dict[str, str] | None = None, + native_codec: NativeCodec | None = None, ): """ Async HTTP Client using aiohttp. Initialization is handled via _initialize(). @@ -252,13 +255,16 @@ def __init__( connector_kwargs["enable_cleanup_closed"] = True self._write_format = "Native" - self._transform = NativeTransform() + self._transform: Transform = _make_native_transform(native_codec) self._client_settings: dict[str, str] = {} self._initialized = False self._reported_libs: set[str] = set() self.headers["User-Agent"] = self.headers["User-Agent"].replace("mode:sync;", "mode:async;") if headers: self.headers.update(headers) + if not isinstance(self._transform, NativeTransform): + # The codec is a client-level choice, so the tag is applied at construction rather than per call. + add_integration_tag(self.headers, self._reported_libs, "clickhouse-connect-core") # Store aiohttp-specific params for deferred initialization self._compress_param = compress @@ -406,7 +412,9 @@ async def _execute_operation(self, operation: Operation) -> object: if isinstance(operation, CommandOp): return await self.command(operation.text, settings=settings, use_database=operation.use_database) if isinstance(operation, QueryOp): - return await self.query(operation.text, settings=settings, query_formats=dict(_INTERNAL_QUERY_FORMATS)) + context = self.create_query_context(query=operation.text, settings=settings, query_formats=dict(_INTERNAL_QUERY_FORMATS)) + context.internal = True + return await self._query_with_context(context) if isinstance(operation, RawQueryOp): return await self.raw_query(operation.text, settings=settings, fmt=operation.fmt) raise TypeError(f"Unsupported operation type: {type(operation).__name__}") @@ -492,8 +500,11 @@ def parse_streaming(): query_result.summary = execution.summary # Attach streaming_source to query_result.source to ensure it gets closed - # when the query result is closed (e.g. by StreamContext.__exit__) - query_result.source = streaming_source + # when the query result is closed (e.g. by StreamContext.__exit__). The rust codec wraps the + # byte source in a ReadAheadSource whose close() stops the read-ahead thread and chains through + # to streaming_source, so do not clobber it. + if not isinstance(query_result.source, ReadAheadSource): + query_result.source = streaming_source return query_result @@ -1192,6 +1203,13 @@ async def data_insert(self, context: InsertContext) -> QuerySummary: # type: ig async def rebuild_body(): nonlocal active_source + recorded = context.insert_exception + if isinstance(recorded, Error): + # Deterministic client-side refusal; a rebuilt insert would fail identically. + context.insert_exception = None + raise recorded + # Reset so a failure on the rebuilt attempt is not masked by the first attempt's error. + context.insert_exception = None await active_source.close(timeout=None) context.current_row = 0 context.current_block = 0 diff --git a/clickhouse_connect/driver/client.py b/clickhouse_connect/driver/client.py index 7da263b6..65db3991 100644 --- a/clickhouse_connect/driver/client.py +++ b/clickhouse_connect/driver/client.py @@ -146,7 +146,7 @@ class Client(ABC): compression: str | None = None write_compression: str | None = None - protocol_version = 0 + protocol_version: int = 0 # User-supplied initial ClickHouse settings, set by subclasses before # initialization so generated setting defaults never overwrite them _initial_settings: dict[str, Any] | None = None @@ -158,8 +158,8 @@ class Client(ABC): # setting and silently corrupt the request. _reserved_setting_names: set[str] = set() _reserved_setting_prefixes: tuple[str, ...] = () - database = None - max_error_message = 0 + database: str | None = None + max_error_message: int = 0 _tz_source: TzSource = "auto" _apply_server_tz = False tz_mode: TzMode = "naive_utc" @@ -170,7 +170,7 @@ def tz_source(self) -> TzSource: return self._tz_source @tz_source.setter - def tz_source(self, value: TzSource): + def tz_source(self, value: TzSource) -> None: if value not in _VALID_TZ_SOURCES: raise ProgrammingError(f'tz_source must be "auto", "server", or "local", got "{value}"') self._tz_source = value @@ -250,7 +250,9 @@ def _execute_operation(self, operation: Operation) -> object: if isinstance(operation, CommandOp): return self.command(operation.text, settings=settings, use_database=operation.use_database) if isinstance(operation, QueryOp): - return self.query(operation.text, settings=settings, query_formats=dict(_INTERNAL_QUERY_FORMATS)) + context = self.create_query_context(query=operation.text, settings=settings, query_formats=dict(_INTERNAL_QUERY_FORMATS)) + context.internal = True + return self._query_with_context(context) if isinstance(operation, RawQueryOp): return self.raw_query(operation.text, settings=settings, fmt=operation.fmt) raise TypeError(f"Unsupported operation type: {type(operation).__name__}") diff --git a/clickhouse_connect/driver/httpclient.py b/clickhouse_connect/driver/httpclient.py index 97d8fba2..ec48d7ea 100644 --- a/clickhouse_connect/driver/httpclient.py +++ b/clickhouse_connect/driver/httpclient.py @@ -41,6 +41,7 @@ get_proxy_manager, ) from clickhouse_connect.driver.query import TzMode, TzSource +from clickhouse_connect.driver.rustcodec import NativeCodec, _make_native_transform from clickhouse_connect.driver.transform import NativeTransform logger = logging.getLogger(__name__) @@ -106,6 +107,7 @@ def __init__( form_encode_query_params: bool = False, rename_response_column: str | None = None, headers: dict[str, str] | None = None, + native_codec: NativeCodec | None = None, ): """ Create an HTTP ClickHouse Connect client @@ -163,7 +165,10 @@ def __init__( if headers: client_headers.update(headers) self._write_format = "Native" - self._transform = NativeTransform() + self._transform = _make_native_transform(native_codec) + if not isinstance(self._transform, NativeTransform): + # The codec is a client-level choice, so the tag is applied at construction rather than per call. + add_integration_tag(client_headers, self._reported_libs, "clickhouse-connect-core") # There are use cases when the client needs to disable timeouts. if connect_timeout is not None: diff --git a/clickhouse_connect/driver/insert.py b/clickhouse_connect/driver/insert.py index 04cc0067..2c75c397 100644 --- a/clickhouse_connect/driver/insert.py +++ b/clickhouse_connect/driver/insert.py @@ -56,7 +56,7 @@ def __init__( self.req_block_size = block_size self.block_row_count = DEFAULT_BLOCK_BYTES self.data = data - self.insert_exception = None + self.insert_exception: Exception | None = None @property def empty(self) -> bool: @@ -154,6 +154,28 @@ def _convert_pandas(self, df): for df_col_name, col_name, ch_type in zip(df.columns, self.column_names, self.column_types): df_col = df[df_col_name] d_type_kind = df_col.dtype.kind + if ch_type.base_type in ("Enum8", "Enum16") and (d_type_kind in ("f", "O") or df_col.hasnans): + # Only dtypes that can hold missing values need the per-row NA and float-code handling. + enum_values = [] + for row, value in enumerate(df_col): + if options.pd.isna(value): + enum_values.append(None if ch_type.nullable else 0) + elif isinstance(value, (float, options.np.floating)): + if not options.np.isfinite(value): + raise DataError( + f"Column {col_name!r} row {row} has non-finite enum code {value!r}; " + "use a valid enum label or finite integral code" + ) + if not float(value).is_integer(): + raise DataError( + f"Column {col_name!r} row {row} has enum code {value!r} that would lose fractional data; " + "use a valid enum label or integral code" + ) + enum_values.append(int(value)) + else: + enum_values.append(value) + data.append(enum_values) + continue if ch_type.python_type is int: if d_type_kind == "f": df_col = df_col.round().astype(ch_type.base_type) diff --git a/clickhouse_connect/driver/query.py b/clickhouse_connect/driver/query.py index ca24fffa..9bce7f19 100644 --- a/clickhouse_connect/driver/query.py +++ b/clickhouse_connect/driver/query.py @@ -50,6 +50,9 @@ class QueryContext(BaseQueryContext): Argument/parameter object for queries. This context is used to set thread/query specific formats """ + bind_params: dict[str, str] + uncommented_query: str + def __init__( self, query: str | bytes = "", @@ -116,10 +119,10 @@ def __init__( ) self.query = query self.parameters = parameters or {} - self.use_none = True if use_none is None else use_none - self.column_oriented = False if column_oriented is None else column_oriented - self.use_numpy = use_numpy if use_numpy is not None else False - self.max_str_len = 0 if max_str_len is None else max_str_len + self.use_none: bool = True if use_none is None else use_none + self.column_oriented: bool = False if column_oriented is None else column_oriented + self.use_numpy: bool = use_numpy if use_numpy is not None else False + self.max_str_len: int = 0 if max_str_len is None else max_str_len self.server_tz = server_tz self.apply_server_tz = apply_server_tz self.external_data = external_data @@ -147,6 +150,8 @@ def __init__( self.column_tz: str | tzinfo | None = None self.response_tz: tzinfo | None = None self.block_info = False + # Marks driver-internal metadata queries, which always decode with the Python codec + self.internal = False self.as_pandas = as_pandas self.streaming = streaming self.show_clickhouse_errors: ShowClickHouseErrors = True diff --git a/clickhouse_connect/driver/rustcodec.py b/clickhouse_connect/driver/rustcodec.py new file mode 100644 index 00000000..e22c613a --- /dev/null +++ b/clickhouse_connect/driver/rustcodec.py @@ -0,0 +1,527 @@ +import logging +from collections.abc import Generator +from typing import Any, Literal, cast + +from clickhouse_connect import common +from clickhouse_connect.datatypes import registry +from clickhouse_connect.datatypes.base import ClickHouseType, ch_read_formats, ch_write_formats +from clickhouse_connect.datatypes.container import Tuple +from clickhouse_connect.datatypes.network import IPv4 +from clickhouse_connect.datatypes.special import UUID, SimpleAggregateFunction +from clickhouse_connect.datatypes.string import FixedString, String +from clickhouse_connect.datatypes.temporal import Date, DateTimeBase +from clickhouse_connect.driver import options +from clickhouse_connect.driver.compression import get_compressor +from clickhouse_connect.driver.exceptions import DataError, Error, NotSupportedError, ProgrammingError, StreamFailureError +from clickhouse_connect.driver.insert import InsertContext +from clickhouse_connect.driver.npquery import NumpyResult +from clickhouse_connect.driver.query import QueryContext, QueryResult +from clickhouse_connect.driver.rustnumpy import ( + _build_converters, + _contains_json_type, + _convert_block, + _normalize_json_shared_column, +) +from clickhouse_connect.driver.streaming import ReadAheadSource +from clickhouse_connect.driver.transform import NativeTransform, Transform, extract_error_message, extract_exception_with_tag +from clickhouse_connect.driver.types import ByteSource, Closable + +logger: logging.Logger = logging.getLogger(__name__) + +NativeCodec = Literal["python", "rust", "rust_strict"] + +_VALID_CODECS = ("python", "rust", "rust_strict") + +REQUIRED_BINDING_API_VERSION = 2 + +_versions_logged = False + + +def _ch_core_module() -> Any: + """Return the imported _ch_core module, or None if it is not available.""" + try: + import _ch_core + except ImportError: + return None + return _ch_core + + +def _log_versions_once(resolved: str, core: Any) -> None: + global _versions_logged + if _versions_logged: + return + _versions_logged = True + logger.info( + "native_codec=%s using clickhouse-connect-core %s (binding API %s), clickhouse-connect %s", + resolved, + getattr(core, "__version__", "unknown"), + getattr(core, "BINDING_API_VERSION", 0), + common.version(), + ) + + +def resolve_native_codec(native_codec: str | None) -> str: + """Resolve the effective codec name and verify a compatible compiled _ch_core module is available.""" + if native_codec is None: + resolved = common.get_setting("native_codec") + else: + resolved = native_codec.strip().lower() + if resolved not in _VALID_CODECS: + raise ProgrammingError(f"Invalid native_codec {native_codec!r}; expected one of {', '.join(_VALID_CODECS)}") + if resolved == "python": + return resolved + core = _ch_core_module() + if core is None: + raise NotSupportedError( + f'native_codec="{resolved}" requires the compiled _ch_core extension module, which ships separately as ' + "the clickhouse-connect-core wheel and is not installed in this environment. Install it with " + 'pip install "clickhouse-connect[rust]", or use native_codec="python" to run without it.' + ) + api_version = getattr(core, "BINDING_API_VERSION", 0) + if api_version < REQUIRED_BINDING_API_VERSION: + raise NotSupportedError( + f"The installed clickhouse-connect-core version {getattr(core, '__version__', 'unknown')} provides " + f"binding API {api_version}, but this version of clickhouse-connect requires binding API " + f"{REQUIRED_BINDING_API_VERSION} or newer. Upgrade it with pip install --upgrade clickhouse-connect-core." + ) + _log_versions_once(resolved, core) + return resolved + + +def _make_native_transform(native_codec: NativeCodec | None = None) -> Transform: + """Build the Transform for the resolved codec.""" + resolved = resolve_native_codec(native_codec) + if resolved == "python": + return NativeTransform() + return _RustNativeTransform(strict=(resolved == "rust_strict")) + + +def _rust_query_ineligible_reason(context: QueryContext) -> str | None: + """Return a short reason the rust decoder cannot serve this query, or None if it can. + + column_oriented, streaming, block_info, column_renamer, numpy/pandas output, and + settings/transport_settings/external_data are honored and not listed here. The columns-only LIMIT 0 branch is + answered from FORMAT JSON metadata in both clients before any transform runs. + """ + if context.use_numpy and options.arrow is None: + # numpy and pandas output route through the zero-copy Arrow exit, which needs pyarrow. + return "pyarrow not installed" + if context.query_formats: + return "query_formats" + if context.column_formats: + return "column_formats" + if ch_read_formats: + return "global read format override" + if not context.use_none: + return "use_none=False" + if context.encoding: + return "custom encoding" + if context.query_tz is not None: + return "query_tz" + if context.column_tzs: + return "column_tzs" + if context.response_tz is not None: + return "server timezone header" + if context.tz_mode != "naive_utc": + return "tz_mode" + # With the checks above excluded, active_tz(None) is None exactly when bare DateTime renders naive-UTC + # (query.py active_tz), the only ambient timezone behavior the rust decoder reproduces. + if context.active_tz(None) is not None: + return "ambient timezone" + return None + + +class _ExceptionTagScanner: + """Scans raw stream chunks for a complete tagged exception block, mirroring ResponseBuffer._check_for_exception.""" + + def __init__(self, exception_tag: str): + tag_bytes = exception_tag.encode() + self._open_marker = b"__exception__\r\n" + tag_bytes + self._close_marker = tag_bytes + b"\r\n__exception__" + self._carryover = b"" + self._exception_buf: bytearray | None = None + + def push(self, chunk: bytes) -> bytes | None: + if self._exception_buf is not None: + self._exception_buf += chunk + if self._close_marker in self._exception_buf: + return bytes(self._exception_buf) + return None + search_data = self._carryover + chunk + marker_pos = search_data.find(self._open_marker) + if marker_pos != -1: + self._exception_buf = bytearray(search_data[marker_pos:]) + if self._close_marker in self._exception_buf: + return bytes(self._exception_buf) + return None + carry_size = len(self._open_marker) - 1 + if len(search_data) >= carry_size: + self._carryover = search_data[-carry_size:] + else: + self._carryover = search_data + return None + + +def _unsupported_decode_error(ex: Exception) -> NotSupportedError: + return NotSupportedError( + f'The rust native codec cannot decode this response: {ex}. Use native_codec="python" to fall back to the Python codec' + ) + + +def _binding_value_error(ex: ValueError) -> Error: + # The binding raises NotImplementedError for unsupported input (handled separately at + # each call site), so any ValueError from it is malformed data. + return DataError(str(ex)) + + +def _contains_datetime_type(ch_type: object) -> bool: + """Return whether a ClickHouse type contains DateTime or DateTime64.""" + if not isinstance(ch_type, ClickHouseType): + return False + if isinstance(ch_type, DateTimeBase): + return True + for attr in ("element_type", "key_type", "value_type"): + child = getattr(ch_type, attr, None) + if child is not None and _contains_datetime_type(child): + return True + return any(_contains_datetime_type(child) for child in getattr(ch_type, "element_types", ())) + + +def _python_codec_tuple_container(ch_type: ClickHouseType) -> bool: + """Whether the Python codec's C-accelerated column readers return this column as a tuple. + + Streamed column blocks match those containers exactly: read_str_col renders String (nullable + included) and FixedString as tuples, the dataconv date/datetime/uuid/ipv4 readers render their + non-nullable columns as tuples, and unnamed Tuple columns are tuples of row tuples. + """ + if isinstance(ch_type, SimpleAggregateFunction): + ch_type = ch_type.element_type + if ch_type.low_card: + return False + if isinstance(ch_type, String): + return True + if isinstance(ch_type, Tuple): + return not ch_type.element_names + if ch_type.nullable: + return False + return isinstance(ch_type, (FixedString, Date, DateTimeBase, UUID, IPv4)) + + +def _chunk_has_server_error(chunk: bytes, exception_tag: str | None) -> bool: + if not chunk: + return False + # On tag-capable servers (v25.11+) a real server error is caught by the scanner before decoding, so a + # tagged block reaching here has no error. Only apply the byte heuristic for untagged servers, otherwise + # block data containing b"Code: " would misclassify an unsupported-type error as a stream failure. + if exception_tag: + return bool(extract_exception_with_tag(chunk, exception_tag)) + return b"Code: " in chunk + + +class _BufferedQueryResult(QueryResult): + """QueryResult over per-batch Python column lists materialized at parse time. + + The stream is fully consumed, decoded, and closed inside parse_response, so decode errors raise from + the query call itself. result_rows chains zip(*batch) per batch, exactly how the Python codec generates + rows from its column blocks, and result_columns concatenates per column across batches. + """ + + def __init__(self, batch_columns: list, names: tuple, col_types: tuple, column_oriented: bool, source: Closable): + super().__init__(None, None, names, col_types, column_oriented, source) + self._block_gen = None + self._batch_columns = batch_columns + + @property + def result_rows(self) -> Any: + if self._result_rows is None: + rows: list = [] + for columns in self._batch_columns: + rows.extend(zip(*columns)) + self._result_rows = rows + return self._result_rows + + @property + def result_columns(self) -> Any: + if self._result_columns is None: + columns = self._batch_columns[0] + for batch in self._batch_columns[1:]: + for base, added in zip(columns, batch): + base.extend(added) + # The first batch now holds the concatenated columns; collapse so a later + # result_rows access does not see the trailing batches twice. + self._batch_columns = [columns] + self._result_columns = columns + return self._result_columns + + +class _RustNativeTransform: + """FORMAT Native codec backed by the compiled _ch_core module.""" + + threaded_insert = True + + def __init__(self, strict: bool = False): + self.strict = strict + + def parse_response(self, source: ByteSource, context: QueryContext) -> NumpyResult | QueryResult: + if context.internal: + # Driver-internal metadata queries pin read formats and always use the Python codec, even under strict. + return NativeTransform.parse_response(source, context) + reason = _rust_query_ineligible_reason(context) + if reason is not None: + if self.strict: + source.close() + raise NotSupportedError(f'native_codec="rust_strict" does not support {reason}; use native_codec="python" or "rust"') + logger.info("Native codec fallback to Python for query: %s", reason) + return NativeTransform.parse_response(source, context) + + core = _ch_core_module() + if core is None: + source.close() + raise NotSupportedError('The rust native codec is unavailable (_ch_core not importable); use native_codec="python"') + + # Read-ahead: a daemon thread pulls transport chunks into a bounded queue so the network read overlaps + # decode. The consumer generator re-raises producer errors verbatim, in stream order, so the exception + # mapping, tag scanning, and last-chunk heuristics below run unchanged on the consuming thread. + read_source = ReadAheadSource(source) + decoder = core.StreamDecoder(has_block_info=context.block_info) + exception_tag = read_source.exception_tag + scanner = _ExceptionTagScanner(exception_tag) if exception_tag else None + last_chunk = b"" + + def raw_blocks() -> Generator[Any, None, None]: + # Decode blocks lazily so streaming queries keep bounded memory. Errors surface here on the + # consuming thread with the same mapping the Python codec uses; the source is closed before every raise. + nonlocal last_chunk + try: + for chunk in read_source.gen: + last_chunk = chunk + if scanner is not None: + hit = scanner.push(chunk) + if hit is not None: + read_source.close() + raise StreamFailureError( + extract_exception_with_tag(hit, cast(str, exception_tag)) or extract_error_message(hit) + ) + yield from decoder.feed(chunk) + yield from decoder.finish() + except StreamFailureError: + raise + except EOFError as ex: + read_source.close() + if _chunk_has_server_error(last_chunk, exception_tag): + raise StreamFailureError(extract_error_message(last_chunk)) from ex + raise StreamFailureError("Stream ended unexpectedly (connection closed by server)") from ex + except NotImplementedError as ex: + read_source.close() + if _chunk_has_server_error(last_chunk, exception_tag): + raise StreamFailureError(extract_error_message(last_chunk)) from ex + raise _unsupported_decode_error(ex) from ex + except ValueError as ex: + read_source.close() + if _chunk_has_server_error(last_chunk, exception_tag): + raise StreamFailureError(extract_error_message(last_chunk)) from ex + raise _binding_value_error(ex) from ex + except Exception as ex: + read_source.close() + if _chunk_has_server_error(last_chunk, exception_tag): + raise StreamFailureError(extract_error_message(last_chunk)) from ex + if ex.__class__.__name__ == "ClientPayloadError": + raise StreamFailureError("Stream failed during read (connection closed by server)") from ex + raise + + blocks = raw_blocks() + try: + first = next(blocks) + except StopIteration: + read_source.close() + return NumpyResult() if context.use_numpy else QueryResult([]) + + renamer = context.column_renamer + try: + names = tuple(renamer(name) if renamer is not None else name for name in first.column_names) + col_types = tuple(registry.get_from_name(type_name) for type_name in first.column_type_names) + except Exception: + read_source.close() + raise + + json_flags = [_contains_json_type(col_type) for col_type in col_types] + needs_post = any(json_flags) + + def normalize_json_columns(columns: list) -> list: + return [_normalize_json_shared_column(col_type, column, context) for col_type, column in zip(col_types, columns)] + + if not context.use_numpy and not context.streaming: + # Buffered query: materialize each batch to Python columns as it arrives so object building + # on this thread overlaps the producer's transport drain, and any decode error surfaces here + # rather than on a later result access. + try: + if needs_post: + batch_columns = [normalize_json_columns(first.to_python_columns())] + for batch in blocks: + batch_columns.append(normalize_json_columns(batch.to_python_columns())) + else: + batch_columns = [first.to_python_columns()] + for batch in blocks: + batch_columns.append(batch.to_python_columns()) + except NotImplementedError as ex: + read_source.close() + raise _unsupported_decode_error(ex) from ex + except ValueError as ex: + read_source.close() + raise _binding_value_error(ex) from ex + except Exception: + read_source.close() + raise + read_source.close() + return _BufferedQueryResult(batch_columns, names, col_types, context.column_oriented, read_source) + + tuple_flags = [_python_codec_tuple_container(col_type) for col_type in col_types] + # JSON-containing columns normalize in Python below, so their tuple wrap stays there too. + # None when no flag is set skips the per-call flag extraction in the binding. + flags = [tf and not jf for tf, jf in zip(tuple_flags, json_flags)] + rust_tuple_flags = flags if any(flags) else None + + def post_process_json(columns: list) -> list: + # Streamed blocks match the Python codec's per-column containers. + for i, json_flag in enumerate(json_flags): + if json_flag: + normalized = _normalize_json_shared_column(col_types[i], columns[i], context) + columns[i] = tuple(normalized) if tuple_flags[i] else normalized + return columns + + try: + if context.use_numpy: + converters = _build_converters(col_types, context) + first_columns = _convert_block(first, converters) + else: + first_columns = first.to_python_columns(typed_numeric=True, tuple_columns=rust_tuple_flags) + if needs_post: + first_columns = post_process_json(first_columns) + except NotImplementedError as ex: + read_source.close() + raise _unsupported_decode_error(ex) from ex + except ValueError as ex: + read_source.close() + raise _binding_value_error(ex) from ex + except Exception: + read_source.close() + raise + + if context.use_numpy: + d_types = [col.dtype if hasattr(col, "dtype") else "O" for col in first_columns] + + def np_block_gen() -> Generator[list, None, None]: + yield first_columns + for batch in blocks: + try: + columns = _convert_block(batch, converters) + except NotImplementedError as ex: + read_source.close() + raise _unsupported_decode_error(ex) from ex + except ValueError as ex: + read_source.close() + raise _binding_value_error(ex) from ex + except Exception: + read_source.close() + raise + yield columns + + return NumpyResult(np_block_gen(), names, col_types, d_types, read_source) + + def block_gen() -> Generator[list, None, None]: + yield first_columns + for batch in blocks: + try: + columns = batch.to_python_columns(typed_numeric=True, tuple_columns=rust_tuple_flags) + if needs_post: + columns = post_process_json(columns) + except NotImplementedError as ex: + read_source.close() + raise _unsupported_decode_error(ex) from ex + except ValueError as ex: + read_source.close() + raise _binding_value_error(ex) from ex + except Exception: + read_source.close() + raise + yield columns + + return QueryResult(None, block_gen(), names, col_types, context.column_oriented, read_source) + + def build_insert(self, context: InsertContext) -> Generator[bytes, None, None]: + core = _ch_core_module() + if core is None: + raise NotSupportedError('The rust native codec is unavailable (_ch_core not importable); use native_codec="python"') + + if common.get_setting("naive_datetime_insert") == "server" and any( + _contains_datetime_type(ch_type) for ch_type in context.column_types + ): + if self.strict: + raise NotSupportedError( + 'native_codec="rust_strict" does not support naive_datetime_insert="server" for DateTime inserts; ' + 'use native_codec="python" or "rust"' + ) + logger.info('Native codec fallback to Python for insert: naive_datetime_insert="server"') + return NativeTransform.build_insert(context) + + if ch_write_formats: + # The rust encoder does not consult the global write-format registry, so per-value conversions + # (e.g. set_write_format) would be ignored. Route these to the Python encoder. + if self.strict: + raise NotSupportedError( + 'native_codec="rust_strict" does not support global write format overrides; use native_codec="python" or "rust"' + ) + logger.info("Native codec fallback to Python for insert: global write format override") + return NativeTransform.build_insert(context) + + if context.col_simple_formats or context.col_type_formats or context.type_formats: + # The rust encoder ignores user column/query formats. Gate on the compiled format dicts, which are + # built from the user dict at init. _convert_pandas injects a harmless column_formats["int"] hint + # for datetime columns post-init, and the rust encoder already accepts the raw int values it feeds. + if self.strict: + raise NotSupportedError( + 'native_codec="rust_strict" does not support per-column or per-type write formats; use native_codec="python" or "rust"' + ) + logger.info("Native codec fallback to Python for insert: column/type write format") + return NativeTransform.build_insert(context) + + column_names = list(context.column_names) + type_names = [col_type.name for col_type in context.column_types] + try: + core.encode_native_block(column_names, type_names, [[] for _ in column_names], 0, None) + except NotImplementedError as ex: + if self.strict: + raise NotSupportedError( + f'native_codec="rust_strict" cannot insert unsupported column type: {ex}; use native_codec="python" or "rust"' + ) from ex + logger.info("Native codec fallback to Python for insert: unsupported type (%s)", ex) + return NativeTransform.build_insert(context) + + compression = context.compression if isinstance(context.compression, str) else None + compressor = get_compressor(compression) + + def chunk_gen(): + for block in context.next_block(): + try: + output = core.encode_native_block( + list(block.column_names), + [col_type.name for col_type in block.column_types], + list(block.column_data), + block.row_count, + block.prefix, + ) + except Exception as ex: + logger.error("Error serializing insert with Rust Native encoder", exc_info=True) + if not isinstance(ex, Error): + wrapped = DataError(str(ex)) + wrapped.__cause__ = ex + ex = wrapped + context.insert_exception = ex + yield b"INTERNAL EXCEPTION WHILE SERIALIZING" + return + yield compressor.compress_block(output) + footer = compressor.flush() + if footer: + yield footer + + return chunk_gen() diff --git a/clickhouse_connect/driver/rustnumpy.py b/clickhouse_connect/driver/rustnumpy.py new file mode 100644 index 00000000..1427ce07 --- /dev/null +++ b/clickhouse_connect/driver/rustnumpy.py @@ -0,0 +1,616 @@ +"""Numpy/pandas column converters for the rust native codec. + +The rust Arrow export is raw: Date is uint16 days, DateTime is uint32 seconds with the timezone dropped, +Enum is raw ints. A naive to_pandas() therefore yields wrong dtypes. These converters are resolved once +per query from the driver's own ClickHouseType (np_type, tzinfo, nullability) so the produced columns match +the Python codec by construction. Non-nullable numeric and temporal columns take the Arrow exit. Time and +Time64 keep their declared duration units through NumPy and extended pandas output. Strings, enums, +and remaining nullable columns take the rust python-object exit and are finalized through the driver's own +_finalize_column. +""" + +import logging +from collections.abc import Callable, Sequence +from typing import Any, cast + +from clickhouse_connect.datatypes import dynamic as dynamic_module +from clickhouse_connect.datatypes.base import ClickHouseType +from clickhouse_connect.datatypes.container import Array, Map, Nested, Tuple +from clickhouse_connect.datatypes.numeric import BFloat16, Interval +from clickhouse_connect.datatypes.special import SimpleAggregateFunction +from clickhouse_connect.datatypes.temporal import Date, DateTime, DateTime64, DateTimeBase, Time, Time64 +from clickhouse_connect.driver import options +from clickhouse_connect.driver.common import first_value +from clickhouse_connect.driver.exceptions import NotSupportedError +from clickhouse_connect.driver.query import QueryContext + +logger: logging.Logger = logging.getLogger(__name__) + +BlockConverter = Callable[[Any, Any, int], Any] + +_TIME64_UNITS = {3: "ms", 6: "us", 9: "ns"} +_COMPOUND_JSON_BINARY_TYPE_INDEXES = frozenset({0x1E, 0x1F, 0x20, 0x23, 0x26, 0x27, 0x2B, 0x30}) + + +def _contains_json_type(ch_type: object) -> bool: + """Return whether a ClickHouse type contains JSON through a supported container.""" + if not isinstance(ch_type, ClickHouseType): + return False + if ch_type.base_type == "JSON": + return True + for attr in ("element_type", "key_type", "value_type"): + if _contains_json_type(getattr(ch_type, attr, None)): + return True + return any(_contains_json_type(child) for child in getattr(ch_type, "element_types", ())) + + +def _decode_json_tree(value: Any, context: QueryContext) -> Any: + """Decode raw shared-data cells inside one materialized JSON object.""" + if isinstance(value, (bytes, bytearray, memoryview)): + binary_value = bytes(value) + if not binary_value or binary_value[0] not in _COMPOUND_JSON_BINARY_TYPE_INDEXES: + return value + decoded = dynamic_module.decode_shared_data_value(binary_value, context) + return value if decoded == binary_value else decoded + if isinstance(value, dict): + return {key: _decode_json_tree(item, context) for key, item in value.items()} + if isinstance(value, list): + return [_decode_json_tree(item, context) for item in value] + if isinstance(value, tuple): + return tuple(_decode_json_tree(item, context) for item in value) + return value + + +def _variant_materialized_type(ch_type: ClickHouseType) -> type | None: + if ch_type.base_type == "JSON" or isinstance(ch_type, Map): + return dict + if isinstance(ch_type, Nested): + return list + if isinstance(ch_type, Tuple) and ch_type.element_names: + return dict + return ch_type.python_type + + +def _normalize_json_shared_value(ch_type: ClickHouseType, value: Any, context: QueryContext) -> Any: + """Apply JSON shared-data decoding without changing non-JSON sibling values.""" + if value is None: + return None + if isinstance(ch_type, SimpleAggregateFunction): + return _normalize_json_shared_value(ch_type.element_type, value, context) + if ch_type.base_type == "JSON": + return _decode_json_tree(value, context) + if isinstance(ch_type, Array): + return [_normalize_json_shared_value(ch_type.element_type, item, context) for item in value] + if isinstance(ch_type, Tuple): + if ch_type.element_names and isinstance(value, dict): + result = dict(value) + for name, element_type in zip(ch_type.element_names, ch_type.element_types): + result[name] = _normalize_json_shared_value(element_type, value[name], context) + return result + return tuple(_normalize_json_shared_value(element_type, item, context) for element_type, item in zip(ch_type.element_types, value)) + if isinstance(ch_type, Map): + return { + _normalize_json_shared_value(ch_type.key_type, key, context): _normalize_json_shared_value(ch_type.value_type, item, context) + for key, item in value.items() + } + if isinstance(ch_type, Nested): + return [ + { + name: _normalize_json_shared_value(element_type, item[name], context) + for name, element_type in zip(ch_type.element_names, ch_type.element_types) + } + for item in value + ] + if isinstance(ch_type, dynamic_module.Variant): + if isinstance(value, dynamic_module.TypedVariant): + for element_type in ch_type.element_types: + if element_type.name == value.type_name: + normalized = _normalize_json_shared_value(element_type, value.value, context) + return dynamic_module.TypedVariant(normalized, value.type_name) + return value + candidates = [element_type for element_type in ch_type.element_types if _variant_materialized_type(element_type) is type(value)] + if len(candidates) == 1: + return _normalize_json_shared_value(candidates[0], value, context) + return value + + +def _normalize_json_shared_column(ch_type: ClickHouseType, column: Sequence[Any], context: QueryContext) -> Sequence[Any]: + if not _contains_json_type(ch_type): + return column + return [_normalize_json_shared_value(ch_type, value, context) for value in column] + + +class _Converter: + """One column's converter. ``needs_arrow`` decides whether the block Arrow table is built.""" + + __slots__ = ("needs_arrow", "_convert") + + def __init__(self, needs_arrow: bool, convert: BlockConverter): + self.needs_arrow = needs_arrow + self._convert = convert + + def __call__(self, arrow_table: Any, col_batch: Any, index: int) -> Any: + return self._convert(arrow_table, col_batch, index) + + +def _arrow_column(arrow_table: Any, index: int) -> Any: + return arrow_table.column(index).combine_chunks() + + +def _numeric_convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + return _arrow_column(arrow_table, index).to_numpy(zero_copy_only=False) + + +def _make_bfloat16_convert(as_extended_pandas: bool) -> BlockConverter: + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + column = _arrow_column(arrow_table, index) + data = column.buffers()[1] + if data is None: + values = options.np.zeros(len(column), dtype=options.np.float32) + else: + words = options.np.frombuffer(data, dtype=" Any: + column = _arrow_column(arrow_table, index) + return column.cast(options.arrow.int64()).to_numpy(zero_copy_only=False) + + +def _make_date_convert(as_pandas: bool) -> BlockConverter: + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + days = _arrow_column(arrow_table, index).to_numpy(zero_copy_only=False).astype("datetime64[D]") + return days.astype("datetime64[s]") if as_pandas else days + + return convert + + +def _make_datetime_convert(as_pandas: bool, active_tz: Any) -> BlockConverter: + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + naive = _arrow_column(arrow_table, index).to_numpy(zero_copy_only=False).astype("datetime64[s]") + if as_pandas and active_tz is not None: + return options.pd.DatetimeIndex(naive, tz="UTC").tz_convert(active_tz) + return naive + + return convert + + +def _make_datetime64_convert(as_pandas: bool, active_tz: Any) -> BlockConverter: + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + # Arrow timestamp[unit] -> datetime64[unit], tz metadata dropped to UTC instants. + column = _arrow_column(arrow_table, index).to_numpy(zero_copy_only=False) + if as_pandas and active_tz is not None: + return options.pd.DatetimeIndex(column, tz="UTC").tz_convert(active_tz) + return column + + return convert + + +def _pandas_infers_ns_timedeltas() -> bool: + """pandas < 3 infers timedelta64[ns] for the object arrays the Python codec's nullable + temporal path hands to the DataFrame constructor. pandas >= 3 keeps the scalar unit.""" + return int(options.pd.__version__.split(".", 1)[0]) < 3 + + +def _make_time_convert( + ch_type: Time | Time64, + as_pandas: bool = False, + use_extended_dtypes: bool = False, +) -> BlockConverter: + unit = "s" if isinstance(ch_type, Time) else _TIME64_UNITS[ch_type.scale] + nullable_pandas_ns = as_pandas and not use_extended_dtypes and _pandas_infers_ns_timedeltas() + + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + column = _arrow_column(arrow_table, index) + if ch_type.nullable: + null_count = column.null_count + if isinstance(ch_type, Time): + column = column.cast(options.arrow.int64()) + values = column.cast(options.arrow.duration(unit)).to_numpy(zero_copy_only=False) + if as_pandas: + return values.astype("timedelta64[ns]") if nullable_pandas_ns else values + # The Python codec's query_np contract for nullable temporal columns + # is an object array of numpy.timedelta64 scalars and None. Assign a + # list here because direct ndarray assignment coerces the scalars to + # datetime.timedelta at microsecond precision. + result = list(values) + if null_count: + for null_index in options.np.flatnonzero(options.np.isnat(values)): + result[null_index] = None + return result + values = column.to_numpy(zero_copy_only=False) + if isinstance(ch_type, Time64): + # The core exports Time64 as its raw signed Int64 tick buffer because + # Arrow time types cannot represent negative or >=24-hour values. + # NumPy timedelta64 has the same 64-bit layout, so only reinterpret + # the dtype here. No values or validity data are copied. + return values.view(ch_type.np_type) + # Time is Int32 on the wire while NumPy timedelta64 uses Int64. Widening + # requires one allocation, but still avoids one Python timedelta object + # per cell and lets NumPy perform the conversion in bulk. + return values.astype(ch_type.np_type, copy=False) + + return convert + + +def _any_leaf(ch_type: ClickHouseType, predicate: Callable[[ClickHouseType], bool]) -> bool: + if isinstance(ch_type, Array): + return _any_leaf(ch_type.element_type, predicate) + if isinstance(ch_type, Tuple): + return any(_any_leaf(elem, predicate) for elem in ch_type.element_types) + if isinstance(ch_type, Map): + return _any_leaf(ch_type.key_type, predicate) or _any_leaf(ch_type.value_type, predicate) + return predicate(ch_type) + + +def _contains_nested_time(ch_type: ClickHouseType) -> bool: + # Container-only: bare Time/Time64 variants keep their dedicated or object-exit converters. + return isinstance(ch_type, (Array, Tuple, Map)) and _any_leaf(ch_type, lambda leaf: isinstance(leaf, (Time, Time64))) + + +def _materialize_raw_times(ch_type: ClickHouseType, raw: Any, extended_time_null: bool = False) -> Any: + """Replace raw Time ticks in an otherwise materialized Python object tree.""" + if raw is None: + if extended_time_null and isinstance(ch_type, Time): + return options.np.timedelta64("NaT", "s") + if extended_time_null and isinstance(ch_type, Time64): + return options.np.timedelta64("NaT", _TIME64_UNITS[ch_type.scale]) + return None + if isinstance(ch_type, Time): + return options.np.timedelta64(raw, "s") + if isinstance(ch_type, Time64): + return options.np.timedelta64(raw, _TIME64_UNITS[ch_type.scale]) + if isinstance(ch_type, Array): + for index, value in enumerate(raw): + raw[index] = _materialize_raw_times(ch_type.element_type, value, extended_time_null) + return raw + if isinstance(ch_type, Tuple): + if isinstance(raw, dict): + return { + name: _materialize_raw_times(elem_type, raw[name], extended_time_null) + for name, elem_type in zip(ch_type.element_names, ch_type.element_types) + } + return tuple(_materialize_raw_times(elem_type, value, extended_time_null) for elem_type, value in zip(ch_type.element_types, raw)) + if isinstance(ch_type, Map): + return { + _materialize_raw_times(ch_type.key_type, key, extended_time_null): _materialize_raw_times( + ch_type.value_type, value, extended_time_null + ) + for key, value in raw.items() + } + return raw + + +def _make_nested_time_convert(ch_type: ClickHouseType, context: QueryContext) -> BlockConverter: + extended_time_null = context.as_pandas and context.use_extended_dtypes + leaf_predicate = _refinalize_predicate(ch_type, context) + + def convert(_arrow_table: Any, col_batch: Any, index: int) -> Any: + try: + column = col_batch.column_data(index, raw_time_ticks=True) + except NotImplementedError as ex: + raise NotSupportedError( + f"The rust native codec cannot decode this column for numpy/pandas output: {ex}. " + 'Use native_codec="python" to fall back to the Python codec' + ) from ex + for row, value in enumerate(column): + column[row] = _materialize_raw_times(ch_type, value, extended_time_null) + if leaf_predicate is not None: + column = _refinalize_leaves(ch_type, column, context, leaf_predicate) + return column + + return convert + + +def _array_time_leaf(ch_type: ClickHouseType) -> tuple[int, Time | Time64] | None: + """Return (nesting depth, leaf type) for pure Array(...(Time/Time64)) columns, else None.""" + depth = 0 + while isinstance(ch_type, Array): + depth += 1 + ch_type = ch_type.element_type + if depth and isinstance(ch_type, (Time, Time64)) and not ch_type.low_card: + return depth, ch_type + return None + + +def _make_array_time_convert(leaf: Time | Time64, depth: int, context: QueryContext) -> BlockConverter: + unit = "s" if isinstance(leaf, Time) else _TIME64_UNITS[leaf.scale] + extended_time_null = context.as_pandas and context.use_extended_dtypes + + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + column = _arrow_column(arrow_table, index) + offset_levels = [] + for _ in range(depth): + offset_levels.append(column.offsets.to_numpy().tolist()) + column = column.values + if leaf.nullable: + null_count = column.null_count + if isinstance(leaf, Time): + column = column.cast(options.arrow.int64()) + values = column.cast(options.arrow.duration(unit)).to_numpy(zero_copy_only=False) + # list() of a timedelta64 array yields numpy scalars, NaT included, + # which is the extended-dtypes null representation already. + cells = list(values) + if null_count and not extended_time_null: + for null_index in options.np.flatnonzero(options.np.isnat(values)): + cells[null_index] = None + else: + values = column.to_numpy(zero_copy_only=False) + values = values.view(leaf.np_type) if isinstance(leaf, Time64) else values.astype(leaf.np_type, copy=False) + cells = list(values) + for offsets in reversed(offset_levels): + cells = [cells[start:stop] for start, stop in zip(offsets, offsets[1:])] + return cells + + return convert + + +def _make_low_card_time_convert(ch_type: Time, as_pandas: bool) -> BlockConverter: + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + # The core exports LowCardinality(Time) as dictionary with + # nulls in the indices. Decoding then casting stays fully vectorized. + column = _arrow_column(arrow_table, index).dictionary_decode() + values = column.cast(options.arrow.int64()).cast(options.arrow.duration("s")).to_numpy(zero_copy_only=False) + if as_pandas or not ch_type.nullable: + return values + result = list(values) + if column.null_count: + for null_index in options.np.flatnonzero(options.np.isnat(values)): + result[null_index] = None + return result + + return convert + + +def _extended_refinalize_leaf(leaf: ClickHouseType) -> bool: + # Nullable leaves gain pd.NA/NaN/NaT, and Date leaves are converted from rust's datetime.date objects + # to the numpy datetime64 values produced by the Python codec. Time leaves are excluded: + # _materialize_raw_times fully renders them before refinalize runs. + return (leaf.nullable and not isinstance(leaf, (Time, Time64))) or isinstance(leaf, Date) + + +def _bfloat16_refinalize_leaf(leaf: ClickHouseType) -> bool: + # Non-extended numpy output only densifies nullable BFloat16 leaves. Other nullable leaves keep + # python None, matching the Python codec. + return isinstance(leaf, BFloat16) and leaf.nullable and not leaf.low_card + + +LeafPredicate = Callable[[ClickHouseType], bool] + + +def _needs_refinalize(ch_type: ClickHouseType, leaf_predicate: LeafPredicate = _extended_refinalize_leaf) -> bool: + # Only leaves need rewriting. + return _any_leaf(ch_type, leaf_predicate) + + +def _refinalize_predicate(ch_type: ClickHouseType, context: QueryContext) -> LeafPredicate | None: + """Select the leaf predicate for nested refinalize, or None when no leaf needs it.""" + if not isinstance(ch_type, (Array, Tuple, Map)): + return None + if context.as_pandas and context.use_extended_dtypes: + predicate: LeafPredicate = _extended_refinalize_leaf + elif context.use_numpy: + predicate = _bfloat16_refinalize_leaf + else: + return None + return predicate if _needs_refinalize(ch_type, predicate) else None + + +def _refinalize_leaves( + ch_type: ClickHouseType, + column: list, + context: QueryContext, + leaf_predicate: LeafPredicate = _extended_refinalize_leaf, +) -> list: + """Rebuild a rust-decoded object column so nested leaves match the Python codec. + + The rust object exit materializes nulls as python None. The Python codec finalizes each flat leaf + column, so pandas/numpy leaves render nulls as pd.NA/NaT and values as numpy scalars. Each affected + leaf is flattened, run through its own _finalize_column, and resliced. Unaffected sibling leaves keep + their value-equal rust-native scalars. + """ + if isinstance(ch_type, Array): + lengths = [None if cell is None else len(cell) for cell in column] + flat = _refinalize_leaves( + ch_type.element_type, + [item for cell in column if cell is not None for item in cell], + context, + leaf_predicate, + ) + out: list = [] + pos = 0 + for length in lengths: + if length is None: + out.append(None) + continue + out.append(flat[pos : pos + length]) + pos += length + return out + if isinstance(ch_type, Tuple): + refinalize_indexes = [index for index, elem in enumerate(ch_type.element_types) if _needs_refinalize(elem, leaf_predicate)] + if not refinalize_indexes: + return column + keyed = bool(ch_type.element_names) and isinstance(first_value(column), dict) + keys: list = [ch_type.element_names[index] if keyed else index for index in refinalize_indexes] + columns: list[list] = [[] for _ in refinalize_indexes] + for row in column: + if row is None: + continue + for slot, key in zip(columns, keys): + slot.append(row[key]) + columns = [ + _refinalize_leaves(ch_type.element_types[index], slot, context, leaf_predicate) + for index, slot in zip(refinalize_indexes, columns) + ] + rows_iter = iter(zip(*columns)) + out = [] + for row in column: + if row is None: + out.append(None) + continue + replaced = next(rows_iter) + if keyed: + updated = dict(row) + updated.update(zip(keys, replaced)) + out.append(updated) + else: + values = list(row) + for index, value in zip(refinalize_indexes, replaced): + values[index] = value + out.append(tuple(values)) + return out + if isinstance(ch_type, Map): + refinalize_keys = _needs_refinalize(ch_type.key_type, leaf_predicate) + refinalize_values = _needs_refinalize(ch_type.value_type, leaf_predicate) + if not refinalize_keys and not refinalize_values: + return column + key_iter = None + if refinalize_keys: + key_iter = iter(_refinalize_leaves(ch_type.key_type, [key for row in column if row for key in row], context, leaf_predicate)) + value_iter = None + if refinalize_values: + value_iter = iter( + _refinalize_leaves(ch_type.value_type, [value for row in column if row for value in row.values()], context, leaf_predicate) + ) + out = [] + for row in column: + if row is None: + out.append(None) + continue + mapped = {} + for key, value in row.items(): + mapped[next(key_iter) if key_iter is not None else key] = next(value_iter) if value_iter is not None else value + out.append(mapped) + return out + # Float leaf nulls become numpy NaN via the Python codec's numpy read rather than _finalize_column. + # BFloat16 leaves fall through to _finalize_column, which produces pd.array Float32 in extended mode. + if _np_kind(ch_type) == "f" and not (isinstance(ch_type, BFloat16) and not ch_type.low_card): + return list(options.np.array(column, dtype=ch_type.np_type)) + # Aware stdlib datetimes from the rust exit become Timestamps so _finalize_column takes its tz-aware path. + if isinstance(ch_type, DateTimeBase) and getattr(first_value(column), "tzinfo", None) is not None: + column = [None if value is None else options.pd.Timestamp(value) for value in column] + # _finalize_column returns a pandas/numpy container for the types that diverge. + finalized = ch_type._finalize_column(column, context) + return finalized if isinstance(finalized, list) else list(finalized) + + +def _make_object_convert(ch_type: ClickHouseType, context: QueryContext) -> BlockConverter: + leaf_predicate = _refinalize_predicate(ch_type, context) + extended_pandas = context.as_pandas and context.use_extended_dtypes + + def convert(_arrow_table: Any, col_batch: Any, index: int) -> Any: + try: + column = col_batch.column_data(index) + except NotImplementedError as ex: + raise NotSupportedError( + f"The rust native codec cannot decode this column for numpy/pandas output: {ex}. " + 'Use native_codec="python" to fall back to the Python codec' + ) from ex + column = cast(list[Any], _normalize_json_shared_column(ch_type, column, context)) + if leaf_predicate is not None: + column = _refinalize_leaves(ch_type, column, context, leaf_predicate) + if extended_pandas and isinstance(ch_type, DateTimeBase) and getattr(first_value(column), "tzinfo", None) is not None: + # DateTimeBase._finalize_column's extended-dtype path recognizes timezone-aware pandas values + # by their `.tz` attribute. The rust object exit produces stdlib datetime values, whose + # equivalent attribute is `.tzinfo`, so wrap them before nullable pandas finalization selects + # a naive dtype. Non-pandas output keeps the stdlib datetimes the Python codec returns. + column = [None if value is None else options.pd.Timestamp(value) for value in column] + return ch_type._finalize_column(column, context) + + return convert + + +def _make_nullable_int_convert(pd_dtype: Any) -> BlockConverter: + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + return pd_dtype.__from_arrow__(_arrow_column(arrow_table, index)) + + return convert + + +def _make_nullable_interval_convert(pd_dtype: Any) -> BlockConverter: + def convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + column = _arrow_column(arrow_table, index).cast(options.arrow.int64()) + return pd_dtype.__from_arrow__(column) + + return convert + + +def _nullable_float_convert(arrow_table: Any, _col_batch: Any, index: int) -> Any: + # The Python codec renders nullable Float32/Float64 as a plain float64 array with NaN in null positions. + return _arrow_column(arrow_table, index).to_numpy(zero_copy_only=False).astype("float64") + + +def _np_kind(ch_type: ClickHouseType) -> str | None: + try: + return str(options.np.dtype(ch_type.np_type).kind) + except Exception: # noqa: BLE001 - any non-numpy np_type is not an Arrow-numeric column + return None + + +def _build_converter(ch_type: ClickHouseType, context: QueryContext) -> _Converter: + # SimpleAggregateFunction is a name-decoration alias: convert as the element type, matching both the + # rust core's physical_delegate expansion and the Python codec's delegated read. + if isinstance(ch_type, SimpleAggregateFunction) and not ch_type.low_card: + ch_type = ch_type.element_type + if isinstance(ch_type, BFloat16) and not ch_type.low_card: + extended = ch_type.nullable and context.as_pandas and context.use_extended_dtypes + return _Converter(True, _make_bfloat16_convert(extended)) + # LowCardinality(T) routes through the object exit regardless of inner type. Its values are correct there, + # and the Python codec's own LowCardinality numpy handling is inconsistent per inner type (and truncates + # LowCardinality(numeric)), so there is no clean parity target for an Arrow dictionary fast path. + if isinstance(ch_type, (Time, Time64)) and not ch_type.low_card: + return _Converter(True, _make_time_convert(ch_type, context.as_pandas, context.use_extended_dtypes)) + if isinstance(ch_type, Time) and ch_type.low_card: + return _Converter(True, _make_low_card_time_convert(ch_type, context.as_pandas)) + if isinstance(ch_type, Interval) and not ch_type.low_card: + if not ch_type.nullable: + return _Converter(True, _interval_convert) + if context.as_pandas and context.use_extended_dtypes: + return _Converter(True, _make_nullable_interval_convert(options.pd.Int64Dtype())) + array_time = _array_time_leaf(ch_type) + if array_time is not None: + depth, leaf = array_time + return _Converter(True, _make_array_time_convert(leaf, depth, context)) + if _contains_nested_time(ch_type): + return _Converter(False, _make_nested_time_convert(ch_type, context)) + if not ch_type.nullable and not ch_type.low_card: + if isinstance(ch_type, DateTime64): + _ = ch_type.np_type # ProgrammingError for precisions outside {0,3,6,9}, matching the Python codec + return _Converter(True, _make_datetime64_convert(context.as_pandas, context.active_tz(ch_type.tzinfo))) + if isinstance(ch_type, DateTime): + return _Converter(True, _make_datetime_convert(context.as_pandas, context.active_tz(ch_type.tzinfo))) + if isinstance(ch_type, Date): # Date32 subclasses Date + return _Converter(True, _make_date_convert(context.as_pandas)) + if _np_kind(ch_type) in ("i", "u", "f", "b"): + return _Converter(True, _numeric_convert) + elif ch_type.nullable and not ch_type.low_card and context.as_pandas and context.use_extended_dtypes: + # query_df renders nullable numeric via zero-copy pandas extension arrays. Building them from the Arrow + # validity+values buffers skips the per-value Python object list the object exit would otherwise create. + kind = _np_kind(ch_type) + if kind in ("i", "u"): + return _Converter(True, _make_nullable_int_convert(options.pd.api.types.pandas_dtype(ch_type.base_type))) + if kind == "f": + return _Converter(True, _nullable_float_convert) + return _Converter(False, _make_object_convert(ch_type, context)) + + +def _build_converters(column_types: Sequence[ClickHouseType], context: QueryContext) -> list[_Converter]: + """Resolve one converter per column from the driver's own ClickHouseType metadata.""" + return [_build_converter(ch_type, context) for ch_type in column_types] + + +def _convert_block(col_batch: Any, converters: Sequence[_Converter]) -> list: + """Convert one decoded ColBatch into a list of numpy arrays, pandas arrays, or object lists.""" + arrow_table = None + if any(conv.needs_arrow for conv in converters): + arrow_table = options.arrow.RecordBatchReader.from_stream(col_batch).read_all() + return [conv(arrow_table, col_batch, index) for index, conv in enumerate(converters)] diff --git a/clickhouse_connect/driver/streaming.py b/clickhouse_connect/driver/streaming.py index 31d4e2fb..0007029a 100644 --- a/clickhouse_connect/driver/streaming.py +++ b/clickhouse_connect/driver/streaming.py @@ -1,22 +1,27 @@ import asyncio import logging +import queue import threading import zlib from collections.abc import Callable, Iterable, Iterator +from typing import cast import lz4.frame from clickhouse_connect.driver.asyncqueue import EOF_SENTINEL, AsyncSyncQueue +from clickhouse_connect.driver.asyncqueue import Full as AsyncQueueFull from clickhouse_connect.driver.compression import _zstd_decompressor, available_compression -from clickhouse_connect.driver.exceptions import OperationalError -from clickhouse_connect.driver.types import Closable +from clickhouse_connect.driver.exceptions import Error, OperationalError +from clickhouse_connect.driver.transform import Transform +from clickhouse_connect.driver.types import ByteSource, Closable -logger = logging.getLogger(__name__) +logger: logging.Logger = logging.getLogger(__name__) __all__ = [ "StreamingResponseSource", "StreamingFileAdapter", "StreamingInsertSource", + "ReadAheadSource", "QueuedStreamSource", "start_streaming_response", ] @@ -336,11 +341,12 @@ def close(self): class StreamingInsertSource: """Streaming source for async inserts (reverse bridge)""" - def __init__(self, transform, context, loop: asyncio.AbstractEventLoop, maxsize: int = 10): + def __init__(self, transform: Transform, context, loop: asyncio.AbstractEventLoop, maxsize: int = 10): self.transform = transform self.context = context self.loop = loop self.queue: AsyncSyncQueue[bytes | bytearray | Exception] = AsyncSyncQueue(maxsize=maxsize) + self._stop_event = threading.Event() self._producer_future = None self._started = False @@ -351,17 +357,27 @@ def start_producer(self): def producer(): try: - for block in self.transform.build_insert(self.context): - self.queue.sync_q.put(block) + block_gen = self.transform.build_insert(self.context) + while not self._stop_event.is_set(): + try: + block = next(block_gen) + except StopIteration: + self._put(EOF_SENTINEL) + return - self.queue.sync_q.put(EOF_SENTINEL) + if not self._put(block): + return except Exception as e: - logger.error("Insert producer error: %s", e, exc_info=True) - try: - self.queue.sync_q.put(e) - except Exception: - pass + # Driver errors are deterministic client-side refusals, not operational failures. + if isinstance(e, Error): + logger.debug("Insert producer error: %s", e) + else: + logger.error("Insert producer error: %s", e, exc_info=True) + if getattr(self.context, "insert_exception", None) is None: + self.context.insert_exception = e + if not self._stop_event.is_set(): + self._put(e) finally: self.queue.shutdown() @@ -385,9 +401,14 @@ async def async_generator(self): yield chunk except Exception as e: - logger.error("Insert consumer error: %s", e, exc_info=True) + if isinstance(e, Error): + logger.debug("Insert consumer error: %s", e) + else: + logger.error("Insert consumer error: %s", e, exc_info=True) raise finally: + self._stop_event.set() + self.queue.shutdown() if self._producer_future and not self._producer_future.done(): try: await self._producer_future @@ -396,6 +417,7 @@ async def async_generator(self): async def close(self, timeout: float | None = 1.0): """Shut down the queue and wait for the producer thread to terminate. Pass ``timeout=None`` to wait without a deadline.""" + self._stop_event.set() self.queue.shutdown() if self._producer_future and not self._producer_future.done(): try: @@ -407,3 +429,172 @@ async def close(self, timeout: float | None = 1.0): logger.warning("Insert producer did not finish within timeout") except Exception: pass + + def _put(self, item: bytes | bytearray | Exception) -> bool: + while not self._stop_event.is_set(): + try: + self.queue.sync_q.put(item, timeout=0.1) + return True + except AsyncQueueFull: + continue + except RuntimeError: + return False + return False + + +class ReadAheadSource(Closable): + """Reads chunks from a byte source on a daemon thread into a bounded queue so transport overlaps decode. + + The consumer generator re-raises any producer-side exception verbatim, in stream order, so the wrapping + codec's error mapping, exception-tag scanning, and last-chunk heuristics run unchanged on the consumer thread. + """ + + def __init__(self, source: ByteSource, maxsize: int = 16): + self.source: ByteSource | None = source + self.exception_tag: str | None = getattr(source, "exception_tag", None) + self.queue: queue.Queue[tuple[str, object]] = queue.Queue(maxsize=maxsize) + self._stop_event = threading.Event() + self._gen_cache: Iterator[bytes] | None = None + self._thread = threading.Thread(target=self._producer, name="clickhouse-read-ahead", daemon=True) + self._thread.start() + + @property + def gen(self) -> Iterator[bytes]: + if self._gen_cache is None: + self._gen_cache = self._consume() + return self._gen_cache + + def _consume(self) -> Iterator[bytes]: + while True: + tag, payload = self.queue.get() + if tag == "data": + yield cast(bytes, payload) + elif tag == "error": + raise cast(BaseException, payload) + else: # eof + return + + def _producer(self): + source = self.source + if source is None: + return + try: + for chunk in source.gen: + if not self._put(("data", chunk)): + return + except BaseException as ex: # noqa: BLE001 - forwarded to the consumer thread verbatim + self._put(("error", ex)) + finally: + self._put(("eof", None)) + + def _drain(self): + try: + while True: + self.queue.get_nowait() + except queue.Empty: + pass + + def _release_source(self): + source, self.source = self.source, None + if source is not None: + source.close() + + def close(self) -> None: + # Join the producer before closing the source. A _put-blocked producer returns within one _put + # timeout of the stop event; a read-blocked producer exits after its in-flight read returns. Closing + # the source only after the join keeps the transport single-reader: the sync source drains on close, + # which would race a producer still reading it. + self._stop_event.set() + if self._thread.is_alive(): + self._thread.join(timeout=1.0) + self._drain() + self._release_source() + + async def aclose(self) -> None: + self._stop_event.set() + if self._thread.is_alive(): + # Join off the event loop so the worst-case wait never blocks it. + await asyncio.get_running_loop().run_in_executor(None, self._thread.join, 1.0) + self._drain() + # Release on the loop thread: the async source's close cancels its producer task, which must not + # run from an executor thread. + self._release_source() + + def _put(self, item: tuple[str, object]) -> bool: + while not self._stop_event.is_set(): + try: + self.queue.put(item, timeout=0.1) + return True + except queue.Full: + continue + return False + + +class _SyncStreamingInsertSource: + """Bounded producer/consumer source for sync inserts.""" + + def __init__(self, transform: Transform, context, maxsize: int = 10): + self.transform = transform + self.context = context + self.queue: queue.Queue[bytes | Exception | object] = queue.Queue(maxsize=maxsize) + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + self._started = False + + def start_producer(self): + if self._started: + raise RuntimeError("Producer already started") + self._started = True + self._thread = threading.Thread(target=self._producer, name="clickhouse-insert-producer", daemon=True) + self._thread.start() + + @property + def gen(self) -> Iterator[bytes]: + if not self._started: + raise RuntimeError("Producer not started, call start_producer() first") + try: + while True: + chunk = self.queue.get() + if chunk is EOF_SENTINEL: + break + if isinstance(chunk, Exception): + raise chunk + yield cast(bytes, chunk) + finally: + self.close() + + def close(self, timeout: float | None = 1.0): + self._stop_event.set() + if self._thread and self._thread.is_alive(): + self._thread.join(timeout=timeout) + + def _producer(self): + try: + block_gen = self.transform.build_insert(self.context) + while not self._stop_event.is_set(): + try: + block = next(block_gen) + except StopIteration: + self._put(EOF_SENTINEL) + return + if not self._put(block): + return + except Exception as ex: + # Driver errors are deterministic client-side refusals, not operational failures. + if isinstance(ex, Error): + logger.debug("Insert producer error: %s", ex) + else: + logger.error("Insert producer error: %s", ex, exc_info=True) + if getattr(self.context, "insert_exception", None) is None: + self.context.insert_exception = ex + if not self._stop_event.is_set(): + self._put(ex) + + def _put(self, item: bytes | Exception | object) -> bool: + while not self._stop_event.is_set(): + try: + self.queue.put(item, timeout=0.1) + return True + except queue.Full: + continue + return False diff --git a/clickhouse_connect/driver/transform.py b/clickhouse_connect/driver/transform.py index 0ec94a5f..dc14acca 100644 --- a/clickhouse_connect/driver/transform.py +++ b/clickhouse_connect/driver/transform.py @@ -1,4 +1,6 @@ import logging +from collections.abc import Generator +from typing import Protocol from clickhouse_connect.datatypes import registry from clickhouse_connect.driver.common import write_leb128 @@ -17,10 +19,22 @@ _EMPTY_CTX = QueryContext() -logger = logging.getLogger(__name__) +logger: logging.Logger = logging.getLogger(__name__) + + +class Transform(Protocol): + """Codec contract for FORMAT Native query decode and insert encode.""" + + threaded_insert: bool + + def parse_response(self, source: ByteSource, context: QueryContext) -> NumpyResult | QueryResult: ... + + def build_insert(self, context: InsertContext) -> Generator[bytes, None, None]: ... class NativeTransform: + threaded_insert: bool = False + @staticmethod def parse_response(source: ByteSource, context: QueryContext = _EMPTY_CTX) -> NumpyResult | QueryResult: names = [] @@ -121,7 +135,7 @@ def gen(): return QueryResult(None, gen(), tuple(names), tuple(col_types), context.column_oriented, source) @staticmethod - def build_insert(context: InsertContext): + def build_insert(context: InsertContext) -> Generator[bytes, None, None]: compression = context.compression if isinstance(context.compression, str) else None compressor = get_compressor(compression) diff --git a/clickhouse_connect/driver/types.py b/clickhouse_connect/driver/types.py index 7dadbac3..d154925f 100644 --- a/clickhouse_connect/driver/types.py +++ b/clickhouse_connect/driver/types.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from typing import Any Matrix = Sequence[Sequence[Any]] @@ -7,12 +7,14 @@ class Closable(ABC): @abstractmethod - def close(self): + def close(self) -> None: pass class ByteSource(Closable): last_message: bytes | None = None + gen: Iterator[bytes] + exception_tag: str | None = None @abstractmethod def read_leb128(self) -> int: @@ -31,19 +33,19 @@ def read_bytes(self, sz: int) -> bytes: pass @abstractmethod - def read_str_col(self, num_rows: int, encoding: str | None, nullable: bool = False, null_obj: Any = None): + def read_str_col(self, num_rows: int, encoding: str | None, nullable: bool = False, null_obj: Any = None) -> Any: pass @abstractmethod - def read_bytes_col(self, sz: int, num_rows: int): + def read_bytes_col(self, sz: int, num_rows: int) -> Any: pass @abstractmethod - def read_fixed_str_col(self, sz: int, num_rows: int, encoding: str): + def read_fixed_str_col(self, sz: int, num_rows: int, encoding: str) -> Any: pass @abstractmethod - def read_array(self, array_type: str, num_rows: int): + def read_array(self, array_type: str, num_rows: int) -> Any: pass @abstractmethod diff --git a/clickhouse_connect/driverc/buffer.pxd b/clickhouse_connect/driverc/buffer.pxd index ef65c0e9..d698ce75 100644 --- a/clickhouse_connect/driverc/buffer.pxd +++ b/clickhouse_connect/driverc/buffer.pxd @@ -2,7 +2,8 @@ cdef class ResponseBuffer: cdef: unsigned long long buf_loc, buf_sz, slice_sz signed long long slice_start - object gen, source, _exception_tag, open_marker, close_marker, carryover, exception_buf, last_message_data, current_chunk + readonly object gen + object source, _exception_tag, open_marker, close_marker, carryover, exception_buf, last_message_data, current_chunk char* buffer char* slice unsigned char _read_byte_load(self) except ?255 diff --git a/docs/additional-options.mdx b/docs/additional-options.mdx index 03a4c2a5..3e659bf4 100644 --- a/docs/additional-options.mdx +++ b/docs/additional-options.mdx @@ -36,6 +36,7 @@ The following global settings are currently defined: | `naive_datetime_binding` | `"wall"` | `"wall"`, `"legacy"` | Controls binding for naive `datetime` query parameters. `wall` formats naive datetimes verbatim. `legacy` restores the older host-local conversion behavior. Attach `tzinfo` to preserve an instant. | | `naive_datetime_insert` | `"local"` | `"local"`, `"server"` | Controls Python object inserts of naive `datetime` values and naive ISO strings accepted by `DateTime64`. `local` uses the process timezone for compatibility. `server` uses the declared column timezone, then the server timezone. `datetime64`-dtype NumPy and Pandas columns are unchanged. | | `max_connection_age` | `600` | Any number of seconds | Maximum age for a reused HTTP keep-alive connection. Rotation helps distribute connections across nodes behind a load balancer. | +| `native_codec` | `"python"` | `"python"`, `"rust"`, `"rust_strict"` | Default codec for client-managed Native format traffic, overridable per client. The `CLICKHOUSE_CONNECT_NATIVE_CODEC` environment variable seeds this setting at import. See [Rust codec](/integrations/language-clients/python/rust-codec). | | `product_name` | `""` | Any string | Product identifier added to client information. Use a value such as `"my-product/1.0"`. | | `readonly` | `0` | `0`, `1` | Deprecated no-op retained for 1.x compatibility. The client reads the server's `readonly` setting directly. | | `send_os_user` | `True` | `True`, `False` | Include the detected operating system user in client information. | diff --git a/docs/driver-api.mdx b/docs/driver-api.mdx index c92c4269..543f8613 100644 --- a/docs/driver-api.mdx +++ b/docs/driver-api.mdx @@ -50,6 +50,7 @@ Use `clickhouse_connect.get_client` to create a synchronous `Client`, or install | `show_clickhouse_errors` | bool, boolean string, `"scrub"`, or None | `True` | Controls `str(exc)` for server errors, transport errors, and mid-stream `StreamFailureError`. `True` includes the request URL and server version trailer. `"scrub"` keeps the SQL error text and symbolic name but strips the host/URL and `(version ...)` trailer. `False` returns a generic message (`code` is still set for server errors). Boolean strings are accepted. Other strings raise `ProgrammingError`. For transport errors, `__cause__` and tracebacks still contain the original transport exception. | | `proxy_path` | str | `""` | Path prefix added to the server URL when routing through a proxy. | | `form_encode_query_params` | bool | `False` | Always place query parameters in the form-encoded request body. Large non-binary parameter payloads are moved automatically even when this is false. | +| `native_codec` | str or None | Global setting, `"python"` | Experimental codec for client-managed Native format traffic: `"python"`, `"rust"`, or `"rust_strict"`. The Rust values require the `clickhouse-connect-core` wheel. See [Rust codec](/integrations/language-clients/python/rust-codec). | | `rename_response_column` | str or None | `None` | Column renaming strategy: `"remove_prefix"`, `"to_camelcase"`, `"to_camelcase_without_prefix"`, `"to_underscore"`, or `"to_underscore_without_prefix"`. | The async factory also accepts `connector_limit=100`, `connector_limit_per_host=20`, and `keepalive_timeout=30.0` to configure its aiohttp connection pool. It does not accept `pool_mgr`. The synchronous chDB backend accepts `path` and `chdb_options`; see [Embedded chDB backend](#embedded-chdb-backend). diff --git a/docs/index.mdx b/docs/index.mdx index 46aaeca7..bb707744 100644 --- a/docs/index.mdx +++ b/docs/index.mdx @@ -17,7 +17,7 @@ ClickHouse Connect is a core database driver providing interoperability with a w - The main interfaces are the synchronous `Client` and native aiohttp-based `AsyncClient` in `clickhouse_connect.driver`. The driver package also provides query and insert contexts, streaming helpers, DB-API support, and lower-level HTTP methods. - The `clickhouse_connect.datatypes` package serializes and deserializes ClickHouse types using the ClickHouse Native binary columnar format. -- The optional Cython extensions in `clickhouse_connect.driverc` accelerate common serialization, conversion, and buffering paths. A pure Python path remains available on platforms where the extensions cannot be built. +- The optional Cython extensions in `clickhouse_connect.driverc` accelerate common serialization, conversion, and buffering paths. A pure Python path remains available on platforms where the extensions cannot be built. An experimental opt-in [Rust codec](/integrations/language-clients/python/rust-codec) can replace Native format processing entirely. - The package ships PEP 561 type information, so downstream type checkers consume annotations for the public driver, DB-API, and SQLAlchemy surfaces. - The [SQLAlchemy](https://www.sqlalchemy.org/) dialect in `clickhouse_connect.cc_sqlalchemy` supports SQLAlchemy Core, schema reflection, ClickHouse-specific query clauses and table engines, and Alembic migrations. Basic ORM reads and inserts work, but the dialect is designed for analytical workloads rather than full unit-of-work ORM behavior. - The core driver and [ClickHouse Connect SQLAlchemy](/integrations/language-clients/python/sqlalchemy) implementation are the preferred method for connecting ClickHouse to Apache Superset. Use the `ClickHouse Connect` database connection, or `clickhousedb` SQLAlchemy dialect connection string. @@ -60,6 +60,7 @@ pip install "clickhouse-connect[polars]" # Polars pip install "clickhouse-connect[sqlalchemy]" # SQLAlchemy dialect pip install "clickhouse-connect[alembic]" # SQLAlchemy and Alembic pip install "clickhouse-connect[chdb]" # Embedded chDB backend +pip install "clickhouse-connect[rust]" # Experimental compiled Rust codec pip install "clickhouse-connect[tzdata]" # IANA time zones on minimal systems ``` diff --git a/docs/navigation.json b/docs/navigation.json index 055c2fb3..19acdfc9 100644 --- a/docs/navigation.json +++ b/docs/navigation.json @@ -9,6 +9,7 @@ "integrations/language-clients/python/advanced-querying", "integrations/language-clients/python/advanced-inserting", "integrations/language-clients/python/advanced-usage", + "integrations/language-clients/python/rust-codec", "integrations/language-clients/python/sqlalchemy" ] } \ No newline at end of file diff --git a/docs/rust-codec.mdx b/docs/rust-codec.mdx new file mode 100644 index 00000000..82c65c7e --- /dev/null +++ b/docs/rust-codec.mdx @@ -0,0 +1,73 @@ +--- +sidebarTitle: 'Rust codec' +keywords: ['clickhouse', 'python', 'rust', 'native', 'codec', 'performance'] +description: 'Opt-in compiled Rust codec for ClickHouse Connect' +slug: /integrations/language-clients/python/rust-codec +title: 'Rust native codec' +doc_type: 'reference' +--- + +ClickHouse Connect can decode query results and encode inserts with a compiled Rust codec instead of the default Python and Cython implementation. The Rust codec applies only to client-managed `FORMAT Native` traffic, which covers `query`, `query_np`, `query_df`, their block and row streaming variants, and inserts including `insert_df`. The Arrow methods use `FORMAT Arrow` and are unaffected, as are raw queries, raw inserts, and non-Native formats. + +The codec is experimental and opt in. The Python codec remains the default. + +## Installation {#installation} + +The compiled codec ships as a separate wheel named `clickhouse-connect-core`, which provides the `_ch_core` extension module. Install it through the `rust` extra: + +```bash +pip install "clickhouse-connect[rust]" +``` + +If a Rust codec is selected and the module is not installed, client creation raises a `NotSupportedError` naming this install command. + +## Enabling the codec {#enabling-the-codec} + +Select the codec with the `native_codec` client option: + +```python +import clickhouse_connect + +client = clickhouse_connect.get_client(host="localhost", native_codec="rust") +``` + +Accepted values: + +| Value | Behavior | +|---|---| +| `python` | Default. The existing Python and Cython codec. | +| `rust` | Prefer the Rust codec. Queries with unsupported options and inserts with unsupported types route to the Python codec. | +| `rust_strict` | Require the Rust codec. Unsupported options and types raise instead of routing. | + +The default can also be seeded with the `native_codec` common setting or the `CLICKHOUSE_CONNECT_NATIVE_CODEC` environment variable. Precedence is the client keyword argument, then the common setting, then the environment variable. + +The option is ignored for `interface="chdb"` clients, which always use the Python codec. + +## Fallback rules {#fallback-rules} + +Fallback decisions are made before any bytes are read or sent, so there is never a mid-stream codec switch. For queries the choice happens before the response body is consumed. For inserts the Rust encoder is only selected when every column type is supported, otherwise the whole insert runs on the Python codec. + +When `naive_datetime_insert="server"` is active, `rust` routes an insert containing any `DateTime` or `DateTime64` column to the Python codec so the declared column timezone or server timezone is applied. `rust_strict` rejects that combination. The default `naive_datetime_insert="local"` mode continues to use the Rust encoder. + +Driver-internal metadata queries always use the Python codec, silently, in every mode. + +Malformed Native payloads detected by the Rust codec raise `DataError`. + +## Versioning {#versioning} + +`clickhouse-connect-core` versions independently of `clickhouse-connect`. The driver declares a compatible range through the `rust` extra, and the module exports a binding API version that the driver checks when a Rust codec is selected. If the installed wheel is too old for the driver, client creation raises a `NotSupportedError` naming the upgrade command: + +```bash +pip install --upgrade clickhouse-connect-core +``` + +Codec fixes and performance improvements ship as `clickhouse-connect-core` releases and can be picked up with a wheel upgrade alone, without waiting for a `clickhouse-connect` release. + +## Known behavior differences {#known-behavior-differences} + +The Rust codec targets cell for cell parity with the Python codec. The deliberate differences: + +- `query_np` and `query_df` results for `Variant` columns contain plain Python objects rather than numpy scalar values. The values are equal, the cell types differ. +- A `LowCardinality` alternative inside a container that materializes per cell, such as `Array(Variant(...))`, produces value-equal cells that do not share the per-dictionary-slot object identity the Python codec exhibits. +- `Nullable(Tuple)` columns decode correctly on the Rust codec. The Python codec misreads this layout and the Rust result is the reference behavior. +- `rust_strict` rejects query options the Rust path does not implement, such as custom per-query `query_formats`, rather than silently changing behavior. diff --git a/pyproject.toml b/pyproject.toml index 0a96625e..5926ab94 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ module = [ "pyarrow.*", "tzlocal.*", "ujson.*", + "_ch_core", "clickhouse_connect.driverc.*", ] ignore_missing_imports = true diff --git a/rust/BINDING_ARCHITECTURE.md b/rust/BINDING_ARCHITECTURE.md new file mode 100644 index 00000000..ebb33c48 --- /dev/null +++ b/rust/BINDING_ARCHITECTURE.md @@ -0,0 +1,178 @@ +# The _ch_core binding + +How the PyO3 binding layer works and how to consume it efficiently. The +core's wire and Arrow contract is documented in ch-core-rs +`DECODER_CONTRACT.md`. This document covers the binding crate `ch-core-py` +and the integration pattern above it. + +## Design + +The core turns ClickHouse Native wire bytes into immutable columnar memory, +one decoded chunk per Native block, with the type system resolved into a +schema. The binding moves that memory into Python along the cheapest legal +path for what the caller wants. It contains no decode logic and never +inspects wire bytes. + +``` + network bytes + | + [intake: decode_native or StreamDecoder.feed] + GIL released while the core decodes + | + ColBatch + Arc-shared, immutable, chunked columnar memory + / \ + [Arrow exit] [Python object exit] + __arrow_c_stream__ to_python_rows / columns + pointer handoff, no copy one PyObject per cell + | | + pyarrow / polars / pandas list of tuples / lists +``` + +Three ideas carry the design: + +1. **One result object, many views.** `ColBatch` wraps an `Arc` over the + decoded chunks. Every consumption path is a view or conversion of the + same memory. Nothing is re-decoded, chunks are never concatenated, and + merging streamed batches with `ColBatch.from_batches` is reference + counting only. + +2. **You pay only at the exit you choose.** The Arrow exit hands consumers + raw buffer pointers and costs near zero at any row count. The Python + object exit allocates one object per cell and dominates decode itself. + Going to pandas, polars, numpy, or any Arrow consumer, use the Arrow + exit. Never round-trip through Python objects to reach a dataframe. + +3. **The GIL is released wherever Python memory is not touched.** Intake + copies each fed chunk out of Python-owned memory, then decodes with the + GIL released. That is what lets a producer thread keep reading the + socket while Rust decodes. + +The division of labor is strict: + +| layer | owns | examples | +|---|---|---| +| core (ch-core-rs) | ClickHouse knowledge | wire framing, type parsing, schema rules, Arrow layout | +| binding (ch-core-py) | Python value policy | what a `DateTime64(6,'America/New_York')` becomes, GIL rules, exception types | +| integration (driver `rustcodec.py`, POC `rust_client.py`) | transport | HTTP, decompression, threads, queues, connection cleanup | + +The core knows nothing about Python, the binding knows nothing about HTTP, +and the integration layer never sees a wire byte. The binding is the layer +rewritten per language. + +## Intake + +Four entry points, all funneling into the same core decode: + +| surface | shape | use | +|---|---|---| +| `ColBatch.decode_native(data)` | whole payload to whole result | buffered fetch | +| `StreamDecoder.feed/finish` | push chunks, get completed blocks | async or thread streaming | +| `BlockDecoder(data)` | iterate blocks of an in-memory buffer | block at a time over buffered bytes | +| `PipeDecoder(read_fd)` | iterate blocks read from a pipe fd | producer writes to a pipe | + +`StreamDecoder` copies each fed chunk before releasing the GIL. The copy is +mandatory for correctness: with the GIL released, another thread may mutate +or free a `bytearray` mid-decode. `decode_native` makes the one exception, +a `bytes` input is immutable so it decodes from the borrowed buffer with no +copy at the price of holding the GIL, the right trade for a one-shot +buffered call. + +Pass `has_block_info=True` whenever the request pinned +`client_protocol_version`, which preserves bare DateTime timezones. + +Schema uniformity across blocks is enforced in the core for every surface. +A mid-stream schema change raises `ValueError` rather than producing a +result that would corrupt row materialization. `from_batches` validates +schema equality across inputs and keeps working column names and types even +when every chunk is empty. + +## Exits + +### Arrow + +`__arrow_c_stream__` exports one Arrow record batch per decoded chunk. The +exported buffers point straight into Rust-owned column memory, kept alive +by `Arc` references held in the stream's private data until the consumer +invokes the Arrow `release` callback. The capsule follows the Arrow +PyCapsule protocol: one capsule, one consumer, and the capsule destructor +releases the stream if it is dropped unconsumed. + +Date exports as raw `uint16` days and DateTime as raw `uint32` epoch +seconds, the core's documented zero-copy choice. DateTime64 and Date32 +carry real Arrow temporal types. Consumers needing temporal dtypes for Date +and DateTime convert after import or use the object exit. + +### Python objects + +Built directly on the CPython C API: result lists are preallocated once and +filled with `PyList_SET_ITEM`/`PyTuple_SET_ITEM` ownership transfer, the +raw list pointer is bound into a managed reference immediately so a panic +mid-fill drops a partially filled list instead of leaking it, and the +validity branch is hoisted out of the per-cell loop so non-nullable columns +run with no per-cell null check. Rows, columns, and single-column paths +share one cell constructor, so value policy cannot drift between them. + +## Value policy + +The binding's reason to exist. The values match the Python codec cell for +cell except where documented in `docs/rust-codec.mdx`. + +- **Strings.** CPython's own UTF-8 decode is the validation. Only on + failure does the fallback render the raw bytes as lowercase hex, matching + the driver's String policy. +- **Temporals.** A per-column context resolves timezone policy once, not + per cell. Naive columns are built by pure epoch arithmetic with no + datetime parsing in the loop. Named non-UTC zones go through + `datetime.fromtimestamp` for DST correctness. +- **Aggregate states.** Exact serialized state `bytes`, uninterpreted. The + Arrow exit exports them zero copy as LargeBinary. +- **Variant and Dynamic.** Intrinsic NULL becomes `None`, unambiguous + alternatives use ordinary Python values, and alternatives sharing a + Python type use `typed_variant`. The Arrow exit is the core's zero-copy + dense union. +- **Everything else** is a direct `PyLong`, `PyFloat`, `PyBool`, or + `PyBytes` constructor call, with wide integers built exactly from their + little-endian buffers. + +## Errors + +One shared mapping from core errors: + +| core error | Python exception | +|---|---| +| unsupported type, invalid block, schema mismatch, corruption | `ValueError`, named column where known | +| unexpected EOF | `EOFError`, truncation is distinct and retryable | +| other IO | `RuntimeError` | + +The driver layer translates these into `DataError` for its public surface. + +## Streaming pattern + +A producer thread reads and decompresses the socket while a consumer thread +calls `StreamDecoder.feed`. Both make progress at once because neither +holds the GIL during its expensive part. The reference implementation is +the driver's `rustcodec.py`, with `rust_client.py` as the standalone POC +version. The essentials of the protocol: + +- Queue items are tagged data, error, or EOF, and the producer uses a + timeout put-loop with a stop flag so it can never block forever on a full + queue. +- `finish()` runs only after a clean EOF, so a transport error is never + misread as a truncated stream. +- The response is closed rather than pooled on any failure, so a half-read + connection never returns to the pool. +- Per-block batches accumulate as `Arc` wrappers and one `from_batches` + merge at the end produces a result shape-identical to `decode_native`, + so every exit works unchanged on streamed data. + +## Compression + +HTTP `Accept-Encoding` compression lives entirely in the integration layer, +upstream of the decoder. The core and the binding never see a compressed +byte. lz4 frames and zstd streams are generic transport formats, not +ClickHouse knowledge, so they do not belong in the zero-dependency core. +ClickHouse's own compressed-block framing, used by the TCP protocol and the +`compress=1` HTTP mode, is protocol knowledge and would belong in the core +behind a cargo feature if ever needed. clickhouse-connect does not use that +mode. diff --git a/rust/COMPLETENESS.md b/rust/COMPLETENESS.md new file mode 100644 index 00000000..4c78b9e8 --- /dev/null +++ b/rust/COMPLETENESS.md @@ -0,0 +1,325 @@ +# ch-core-rs Binding Completeness + +This document tracks what the `clickhouse-connect` binding has exposed from +`ch-core-rs`. The upstream core tracker lives in +`/Users/joe/dev/ch-core-rs/COMPLETENESS.md`; this file is the binding-side +handoff for Python integration work. + +## Context Handoff + +- **Current state:** `_ch_core` has the query decode path bound through + `ColBatch`, `StreamDecoder`, `BlockDecoder`, and Unix `PipeDecoder`. Both the + query decode and insert encode paths are now wired into the driver behind the + public `native_codec` selector. The Python serializer remains the default + compatibility path. The temporary `transport_settings={"rust_insert": ...}` + hook has been removed. +- **Codec selector:** clients take `native_codec` with values `python` + (default), `rust`, and `rust_strict`. There is a matching common setting and + the `CLICKHOUSE_CONNECT_NATIVE_CODEC` environment variable that seeds the + default. The seam lives in `clickhouse_connect/driver/rustcodec.py`. `rust` + routes ineligible query contexts and unsupported insert types to Python; + `rust_strict` raises. Query decode makes the Rust vs Python choice before + reading the response body, so there is no mid-stream query fallback. +- **Build note:** the streamed Rust decode path reads decompressed chunks + straight from the response buffer, so the compiled Cython `ResponseBuffer` + must expose `.gen`. `driverc/buffer.pxd` now declares `readonly object gen`, + which requires a rebuild (`python setup.py build_ext --inplace`). Environments + running the Rust query path against a stale `.so` will hit `AttributeError` on + `.gen`. +- **Upstream pin:** `ch-core-rs` currently tracks encode against + ClickHouse `v26.6.1.1193-stable`, protocol revision `54485`. HTTP Native + inserts use `EncodeOptions { protocol_revision: 0 }`, so there is no + `BlockInfo` preamble and no per-column custom-serialization marker. +- **Current type scope:** the first binding encoder targets the upstream + encodable set: `Nothing`, `Bool`, fixed-width numerics, floats, `String`, + `FixedString(N)`, `Date`, `Date32`, `DateTime`, `DateTime64(P[, tz])`, `UUID`, + `IPv4`, `IPv6`, `Enum8`/`Enum16`, `Decimal(P, S)`, `QBit(T, D)`, `Array(T)`, `Tuple(...)` + (named or unnamed, including `Nullable(Tuple)`), `Map(K, V)`, `Variant(...)`, + `Dynamic`, `JSON`, `Nullable(T)`, and `LowCardinality(T)` where the upstream + core permits it. It also covers the three + name-decoration alias families the core resolves through `ChType::physical_delegate`: + `SimpleAggregateFunction(func, T)`, the six geo types (`Point`, `Ring`, `LineString`, + `MultiLineString`, `Polygon`, `MultiPolygon`), and `Nested(...)`. Registered + `AggregateFunction(...)` signatures use the core's function-specific state + framing and validation. +- **numpy/pandas:** `query_np`, `query_df`, and their block and row stream variants + now route through the Rust codec via the zero-copy Arrow exit + (`clickhouse_connect/driver/rustnumpy.py`). Per-column dtype conversion is driven by + the driver's own `ClickHouseType` (np_type, tzinfo, nullability) so dtypes match the + Python codec. The Arrow export is raw (Date is uint16 days, DateTime is uint32 seconds + with the timezone dropped, Enum is raw ints), so the converters, not a naive + `to_pandas()`, produce the final columns. Nullable numeric columns in `query_df` build + pandas extension arrays directly from the Arrow buffers. LowCardinality columns route through the object + exit regardless of inner type (values are correct there). Note: the Python codec truncates + `LowCardinality()` numpy/pandas output to the dictionary length (`ArrayType._build_lc_column` in + `datatypes/base.py` passes `count=len(index)` where it should be `count=len(keys)`), so there is no clean + parity target for those rare suspicious types; the rust object exit returns the full column. +- **Type coverage: COMPLETE.** The Rust codec now covers every scalar type in the + supported set plus `Array(T)`, `Tuple(...)`, and `Map(K, V)` over any supported + element type in BOTH directions. Decode policy: an unnamed tuple materializes as + a Python `tuple`, a named tuple as a `dict` keyed by the element names, and a map + as a `dict` (wire order, last duplicate key wins), matching the Python codec's + default read format. Top-level Tuple and Map columns decode through hoisted + column-major fills (`fill_tuple`/`fill_map` in + `ch-core-py/src/pyval/containers.rs`) with per-chunk LowCardinality slot caching, + so LC fields share object identity like top-level LC columns. Encode policy mirrors + the Python codec's acceptance: Tuple rows are any + positional iterable, or dicts for named tuples (read via element name, missing + keys become `None`); Map rows must be dict-like. Intentional divergences from the + Python codec, where the Rust path is the reference: a wrong-arity Tuple row raises + a clear error where Python silently truncates the block, `Nullable(Tuple(...))` + works in both directions where the Python codec misdecodes reads and cannot insert + at all, and the zero-element `Tuple()` decodes correctly at the binding level (the + driver's Python type registry still cannot parse `Tuple()`, so it does not + round-trip through `get_client` under either codec). The rustcodec insert seam + passes `col_type.name` (not `insert_name`) to the encoder because + `Tuple.insert_name` drops the `Nullable` wrapper. Query decode picks Rust or + Python before reading the response body, so once a query is eligible there is no + mid-stream fallback: an unsupported type raises mid-stream in both `rust` and + `rust_strict` mode rather than routing to Python. +- **Nothing.** Decode materializes both bare `Nothing` and `Nullable(Nothing)` + as Python `None` regardless of the latter's structural null map. The flat + object exit uses a column-wide fill, while recursive Array/Tuple/Map paths use + the same per-cell policy. Arrow remains the core's zero-copy Null export, and + the pandas exits ride the object path via a `Nothing._finalize_column` + override in the Python datatypes (`base_type` is not a pandas dtype). + Encode builds a length-only column: bare `Nothing` ignores Python placeholder + values, matching the Python codec, and `Nullable(Nothing)` scans them only to + retain the Native null map before the canonical one-byte marker run. + `LowCardinality(Nothing)` remains invalid, as required by ClickHouse. +- **QBit.** Decode materializes one fixed-length Python float list per row, + preserves parent nullability, and composes through Array, Tuple, Map, and + Variant. Insert dispatches the element width once and fills one final typed + child vector. Exact lists and tuples use CPython fast paths, while contiguous + two-dimensional native-endian PEP 3118 buffers use a bulk copy or conversion + path with exact shape validation. The core performs the bit-plane transpose + directly between that row-major child and Native bytes. Arrow remains the + core's zero-copy FixedSizeList export. +- **Variant.** Decode preserves the server's intrinsic NULL as Python `None` + and scatters each dense alternative directly into its logical row positions, + including nested Array, Tuple, and Map shapes. Insert scans logical rows once, + stores one discriminator byte per row, and builds each dense child once through + the existing per-type fast paths. Automatic selection matches the Python + codec's exact `type(value)` policy; ambiguous alternatives use + `typed_variant(value, type_name)`. Arrow export remains the core's zero-copy + dense union. Two deliberate divergences from the Python codec: `query_np`/ + `query_df` Variant cells are plain Python objects rather than value-equal + numpy scalars, and a `LowCardinality` alternative inside a per-cell- + materialized container (`Array(Variant(...))`) does not share + per-dictionary-slot objects (the per-cell Tuple/Map arms behave the same). +- **Dynamic.** Decode materializes each block-local typed child through the + same value machinery as a standalone column, including Dynamic nested under + Array, Tuple, and Map. Intrinsic NULL rows become `None`. The Python object + exits decode SharedVariant cells (binary type descriptor plus single-value + payload) to the same typed values as ordinary columns for every supported + type; AggregateFunction states and unsupported descriptors stay exact Python + `bytes`, and malformed cells raise a column-named `ValueError`. The Python + codec only heuristically decodes int/float/str/bool shared cells, so the rust + path is more complete there (FINDINGS.md finding 4). Arrow export is the + core's result-wide, name-sorted dense union, including its union-of-unions + representation and clean failure above 16,256 children; shared cells stay + `bytes` in Arrow because the schema needs a stable child type. Dynamic object + cells in np/df output carry the general np-scalar residue: rust yields + python-native scalars where the Python codec yields value-equal numpy + scalars. Insert builds the driver's established String input column natively + with exact `str(value)` parity (`None` becomes the literal `"NULL"`), + producing the same wire bytes as the Python codec, so the server keeps its + setting-dependent text inference and both `native_codec="rust"` and + `"rust_strict"` insert Dynamic without a fallback. +- **JSON.** Decode supports the core's V1, V2, flattened, and String Native + layouts. Structured Python exits allocate one dictionary per valid row and + fill typed and dynamic paths column-major, then decode shared-data cells with + cached path and type-descriptor metadata. This avoids cloning Dynamic columns + or performing a Python mapping lookup for every path segment. Escaped dots in + path names remain literal keys. Arrow export delegates to the core's + structured arrays and dynamic union representation; String-layout output is + UTF-8. Insert accepts the same dictionary and JSON object string inputs as the + Python codec, serializes each Python object once with the driver's configured + JSON serializer, and hands a contiguous UTF-8 column to the core's JSON text + encoder. Nullable and Array/Tuple/Map/Variant nesting compose through the + ordinary container paths. ClickHouse 24.8-24.9 requires the legacy String + column header for JSON inserts, so non-strict mode selects the Python encoder + before producing bytes and strict mode raises a targeted compatibility error. +- **AggregateFunction.** The Python object exit returns each supported + function's exact serialized state as `bytes`; Arrow remains the core's + zero-copy LargeBinary export. Insert accepts bytes-like state values and + builds one i64 offsets run plus one contiguous data buffer, then delegates + signature and per-state validation to the core encoder. This includes base + `sum` over a nullable numeric, Decimal, BFloat16, or Enum argument, whose + one-byte presence flag and conditional accumulator remain part of the opaque + state slice rather than becoming Python or Arrow nullability. Scalar, Array, + Tuple, Array(Tuple), and Map-value shapes compose through the existing + container machinery. Direct Nullable and LowCardinality wrappers remain + server-illegal, and aggregate signatures without a registered core boundary + codec remain unsupported because Native has no generic per-state length + framing. Null rows in a legal `Nullable(Tuple(...))` containing an aggregate + state are rejected explicitly on insert until the core exposes a canonical + valid placeholder state. Decode of such null rows works, the server writes + placeholder states under the null mask. +- **Name-decoration aliases (`SimpleAggregateFunction`, geo, `Nested`).** These add + no new `Column` variant. The core resolves each to a physical type via + `ChType::physical_delegate` (`Point` -> unnamed `Tuple(Float64, Float64)`, the five + array-shaped geo kinds -> nested `Array(...(Point))`, `Nested(names...)` -> + `Array(Tuple(named...))`, `SimpleAggregateFunction(func, T)` -> `T`), and the binding + calls the same `physical_delegate` at each ChType-dispatch point (`prepare_column_ctx` + on decode; `build_column`/`build_element_column`/`build_nullable_column`/ + `default_pyobject` on encode) and recurses into the existing Tuple/Array/scalar + machinery. Value shapes therefore fall out for free and match the Python codec: a + `Point` reads as a 2-`tuple`, a `Ring` as a list of 2-tuples, a `Nested` as a list of + `dict`s keyed by the field names, and a `SimpleAggregateFunction` as its inner type's + value. `Nullable(Point)` works both directions. No driver-side (`rustcodec.py`, + `rustnumpy.py`) change was needed: query eligibility is context-level and insert + support is probed by a zero-row `encode_native_block`, so these types route to Rust + automatically once the binding supports them. +- **Deferred (LowCardinality object identity below Array-nested containers):** a + `LowCardinality` value inside a container that is itself nested under `Array` + (for example `Array(Tuple(a LowCardinality(String)))`) is rebuilt as a fresh + Python object per occurrence. Top-level `Tuple`/`Map` fields and + `Array(LowCardinality)` share identity per chunk; the residual gap is only on + the per-cell path for deeper nesting, where each LC sub-column would need its + own slot cache threaded through the recursion. Values compare equal; this is a + memory/allocation gap only. Track before flipping the default; acceptable while + the codec is opt-in. +- **insert_df bulk encode (follow-up, ch-core-py workstream):** `insert_df` is correct + under the Rust codec but not yet faster than Python. `encode_native_block` takes Python + columnar values one value at a time. A buffer-protocol or `ArrowArrayStream` import entry + for `encode_native_block` would let the encoder consume numpy/Arrow buffers without the + per-value Python round-trip. Track this in `ch-core-py`. + +## Path Forward + +Where the effort stands and what to do next, in order. Type coverage and +per-type performance work are done: every supported type is measured in both +directions on `rust/profile_codecs.py` (5M-row single-column workloads, object +rows/columns plus df/np exits plus encode and e2e insert), every loser found was +fixed, and the Rust path now wins or ties everywhere on uncompressed localhost. +More depth on that micro matrix is low value. The remaining risk is breadth and +realism, then productionization. + +1. **Realistic benchmark pass (next).** Everything measured so far is + uncompressed localhost, which hides transport latency and therefore has never + actually tested the read-ahead/producer-thread overlap that motivates the + Rust architecture (localhost A/B showed read-ahead as a wash; its value + thesis is latency hiding). Run the existing workloads against a real network: + the Cloud staging instance wired up in `run_cloud_tests.sh` (TLS, compression + on) is the target. Prereq: `profile_codecs.py` hardcodes `localhost:8123` in + three places and needs host/port/TLS/password taken from the + `CLICKHOUSE_CONNECT_TEST_*` env vars. Measure at least: wide mixed reads, + nullable df (the big localhost win, does it survive TLS+lz4), streaming vs + buffered, insert paths, and 2-4 concurrent clients (GIL contention with + producer threads is unmeasured). Also flip `--compress` on locally for one + sweep since compression shifts the transport/decode balance. +2. **CI leg.** Nothing in CI builds `_ch_core` or runs the suite under + `rust`/`rust_strict`. Add a job that builds the extension (needs the + `ch-core-rs` checkout or a published crate, see next item) and runs + `ch-core-py/tests/` plus the integration suite with + `CLICKHOUSE_CONNECT_NATIVE_CODEC=rust_strict`. Until this exists every + refactor risks silently breaking the opt-in path. +3. **Distribution decisions.** The build is cp312-specific by design + (`_PyDict_NewPresized`, slot-offset reads, and the pyo3-ffi + `PyMemberDescrObject` cast all preclude abi3), so shipping means per-version + wheels. The core crate is a path dependency on a private working tree and + needs a publish-or-vendor decision before any wheel can build outside this + machine. +4. **Hardening before any default flip.** Malformed-block fuzzing against the + decoder (the encode side is validated by the core pre-write), a soak test for + the producer-thread insert path, and the open maintainer decisions below. +5. **Open maintainer decisions (queued, not blockers for opt-in):** + df/np element scalars inside Array/Tuple/Map cells are python-native under + Rust vs numpy scalars under Python (value-equal); `column_block_stream` + yields lists under Rust vs `array.array` under Python (exit-format + asymmetry, 0.17x on ints but zero-object python exit is the anomaly); + `row_block_stream` ints at 0.87x (needs a per-block `to_python_rows` through + a QueryResult shape change); `insert_df` bulk encode via a buffer/Arrow + import entry; file the Python-codec `Nullable(Tuple)` misdecode upstream + (it returns garbage values, server-verified); the Python type registry + cannot parse `Tuple()` under either codec. + +## Future Public Opt-In Design + +Status: implemented as `native_codec` (see Context Handoff above). This section is +retained as the design rationale. + +The release-facing opt-in should be a driver-level Native codec selector rather +than a transport setting. `transport_settings` reads as HTTP headers and +transport behavior, so it is not a good public home for "choose the Python or +Rust Native implementation". + +Recommended API shape: + +- Add a client/common setting named `native_codec`. +- Add an environment variable, for example + `CLICKHOUSE_CONNECT_NATIVE_CODEC=rust`. +- Support initial values: + - `python`: current behavior. This should remain the first public default. + - `rust`: prefer Rust for eligible client-managed `FORMAT Native` query and + insert paths, with Python fallback only where fallback is safe before bytes + are consumed or sent. + - `rust_strict` or `rust-only`: require Rust and fail fast when Rust is not + available or the type/path is unsupported. + +Suggested precedence: + +1. Per-call override, if a public override is later added. +2. Client constructor option, for example `get_client(native_codec="rust")`. +3. Environment variable. +4. Library default. + +Scope should stay precise: `native_codec` applies only to client-managed +`FORMAT Native` encode/decode. It should not affect Arrow, raw inserts, raw +queries, non-Native JSON formats, Parquet, or caller-provided byte streams. + +Fallback policy differs by direction: + +- Inserts can fall back from Rust to Python before the first encoded chunk is + yielded or sent. After bytes have been sent, local failures should surface + rather than silently switching encoders. +- Query decode cannot generally rewind the HTTP response once bytes have been + consumed. The client should choose Rust or Python before reading a response + stream. Mid-stream Rust decode failures should surface as errors, especially + in strict mode. + +Deprecation path: ship `native_codec="python"` as the default, encourage early +adopters and CI/performance testing to use `native_codec="rust"` or the env var, +then flip the default only after query and insert parity are stable. Keep a +temporary `python` escape hatch when the default changes. + +## Binding Checklist + +- [x] Decode bindings for buffered and streamed Native query results. +- [x] Arrow C Data Interface exit for decoded query results. +- [x] Python object exits for decoded rows and columns. +- [x] Insert block encoder binding from Python columnar values to + `ch_core_rs::ColBatch`. +- [x] Streaming sync insert source with producer-thread encoding and bounded + queue backpressure. +- [x] Async insert source selection for the Rust transform path. +- [x] Client integration with opt-in selector and Python fallback. +- [x] Binding unit coverage for supported inserts, malformed values, and + unsupported types. +- [x] Driver unit coverage for framing parity, fallback behavior, and producer + errors. +- [x] Sync and async integration coverage against a live ClickHouse server. + +## Notes + +- The binding owns Python value policy. The upstream core owns Native framing, + type metadata, and encoded block bytes. +- Python values must be copied into Rust-owned buffers before the GIL is + released. Encoding with `ch_core_rs::native::encode::encode_block` can then run + without the GIL. +- Unsupported binding-side conversion should raise a clear local exception. In + non-strict opt-in mode the driver falls back to the Python serializer when + Rust fails before any Rust-encoded bytes are sent. Single-block inserts are + fully encoded before the first yield so fallback remains safe. Multi-block + inserts keep the streaming overlap; once a Rust chunk has been yielded, later + Rust failures surface as local insert errors instead of switching encoders + mid-body. In strict mode it should surface the Rust-path failure. +- Server-source review at ClickHouse `v26.3.9.8-lts` confirmed that Native + `DateTime64` payloads are signed Int64 ticks at declared scale and that Native + Decimal deserialization does not enforce declared precision. The binding + therefore floor/euclidean-scales pre-epoch DateTime64 objects and rejects + over-precision Decimal values before encoding. +- `LowCardinality(Decimal...)`, `LowCardinality(Enum...)`, and + `LowCardinality(DateTime64...)` are not part of the current upstream LC + encodable set because the ClickHouse server forbids those inner types. diff --git a/rust/Cargo.lock b/rust/Cargo.lock new file mode 100644 index 00000000..44ecbf7f --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,139 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ch-core-py" +version = "0.1.0" +dependencies = [ + "ch-core-rs", + "pyo3", + "pyo3-build-config", +] + +[[package]] +name = "ch-core-rs" +version = "0.1.1" +source = "git+https://github.com/ClickHouse/ch-core-rs?tag=v0.1.1#303c09fc5c064505ade7e2c7bd5a12cd308de549" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "libc" +version = "0.2.185" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f" + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 00000000..cd6fb496 --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,7 @@ +[workspace] +resolver = "2" +members = ["ch-core-py"] + +[profile.release] +codegen-units = 1 +lto = "fat" diff --git a/rust/README.md b/rust/README.md new file mode 100644 index 00000000..3768d70d --- /dev/null +++ b/rust/README.md @@ -0,0 +1,170 @@ +# Rust core binding POC + +Every ClickHouse client reimplements wire decoding and the type system +today. `ch-core-rs` is a shared Rust core that does it once: it decodes +`FORMAT Native` bytes into Arrow-shaped columnar memory, and each language +client wraps it with a thin binding. This directory is the Python proof of +concept: a PyO3 binding (`_ch_core`), an end-to-end `query_rust` path over +the existing clickhouse-connect transport, and the verification and +benchmark harnesses behind the measured results. + +The mental model in five lines: + +``` +network bytes -> [decode in Rust, GIL released] -> ColBatch + | + +------------------------------------+ + | | + Arrow C stream capsule Python rows / columns + zero copy, ~free one object per cell + pyarrow / pandas / polars list of tuples / lists +``` + +## Quick look + +```python +import _ch_core +import pyarrow as pa + +batch = _ch_core.ColBatch.decode_native(native_bytes, has_block_info=True) +batch.column_names, batch.num_rows +rows = batch.to_python_rows() +table = pa.RecordBatchReader.from_stream(batch).read_all() +``` + +Or end to end on an existing clickhouse-connect client, streamed with +transport and decode overlapped: + +```python +from rust_client import query_rust + +result = query_rust(client, "SELECT id, name FROM events ORDER BY id") +result.result_rows +result.arrow_table() +result.to_pandas() +``` + +## Supported types + +Nothing, Bool, Int8 through Int64 plus Int128/256, UInt8 through UInt64 plus +UInt128/256, Float32/64, BFloat16, QBit, String, +FixedString, Date, Date32, DateTime, DateTime64 with precision and +timezone, Nullable, LowCardinality where ClickHouse permits it, Array, Tuple, +Map, Variant, the supported name-decoration aliases, and the function +signatures for `AggregateFunction` registered by the core. Dynamic query +decode is also supported, including nested shapes and Arrow's result-wide dense +union. Typed Dynamic children use their ordinary Python values, intrinsic NULL +uses `None`, and the Python object exits decode SharedVariant cells, including +compound JSON shared-data values, to the same typed values as ordinary columns; +AggregateFunction states and unsupported +descriptors stay exact Python `bytes`, and the Arrow exit keeps every shared +cell as `bytes` for schema stability. Variant uses `None` +for its intrinsic NULL, ordinary Python values for unambiguous alternatives, +and `typed_variant` for alternatives that share a Python type. Its Arrow exit +is the core's zero-copy dense union. Aggregate states materialize as exact +Python `bytes` and export zero-copy as Arrow LargeBinary. Dynamic insert +builds the driver's established String input column natively, with exact +`str(value)` parity (`None` becomes the literal `"NULL"`), so the server keeps +its setting-dependent text inference and both `native_codec="rust"` and +`"rust_strict"` insert Dynamic without a fallback. JSON query decode supports +the core's structured and text Native layouts, including typed, dynamic, and +shared paths. Python object exits reconstruct dictionaries with escaped path +segments preserved, while Arrow uses the core's structured zero-copy export. +JSON inserts accept dictionaries or JSON object strings and use the core's +text Native encoder; JSON also composes under Nullable, Array, Tuple, Map, and +Variant. Other unsupported types, plus +unsupported aggregate signatures, are rejected at decode time with a clean +`ValueError` naming the column; malformed payloads raise a column-named +`ValueError` as well. Type coverage lives in the core, so new types land there +once and every binding gets them. + +QBit object results materialize as fixed-length Python lists of floats, with +`None` at the parent level for nullable rows. Inserts accept row containers and +contiguous two-dimensional PEP 3118 buffers such as NumPy float32/float64 +matrices. The buffer path builds one typed child allocation, and Arrow exports +the core's zero-copy FixedSizeList representation. + +## Prerequisite: the core crate + +The binding depends on `ch-core-rs`, pinned by release tag in +`ch-core-py/Cargo.toml`: + +```toml +ch-core-rs = { git = "https://github.com/ClickHouse/ch-core-rs", tag = "v0.1.0" } +``` + +Release and CI builds resolve the tag directly, so the repository must be +reachable from the build environment. See `DISTRIBUTION_PLAN.md` at the repo +root for how the binding ships as the `clickhouse-connect-core` wheel. + +To develop against a local core checkout, add an untracked +`.cargo/config.toml` at this repository's root that patches the git source to +your working tree: + +```toml +[patch."https://github.com/ClickHouse/ch-core-rs"] +ch-core-rs = { path = "/path/to/ch-core-rs" } +``` + +With the patch in place, builds compile the core working tree as-is, which is +the intended inner loop for core changes. Remove or ignore the patch to build +what the tag pins. + +The committed `Cargo.lock` must record the git source and revision for +`ch-core-rs`. Any cargo resolution while the patch is active rewrites that +entry to the path form in your working tree. Discard those lock changes and +never commit them. To regenerate the lock legitimately, for example after a +repin, move the patch aside first: + +```sh +mv .cargo/config.toml .cargo/config.toml.disabled +cd rust && CARGO_NET_GIT_FETCH_WITH_CLI=true cargo fetch +mv ../.cargo/config.toml.disabled ../.cargo/config.toml +``` + +## Build and test + +```sh +pip install maturin +maturin develop --release -m rust/ch-core-py/Cargo.toml +python -m pytest rust/ch-core-py/tests/ +``` + +Test extras: `pytest`, and `pyarrow` for the Arrow round-trip tests. + +If the shell exports `VIRTUAL_ENV` pointing at a different environment than +the one on `PATH`, maturin installs into the exported one. Set +`VIRTUAL_ENV=/path/to/repo/.venv` explicitly when in doubt. + +## What is here + +| file | purpose | +|---------------------------------|---------------------------------------------------------------------------| +| `ch-core-py/` | the PyO3 binding crate, module `_ch_core` | +| `rust_client.py` | end-to-end `query_rust` over the existing clickhouse-connect transport | +| `e2e_query_check.py` | parity gate vs `client.query()` on a live server, run before benchmarking | +| `temporal_e2e_check.py` | cell-by-cell temporal value policy check incl. named timezones | +| `bench_native_decode_strict.py` | decode-isolated benchmark, value-parity gated | +| `bench_query_e2e.py` | start-to-finish query vs query benchmark | +| `streaming_demo.py` | sync and async streaming overlap demonstration | +| `BINDING_ARCHITECTURE.md` | how the binding works and how to use it for best performance | + +The live-server scripts expect ClickHouse on `localhost:8123` and a +clickhouse-connect checkout with compiled C extensions for the comparison +side, configurable via `CHC_BASELINE_PATH`. + +## Where to read next + +- `BINDING_ARCHITECTURE.md`: the layered design, the intake and exit paths, + GIL rules, the streaming pattern, and the practical performance guidance. +- The `ch-core-rs` repository: the core's own README, `ARCHITECTURE.md`, + and `DECODER_CONTRACT.md`, the per-type wire and Arrow contract. + +## Status + +The Native codec supports query decode, streaming, Arrow export, Python +materialization, and insert encoding through the `clickhouse_connect` +`native_codec="rust"` and `native_codec="rust_strict"` client options. It +does not implement the TCP protocol. The binding still builds against a +local or git checkout of the private core and is not part of the default +package build. diff --git a/rust/bench_native_decode_strict.py b/rust/bench_native_decode_strict.py new file mode 100644 index 00000000..c4ba868e --- /dev/null +++ b/rust/bench_native_decode_strict.py @@ -0,0 +1,339 @@ +"""Strict Native decode benchmark for clickhouse-connect vs _ch_core. + +This benchmark fetches ClickHouse FORMAT Native bytes once per workload, then +feeds the exact same bytes to both decoders. That avoids estimating decode CPU +as "full query time - raw_query time", which is useful but noisy on localhost. + +The comparison is intentionally destination-specific: + +* Python rows/columns measure Python-native materialization. +* NumPy/pandas measure the existing clickhouse-connect Native path. +* Arrow/polars/pandas Arrow measure the Rust columnar export path. + +Set CHC_BASELINE_PATH to a clickhouse-connect checkout with compiled C +extensions. Defaults to this repo's root, which only works if the extensions +are built there. +""" + +from __future__ import annotations + +import os +import statistics +import sys +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +BASELINE_PATH = os.environ.get("CHC_BASELINE_PATH", REPO_ROOT) +sys.path.insert(0, BASELINE_PATH) + +import _ch_core # noqa: E402 +import pandas as pd # noqa: E402 +import polars as pl # noqa: E402 +import pyarrow as pa # noqa: E402 + +import clickhouse_connect # noqa: E402 +from clickhouse_connect.driver.query import QueryContext # noqa: E402 +from clickhouse_connect.driver.transform import NativeTransform # noqa: E402 +from clickhouse_connect.driverc import dataconv # noqa: E402,F401 +from tests.helpers import bytes_source # noqa: E402 + +ITERS = int(os.environ.get("BENCH_ITERS", "5")) +ROW_LIMIT = int(os.environ.get("BENCH_ROW_LIMIT", "1000000")) + + +@dataclass(frozen=True) +class Workload: + name: str + query: str + row_materialization: bool = True + pandas_numpy: bool = True + + +WORKLOADS = ( + Workload( + "mixed_6col_100k", + "SELECT id, val, name, flag, small_int, big_uint FROM bench_types LIMIT 100000", + ), + Workload( + "mixed_6col_1M", + "SELECT id, val, name, flag, small_int, big_uint FROM bench_types LIMIT 1000000", + ), + Workload( + "int_3col_10M", + "SELECT id, small_int, big_uint FROM bench_types", + row_materialization=False, + pandas_numpy=False, + ), + Workload( + "string_1col_10M", + "SELECT name FROM bench_types", + row_materialization=False, + pandas_numpy=False, + ), + # Temporal: dt is a bare DateTime (revision-0 Native drops a DateTime's + # zone, so both decoders see naive UTC); dt64_ny keeps its zone in the + # type string at revision 0 and exercises the tz-aware path. + Workload( + "temporal_3col_1M", + "SELECT d, dt, dt64 FROM bench_temporal LIMIT 1000000", + ), + Workload( + "temporal_3col_10M", + "SELECT d, dt, dt64 FROM bench_temporal", + row_materialization=False, + pandas_numpy=False, + ), + Workload( + "dt64_tz_1M", + "SELECT dt64_ny FROM bench_nullable LIMIT 1000000", + ), + Workload( + "nullable_3col_1M", + "SELECT n_int, n_float, n_str FROM bench_nullable LIMIT 1000000", + ), + Workload( + "nullable_3col_10M", + "SELECT n_int, n_float, n_str FROM bench_nullable", + row_materialization=False, + pandas_numpy=False, + ), +) + +TABLE_ROWS = 10_000_000 + + +def ensure_tables(ch_client): + """Create the benchmark tables on first run. Deterministic data.""" + specs = { + "bench_temporal": f""" + CREATE TABLE bench_temporal ENGINE = MergeTree ORDER BY id AS + SELECT + number AS id, + toDate('2000-01-01') + (number % 30000) AS d, + toDateTime('2000-01-01 00:00:00') + number AS dt, + addMilliseconds(toDateTime64('2000-01-01 00:00:00.000', 3), number) AS dt64, + toDateTime('2000-01-01 00:00:00', 'America/New_York') + number AS dt_ny + FROM numbers({TABLE_ROWS}) + """, + "bench_nullable": f""" + CREATE TABLE bench_nullable ENGINE = MergeTree ORDER BY id AS + SELECT + number AS id, + IF(number % 10 = 0, NULL, toInt64(number)) AS n_int, + IF(number % 10 = 3, NULL, number / 7) AS n_float, + IF(number % 10 = 7, NULL, concat('user_', toString(number % 100000))) AS n_str, + addMicroseconds( + toDateTime64('2020-06-01 00:00:00.000000', 6, 'America/New_York'), + number + ) AS dt64_ny + FROM numbers({TABLE_ROWS}) + """, + } + for table, ddl in specs.items(): + if ch_client.command(f"EXISTS TABLE {table}"): + if int(ch_client.command(f"SELECT count() FROM {table}")) == TABLE_ROWS: + continue + ch_client.command(f"DROP TABLE {table}") + print(f"creating {table} ({TABLE_ROWS:,} rows) ...") + ch_client.command(ddl) + + +client = clickhouse_connect.get_client(host="localhost", port=8123, compress=False) +transform = NativeTransform() + + +def timeit(fn: Callable[[], Any], iters: int = ITERS) -> float: + fn() + fn() + samples = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + samples.append(time.perf_counter() - t0) + return statistics.median(samples) + + +def ms(seconds: float | None) -> str: + if seconds is None: + return " n/a" + return f"{seconds * 1000:9.1f} ms" + + +def speedup(existing: float | None, rust: float | None) -> str: + if existing is None or rust is None or rust == 0: + return " n/a" + return f"{existing / rust:6.2f}x" + + +def source(raw: bytes): + return bytes_source(raw, chunk_size=len(raw)) + + +# apply_server_tz=True replicates client.query() against a UTC server: naive +# UTC datetimes unless the column type carries its own zone. A bare +# QueryContext would fall through to the local zone and skew temporal parity. +def v1_rows(raw: bytes): + return transform.parse_response(source(raw), QueryContext(apply_server_tz=True)).result_rows + + +def v1_columns(raw: bytes): + return transform.parse_response(source(raw), QueryContext(apply_server_tz=True)).result_columns + + +def v1_numpy(raw: bytes): + ctx = QueryContext(use_numpy=True, apply_server_tz=True) + return transform.parse_response(source(raw), ctx).np_result + + +def v1_pandas(raw: bytes): + ctx = QueryContext(use_numpy=True, as_pandas=True, use_extended_dtypes=True, apply_server_tz=True) + return transform.parse_response(source(raw), ctx).df_result + + +def rust_batch(raw: bytes): + return _ch_core.ColBatch.decode_native(raw, has_block_info=False) + + +def rust_arrow(raw: bytes): + return pa.table(rust_batch(raw)) + + +def rust_rows(raw: bytes): + return rust_batch(raw).to_python_rows() + + +def rust_columns(raw: bytes): + return rust_batch(raw).to_python_columns() + + +def rust_numpy_dict(raw: bytes): + table = rust_arrow(raw) + return {field.name: column.to_numpy(zero_copy_only=False) for field, column in zip(table.schema, table.columns)} + + +def rust_pandas_numpy(raw: bytes): + return rust_arrow(raw).to_pandas() + + +def rust_pandas_arrow(raw: bytes): + return rust_arrow(raw).to_pandas(types_mapper=pd.ArrowDtype) + + +def rust_polars(raw: bytes): + return pl.from_arrow(rust_arrow(raw)) + + +def report_line(label: str, rust: float | None, existing: float | None): + print(f" {label:18s} rust {ms(rust)} existing {ms(existing)} speedup {speedup(existing, rust)}") + + +def native_bytes(query: str): + return client.raw_query(query, fmt="Native") + + +def rust_arrow_total(query: str): + return rust_arrow(native_bytes(query)) + + +def rust_polars_total(query: str): + return rust_polars(native_bytes(query)) + + +def rust_pandas_arrow_total(query: str): + return rust_pandas_arrow(native_bytes(query)) + + +def rust_pandas_numpy_total(query: str): + return rust_pandas_numpy(native_bytes(query)) + + +def rust_rows_total(query: str): + return rust_rows(native_bytes(query)) + + +def check_parity(raw: bytes, name: str): + """Apples-to-apples gate: both decoders must produce identical row values + from the same Native bytes before any timing is trusted.""" + rust = rust_rows(raw) + v1 = v1_rows(raw) + assert len(rust) == len(v1), f"{name}: row count {len(rust)} != {len(v1)}" + for i, (r, v) in enumerate(zip(rust, v1)): + assert r == tuple(v), f"{name}: row {i} mismatch: rust={r!r} v1={v!r}" + print(f" parity OK ({len(rust):,} rows identical)") + + +def run_workload(workload: Workload): + raw = native_bytes(workload.query) + batch = rust_batch(raw) + rows = batch.num_rows + nbytes = len(raw) + if rows > ROW_LIMIT: + row_enabled = False + else: + row_enabled = workload.row_materialization + + print(f"\n### {workload.name} ({rows:,} rows, {nbytes / 1e6:.1f} MB Native)") + if row_enabled: + check_parity(raw, workload.name) + + rust_decode_t = timeit(lambda: rust_batch(raw)) + rust_arrow_t = timeit(lambda: rust_arrow(raw)) + rust_polars_t = timeit(lambda: rust_polars(raw)) + rust_pandas_arrow_t = timeit(lambda: rust_pandas_arrow(raw)) + + print(f" rust decode floor {ms(rust_decode_t)}") + report_line("Arrow table", rust_arrow_t, None) + report_line("Polars", rust_polars_t, None) + report_line("pandas ArrowDtype", rust_pandas_arrow_t, None) + + if row_enabled: + report_line("Python columns", timeit(lambda: rust_columns(raw)), timeit(lambda: v1_columns(raw))) + report_line("Python rows", timeit(lambda: rust_rows(raw)), timeit(lambda: v1_rows(raw))) + else: + print(" Python columns skipped for this row count") + print(" Python rows skipped for this row count") + + if workload.pandas_numpy: + report_line("NumPy dict/array", timeit(lambda: rust_numpy_dict(raw)), timeit(lambda: v1_numpy(raw))) + report_line("pandas NumPy", timeit(lambda: rust_pandas_numpy(raw)), timeit(lambda: v1_pandas(raw))) + else: + print(" NumPy/pandas NumPy skipped for this row count") + + if row_enabled: + print(" end-to-end totals, including localhost fetch") + report_line("pyarrow.Table", timeit(lambda: rust_arrow_total(workload.query)), timeit(lambda: client.query_arrow(workload.query))) + report_line( + "Polars", + timeit(lambda: rust_polars_total(workload.query)), + timeit(lambda: client.query_df_arrow(workload.query, dataframe_library="polars")), + ) + report_line( + "pandas ArrowDtype", + timeit(lambda: rust_pandas_arrow_total(workload.query)), + timeit(lambda: client.query_df_arrow(workload.query, dataframe_library="pandas")), + ) + report_line( + "pandas NumPy", timeit(lambda: rust_pandas_numpy_total(workload.query)), timeit(lambda: client.query_df(workload.query)) + ) + report_line( + "Python rows", timeit(lambda: rust_rows_total(workload.query)), timeit(lambda: client.query(workload.query).result_rows) + ) + + +def main(): + print( + f"clickhouse-connect {clickhouse_connect.__version__} from {clickhouse_connect.__file__}\n" + f"C extension {dataconv.__file__}\n" + f"_ch_core {_ch_core.__version__} from {_ch_core.__file__}\n" + f"pyarrow {pa.__version__} pandas {pd.__version__} polars {pl.__version__} iters {ITERS}" + ) + ensure_tables(client) + for workload in WORKLOADS: + run_workload(workload) + + +if __name__ == "__main__": + main() diff --git a/rust/bench_query_e2e.py b/rust/bench_query_e2e.py new file mode 100644 index 00000000..34d1bdf1 --- /dev/null +++ b/rust/bench_query_e2e.py @@ -0,0 +1,229 @@ +"""End-to-end query benchmark: client.query() vs query_rust() on one client. + +Every line is a start-to-finish destination comparison: server execution, +HTTP transport, decompression, decode, and conversion, down two different +paths. The v1 path is clickhouse-connect as users call it. The rust path is +rust_client.query_rust. Decode-isolated numbers live in +bench_native_decode_strict.py. Run e2e_query_check.py first: timing means +nothing without the parity gate. + +Method: one discarded warmup per cell (which also primes the User-Agent +integration tags the destination helpers mutate), then ITERS timed pairs +with AB/BA alternating order, reported as median with min..max. + +Caveats printed with results: the v1 path sends wait_end_of_query=1 and the +rust path streams, by design of each path. Date and DateTime reach Arrow as +raw uint16/uint32 storage (DECODER_CONTRACT.md), so arrow-derived numpy and +pandas destinations carry raw ints for those columns where v1 carries +datetime64. The numbers attribute time as transport+server vs client decode, +not a true client/server split. + +Usage: bench_query_e2e.py full matrix + bench_query_e2e.py --memprobe ENGINE DEST WORKLOAD COMPRESS internal +""" + +from __future__ import annotations + +import os +import platform +import resource +import statistics +import subprocess +import sys +import time +from dataclasses import dataclass +from functools import partial + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +BASELINE_PATH = os.environ.get("CHC_BASELINE_PATH", REPO_ROOT) +sys.path.insert(0, BASELINE_PATH) +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import _ch_core # noqa: E402 +import pandas as pd # noqa: E402 +import pyarrow as pa # noqa: E402 +from rust_client import _decompressed_chunks, _mirrored_request, query_rust # noqa: E402 + +import clickhouse_connect # noqa: E402 +from clickhouse_connect.driverc import dataconv # noqa: E402 + +ITERS = int(os.environ.get("BENCH_ITERS", "7")) + + +@dataclass(frozen=True) +class Workload: + name: str + query: str + row_materialization: bool = True + pandas_numpy: bool = True + + +WORKLOADS = ( + Workload("mixed_6col_1M", "SELECT id, val, name, flag, small_int, big_uint FROM bench_types LIMIT 1000000"), + Workload("int_3col_10M", "SELECT id, small_int, big_uint FROM bench_types", False, False), + Workload("string_1col_10M", "SELECT name FROM bench_types", False, False), + Workload("temporal_3col_1M", "SELECT d, dt, dt64 FROM bench_temporal LIMIT 1000000"), + Workload("temporal_3col_10M", "SELECT d, dt, dt64 FROM bench_temporal", False, False), + Workload("dt64_tz_1M", "SELECT dt64_ny FROM bench_nullable LIMIT 1000000"), + Workload("nullable_3col_1M", "SELECT n_int, n_float, n_str FROM bench_nullable LIMIT 1000000"), + Workload("nullable_3col_10M", "SELECT n_int, n_float, n_str FROM bench_nullable", False, False), +) + +MEM_WORKLOADS = ("mixed_6col_1M", "string_1col_10M") + + +def rust_numpy(client, query): + table = query_rust(client, query).arrow_table() + return {field.name: column.to_numpy(zero_copy_only=False) for field, column in zip(table.schema, table.columns)} + + +DESTINATIONS = ( + ("rows", "rows", lambda c, q: c.query(q).result_rows, lambda c, q: query_rust(c, q).result_rows), + ("columns", "rows", lambda c, q: c.query(q).result_columns, lambda c, q: query_rust(c, q).result_columns), + ("numpy", "numpy", lambda c, q: c.query_np(q), rust_numpy), + ("pandas", "numpy", lambda c, q: c.query_df(q), lambda c, q: query_rust(c, q).arrow_table().to_pandas()), + ("arrow", "always", lambda c, q: c.query_arrow(q), lambda c, q: query_rust(c, q).arrow_table()), +) + + +def summarize(times): + return statistics.median(times), min(times), max(times) + + +def fmt(stats): + med, lo, hi = stats + return f"{med * 1000:8.1f} ms ({lo * 1000:7.1f}..{hi * 1000:7.1f})" + + +def bench_pair(v1_fn, rust_fn): + """One warmup each, then ITERS timed pairs in alternating order.""" + v1_fn() + rust_fn() + v1_times, rust_times = [], [] + for i in range(ITERS): + pair = ((v1_fn, v1_times), (rust_fn, rust_times)) + if i % 2: + pair = pair[::-1] + for fn, times in pair: + start = time.perf_counter() + fn() + times.append(time.perf_counter() - start) + return summarize(v1_times), summarize(rust_times) + + +def rust_streamed_batch(client, query): + return query_rust(client, query).batch + + +def rust_sequential_batch(client, query): + """Fetch fully, then decode: the buffered baseline for the same bytes.""" + response = _mirrored_request(client, query, None) + try: + data = b"".join(_decompressed_chunks(response)) + except BaseException: + response.close() + raise + response.release_conn() + return _ch_core.ColBatch.decode_native(data, has_block_info=bool(client.protocol_version)) + + +def assert_encoding(client, expected): + response = _mirrored_request(client, "SELECT 1", None) + encoding = response.headers.get("content-encoding") + b"".join(_decompressed_chunks(response)) + response.release_conn() + assert encoding == expected, f"content-encoding {encoding!r}, expected {expected!r}" + + +def run_workload(client, workload): + probe = query_rust(client, workload.query) + rows = probe.batch.num_rows if probe.batch is not None else 0 + del probe + print(f"\n### {workload.name} ({rows:,} rows)") + + for label, gate, v1_fn, rust_fn in DESTINATIONS: + if gate == "rows" and not workload.row_materialization: + continue + if gate == "numpy" and not workload.pandas_numpy: + continue + v1_stats, rust_stats = bench_pair(partial(v1_fn, client, workload.query), partial(rust_fn, client, workload.query)) + ratio = v1_stats[0] / rust_stats[0] if rust_stats[0] else float("inf") + print(f" {label:8s} v1 {fmt(v1_stats)} rust {fmt(rust_stats)} speedup {ratio:5.2f}x") + + # Pipeline comparison, not pure overlap isolation: the buffered side + # also pays the full-buffer join and a different decode entry point. + seq_stats, stream_stats = bench_pair( + lambda: rust_sequential_batch(client, workload.query), + lambda: rust_streamed_batch(client, workload.query), + ) + gain = (seq_stats[0] - stream_stats[0]) / seq_stats[0] * 100 if seq_stats[0] else 0.0 + print(f" pipeline buffered {fmt(seq_stats)} streamed {fmt(stream_stats)} streamed gain {gain:4.1f}%") + + +def rss_mb(): + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if sys.platform != "darwin": + peak *= 1024 + return peak / 1e6 + + +def memprobe(engine, dest, workload_name, compress): + query = next(w.query for w in WORKLOADS if w.name == workload_name) + client = clickhouse_connect.get_client(host="localhost", port=8123, compress=compress if compress != "False" else False) + if engine == "v1": + result = client.query(query).result_rows if dest == "rows" else client.query_arrow(query) + else: + r = query_rust(client, query) + result = r.result_rows if dest == "rows" else r.arrow_table() + print(f"{rss_mb():.1f}") + del result + + +def memory_snapshot(): + print("\n### peak RSS snapshot (single uncompressed run per cell, whole-process ru_maxrss in MB)") + for workload_name in MEM_WORKLOADS: + for dest in ("rows", "arrow"): + line = f" {workload_name:16s} {dest:6s}" + for engine in ("v1", "rust"): + proc = subprocess.run( + [sys.executable, os.path.abspath(__file__), "--memprobe", engine, dest, workload_name, "False"], + capture_output=True, + text=True, + check=True, + ) + line += f" {engine} {float(proc.stdout.strip()):8.1f} MB" + print(line) + + +def header(client, compress): + settings = {"client_protocol_version": client.protocol_version} + if client.compression and client._send_comp_setting: + settings["enable_http_compression"] = "1" + print( + f"\n{'=' * 72}\n" + f"compress={compress!r} server {client.server_version} iters {ITERS}\n" + f"clickhouse-connect {clickhouse_connect.__version__} from {clickhouse_connect.__file__}\n" + f"C extension {dataconv.__file__}\n" + f"_ch_core {_ch_core.__version__} from {_ch_core.__file__}\n" + f"pyarrow {pa.__version__} pandas {pd.__version__} python {platform.python_version()}\n" + f"platform {platform.platform()} {platform.machine()}\n" + f"mirrored settings {settings} (v1 also sends wait_end_of_query=1, the rust path streams)" + ) + + +def main(): + for compress in (False, "lz4"): + client = clickhouse_connect.get_client(host="localhost", port=8123, compress=compress) + header(client, compress) + assert_encoding(client, "lz4" if compress else None) + for workload in WORKLOADS: + run_workload(client, workload) + client.close() + memory_snapshot() + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "--memprobe": + memprobe(*sys.argv[2:6]) + else: + main() diff --git a/rust/ch-core-py/Cargo.toml b/rust/ch-core-py/Cargo.toml new file mode 100644 index 00000000..0a6a2984 --- /dev/null +++ b/rust/ch-core-py/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "ch-core-py" +version = "0.1.0" +edition = "2021" +description = "PyO3 binding crate for the clickhouse-connect-core wheel" +license = "Apache-2.0" +repository = "https://github.com/ClickHouse/clickhouse-connect" + +[lib] +name = "_ch_core" +crate-type = ["cdylib"] + +[dependencies] +ch-core-rs = { git = "https://github.com/ClickHouse/ch-core-rs", tag = "v0.1.1" } +# generate-import-lib synthesizes the CPython import library when +# cross-compiling to Windows targets with no local interpreter +pyo3 = { version = "0.29", features = ["extension-module", "generate-import-lib"] } + +[build-dependencies] +pyo3-build-config = { version = "0.29", features = ["resolve-config"] } diff --git a/rust/ch-core-py/README.md b/rust/ch-core-py/README.md new file mode 100644 index 00000000..c702b338 --- /dev/null +++ b/rust/ch-core-py/README.md @@ -0,0 +1,23 @@ +# clickhouse-connect-core + +Compiled Rust codec for [clickhouse-connect](https://pypi.org/project/clickhouse-connect/). This wheel provides the `_ch_core` extension module that decodes ClickHouse `FORMAT Native` query results and encodes inserts. + +This package is not meant to be used directly. Install it through the `rust` extra of the main driver: + +```bash +pip install "clickhouse-connect[rust]" +``` + +Then enable the codec on a client: + +```python +import clickhouse_connect + +client = clickhouse_connect.get_client(host="localhost", native_codec="rust") +``` + +The codec is experimental and opt in. See the [clickhouse-connect documentation](https://clickhouse.com/docs/integrations/language-clients/python/rust-codec) for supported values, fallback rules, and version compatibility. + +The `_ch_core` module has no public API contract. Its interface exists solely for the clickhouse-connect driver, which checks a binding API version at client creation and tells you when this wheel needs an upgrade. + +Issues and source live in the [clickhouse-connect repository](https://github.com/ClickHouse/clickhouse-connect) under `rust/ch-core-py`. The decoding core itself is [ch-core-rs](https://github.com/ClickHouse/ch-core-rs). diff --git a/rust/ch-core-py/build.rs b/rust/ch-core-py/build.rs new file mode 100644 index 00000000..90a1aa32 --- /dev/null +++ b/rust/ch-core-py/build.rs @@ -0,0 +1,4 @@ +fn main() { + pyo3_build_config::use_pyo3_cfgs(); + pyo3_build_config::add_extension_module_link_args(); +} diff --git a/rust/ch-core-py/pyproject.toml b/rust/ch-core-py/pyproject.toml new file mode 100644 index 00000000..0478aafa --- /dev/null +++ b/rust/ch-core-py/pyproject.toml @@ -0,0 +1,29 @@ +[build-system] +requires = ["maturin>=1.7,<2.0"] +build-backend = "maturin" + +[project] +name = "clickhouse-connect-core" +dynamic = ["version"] +description = "Compiled Rust native-format codec for clickhouse-connect" +readme = "README.md" +license = "Apache-2.0" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Programming Language :: Rust", +] + +[project.urls] +Homepage = "https://github.com/ClickHouse/clickhouse-connect" +Repository = "https://github.com/ClickHouse/clickhouse-connect" +Issues = "https://github.com/ClickHouse/clickhouse-connect/issues" +Documentation = "https://clickhouse.com/docs/integrations/language-clients/python/rust-codec" + +[tool.maturin] +module-name = "_ch_core" +features = ["pyo3/extension-module"] diff --git a/rust/ch-core-py/src/batch.rs b/rust/ch-core-py/src/batch.rs new file mode 100644 index 00000000..07f13e80 --- /dev/null +++ b/rust/ch-core-py/src/batch.rs @@ -0,0 +1,518 @@ +use std::sync::Arc; + +use pyo3::exceptions::PyValueError; +use pyo3::ffi; +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyBytes, PyCapsule, PyList, PyModule, PyTuple}; + +use ch_core_rs::batch::{ChunkedBatch, ColBatch as RustColBatch}; +use ch_core_rs::column::Column; +use ch_core_rs::ffi as core_ffi; +use ch_core_rs::native::decode::decode_all_bytes; +use ch_core_rs::schema::ChType; + +use crate::decoder::{buffer_to_vec, decode_err, decode_options}; +use crate::pyval::{fill_column, prepare_column_ctx, ColumnCtx}; + +/// Wrapper to make ArrowArrayStream Send-safe for PyCapsule. +#[repr(transparent)] +struct SendableStream(core_ffi::ArrowArrayStream); +// Safety: the stream's private_data is a Box owning only +// Send + Sync data (Schema, Arc chunks, CString), no Python objects +// or thread-affine state, so the capsule destructor may drop it on any thread. +unsafe impl Send for SendableStream {} + +/// A query result: a sequence of decoded columnar chunks (one per Native +/// block), never concatenated. Exposes the chunks as a single Arrow C Stream +/// or materializes them into Python objects on demand. +#[pyclass(name = "ColBatch")] +pub struct ColBatch { + inner: Arc, +} + +impl ColBatch { + pub(crate) fn from_chunked(inner: ChunkedBatch) -> Self { + Self { + inner: Arc::new(inner), + } + } + + /// Wrap a single decoded block as a one-chunk result. Used by the + /// incremental decoders, which emit one block at a time. + pub(crate) fn from_block(block: RustColBatch) -> Self { + let schema = block.schema.clone(); + Self { + inner: Arc::new(ChunkedBatch { + schema, + chunks: vec![Arc::new(block)], + }), + } + } +} + +#[pymethods] +impl ColBatch { + /// Decode a complete Native payload from bytes, bytearray, or memoryview. + #[staticmethod] + #[pyo3(signature = (data, has_block_info = false))] + fn decode_native(data: &Bound<'_, PyAny>, has_block_info: bool) -> PyResult { + let options = decode_options(has_block_info); + let chunked = if let Ok(bytes) = data.cast::() { + // bytes: decode straight from the borrowed buffer, no copy. The + // borrow ties the decode to the GIL, so it cannot be released. + decode_all_bytes(bytes.as_bytes(), &options) + } else { + // bytearray, memoryview: copy out, then decode without the GIL. + let owned = buffer_to_vec(data)?; + data.py().detach(|| decode_all_bytes(&owned, &options)) + } + .map_err(decode_err)?; + Ok(Self::from_chunked(chunked)) + } + + /// Merge already-decoded batches into one. The schema comes from the + /// first batch and every later batch must match it. Chunks are shared via + /// Arc clones; zero-row chunks are dropped, as in `decode_all_bytes`. + #[staticmethod] + fn from_batches(py: Python<'_>, batches: Vec>) -> PyResult { + let first = batches + .first() + .ok_or_else(|| PyValueError::new_err("from_batches requires at least one batch"))?; + let schema = first.borrow(py).inner.schema.clone(); + let mut chunks = Vec::new(); + for (i, batch) in batches.iter().enumerate() { + let guard = batch.borrow(py); + if guard.inner.schema != schema { + return Err(PyValueError::new_err(format!( + "Batch {i} schema differs from the first batch" + ))); + } + chunks.extend( + guard + .inner + .chunks + .iter() + .filter(|c| c.num_rows > 0) + .cloned(), + ); + } + Ok(Self::from_chunked(ChunkedBatch { schema, chunks })) + } + + #[getter] + fn num_rows(&self) -> usize { + self.inner.num_rows() + } + + #[getter] + fn num_columns(&self) -> usize { + self.inner.num_columns() + } + + #[getter] + fn num_chunks(&self) -> usize { + self.inner.num_chunks() + } + + #[getter] + fn column_names(&self) -> Vec { + self.inner + .schema + .fields + .iter() + .map(|f| f.name.clone()) + .collect() + } + + #[getter] + fn column_type_names(&self) -> Vec { + self.inner + .schema + .fields + .iter() + .map(|f| f.ch_type.to_string()) + .collect() + } + + /// Export all chunks as an Arrow C Stream capsule. `requested_schema` is + /// accepted per the Arrow PyCapsule interface and ignored; the stream + /// always carries its native schema. + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_stream__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyResult> { + let _ = requested_schema; + + // Safety: an all-zero ArrowArrayStream is a valid initial value: every + // field is a raw pointer (null) or an Option whose None is + // the all-zero bit pattern. export_chunks_to_stream overwrites every + // field before the struct is observed. + let mut stream: core_ffi::ArrowArrayStream = unsafe { std::mem::zeroed() }; + unsafe { + // Safety: `stream` is a valid, writable, zeroed ArrowArrayStream. + // Ownership of the exported data passes to its release callback. + // Cheap: clones the chunk Vec (Arc clones), no column data copied. + core_ffi::export_chunks_to_stream( + self.inner.schema.clone(), + self.inner.chunks.clone(), + &mut stream, + ); + } + + // The destructor frees the stream if the capsule is dropped unconsumed. + // A consumer that imports the stream moves it out and clears `release` + // in place, so the destructor sees `None` and does nothing. + let capsule = PyCapsule::new_with_value_and_destructor( + py, + SendableStream(stream), + c"arrow_array_stream", + |mut stream: SendableStream, _context| { + // Safety: `stream` was initialized by `export_chunks_to_stream` + // and `release_if_set` is a no-op once a consumer cleared + // `release`. Must not panic: pyo3's capsule destructor + // trampoline has no unwind guard. + unsafe { stream.0.release_if_set() } + }, + )?; + Ok(capsule) + } + + /// Get column `index` as a single Python list, concatenated across chunks. + #[pyo3(signature = (index, raw_time_ticks = false))] + fn column_data<'py>( + &self, + py: Python<'py>, + index: usize, + raw_time_ticks: bool, + ) -> PyResult> { + if index >= self.inner.num_columns() { + return Err(PyValueError::new_err(format!( + "Column index {index} out of range (0..{})", + self.inner.num_columns() + ))); + } + + let ctx = prepare_column_ctx(py, &self.inner.schema.fields[index].ch_type, raw_time_ticks)?; + column_to_pylist(py, &self.inner.chunks, index, &ctx) + } + + /// Get all rows as a list of tuples, across all chunks. + fn to_python_rows<'py>(&self, py: Python<'py>) -> PyResult> { + let total_rows = self.inner.num_rows(); + let num_cols = self.inner.num_columns(); + // Resolve each column's temporal context once, not per cell, so a + // tz-aware column imports its zoneinfo a single time for the whole table. + let ctxs: Vec = self + .inner + .schema + .fields + .iter() + .map(|f| prepare_column_ctx(py, &f.ch_type, false)) + .collect::>()?; + for chunk in &self.inner.chunks { + check_chunk_shape(chunk, num_cols)?; + } + unsafe { + let list_ptr = ffi::PyList_New(total_rows as ffi::Py_ssize_t); + if list_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: list_ptr came from PyList_New, so it is a list and this + // is the sole owned reference. Binding it makes the error and + // panic paths drop the list; list_dealloc tolerates NULL slots. + let list = Bound::from_owned_ptr(py, list_ptr).cast_into_unchecked::(); + + // Allocate every row tuple up front, moved into the list at once, + // so the error path drops them through list_dealloc. Column-major + // fill leaves NULL slots at later column indexes while an error + // unwinds; tuple_dealloc tolerates them. + for out_row in 0..total_rows { + let tuple_ptr = ffi::PyTuple_New(num_cols as ffi::Py_ssize_t); + if tuple_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: out_row < total_rows, the list's allocated length; + // the list takes over the owned tuple. + ffi::PyList_SET_ITEM(list.as_ptr(), out_row as ffi::Py_ssize_t, tuple_ptr); + } + + // Column-major fill: dispatch each column's type once per chunk + // and write into slot col_idx of each row's tuple. + let mut base: usize = 0; + for chunk in &self.inner.chunks { + let rows = chunk.num_rows; + for (col_idx, (col, ctx)) in chunk.columns.iter().zip(&ctxs).enumerate() { + let mut sink = |i: usize, item: *mut ffi::PyObject| { + // Safety: base + i < total_rows, the chunk-row sum the + // list was filled to, so GET_ITEM borrows a live tuple; + // col_idx < num_cols, its allocated length; the tuple + // takes over the owned item. + let tuple = + ffi::PyList_GET_ITEM(list.as_ptr(), (base + i) as ffi::Py_ssize_t); + ffi::PyTuple_SET_ITEM(tuple, col_idx as ffi::Py_ssize_t, item); + }; + fill_column(py, col, ctx, rows, &mut sink)?; + } + base += rows; + } + + Ok(list) + } + } + + /// Get all columns as Python containers, each concatenated across chunks. + /// + /// With `typed_numeric`, top-level non-nullable fixed-width numeric columns + /// use `array.array`, matching clickhouse-connect's Python Native decoder. + /// A true entry in `tuple_columns` renders that column as a tuple instead + /// of a list. The tuple flag is checked first, so it takes precedence over + /// `typed_numeric` for any column carrying both. + #[pyo3(signature = (*, typed_numeric = false, tuple_columns = None))] + fn to_python_columns<'py>( + &self, + py: Python<'py>, + typed_numeric: bool, + tuple_columns: Option>, + ) -> PyResult> { + if let Some(flags) = &tuple_columns { + if flags.len() != self.inner.num_columns() { + return Err(PyValueError::new_err(format!( + "tuple_columns has {} entries, expected {}", + flags.len(), + self.inner.num_columns() + ))); + } + } + for chunk in &self.inner.chunks { + check_chunk_shape(chunk, self.inner.num_columns())?; + } + let array_ctor = if typed_numeric { + Some(PyModule::import(py, "array")?.getattr("array")?) + } else { + None + }; + let cols: Vec> = (0..self.inner.num_columns()) + .map(|ci| { + let ch_type = &self.inner.schema.fields[ci].ch_type; + if tuple_columns.as_ref().is_some_and(|flags| flags[ci]) { + let ctx = prepare_column_ctx(py, ch_type, false)?; + return Ok(column_to_pytuple(py, &self.inner.chunks, ci, &ctx)?.into_any()); + } + if let Some(ctor) = &array_ctor { + if let Some(column) = + typed_numeric_column(py, &self.inner.chunks, ci, ch_type, ctor)? + { + return Ok(column); + } + } + let ctx = prepare_column_ctx(py, ch_type, false)?; + Ok(column_to_pylist(py, &self.inner.chunks, ci, &ctx)?.into_any()) + }) + .collect::>()?; + PyList::new(py, &cols) + } +} + +/// Build a Python `array.array` directly from native primitive buffers. +/// +/// The Rust vectors and `array.array` both use host-native byte order. The +/// binding only ships on the mainstream CPython platforms where the selected +/// typecodes have the same widths as their Rust primitives. +fn typed_numeric_column<'py>( + py: Python<'py>, + chunks: &[Arc], + col_idx: usize, + ch_type: &ChType, + array_ctor: &Bound<'py, PyAny>, +) -> PyResult>> { + let Some(first) = chunks.first() else { + // Zero-row result: no chunks, so the schema type picks the typecode. + // Aliases (SimpleAggregateFunction) resolve to the physical type the + // non-empty case would have matched on. + let resolved; + let ch_type = match ch_type.physical_delegate() { + Some(delegate) => { + resolved = delegate; + &resolved + } + None => ch_type, + }; + let typecode = match ch_type { + ChType::Int8 => "b", + ChType::Int16 => "h", + ChType::Int32 => "i", + ChType::Int64 => "q", + ChType::UInt8 => "B", + ChType::UInt16 => "H", + ChType::UInt32 => "I", + ChType::UInt64 => "Q", + ChType::Float32 => "f", + ChType::Float64 => "d", + _ => return Ok(None), + }; + return Ok(Some(array_ctor.call1((typecode,))?)); + }; + + macro_rules! build_array { + ($variant:ident, $typecode:literal, $ty:ty) => {{ + let output = array_ctor.call1(($typecode,))?; + debug_assert_eq!( + output.getattr("itemsize")?.extract::()?, + std::mem::size_of::<$ty>(), + "array.array typecode width differs from the Rust primitive" + ); + for chunk in chunks { + let Column::$variant(column) = &chunk.columns[col_idx] else { + return Ok(None); + }; + if column.validity.is_some() { + return Ok(None); + } + let byte_len = std::mem::size_of_val(column.values.as_slice()); + // Safety: primitive Vec storage is contiguous and live for this + // call. Reinterpreting it as bytes preserves its native layout. + let bytes = unsafe { + std::slice::from_raw_parts(column.values.as_ptr().cast::(), byte_len) + }; + output.call_method1("frombytes", (PyBytes::new(py, bytes),))?; + } + Ok(Some(output)) + }}; + } + + match &first.columns[col_idx] { + Column::Int8(column) if column.validity.is_none() => build_array!(Int8, "b", i8), + Column::Int16(column) if column.validity.is_none() => build_array!(Int16, "h", i16), + Column::Int32(column) if column.validity.is_none() => build_array!(Int32, "i", i32), + Column::Int64(column) if column.validity.is_none() => build_array!(Int64, "q", i64), + Column::UInt8(column) if column.validity.is_none() => build_array!(UInt8, "B", u8), + Column::UInt16(column) if column.validity.is_none() => build_array!(UInt16, "H", u16), + Column::UInt32(column) if column.validity.is_none() => build_array!(UInt32, "I", u32), + Column::UInt64(column) if column.validity.is_none() => build_array!(UInt64, "Q", u64), + Column::Float32(column) if column.validity.is_none() => build_array!(Float32, "f", f32), + Column::Float64(column) if column.validity.is_none() => build_array!(Float64, "d", f64), + _ => Ok(None), + } +} + +/// Reject a chunk whose column count differs from the schema or whose column +/// lengths differ from the chunk row count. The core only debug_asserts these +/// invariants, so malformed payloads must fail here before any fill loop reads +/// a short buffer or Bool padding bits. +fn check_chunk_shape(chunk: &RustColBatch, num_cols: usize) -> PyResult<()> { + if chunk.columns.len() != num_cols { + return Err(PyValueError::new_err(format!( + "Malformed payload: chunk has {} columns, expected {num_cols}", + chunk.columns.len() + ))); + } + for (idx, col) in chunk.columns.iter().enumerate() { + if col.len() != chunk.num_rows { + return Err(PyValueError::new_err(format!( + "Malformed payload: column {idx} has {} rows, chunk expects {}", + col.len(), + chunk.num_rows + ))); + } + } + Ok(()) +} + +/// Reject chunks missing column `col_idx` or whose column length differs from +/// the chunk row count. Shared by the per-column materializers. +fn check_column_shape(chunks: &[Arc], col_idx: usize) -> PyResult<()> { + for chunk in chunks { + if col_idx >= chunk.columns.len() { + return Err(PyValueError::new_err(format!( + "Malformed payload: chunk has {} columns, expected at least {}", + chunk.columns.len(), + col_idx + 1 + ))); + } + if chunk.columns[col_idx].len() != chunk.num_rows { + return Err(PyValueError::new_err(format!( + "Malformed payload: column {col_idx} has {} rows, chunk expects {}", + chunk.columns[col_idx].len(), + chunk.num_rows + ))); + } + } + Ok(()) +} + +/// Build one column as a Python list across `chunks`, one owned pointer per +/// cell. Shared by `column_data` and `to_python_columns` so every path applies +/// one host-value policy. +fn column_to_pylist<'py>( + py: Python<'py>, + chunks: &[Arc], + col_idx: usize, + ctx: &ColumnCtx<'_>, +) -> PyResult> { + check_column_shape(chunks, col_idx)?; + let total_rows: usize = chunks.iter().map(|c| c.num_rows).sum(); + unsafe { + let list_ptr = ffi::PyList_New(total_rows as ffi::Py_ssize_t); + if list_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: list_ptr came from PyList_New, so it is a list and this is + // the sole owned reference. Binding it makes the error and panic paths + // drop the list; list_dealloc tolerates the NULL slots not yet filled. + let list = Bound::from_owned_ptr(py, list_ptr).cast_into_unchecked::(); + + let mut out_row: usize = 0; + for chunk in chunks { + let col = &chunk.columns[col_idx]; + let base = out_row; + let mut sink = |i: usize, item: *mut ffi::PyObject| { + // Safety: base + i < total_rows, the chunk-row sum the list + // was allocated with, and the list takes over the owned item. + ffi::PyList_SET_ITEM(list.as_ptr(), (base + i) as ffi::Py_ssize_t, item); + }; + fill_column(py, col, ctx, chunk.num_rows, &mut sink)?; + out_row += chunk.num_rows; + } + + Ok(list) + } +} + +/// Tuple twin of `column_to_pylist`, for columns the Python codec's readers +/// return as tuples. +fn column_to_pytuple<'py>( + py: Python<'py>, + chunks: &[Arc], + col_idx: usize, + ctx: &ColumnCtx<'_>, +) -> PyResult> { + check_column_shape(chunks, col_idx)?; + let total_rows: usize = chunks.iter().map(|c| c.num_rows).sum(); + unsafe { + let tuple_ptr = ffi::PyTuple_New(total_rows as ffi::Py_ssize_t); + if tuple_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: tuple_ptr came from PyTuple_New, so it is a tuple and this is + // the sole owned reference. Binding it makes the error and panic paths + // drop the tuple; tuple_dealloc tolerates the NULL slots not yet filled. + let tuple = Bound::from_owned_ptr(py, tuple_ptr).cast_into_unchecked::(); + + let mut out_row: usize = 0; + for chunk in chunks { + let col = &chunk.columns[col_idx]; + let base = out_row; + let mut sink = |i: usize, item: *mut ffi::PyObject| { + // Safety: base + i < total_rows, the chunk-row sum the tuple + // was allocated with, and the tuple takes over the owned item. + ffi::PyTuple_SET_ITEM(tuple.as_ptr(), (base + i) as ffi::Py_ssize_t, item); + }; + fill_column(py, col, ctx, chunk.num_rows, &mut sink)?; + out_row += chunk.num_rows; + } + + Ok(tuple) + } +} diff --git a/rust/ch-core-py/src/decoder.rs b/rust/ch-core-py/src/decoder.rs new file mode 100644 index 00000000..c19209df --- /dev/null +++ b/rust/ch-core-py/src/decoder.rs @@ -0,0 +1,330 @@ +use pyo3::buffer::PyBuffer; +use pyo3::exceptions::{ + PyEOFError, PyNotImplementedError, PyOSError, PyRuntimeError, PyStopIteration, PyValueError, +}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyList}; + +use ch_core_rs::native::decode::{decode_next_block, DecodeError, DecodeOptions}; +use ch_core_rs::native::stream_decoder::StreamDecoder as RustStreamDecoder; +use ch_core_rs::native::varint::ByteReader; + +#[cfg(unix)] +use std::collections::VecDeque; + +#[cfg(unix)] +use ch_core_rs::batch::ColBatch as RustColBatch; + +use crate::batch::ColBatch; + +/// Map a core decode error to the Python exception for its failure class. +/// Unsupported input becomes NotImplementedError, matching the insert side. +/// Truncation (UnexpectedEof) becomes EOFError so callers can tell incomplete +/// data apart from corruption. +pub(crate) fn decode_err(e: DecodeError) -> PyErr { + match e { + DecodeError::UnsupportedType { column, type_name } => PyNotImplementedError::new_err( + format!("Unsupported ClickHouse type '{type_name}' for column '{column}'"), + ), + DecodeError::InvalidBlockInfo { field_num } => { + PyValueError::new_err(format!("Unknown BlockInfo field number {field_num}")) + } + DecodeError::UnsupportedSerialization { + column, + serialization_byte, + } => PyNotImplementedError::new_err(format!( + "Unsupported custom serialization (marker {serialization_byte}) for column '{column}'" + )), + DecodeError::BlockSchemaMismatch { block_index } => PyValueError::new_err(format!( + "Block {block_index} schema differs from the first block" + )), + DecodeError::InvalidLowCardinality { column, reason } => PyValueError::new_err(format!( + "Invalid LowCardinality layout for column '{column}': {reason}" + )), + DecodeError::InvalidArray { column, reason } => PyValueError::new_err(format!( + "Invalid Array layout for column '{column}': {reason}" + )), + DecodeError::InvalidTuple { column, reason } => PyValueError::new_err(format!( + "Invalid Tuple layout for column '{column}': {reason}" + )), + DecodeError::InvalidVariant { column, reason } => PyValueError::new_err(format!( + "Invalid Variant layout for column '{column}': {reason}" + )), + DecodeError::InvalidDynamic { column, reason } => PyValueError::new_err(format!( + "Invalid Dynamic layout for column '{column}': {reason}" + )), + DecodeError::InvalidJson { column, reason } => PyValueError::new_err(format!( + "Invalid JSON layout for column '{column}': {reason}" + )), + DecodeError::ResourceLimit { + limit, + requested, + what, + } => PyValueError::new_err(format!( + "Resource limit exceeded for {what}: requested {requested} bytes cumulatively, limit is {limit} bytes" + )), + DecodeError::Io(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { + PyEOFError::new_err(format!("Truncated Native data: {e}")) + } + DecodeError::Io(e) => PyRuntimeError::new_err(format!("Decode error: {e}")), + _ => PyRuntimeError::new_err(format!("Decode error: {e}")), + } +} + +pub(crate) fn decode_options(has_block_info: bool) -> DecodeOptions { + let mut options = DecodeOptions::default(); + options.protocol_revision = if has_block_info { 1 } else { 0 }; + options +} + +/// Copy the bytes out of any u8 buffer object (bytes, bytearray, memoryview). +pub(crate) fn buffer_to_vec(data: &Bound<'_, PyAny>) -> PyResult> { + if let Ok(bytes) = data.cast::() { + Ok(bytes.as_bytes().to_vec()) + } else { + PyBuffer::::get(data)?.to_vec(data.py()) + } +} + +/// Incremental block decoder over a byte buffer. +/// +/// Yields ColBatch objects one block at a time from an in-memory buffer. +/// Used for the materialized streaming path (buffer first, iterate blocks). +/// Every block must carry the first block's schema, the same invariant the +/// core's `decode_all_bytes` and `StreamDecoder` enforce; `decode_next_block` +/// itself is the stateless single-block primitive, so this loop applies the +/// check. +#[pyclass] +pub struct BlockDecoder { + data: Vec, + pos: u64, + options: DecodeOptions, + exhausted: bool, + schema: Option, + blocks_seen: usize, +} + +#[pymethods] +impl BlockDecoder { + #[new] + #[pyo3(signature = (data, has_block_info = false))] + fn new(data: &Bound<'_, PyAny>, has_block_info: bool) -> PyResult { + Ok(Self { + data: buffer_to_vec(data)?, + pos: 0, + options: decode_options(has_block_info), + exhausted: false, + schema: None, + blocks_seen: 0, + }) + } + + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __next__(&mut self) -> PyResult { + if self.exhausted { + return Err(PyStopIteration::new_err(())); + } + + let mut reader = ByteReader::new(&self.data[self.pos as usize..]); + + match decode_next_block(&mut reader, &self.options) { + Ok(Some(batch)) => { + match &self.schema { + None => self.schema = Some(batch.schema.clone()), + Some(first) => { + if batch.schema != *first { + self.exhausted = true; + return Err(decode_err(DecodeError::BlockSchemaMismatch { + block_index: self.blocks_seen, + })); + } + } + } + self.blocks_seen += 1; + self.pos += reader.position() as u64; + Ok(ColBatch::from_block(batch)) + } + Ok(None) => { + self.exhausted = true; + Err(PyStopIteration::new_err(())) + } + Err(e) => { + self.exhausted = true; + Err(decode_err(e)) + } + } + } +} + +/// Streaming block decoder that reads from a file descriptor. +/// +/// Designed for true streaming: an async producer writes HTTP response chunks to +/// the write end of a pipe, and this decoder reads from the read end. Each +/// `__next__` reads a chunk from the pipe and feeds it to the push-based +/// [`RustStreamDecoder`], buffering any completed blocks and yielding them one at +/// a time. The pipe read and the decode both run with the GIL released, so the +/// async producer can keep writing concurrently. +/// +/// The core decodes over in-memory slices, so the incremental read-and-feed loop +/// lives here rather than in the core: the stream decoder owns the partial-block +/// buffering and only allocates a block's columns once its last byte arrives. +/// +/// Ownership: the constructor duplicates `read_fd` and reads from its own +/// duplicate, which it closes when dropped. The caller keeps ownership of +/// `read_fd` and is responsible for closing it. An invalid `read_fd` raises +/// `OSError` at construction. +/// +/// Usage:: +/// +/// import os +/// read_fd, write_fd = os.pipe() +/// # async producer writes chunks to write_fd +/// decoder = PipeDecoder(read_fd) +/// for batch in decoder: +/// process(batch) +/// os.close(read_fd) +#[cfg(unix)] +#[pyclass] +pub struct PipeDecoder { + reader: std::fs::File, + decoder: RustStreamDecoder, + /// Blocks already decoded from fed chunks but not yet yielded. + pending: VecDeque, + /// Reused read buffer, refilled from the pipe on each network read. + buf: Vec, + finished: bool, +} + +#[cfg(unix)] +#[pymethods] +impl PipeDecoder { + #[new] + #[pyo3(signature = (read_fd, has_block_info = false))] + fn new(read_fd: i32, has_block_info: bool) -> PyResult { + use std::os::fd::BorrowedFd; + if read_fd < 0 { + return Err(PyOSError::new_err(format!( + "invalid file descriptor: {read_fd}" + ))); + } + // Safety: borrow_raw requires fd != -1, guarded above. Its stays-open + // contract is knowingly relaxed: the borrow is used only for the dup + // call, never for I/O, so a closed fd yields EBADF from dup rather + // than unsafety, and a recycled fd dups the caller's wrong descriptor, + // a caller contract violation, not memory unsafety. + let borrowed = unsafe { BorrowedFd::borrow_raw(read_fd) }; + let owned = borrowed.try_clone_to_owned()?; + Ok(Self { + reader: std::fs::File::from(owned), + decoder: RustStreamDecoder::new(decode_options(has_block_info)), + pending: VecDeque::new(), + buf: vec![0u8; 1 << 16], + finished: false, + }) + } + + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __next__(&mut self, py: Python<'_>) -> PyResult { + loop { + if let Some(block) = self.pending.pop_front() { + return Ok(ColBatch::from_block(block)); + } + if self.finished { + return Err(PyStopIteration::new_err(())); + } + + // Read one chunk from the pipe and feed it to the stream decoder, + // releasing the GIL across the blocking read and the decode so the + // async producer can keep writing to the pipe concurrently. A + // zero-length read is EOF: finish the decoder to flush trailing + // blocks and surface a truncated final block as an error. + let reader = &mut self.reader; + let decoder = &mut self.decoder; + let buf = &mut self.buf; + let result = py.detach(|| -> Result<(bool, Vec), DecodeError> { + use std::io::Read; + let n = reader.read(&mut buf[..])?; + if n == 0 { + Ok((true, decoder.finish()?)) + } else { + Ok((false, decoder.feed(&buf[..n])?)) + } + }); + + match result { + Ok((eof, blocks)) => { + self.finished = eof; + self.pending.extend(blocks); + } + Err(e) => { + self.finished = true; + return Err(decode_err(e)); + } + } + } + } +} + +/// Push-based incremental block decoder. +/// +/// Receives byte chunks via `feed()` and returns any complete blocks. +/// No blocking I/O, no pipes, no threads — designed for async contexts +/// where an event loop pushes chunks as they arrive from the network. +/// +/// Usage:: +/// +/// decoder = StreamDecoder() +/// for chunk in byte_chunks: +/// blocks = decoder.feed(chunk) +/// for block in blocks: +/// process(block) +/// final_blocks = decoder.finish() +#[pyclass] +pub struct StreamDecoder { + inner: RustStreamDecoder, +} + +#[pymethods] +impl StreamDecoder { + #[new] + #[pyo3(signature = (has_block_info = false))] + fn new(has_block_info: bool) -> Self { + let options = decode_options(has_block_info); + Self { + inner: RustStreamDecoder::new(options), + } + } + + /// Feed a chunk of bytes (bytes, bytearray, or memoryview). Returns a + /// list of any complete ColBatch objects that could be decoded from the + /// accumulated data. + /// + /// The GIL is released during the actual decode so a producer on another + /// thread (or the asyncio event loop) can keep pulling network bytes while + /// this call decodes — this is what makes pull and parse overlap. + fn feed<'py>( + &mut self, + py: Python<'py>, + data: &Bound<'py, PyAny>, + ) -> PyResult> { + // Copy out of Python-owned memory so we can drop the GIL during decode. + let owned = buffer_to_vec(data)?; + let batches = py.detach(|| self.inner.feed(&owned)).map_err(decode_err)?; + let py_batches: Vec = batches.into_iter().map(ColBatch::from_block).collect(); + PyList::new(py, py_batches) + } + + /// Signal end of stream. Returns any remaining complete blocks. + /// Raises an error if the stream ends with a truncated block. + fn finish<'py>(&mut self, py: Python<'py>) -> PyResult> { + let batches = py.detach(|| self.inner.finish()).map_err(decode_err)?; + let py_batches: Vec = batches.into_iter().map(ColBatch::from_block).collect(); + PyList::new(py, py_batches) + } +} diff --git a/rust/ch-core-py/src/insert/containers.rs b/rust/ch-core-py/src/insert/containers.rs new file mode 100644 index 00000000..799b2625 --- /dev/null +++ b/rust/ch-core-py/src/insert/containers.rs @@ -0,0 +1,840 @@ +use super::*; + +/// Build an `Array(T)` column: each row is a sequence of elements, flattened +/// into one strong-reference element run with an Arrow LargeList offsets run. +/// The element column is built once over the flat run, so elements hit the +/// same per-type fast paths as a plain column and nested Arrays compose +/// recursively. Arrays are never nullable at the array level, so a None row +/// is an error. +pub(super) fn build_array_column( + py: Python<'_>, + name: &str, + inner: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + let column_values = ColumnValues::new(values, name)?; + check_row_count(name, &column_values, row_count)?; + + if let Ok(list) = values.cast_exact::() { + return array_column_from_seq(py, name, inner, &ListSeq(list), row_count); + } + if let Ok(tuple) = values.cast_exact::() { + return array_column_from_seq(py, name, inner, &TupleSeq(tuple), row_count); + } + + // Generic outer container: safe indexed reads, same per-row flatten. + let mut offsets = Vec::with_capacity(row_count + 1); + offsets.push(0i64); + let mut flat = FlatRefs::default(); + for row in 0..row_count { + let value = column_values.get_item(row)?; + flatten_array_row(py, name, inner, &value, row, &mut flat)?; + offsets.push(flat.end_offset(name, "Array element")?); + } + let element_column = build_element_column(py, name, inner, &flat.ptrs) + .map_err(|err| remap_element_err(py, name, &offsets, err))?; + Ok(Column::Array(ArrayColumn::new(offsets, element_column))) +} + +/// Array flatten over an exact list or tuple of rows: borrowed row reads, one +/// flat element run. A row that is not an exact list or tuple flattens through +/// `list.extend`, which can run Python, so the container size is revalidated +/// before the next borrowed read. +fn array_column_from_seq( + py: Python<'_>, + name: &str, + inner: &ChType, + seq: &S, + row_count: usize, +) -> PyResult { + let mut offsets = Vec::with_capacity(row_count + 1); + offsets.push(0i64); + let mut flat = FlatRefs::default(); + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every fallback revalidates; the strong reference keeps the row + // alive across any Python code the fallback runs. + let value = unsafe { Bound::from_borrowed_ptr(py, seq.get(row)) }; + flatten_array_row(py, name, inner, &value, row, &mut flat)?; + check_not_resized(seq, name, row_count)?; + offsets.push(flat.end_offset(name, "Array element")?); + } + let element_column = build_element_column(py, name, inner, &flat.ptrs) + .map_err(|err| remap_element_err(py, name, &offsets, err))?; + Ok(Column::Array(ArrayColumn::new(offsets, element_column))) +} + +/// Rewrite an element-run error's flat index as the outer row and element +/// index, using the offsets built alongside the flat run. Nested Arrays remap +/// at each level, so the final message leads with the outermost row. An error +/// that is not a ValueError or whose text does not carry the `column "name" +/// row N` prefix passes through unchanged. +fn remap_element_err(py: Python<'_>, name: &str, offsets: &[i64], err: PyErr) -> PyErr { + remap_flat_err(py, name, offsets, "element", err) +} + +/// Shared flat-run error rewrite: `unit` names what the flat index counts +/// ("element" for Array, "key"/"value" for Map). +fn remap_flat_err(py: Python<'_>, name: &str, offsets: &[i64], unit: &str, err: PyErr) -> PyErr { + let Some((flat, prefix, tail)) = parse_row_err(py, name, &err) else { + return err; + }; + let Ok(flat) = i64::try_from(flat) else { + return err; + }; + let row = offsets[1..].partition_point(|&end| end <= flat); + if row + 1 >= offsets.len() { + return err; + } + let element = flat - offsets[row]; + PyValueError::new_err(format!("{prefix}{row} {unit} {element}{tail}")) +} + +/// Parse a ValueError's `column {name:?} row N` prefix into the row index, +/// the prefix through "row ", and the message tail after the digits. `None` +/// for any error that does not carry the prefix. +fn parse_row_err(py: Python<'_>, name: &str, err: &PyErr) -> Option<(usize, String, String)> { + if !err.is_instance_of::(py) { + return None; + } + let text = err.value(py).str().ok()?; + let text = text.to_str().ok()?; + let prefix = format!("column {name:?} row "); + let rest = text.strip_prefix(&prefix)?; + let digits = rest.bytes().take_while(u8::is_ascii_digit).count(); + let row = rest[..digits].parse::().ok()?; + Some((row, prefix, rest[digits..].to_string())) +} + +/// Rewrite a tuple field-run error: the flat index equals the outer row, so +/// only the element label (index, or name for a named tuple) is inserted. +fn remap_tuple_element_err(py: Python<'_>, name: &str, label: &str, err: PyErr) -> PyErr { + let Some((row, prefix, tail)) = parse_row_err(py, name, &err) else { + return err; + }; + PyValueError::new_err(format!("{prefix}{row} element {label}{tail}")) +} + +/// Rewrite an alternative's dense child row index to its logical Variant row. +pub(super) fn remap_variant_child_err( + py: Python<'_>, + name: &str, + discriminators: &[u8], + alternative: usize, + err: PyErr, +) -> PyErr { + let Some((dense_row, prefix, tail)) = parse_row_err(py, name, &err) else { + return err; + }; + let Ok(discriminator) = u8::try_from(alternative) else { + return err; + }; + let Some(logical_row) = discriminators + .iter() + .enumerate() + .filter_map(|(row, &disc)| (disc == discriminator).then_some(row)) + .nth(dense_row) + else { + return err; + }; + PyValueError::new_err(format!("{prefix}{logical_row}{tail}")) +} + +/// Append one Array row's elements to the flat run. Exact list/tuple rows +/// copy borrowed pointers without running Python; anything else keeps the +/// generic path's accepted containers and error messages via `list.extend`. +fn flatten_array_row( + py: Python<'_>, + name: &str, + inner: &ChType, + value: &Bound<'_, PyAny>, + row: usize, + flat: &mut FlatRefs, +) -> PyResult<()> { + if let Ok(list) = value.cast_exact::() { + // SAFETY: copying exact-list items runs no Python code. + unsafe { flat.extend_from_seq(&ListSeq(list)) }; + return Ok(()); + } + if let Ok(tuple) = value.cast_exact::() { + // SAFETY: copying exact-tuple items runs no Python code. + unsafe { flat.extend_from_seq(&TupleSeq(tuple)) }; + return Ok(()); + } + if value.is_none() { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but Array({inner}) is not Nullable" + ))); + } + if value.cast::().is_ok() { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is a str, not an Array sequence" + ))); + } + // bytes-like rows flatten as int elements, matching the python codec's + // `data.extend(row)` iteration semantics. + if let Ok(bytes) = value.cast::() { + return flat.extend_from_byte_run(py, bytes.as_bytes()); + } + if let Ok(bytes) = value.cast::() { + return flat.extend_from_byte_run(py, &bytes.to_vec()); + } + if value.cast::().is_ok() + || value.cast::().is_ok() + || value.cast::().is_ok() + { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is an unordered set/dict, which has no defined Array element order" + ))); + } + let items = PyList::empty(py); + items + .call_method1(intern!(py, "extend"), (value,)) + .map_err(|err| { + let wrapped = PyValueError::new_err(format!( + "column {name:?} row {row} is not a valid Array value" + )); + wrapped.set_cause(py, Some(err)); + wrapped + })?; + // SAFETY: items is an owned exact list; copying its items runs no Python + // code and the strong references outlive the temporary list. + unsafe { flat.extend_from_seq(&ListSeq(&items)) }; + Ok(()) +} + +/// Strong references to flattened Array elements. Holding a reference per +/// element keeps every pointer in the run valid across any Python code the +/// element conversion runs, so the run can be consumed with borrowed reads. +#[derive(Default)] +pub(super) struct FlatRefs { + pub(super) ptrs: Vec<*mut ffi::PyObject>, +} + +impl FlatRefs { + /// Append every element of `seq` as a strong reference. + /// + /// # Safety + /// + /// Requires the GIL; `seq` must allow borrowed reads with no Python code + /// running during the copy (an exact list or tuple). + unsafe fn extend_from_seq(&mut self, seq: &S) { + let len = seq.size(); + self.ptrs.reserve(len); + for index in 0..len { + let item = seq.get(index); + ffi::Py_INCREF(item); + self.ptrs.push(item); + } + } + + /// Append each byte of `bytes` as an owned Python int element. + fn extend_from_byte_run(&mut self, py: Python<'_>, bytes: &[u8]) -> PyResult<()> { + self.ptrs.reserve(bytes.len()); + for &byte in bytes { + // SAFETY: GIL held; PyLong_FromLong returns an owned reference + // (a cached small int here) that Drop releases. + let item = unsafe { ffi::PyLong_FromLong(c_long::from(byte)) }; + if item.is_null() { + return Err(PyErr::fetch(py)); + } + self.ptrs.push(item); + } + Ok(()) + } + + /// Append a strong reference to `obj`. + pub(super) fn push_ref(&mut self, obj: &Bound<'_, PyAny>) { + // SAFETY: GIL held (the Bound proves it); Drop releases the reference. + unsafe { ffi::Py_INCREF(obj.as_ptr()) }; + self.ptrs.push(obj.as_ptr()); + } + + /// Append a strong reference to None. + fn push_none(&mut self, _py: Python<'_>) { + // SAFETY: GIL held; Py_None is a valid object pointer. + unsafe { + let none = ffi::Py_None(); + ffi::Py_INCREF(none); + self.ptrs.push(none); + } + } + + /// Current run length as an i64 offset. `unit` names what the run counts + /// ("Array element", "Map entry"). + fn end_offset(&self, name: &str, unit: &str) -> PyResult { + i64::try_from(self.ptrs.len()).map_err(|_| { + PyValueError::new_err(format!( + "column {name:?} {unit} count exceeds i64 offset capacity" + )) + }) + } +} + +impl Drop for FlatRefs { + fn drop(&mut self) { + // SAFETY: FlatRefs is only built and dropped inside a frame that holds + // a `Python` token, so the GIL is held here. + for &ptr in &self.ptrs { + unsafe { ffi::Py_DECREF(ptr) }; + } + } +} + +/// Build the flattened Array element column. Mirrors `build_column`'s wrapper +/// dispatch with the flat run standing in for the Python container, feeding +/// the same seq fast paths and generic scalar loops. +pub(super) fn build_element_column( + py: Python<'_>, + name: &str, + ch_type: &ChType, + ptrs: &[*mut ffi::PyObject], +) -> PyResult { + let row_count = ptrs.len(); + let seq = PtrSeq(ptrs); + // Expand SimpleAggregateFunction/geo/Nested aliases before dispatch so a + // nested alias (e.g. Array(Point)) builds its physical element column. + if let Some(delegate) = ch_type.physical_delegate() { + return build_element_column(py, name, &delegate, ptrs); + } + match ch_type { + ChType::Nothing => Ok(nothing_column_from_seq(&seq, row_count, false)), + ChType::AggregateFunction { .. } => { + aggregate_state_column_from_seq(py, name, ch_type, &seq, row_count) + } + ChType::Array(inner) => array_column_from_seq(py, name, inner, &seq, row_count), + ChType::QBit { + element_type, + dimension, + } => { + build_qbit_column_from_seq(py, name, *element_type, *dimension, &seq, row_count, false) + } + ChType::Tuple(elements) => { + tuple_column_from_seq(py, name, elements, &seq, row_count, false) + } + ChType::Map(key, value) => map_column_from_seq(py, name, key, value, &seq, row_count), + ChType::Variant(alternatives) => { + variant_column_from_seq(py, name, alternatives, &seq, row_count) + } + ChType::Dynamic { .. } => dynamic_string_column(&PtrRows { py, ptrs }, row_count), + ChType::Json { .. } => { + json_text_column_from_rows(py, name, &PtrRows { py, ptrs }, row_count, false) + } + ChType::Nullable(inner) => { + // Nullable(Point) -> Nullable(Tuple), Nullable(SAF(T)) -> Nullable(T); + // the physical delegate governs the nullable element shape. + if let Some(delegate) = inner.physical_delegate() { + return build_element_column(py, name, &ChType::Nullable(Box::new(delegate)), ptrs); + } + if let ChType::Tuple(elements) = inner.as_ref() { + return tuple_column_from_seq(py, name, elements, &seq, row_count, true); + } + if matches!(inner.as_ref(), ChType::Nothing) { + return Ok(nothing_column_from_seq(&seq, row_count, true)); + } + if matches!(inner.as_ref(), ChType::Json { .. }) { + return json_text_column_from_rows( + py, + name, + &PtrRows { py, ptrs }, + row_count, + true, + ); + } + if let ChType::QBit { + element_type, + dimension, + } = inner.as_ref() + { + return build_qbit_column_from_seq( + py, + name, + *element_type, + *dimension, + &seq, + row_count, + true, + ); + } + if matches!( + inner.as_ref(), + ChType::Nullable(_) + | ChType::LowCardinality(_) + | ChType::Array(_) + | ChType::Map(..) + ) { + return Err(PyNotImplementedError::new_err(format!( + "unsupported Nullable inner type {inner} for column {name:?}" + ))); + } + if let Some(column) = try_fast_column_seq(py, name, inner, &seq, row_count, true)? { + return Ok(column); + } + nullable_scalar_column(py, name, inner, &PtrRows { py, ptrs }, row_count) + } + ChType::LowCardinality(inner) => { + // Resolve the physical dict value type: strip the SAF chain, unwrap an + // optional Nullable, strip again. So LowCardinality(SAF(anyLast, String)) + // and LowCardinality(SAF(anyLast, Nullable(String))) reach the String path. + let (nullable, value_type) = low_cardinality_dict_value_type(inner); + if !is_low_cardinality_inner(value_type) { + return Err(PyNotImplementedError::new_err(format!( + "unsupported LowCardinality inner type {value_type} for column {name:?}" + ))); + } + if matches!(value_type, ChType::String) { + return lc_string_seq(py, name, value_type, &seq, row_count, nullable); + } + if wide_int_layout(value_type).is_some() { + return lc_wide_column( + py, + name, + value_type, + &PtrRows { py, ptrs }, + row_count, + nullable, + ); + } + lc_scalar_column( + py, + name, + value_type, + &PtrRows { py, ptrs }, + row_count, + nullable, + ) + } + _ => { + if let Some(column) = try_fast_column_seq(py, name, ch_type, &seq, row_count, false)? { + return Ok(column); + } + plain_scalar_column(py, name, ch_type, &PtrRows { py, ptrs }, row_count) + } + } +} + +/// Build a `Tuple(T1, ...)` column: each row fans out into one strong-ref run +/// per element, then each element column is built once over its run, hitting +/// the same per-type fast paths as a plain column. `nullable` builds the +/// tuple-level validity of a `Nullable(Tuple)`; a None row keeps the children +/// full length with per-type default placeholders. +pub(super) fn build_tuple_column( + py: Python<'_>, + name: &str, + elements: &[(Option, ChType)], + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult { + let column_values = ColumnValues::new(values, name)?; + check_row_count(name, &column_values, row_count)?; + if let Ok(list) = values.cast_exact::() { + return tuple_column_from_seq(py, name, elements, &ListSeq(list), row_count, nullable); + } + if let Ok(tuple) = values.cast_exact::() { + return tuple_column_from_seq(py, name, elements, &TupleSeq(tuple), row_count, nullable); + } + let mut builder = TupleBuilder::new(py, name, elements, row_count, nullable); + for row in 0..row_count { + let value = column_values.get_item(row)?; + builder.push_row(&value, row)?; + } + builder.finish() +} + +/// Tuple rows over an exact list or tuple of rows, or a flattened element run. +fn tuple_column_from_seq( + py: Python<'_>, + name: &str, + elements: &[(Option, ChType)], + seq: &S, + row_count: usize, + nullable: bool, +) -> PyResult { + let mut builder = TupleBuilder::new(py, name, elements, row_count, nullable); + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every fallback revalidates; the strong reference keeps the row + // alive across any Python code the row read runs. + let value = unsafe { Bound::from_borrowed_ptr(py, seq.get(row)) }; + builder.push_row(&value, row)?; + check_not_resized(seq, name, row_count)?; + } + builder.finish() +} + +/// Streaming Tuple row builder: one strong-ref run per element, an optional +/// tuple-level null map, and a row-read mode decided by the first non-None +/// row (element-name dict reads for a fully named tuple whose first non-None +/// row is a dict, positional reads otherwise). +struct TupleBuilder<'a, 'py> { + py: Python<'py>, + name: &'a str, + elements: &'a [(Option, ChType)], + flats: Vec, + null_map: Option>, + defaults: Option>>, + names: Option>>, + dict_mode: Option, + row_count: usize, +} + +impl<'a, 'py> TupleBuilder<'a, 'py> { + fn new( + py: Python<'py>, + name: &'a str, + elements: &'a [(Option, ChType)], + row_count: usize, + nullable: bool, + ) -> Self { + let names: Option>> = + (!elements.is_empty() && elements.iter().all(|(n, _)| n.is_some())).then(|| { + elements + .iter() + .map(|(n, _)| PyString::new(py, n.as_deref().unwrap_or_default())) + .collect() + }); + Self { + py, + name, + elements, + flats: elements + .iter() + .map(|_| FlatRefs { + ptrs: Vec::with_capacity(row_count), + }) + .collect(), + null_map: nullable.then(|| Vec::with_capacity(row_count)), + defaults: None, + names, + dict_mode: None, + row_count, + } + } + + fn tuple_type(&self) -> ChType { + ChType::Tuple(self.elements.to_vec()) + } + + fn push_row(&mut self, value: &Bound<'_, PyAny>, row: usize) -> PyResult<()> { + if value.is_none() { + if self.null_map.is_none() { + return Err(PyValueError::new_err(format!( + "column {:?} row {row} is None but {} is not Nullable", + self.name, + self.tuple_type() + ))); + } + if self.defaults.is_none() { + self.defaults = Some( + self.elements + .iter() + .map(|(_, element_type)| default_pyobject(self.py, element_type)) + .collect::>()?, + ); + } + if let Some(null_map) = &mut self.null_map { + null_map.push(1); + } + if let Some(defaults) = &self.defaults { + for (flat, default) in self.flats.iter_mut().zip(defaults) { + flat.push_ref(default.bind(self.py)); + } + } + return Ok(()); + } + if let Some(null_map) = &mut self.null_map { + null_map.push(0); + } + let dict_mode = match self.dict_mode { + Some(mode) => mode, + None => { + let mode = self.names.is_some() && value.is_instance_of::(); + self.dict_mode = Some(mode); + mode + } + }; + if dict_mode { + self.push_dict_row(value, row) + } else { + self.push_positional_row(value, row) + } + } + + /// Read one row by element name: missing keys become None, extra keys are + /// ignored. Non-dict rows fall back to a `.get` method read. + fn push_dict_row(&mut self, value: &Bound<'_, PyAny>, row: usize) -> PyResult<()> { + let Some(names) = &self.names else { + return Err(PyValueError::new_err("internal tuple row mode mismatch")); + }; + if let Ok(dict) = value.cast_exact::() { + for (flat, key) in self.flats.iter_mut().zip(names) { + match dict.get_item(key)? { + Some(item) => flat.push_ref(&item), + None => flat.push_none(self.py), + } + } + return Ok(()); + } + let get = value.getattr(intern!(self.py, "get")).map_err(|err| { + let wrapped = PyValueError::new_err(format!( + "column {:?} row {row} cannot be read as a dict for the named Tuple", + self.name + )); + wrapped.set_cause(self.py, Some(err)); + wrapped + })?; + for (flat, key) in self.flats.iter_mut().zip(names) { + let item = get.call1((key,)).map_err(|err| { + let wrapped = PyValueError::new_err(format!( + "column {:?} row {row} element {key:?} cannot be read from the dict-like row", + self.name + )); + wrapped.set_cause(self.py, Some(err)); + wrapped + })?; + flat.push_ref(&item); + } + Ok(()) + } + + /// Read one row positionally. Exact list/tuple rows copy borrowed + /// pointers; other iterables flatten through `list.extend`. + fn push_positional_row(&mut self, value: &Bound<'_, PyAny>, row: usize) -> PyResult<()> { + if let Ok(list) = value.cast_exact::() { + // SAFETY: copying exact-list items runs no Python code. + return unsafe { self.push_positional_seq(&ListSeq(list), row) }; + } + if let Ok(tuple) = value.cast_exact::() { + // SAFETY: copying exact-tuple items runs no Python code. + return unsafe { self.push_positional_seq(&TupleSeq(tuple), row) }; + } + if value.cast::().is_ok() { + return Err(PyValueError::new_err(format!( + "column {:?} row {row} is a str, not a Tuple row", + self.name + ))); + } + if value.is_instance_of::() { + return Err(PyValueError::new_err(format!( + "column {:?} row {row} is a dict but Tuple rows are read positionally", + self.name + ))); + } + if value.cast::().is_ok() || value.cast::().is_ok() { + return Err(PyValueError::new_err(format!( + "column {:?} row {row} is an unordered set, which has no defined Tuple element order", + self.name + ))); + } + let items = PyList::empty(self.py); + items + .call_method1(intern!(self.py, "extend"), (value,)) + .map_err(|err| { + let wrapped = PyValueError::new_err(format!( + "column {:?} row {row} is not a valid Tuple row", + self.name + )); + wrapped.set_cause(self.py, Some(err)); + wrapped + })?; + // SAFETY: items is an owned exact list; copying runs no Python code. + unsafe { self.push_positional_seq(&ListSeq(&items), row) } + } + + /// Copy one arity-checked row into the element runs. + /// + /// # Safety + /// + /// Requires the GIL; `seq` must allow borrowed reads with no Python code + /// running during the copy (an exact list or tuple). + unsafe fn push_positional_seq(&mut self, seq: &S, row: usize) -> PyResult<()> { + let got = seq.size(); + if got != self.elements.len() { + return Err(PyValueError::new_err(format!( + "column {:?} row {row} has {got} elements but the Tuple declares {}", + self.name, + self.elements.len() + ))); + } + for (index, flat) in self.flats.iter_mut().enumerate() { + let item = seq.get(index); + ffi::Py_INCREF(item); + flat.ptrs.push(item); + } + Ok(()) + } + + fn finish(self) -> PyResult { + let mut fields = Vec::with_capacity(self.elements.len()); + for (index, ((element_name, element_type), flat)) in + self.elements.iter().zip(&self.flats).enumerate() + { + let label = match element_name { + Some(n) => format!("{n:?}"), + None => index.to_string(), + }; + let column = build_element_column(self.py, self.name, element_type, &flat.ptrs) + .map_err(|err| remap_tuple_element_err(self.py, self.name, &label, err))?; + fields.push(column); + } + Ok(Column::Tuple(match self.null_map { + Some(nulls) => { + TupleColumn::new_nullable(fields, self.row_count, Bitmap::from_ch_null_map(&nulls)) + } + None => TupleColumn::new(fields, self.row_count), + })) + } +} + +/// Build a `Map(K, V)` column: dict rows flatten into a key run and a value +/// run with an Arrow list offsets run; the two entry columns are built once +/// over their flat runs. Maps are never nullable at the map level, so a None +/// row is an error. +pub(super) fn build_map_column( + py: Python<'_>, + name: &str, + key_type: &ChType, + value_type: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + let column_values = ColumnValues::new(values, name)?; + check_row_count(name, &column_values, row_count)?; + if let Ok(list) = values.cast_exact::() { + return map_column_from_seq(py, name, key_type, value_type, &ListSeq(list), row_count); + } + if let Ok(tuple) = values.cast_exact::() { + return map_column_from_seq(py, name, key_type, value_type, &TupleSeq(tuple), row_count); + } + let mut builder = MapBuilder::new(py, name, key_type, value_type, row_count); + for row in 0..row_count { + let value = column_values.get_item(row)?; + builder.push_row(&value, row)?; + } + builder.finish() +} + +/// Map rows over an exact list or tuple of rows, or a flattened element run. +fn map_column_from_seq( + py: Python<'_>, + name: &str, + key_type: &ChType, + value_type: &ChType, + seq: &S, + row_count: usize, +) -> PyResult { + let mut builder = MapBuilder::new(py, name, key_type, value_type, row_count); + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every fallback revalidates; the strong reference keeps the row + // alive across any Python code the row read runs. + let value = unsafe { Bound::from_borrowed_ptr(py, seq.get(row)) }; + builder.push_row(&value, row)?; + check_not_resized(seq, name, row_count)?; + } + builder.finish() +} + +/// Streaming Map row builder: an offsets run plus parallel key and value +/// strong-ref runs. Exact dict rows copy entries without running Python; +/// dict-like rows read through `.items()`. +struct MapBuilder<'a, 'py> { + py: Python<'py>, + name: &'a str, + key_type: &'a ChType, + value_type: &'a ChType, + offsets: Vec, + keys: FlatRefs, + values: FlatRefs, +} + +impl<'a, 'py> MapBuilder<'a, 'py> { + fn new( + py: Python<'py>, + name: &'a str, + key_type: &'a ChType, + value_type: &'a ChType, + row_count: usize, + ) -> Self { + let mut offsets = Vec::with_capacity(row_count + 1); + offsets.push(0i64); + Self { + py, + name, + key_type, + value_type, + offsets, + keys: FlatRefs::default(), + values: FlatRefs::default(), + } + } + + fn push_row(&mut self, value: &Bound<'_, PyAny>, row: usize) -> PyResult<()> { + if value.is_none() { + return Err(PyValueError::new_err(format!( + "column {:?} row {row} is None but Map({}, {}) is not Nullable", + self.name, self.key_type, self.value_type + ))); + } + if let Ok(dict) = value.cast_exact::() { + for (key, val) in dict.iter() { + self.keys.push_ref(&key); + self.values.push_ref(&val); + } + } else { + self.push_mapping_row(value, row)?; + } + self.offsets + .push(self.keys.end_offset(self.name, "Map entry")?); + Ok(()) + } + + /// Dict-like fallback: rows must expose `.items()`; anything else is not + /// a Map row. Pair iterables are deliberately not accepted. + fn push_mapping_row(&mut self, value: &Bound<'_, PyAny>, row: usize) -> PyResult<()> { + let not_a_dict = |err: PyErr| { + let wrapped = PyValueError::new_err(format!( + "column {:?} row {row} is not a dict for Map({}, {})", + self.name, self.key_type, self.value_type + )); + wrapped.set_cause(self.py, Some(err)); + wrapped + }; + let items = value + .call_method0(intern!(self.py, "items")) + .map_err(not_a_dict)?; + let entries = PyList::empty(self.py); + entries + .call_method1(intern!(self.py, "extend"), (items,)) + .map_err(not_a_dict)?; + for index in 0..entries.len() { + let entry = entries.get_item(index)?; + let (key, val) = entry + .extract::<(Bound<'_, PyAny>, Bound<'_, PyAny>)>() + .map_err(|err| { + let wrapped = PyValueError::new_err(format!( + "column {:?} row {row} Map entry is not a key/value pair", + self.name + )); + wrapped.set_cause(self.py, Some(err)); + wrapped + })?; + self.keys.push_ref(&key); + self.values.push_ref(&val); + } + Ok(()) + } + + fn finish(self) -> PyResult { + let total = self.keys.ptrs.len(); + let keys_column = build_element_column(self.py, self.name, self.key_type, &self.keys.ptrs) + .map_err(|err| remap_flat_err(self.py, self.name, &self.offsets, "key", err))?; + let values_column = + build_element_column(self.py, self.name, self.value_type, &self.values.ptrs) + .map_err(|err| remap_flat_err(self.py, self.name, &self.offsets, "value", err))?; + let entries = Column::Tuple(TupleColumn::new(vec![keys_column, values_column], total)); + Ok(Column::Map(MapColumn::new(self.offsets, entries))) + } +} diff --git a/rust/ch-core-py/src/insert/fastpath.rs b/rust/ch-core-py/src/insert/fastpath.rs new file mode 100644 index 00000000..fd5115f1 --- /dev/null +++ b/rust/ch-core-py/src/insert/fastpath.rs @@ -0,0 +1,1387 @@ +use super::*; + +pub(super) fn build_plain_column( + py: Python<'_>, + name: &str, + ch_type: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + value_column(py, name, ch_type, values, row_count, false) +} + +pub(super) fn plain_scalar_column<'py, R: RowAccess<'py>>( + py: Python<'py>, + name: &str, + ch_type: &ChType, + rows: &R, + row_count: usize, +) -> PyResult { + let mut scalars = Vec::with_capacity(row_count); + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + } + scalars.push(convert_scalar(py, ch_type, &value, name, row)?); + } + column_from_scalars(ch_type, scalars, None) +} + +pub(super) fn build_nullable_column( + py: Python<'_>, + name: &str, + inner: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + // Expand a Nullable over a name-decoration alias by its delegate: + // Nullable(Point) -> Nullable(Tuple), Nullable(SimpleAggregateFunction(T)) + // -> Nullable(T). The physical delegate governs the nullable value shape. + if let Some(delegate) = inner.physical_delegate() { + return build_nullable_column(py, name, &delegate, values, row_count); + } + if let ChType::Tuple(elements) = inner { + return build_tuple_column(py, name, elements, values, row_count, true); + } + if matches!(inner, ChType::Nothing) { + return build_nothing_column(name, values, row_count, true); + } + if matches!(inner, ChType::Json { .. }) { + return build_json_text_column(py, name, values, row_count, true); + } + if let ChType::QBit { + element_type, + dimension, + } = inner + { + return build_qbit_column(py, name, *element_type, *dimension, values, row_count, true); + } + if matches!( + inner, + ChType::Nullable(_) | ChType::LowCardinality(_) | ChType::Array(_) | ChType::Map(..) + ) { + return Err(PyNotImplementedError::new_err(format!( + "unsupported Nullable inner type {inner} for column {name:?}" + ))); + } + + value_column(py, name, inner, values, row_count, true) +} + +/// Shared plain/Nullable value column build over a resolved container. +fn value_column( + py: Python<'_>, + name: &str, + ch_type: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult { + match ColumnSource::resolve(values, name, row_count)? { + ColumnSource::List(list) => value_column_from_seq( + py, + name, + ch_type, + &ListSeq(&list), + values, + row_count, + nullable, + ), + ColumnSource::Tuple(tuple) => value_column_from_seq( + py, + name, + ch_type, + &TupleSeq(&tuple), + values, + row_count, + nullable, + ), + ColumnSource::ObjectArray(seq) => { + if let Some(column) = try_fast_column_seq(py, name, ch_type, &seq, row_count, nullable)? + { + return Ok(column); + } + if nullable { + nullable_scalar_column(py, name, ch_type, &seq, row_count) + } else { + plain_scalar_column(py, name, ch_type, &seq, row_count) + } + } + ColumnSource::Values(column_values) => { + if let Some(column) = + try_numpy_timedelta_column(py, name, ch_type, values, row_count, nullable)? + { + return Ok(column); + } + if let Some(column) = try_buffer_column(py, name, ch_type, values, row_count, nullable)? + { + return Ok(column); + } + if wide_int_layout(ch_type).is_some() { + return wide_column_from_rows( + py, + name, + ch_type, + &column_values, + row_count, + nullable, + ); + } + if matches!(ch_type, ChType::Decimal { .. }) { + return decimal_column_from_rows( + name, + ch_type, + &column_values, + row_count, + nullable, + ); + } + if nullable { + nullable_scalar_column(py, name, ch_type, &column_values, row_count) + } else { + plain_scalar_column(py, name, ch_type, &column_values, row_count) + } + } + } +} + +/// Seq fast paths over an exact list or tuple; a type with no seq fast path +/// falls back to safe indexed reads, matching the generic container path. +fn value_column_from_seq( + py: Python<'_>, + name: &str, + ch_type: &ChType, + seq: &S, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult { + if let Some(column) = try_fast_column_seq(py, name, ch_type, seq, row_count, nullable)? { + return Ok(column); + } + let column_values = ColumnValues::new(values, name)?; + if nullable { + nullable_scalar_column(py, name, ch_type, &column_values, row_count) + } else { + plain_scalar_column(py, name, ch_type, &column_values, row_count) + } +} + +/// Build a wide-integer column over generic row access with one checked data +/// allocation. Each Python value writes directly into its final row slice. +fn wide_column_from_rows<'py, R: RowAccess<'py>>( + py: Python<'py>, + name: &str, + ch_type: &ChType, + rows: &R, + row_count: usize, + nullable: bool, +) -> PyResult { + let (width, signed, type_name) = wide_int_layout(ch_type) + .ok_or_else(|| PyValueError::new_err("internal wide integer type mismatch"))?; + let mut data = wide_data_buffer(name, width, row_count)?; + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + }; + null_map.push(1); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + let start = row * width; + wide_int_into( + py, + &value, + &mut data[start..start + width], + signed, + name, + row, + type_name, + )?; + } + let validity = null_map.map(|nulls| Bitmap::from_ch_null_map(&nulls)); + finish_wide_int_column(ch_type, data, validity) +} + +/// Decimal layout for the direct-into-buffer builders: byte width plus the +/// declared precision and scale. +fn decimal_layout(ch_type: &ChType) -> PyResult<(usize, u8, u8)> { + let ChType::Decimal { + precision, scale, .. + } = ch_type + else { + return Err(PyValueError::new_err("internal decimal type mismatch")); + }; + Ok((decimal_width(*precision)?, *precision, *scale)) +} + +fn decimal_data_buffer(name: &str, width: usize, row_count: usize) -> PyResult> { + let byte_len = width.checked_mul(row_count).ok_or_else(|| { + PyValueError::new_err(format!( + "column {name:?} Decimal byte size exceeds usize capacity" + )) + })?; + Ok(vec![0u8; byte_len]) +} + +fn finish_decimal_column( + data: Vec, + width: usize, + precision: u8, + scale: u8, + null_map: Option>, +) -> Column { + Column::Decimal(match null_map { + Some(nulls) => DecimalColumn::new_nullable( + data, + width, + precision, + scale, + Bitmap::from_ch_null_map(&nulls), + ), + None => DecimalColumn::new(data, width, precision, scale), + }) +} + +/// Build a Decimal column over generic row access with one data allocation. +/// Each value is stringified and its scaled integer written directly into its +/// final row slice. +fn decimal_column_from_rows<'py, R: RowAccess<'py>>( + name: &str, + ch_type: &ChType, + rows: &R, + row_count: usize, + nullable: bool, +) -> PyResult { + let (width, precision, scale) = decimal_layout(ch_type)?; + let mut data = decimal_data_buffer(name, width, row_count)?; + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + }; + null_map.push(1); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + let text = decimal_text(&value, name, row)?; + let start = row * width; + decimal_wire_into(&text, &mut data[start..start + width], precision, scale) + .map_err(|err| decimal_err(err, &text, width, precision, name, row))?; + } + Ok(finish_decimal_column( + data, width, precision, scale, null_map, + )) +} + +/// Decimal equivalent of `wide_column_from_seq`. Stringification runs Python +/// (`str` on the value), so every row holds a strong reference to its item +/// and revalidates the mutable container size before the next borrowed read. +fn decimal_column_from_seq( + py: Python<'_>, + name: &str, + ch_type: &ChType, + seq: &S, + row_count: usize, + nullable: bool, +) -> PyResult { + let (width, precision, scale) = decimal_layout(ch_type)?; + let mut data = decimal_data_buffer(name, width, row_count)?; + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every row revalidates below. + let ptr = unsafe { seq.get(row) }; + if ptr == unsafe { ffi::Py_None() } { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + }; + null_map.push(1); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + // SAFETY: taking a strong reference keeps the current item alive + // across the Python code its stringification runs. + let value = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + let text = decimal_text(&value, name, row)?; + let start = row * width; + decimal_wire_into(&text, &mut data[start..start + width], precision, scale) + .map_err(|err| decimal_err(err, &text, width, precision, name, row))?; + check_not_resized(seq, name, row_count)?; + } + Ok(finish_decimal_column( + data, width, precision, scale, null_map, + )) +} + +pub(super) fn nullable_scalar_column<'py, R: RowAccess<'py>>( + py: Python<'py>, + name: &str, + inner: &ChType, + rows: &R, + row_count: usize, +) -> PyResult { + let mut null_map = Vec::with_capacity(row_count); + let mut scalars = Vec::with_capacity(row_count); + let mut time_probe = TimeScalarProbe::new(inner); + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + null_map.push(1); + scalars.push(default_scalar(inner)?); + continue; + } + if let Some(probe) = time_probe.as_mut() { + match probe.probe(&value, name, row)? { + Some(TimeProbe::Nat) => { + null_map.push(1); + scalars.push(default_scalar(inner)?); + continue; + } + Some(TimeProbe::Ticks(ticks)) => { + null_map.push(0); + scalars.push(time_ticks_scalar(inner, ticks, name, row)?); + continue; + } + None => {} + } + } + if is_enum_nan(inner, &value) { + null_map.push(1); + scalars.push(default_scalar(inner)?); + continue; + } + null_map.push(0); + scalars.push(convert_scalar(py, inner, &value, name, row)?); + } + column_from_scalars(inner, scalars, Some(Bitmap::from_ch_null_map(&null_map))) +} +pub(super) struct ColumnValues<'py> { + values: Bound<'py, PyAny>, + len: usize, + name: String, +} + +impl<'py> ColumnValues<'py> { + pub(super) fn new(values: &Bound<'py, PyAny>, name: &str) -> PyResult { + let len = checked_container_len(values, name)?; + Ok(Self::from_parts(values, name, len)) + } + + fn from_parts(values: &Bound<'py, PyAny>, name: &str, len: usize) -> Self { + Self { + values: positional_source(values), + len, + name: name.to_string(), + } + } + + fn len(&self) -> usize { + self.len + } + + pub(super) fn get_item(&self, row: usize) -> PyResult> { + self.values.get_item(row).map_err(|_| { + PyValueError::new_err(format!( + "column {:?} row {row} could not be read from column container", + self.name + )) + }) + } +} + +fn checked_container_len(values: &Bound<'_, PyAny>, name: &str) -> PyResult { + if is_string_or_bytes_like(values) { + return Err(PyValueError::new_err(format!( + "column {name:?} values must be an indexable column container, not bare str or bytes" + ))); + } + values.len().map_err(|_| { + PyValueError::new_err(format!( + "column {name:?} values must be an indexable column container" + )) + }) +} + +/// Resolved input column container. Exact lists, exact tuples, and 1-D +/// contiguous object ndarrays feed the borrowed-pointer seq fast paths; +/// everything else uses safe indexed `ColumnValues` access. +pub(super) enum ColumnSource<'py> { + List(Bound<'py, PyList>), + Tuple(Bound<'py, PyTuple>), + ObjectArray(ObjectArraySeq<'py>), + Values(ColumnValues<'py>), +} + +impl<'py> ColumnSource<'py> { + /// Classify `values`, validating the container shape and row count. + pub(super) fn resolve( + values: &Bound<'py, PyAny>, + name: &str, + row_count: usize, + ) -> PyResult { + let len = checked_container_len(values, name)?; + if len != row_count { + return Err(PyValueError::new_err(format!( + "column {name:?} has {len} values but row_count is {row_count}" + ))); + } + if let Ok(list) = values.cast_exact::() { + return Ok(Self::List(list.clone())); + } + if let Ok(tuple) = values.cast_exact::() { + return Ok(Self::Tuple(tuple.clone())); + } + if let Some(seq) = ObjectArraySeq::matching(values, row_count) { + return Ok(Self::ObjectArray(seq)); + } + Ok(Self::Values(ColumnValues::from_parts(values, name, len))) + } +} + +/// Borrowed view over the object-pointer buffer of a 1-D C-contiguous +/// `dtype=object` ndarray. The buffer export pins the array length, so reads +/// stay in bounds for the whole build; a NULL slot reads as None. +pub(super) struct ObjectArraySeq<'py> { + py: Python<'py>, + view: RawBuffer<'py>, + len: usize, +} + +impl<'py> ObjectArraySeq<'py> { + fn matching(values: &Bound<'py, PyAny>, row_count: usize) -> Option { + let view = RawBuffer::matching( + values, + b'O', + std::mem::size_of::<*mut ffi::PyObject>(), + row_count, + )?; + Some(Self { + py: values.py(), + view, + len: row_count, + }) + } +} + +impl FastSeq for ObjectArraySeq<'_> { + // The exported buffer cannot be resized, but Python code run from a + // fallback can still replace items in place, so the pointer-identity + // caches keyed on MUTABLE must invalidate like they do for lists. + const MUTABLE: bool = true; + + #[inline] + unsafe fn get(&self, index: usize) -> *mut ffi::PyObject { + debug_assert!(index < self.len); + // SAFETY: the trait contract requires index < size(); each slot holds + // a pointer the array owns, read fresh so in-place replacement by + // earlier fallbacks is observed. + let ptr = *self.view.data().cast::<*mut ffi::PyObject>().add(index); + if ptr.is_null() { + ffi::Py_None() + } else { + ptr + } + } + + fn size(&self) -> usize { + self.len + } +} + +impl<'py> RowAccess<'py> for ObjectArraySeq<'py> { + fn value(&self, row: usize) -> PyResult> { + if row >= self.len { + return Err(PyValueError::new_err(format!( + "row {row} out of bounds for column of {} values", + self.len + ))); + } + // SAFETY: row is in bounds for the exported buffer; the array slot + // holds a strong reference and from_borrowed_ptr takes its own before + // any Python code can run. + Ok(unsafe { Bound::from_borrowed_ptr(self.py, self.get(row)) }) + } +} + +/// Positional row access for the generic scalar loops: a Python column +/// container or a flattened strong-reference run. +pub(super) trait RowAccess<'py> { + fn value(&self, row: usize) -> PyResult>; + + fn validate(&self) -> PyResult<()> { + Ok(()) + } +} + +impl<'py> RowAccess<'py> for ColumnValues<'py> { + fn value(&self, row: usize) -> PyResult> { + self.get_item(row) + } +} + +/// Borrowed exact-list access for JSON serialization. The configured Python +/// serializer can execute arbitrary code, so the list size is revalidated +/// before the next unchecked item read. +pub(super) struct ListRows<'a, 'py> { + pub(super) py: Python<'py>, + pub(super) list: &'a Bound<'py, PyList>, + pub(super) name: &'a str, + pub(super) expected: usize, +} + +impl<'py> RowAccess<'py> for ListRows<'_, 'py> { + fn value(&self, row: usize) -> PyResult> { + // SAFETY: callers iterate below `expected`, which `validate` confirms + // after every operation that may execute Python. + Ok(unsafe { + Bound::from_borrowed_ptr( + self.py, + ffi::PyList_GET_ITEM(self.list.as_ptr(), row as ffi::Py_ssize_t), + ) + }) + } + + fn validate(&self) -> PyResult<()> { + if self.list.len() == self.expected { + Ok(()) + } else { + Err(PyValueError::new_err(format!( + "column {:?} values changed size during JSON serialization", + self.name + ))) + } + } +} + +/// Borrowed exact-tuple access needs no resize guard because tuples are +/// immutable. +pub(super) struct TupleRows<'a, 'py> { + pub(super) py: Python<'py>, + pub(super) tuple: &'a Bound<'py, PyTuple>, +} + +impl<'py> RowAccess<'py> for TupleRows<'_, 'py> { + fn value(&self, row: usize) -> PyResult> { + // SAFETY: row_count was checked against the immutable tuple length. + Ok(unsafe { + Bound::from_borrowed_ptr( + self.py, + ffi::PyTuple_GET_ITEM(self.tuple.as_ptr(), row as ffi::Py_ssize_t), + ) + }) + } +} + +/// Row access over flattened Array element pointers, kept valid by the +/// `FlatRefs` strong references. +pub(super) struct PtrRows<'a, 'py> { + pub(super) py: Python<'py>, + pub(super) ptrs: &'a [*mut ffi::PyObject], +} + +impl<'py> RowAccess<'py> for PtrRows<'_, 'py> { + fn value(&self, row: usize) -> PyResult> { + // SAFETY: FlatRefs holds a strong reference for every pointer in the + // run for the whole build. + Ok(unsafe { Bound::from_borrowed_ptr(self.py, self.ptrs[row]) }) + } +} + +pub(super) fn check_row_count( + name: &str, + values: &ColumnValues<'_>, + row_count: usize, +) -> PyResult<()> { + let len = values.len(); + if len != row_count { + return Err(PyValueError::new_err(format!( + "column {name:?} has {len} values but row_count is {row_count}" + ))); + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// Fast paths for primitive numeric columns. The ChType dispatch runs once per +// column; the per-value loop reads borrowed pointers from an exact list or +// tuple, or a matching buffer-protocol container copies straight into Vec. +// Anything the exact-type checks reject falls back to `convert_scalar` per +// item, so accepted-type and error semantics match the generic path. +// --------------------------------------------------------------------------- + +/// Per-value conversion for the primitive fast paths. +pub(super) trait FastValue: Copy { + const DEFAULT: Self; + + /// Convert an exact-type Python object without running Python code. + /// `Err(())` means "no fast conversion, use the generic fallback" and + /// guarantees no Python exception is left pending. + /// + /// # Safety + /// + /// Requires the GIL; `ptr` must be a valid, non-null object pointer. + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + ch_type: &ChType, + fast_limit: i64, + ) -> Result; + + /// Unwrap the `Scalar` produced by the `convert_scalar` fallback. + fn from_scalar(scalar: Scalar) -> PyResult; + + /// Copy a matching buffer-protocol container. `Ok(None)` for types with + /// no buffer representation or containers that do not match. + fn from_buffer( + py: Python<'_>, + name: &str, + values: &Bound<'_, PyAny>, + row_count: usize, + ) -> PyResult>> { + let _ = (py, name, values, row_count); + Ok(None) + } + + fn into_column(values: Vec, validity: Option) -> Column; +} + +/// Truncate a Float32 to ClickHouse's upper-16-bit BFloat16 representation. +/// NaN inputs get the quiet bit set so truncation cannot yield infinity. +#[inline] +fn f32_to_bfloat16(value: f32) -> [u8; 2] { + let bits = value.to_bits() | if value.is_nan() { 0x0040_0000 } else { 0 }; + ((bits >> 16) as u16).to_le_bytes() +} + +/// Narrow a Float64 without turning a finite out-of-range value into infinity. +#[inline] +pub(super) fn checked_f64_to_bfloat16(value: f64) -> Result<[u8; 2], ()> { + let narrowed = value as f32; + if value.is_finite() && narrowed.is_infinite() { + return Err(()); + } + Ok(f32_to_bfloat16(narrowed)) +} + +/// Read an exact `int` as i64. On overflow the pending exception is cleared +/// and `Err(())` sends the value through the generic fallback, which produces +/// the standard conversion error. +/// +/// # Safety +/// +/// Requires the GIL; `ptr` must be a valid, non-null object pointer. +#[inline] +pub(super) unsafe fn exact_long_as_i64(ptr: *mut ffi::PyObject) -> Result { + if ffi::PyLong_CheckExact(ptr) == 0 { + return Err(()); + } + let value = ffi::PyLong_AsLongLong(ptr); + if value == -1 && !ffi::PyErr_Occurred().is_null() { + ffi::PyErr_Clear(); + return Err(()); + } + Ok(value) +} + +macro_rules! impl_fast_prim_common { + ($ty:ty, $variant:ident) => { + fn from_scalar(scalar: Scalar) -> PyResult { + match scalar { + Scalar::$variant(value) => Ok(value), + _ => Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + + fn from_buffer( + py: Python<'_>, + _name: &str, + values: &Bound<'_, PyAny>, + row_count: usize, + ) -> PyResult>> { + buffer_values::<$ty>(py, values, row_count) + } + + fn into_column(values: Vec, validity: Option) -> Column { + Column::$variant(match validity { + Some(validity) => PrimitiveColumn::new_nullable(values, validity), + None => PrimitiveColumn::new(values), + }) + } + }; +} + +macro_rules! impl_fast_narrow_int { + ($ty:ty, $variant:ident) => { + impl FastValue for $ty { + const DEFAULT: Self = 0; + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + _ch_type: &ChType, + _fast_limit: i64, + ) -> Result { + <$ty>::try_from(exact_long_as_i64(ptr)?).map_err(|_| ()) + } + + impl_fast_prim_common!($ty, $variant); + } + }; +} + +impl_fast_narrow_int!(i8, Int8); +impl_fast_narrow_int!(i16, Int16); +impl_fast_narrow_int!(i32, Int32); +impl_fast_narrow_int!(u8, UInt8); +impl_fast_narrow_int!(u16, UInt16); +impl_fast_narrow_int!(u32, UInt32); + +impl FastValue for i64 { + const DEFAULT: Self = 0; + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + _ch_type: &ChType, + _fast_limit: i64, + ) -> Result { + exact_long_as_i64(ptr) + } + + impl_fast_prim_common!(i64, Int64); +} + +impl FastValue for u64 { + const DEFAULT: Self = 0; + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + _ch_type: &ChType, + _fast_limit: i64, + ) -> Result { + if ffi::PyLong_CheckExact(ptr) == 0 { + return Err(()); + } + let value = ffi::PyLong_AsUnsignedLongLong(ptr); + if value == u64::MAX && !ffi::PyErr_Occurred().is_null() { + ffi::PyErr_Clear(); + return Err(()); + } + Ok(value) + } + + impl_fast_prim_common!(u64, UInt64); +} + +impl FastValue for f64 { + const DEFAULT: Self = 0.0; + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + _ch_type: &ChType, + _fast_limit: i64, + ) -> Result { + if ffi::PyFloat_CheckExact(ptr) != 0 { + return Ok(ffi::PyFloat_AS_DOUBLE(ptr)); + } + // Exact int: same result as extract::'s PyFloat_AsDouble, which + // reaches PyLong_AsDouble through int.__float__. + if ffi::PyLong_CheckExact(ptr) != 0 { + let value = ffi::PyLong_AsDouble(ptr); + if value == -1.0 && !ffi::PyErr_Occurred().is_null() { + ffi::PyErr_Clear(); + return Err(()); + } + return Ok(value); + } + Err(()) + } + + impl_fast_prim_common!(f64, Float64); +} + +impl FastValue for f32 { + const DEFAULT: Self = 0.0; + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + ch_type: &ChType, + fast_limit: i64, + ) -> Result { + // Matches extract::: extract as f64, then `as` cast. + f64::from_exact(ptr, ch_type, fast_limit).map(|value| value as f32) + } + + impl_fast_prim_common!(f32, Float32); +} + +/// BFloat16 wire word; a distinct type so `[u8; 2]` can serve other 2-byte +/// wire types. +#[derive(Clone, Copy)] +#[repr(transparent)] +struct Bf16Word([u8; 2]); + +impl FastValue for Bf16Word { + const DEFAULT: Self = Bf16Word([0; 2]); + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + ch_type: &ChType, + fast_limit: i64, + ) -> Result { + f64::from_exact(ptr, ch_type, fast_limit) + .and_then(checked_f64_to_bfloat16) + .map(Self) + } + + fn from_scalar(scalar: Scalar) -> PyResult { + match scalar { + Scalar::BFloat16(value) => Ok(Self(value)), + _ => Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + + fn from_buffer( + py: Python<'_>, + name: &str, + values: &Bound<'_, PyAny>, + row_count: usize, + ) -> PyResult>> { + if let Some(values) = + map_buffer_values::(py, values, row_count, |_row, value| { + Ok(Self(f32_to_bfloat16(value))) + })? + { + return Ok(Some(values)); + } + map_buffer_values::(py, values, row_count, |row, value| { + checked_f64_to_bfloat16(value) + .map(Self) + .map_err(|_| conversion_error(name, row, "BFloat16")) + }) + } + + fn into_column(values: Vec, validity: Option) -> Column { + // SAFETY: Bf16Word is #[repr(transparent)] over [u8; 2]. + let values = unsafe { cast_vec::(values) }; + Column::BFloat16(match validity { + Some(validity) => PrimitiveColumn::new_nullable(values, validity), + None => PrimitiveColumn::new(values), + }) + } +} + +/// Bool wire byte (0/1); a distinct type so u8 can serve UInt8. +#[derive(Clone, Copy)] +#[repr(transparent)] +struct WireBool(u8); + +impl FastValue for WireBool { + const DEFAULT: Self = WireBool(0); + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + _ch_type: &ChType, + _fast_limit: i64, + ) -> Result { + if ptr == ffi::Py_True() { + Ok(WireBool(1)) + } else if ptr == ffi::Py_False() { + Ok(WireBool(0)) + } else { + Err(()) + } + } + + fn from_scalar(scalar: Scalar) -> PyResult { + match scalar { + Scalar::Bool(value) => Ok(WireBool(u8::from(value))), + _ => Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + + fn into_column(values: Vec, validity: Option) -> Column { + // SAFETY: WireBool is #[repr(transparent)] over u8, so the buffer can + // be viewed as bytes directly. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) }; + Column::Bool(match validity { + Some(validity) => BoolColumn::from_wire_bytes_nullable(bytes, validity), + None => BoolColumn::from_wire_bytes(bytes), + }) + } +} + +/// Copy a one-dimensional buffer whose element type matches `T` exactly +/// (itemsize, signedness, and alignment are validated by `PyBuffer::get`). +/// Non-buffer containers and mismatched dtypes return `Ok(None)`. +/// +/// The format string is revalidated here because `PyBuffer::get` in pyo3 +/// 0.23 accepts b'>' as a matching byte order on little-endian targets and +/// maps format b'c' (char) to a 1-byte unsigned integer; both would change +/// the encoded values versus the generic per-item path. Any such buffer +/// falls through to the generic path instead. +fn matching_buffer(values: &Bound<'_, PyAny>, row_count: usize) -> Option> { + matching_native_buffer(values, &[row_count]) +} + +pub(super) fn buffer_values( + py: Python<'_>, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult>> { + let Some(buffer) = matching_buffer::(values, row_count) else { + return Ok(None); + }; + buffer.to_vec(py).map(Some) +} + +/// Map a matching one-dimensional Python buffer through a per-element +/// conversion. Contiguous buffers are read directly and allocate only the +/// destination; strided buffers use PyO3's safe gather before the same +/// Rust conversion. +fn map_buffer_values PyResult>( + py: Python<'_>, + values: &Bound<'_, PyAny>, + row_count: usize, + map: F, +) -> PyResult>> { + let Some(buffer) = matching_buffer::(values, row_count) else { + return Ok(None); + }; + if let Some(values) = buffer.as_slice(py) { + return values + .iter() + .enumerate() + .map(|(row, value)| map(row, value.get())) + .collect::>>() + .map(Some); + } + buffer + .to_vec(py)? + .into_iter() + .enumerate() + .map(|(row, value)| map(row, value)) + .collect::>>() + .map(Some) +} + +/// Borrowed positional access to an exact list or tuple. +pub(super) trait FastSeq { + /// Whether Python code run from a conversion fallback can resize the + /// container, invalidating later `get` calls. + const MUTABLE: bool; + + /// # Safety + /// + /// Requires the GIL and `index < size()`. The returned pointer is + /// borrowed and must be consumed before any Python code runs. + unsafe fn get(&self, index: usize) -> *mut ffi::PyObject; + + fn size(&self) -> usize; +} + +pub(super) struct ListSeq<'a, 'py>(pub(super) &'a Bound<'py, PyList>); + +impl FastSeq for ListSeq<'_, '_> { + const MUTABLE: bool = true; + + #[inline] + unsafe fn get(&self, index: usize) -> *mut ffi::PyObject { + ffi::PyList_GET_ITEM(self.0.as_ptr(), index as ffi::Py_ssize_t) + } + + fn size(&self) -> usize { + self.0.len() + } +} + +pub(super) struct TupleSeq<'a, 'py>(pub(super) &'a Bound<'py, PyTuple>); + +impl FastSeq for TupleSeq<'_, '_> { + const MUTABLE: bool = false; + + #[inline] + unsafe fn get(&self, index: usize) -> *mut ffi::PyObject { + ffi::PyTuple_GET_ITEM(self.0.as_ptr(), index as ffi::Py_ssize_t) + } + + fn size(&self) -> usize { + self.0.len() + } +} + +/// Flattened Array element run; the `FlatRefs` strong references keep every +/// pointer valid for the whole build and the slice can never be resized, so +/// fallbacks that run Python need no revalidation. +pub(super) struct PtrSeq<'a>(pub(super) &'a [*mut ffi::PyObject]); + +impl FastSeq for PtrSeq<'_> { + const MUTABLE: bool = false; + + #[inline] + unsafe fn get(&self, index: usize) -> *mut ffi::PyObject { + debug_assert!(index < self.0.len()); + // SAFETY: the trait contract requires index < size(). + *self.0.get_unchecked(index) + } + + fn size(&self) -> usize { + self.0.len() + } +} + +/// Convert `row_count` items from an exact list/tuple into a typed vector, +/// plus a validity bitmap when `nullable`. Items that fail the exact-type +/// check fall back to `convert_scalar`. The fallback can run arbitrary Python +/// (`__index__`, `__float__`, ...), so it holds a strong reference to its +/// item and the container size is revalidated afterwards, before the next +/// borrowed read could go out of bounds on a shrunk list. +fn seq_values( + py: Python<'_>, + seq: &S, + ch_type: &ChType, + name: &str, + row_count: usize, + nullable: bool, +) -> PyResult<(Vec, Option)> { + let mut values = Vec::with_capacity(row_count); + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + let mut time_probe = TimeScalarProbe::new(ch_type); + let fast_limit = match ch_type { + ChType::Time => MAX_TIME_SECONDS, + ChType::Time64 { precision } => max_time64_ticks(*precision), + _ => 0, + }; + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every fallback revalidates; the borrowed pointer is consumed before + // any Python code can run. + let ptr = unsafe { seq.get(row) }; + if ptr == unsafe { ffi::Py_None() } { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + }; + null_map.push(1); + values.push(T::DEFAULT); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + // SAFETY: GIL held; ptr is a valid borrowed item pointer. + match unsafe { T::from_exact(ptr, ch_type, fast_limit) } { + Ok(value) => values.push(value), + Err(()) => { + // SAFETY: ptr is valid here; taking a strong reference keeps + // the item alive across any Python code the fallback runs. + let obj = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + if let Some(probe) = time_probe.as_mut() { + if let Some(hit) = probe.probe(&obj, name, row)? { + match hit { + TimeProbe::Nat => { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is NaT but {ch_type} is not Nullable" + ))); + }; + if let Some(entry) = null_map.last_mut() { + *entry = 1; + } + values.push(T::DEFAULT); + } + TimeProbe::Ticks(ticks) => { + let scalar = time_ticks_scalar(ch_type, ticks, name, row)?; + values.push(T::from_scalar(scalar)?); + } + } + check_not_resized(seq, name, row_count)?; + continue; + } + } + let scalar = convert_scalar(py, ch_type, &obj, name, row)?; + values.push(T::from_scalar(scalar)?); + check_not_resized(seq, name, row_count)?; + } + } + } + Ok(( + values, + null_map.map(|nulls| Bitmap::from_ch_null_map(&nulls)), + )) +} + +/// Wide-integer equivalent of `seq_values`: allocate the final byte buffer +/// once, then convert each borrowed list/tuple/container item directly into +/// its row slice. Exact ints within i64 convert in place without running +/// Python; any other conversion may execute Python (`__index__`, or `int()` +/// for a string), so the current item is held strongly and mutable sequence +/// length is revalidated before the next borrowed access. +fn wide_column_from_seq( + py: Python<'_>, + name: &str, + ch_type: &ChType, + seq: &S, + row_count: usize, + nullable: bool, +) -> PyResult { + let (width, signed, type_name) = wide_int_layout(ch_type) + .ok_or_else(|| PyValueError::new_err("internal wide integer type mismatch"))?; + let mut data = wide_data_buffer(name, width, row_count)?; + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + for row in 0..row_count { + // SAFETY: row < row_count, the caller-validated size. A conversion + // that executes Python is followed by the resize check below. + let ptr = unsafe { seq.get(row) }; + if ptr == unsafe { ffi::Py_None() } { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + }; + null_map.push(1); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + let start = row * width; + let slice = &mut data[start..start + width]; + // SAFETY: the borrowed ptr stays valid without a strong ref only + // because the GIL is held across this whole loop (the module keeps + // pyo3's default gil_used = true) and the fast try never executes + // Python, so nothing can mutate the container or drop the item under + // it. A future free-threading (gil_used = false) opt-in invalidates + // this and requires a strong ref before the fast try. + let outcome = unsafe { wide_int_fast_into(ptr, slice, signed, name, row, type_name)? }; + if outcome == WideFast::Done { + continue; + } + // SAFETY: taking a strong reference keeps the current item alive if + // its conversion mutates the source list. + let value = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + wide_int_slow_into( + py, + &value, + outcome == WideFast::WideInt, + slice, + signed, + name, + row, + type_name, + )?; + check_not_resized(seq, name, row_count)?; + } + let validity = null_map.map(|nulls| Bitmap::from_ch_null_map(&nulls)); + finish_wide_int_column(ch_type, data, validity) +} + +/// Per-type dispatch over a borrowed-pointer run; runs once per column. +pub(super) fn try_fast_column_seq( + py: Python<'_>, + name: &str, + ch_type: &ChType, + seq: &S, + row_count: usize, + nullable: bool, +) -> PyResult> { + fn prim( + py: Python<'_>, + name: &str, + ch_type: &ChType, + seq: &S, + row_count: usize, + nullable: bool, + ) -> PyResult> { + let (values, validity) = seq_values::(py, seq, ch_type, name, row_count, nullable)?; + Ok(Some(T::into_column(values, validity))) + } + + match ch_type { + ChType::Bool => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Int8 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Int16 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Int32 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Int64 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::UInt8 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::UInt16 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::UInt32 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::UInt64 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Int128 | ChType::UInt128 | ChType::Int256 | ChType::UInt256 => { + wide_column_from_seq(py, name, ch_type, seq, row_count, nullable).map(Some) + } + ChType::Float32 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Float64 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::BFloat16 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Date => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Date32 => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::DateTime { .. } => { + prim::(py, name, ch_type, seq, row_count, nullable) + } + ChType::DateTime64 { .. } => { + prim::(py, name, ch_type, seq, row_count, nullable) + } + ChType::Time => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Time64 { .. } => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Interval(_) => prim::(py, name, ch_type, seq, row_count, nullable), + ChType::Uuid => uuid_seq(py, seq, ch_type, name, row_count, nullable).map(Some), + ChType::Ipv4 => ipv4_seq(py, seq, ch_type, name, row_count, nullable).map(Some), + ChType::Enum8 { variants } => { + enum_seq(py, seq, ch_type, variants, name, row_count, nullable).map(Some) + } + ChType::Enum16 { variants } => { + enum_seq(py, seq, ch_type, variants, name, row_count, nullable).map(Some) + } + ChType::Decimal { .. } => { + decimal_column_from_seq(py, name, ch_type, seq, row_count, nullable).map(Some) + } + _ => Ok(None), + } +} + +/// Copy from a buffer-protocol container whose element type matches exactly. +fn try_buffer_column( + py: Python<'_>, + name: &str, + ch_type: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult> { + fn buf( + py: Python<'_>, + name: &str, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, + ) -> PyResult> { + Ok(T::from_buffer(py, name, values, row_count)?.map(|vals| { + // A buffer holds no Python objects, so a nullable column is all-valid. + let validity = nullable.then(|| Bitmap::all_valid(row_count)); + T::into_column(vals, validity) + })) + } + + match ch_type { + ChType::Bool => Ok(bool_buffer_column(values, row_count, nullable)), + ChType::Int8 => buf::(py, name, values, row_count, nullable), + ChType::Int16 => buf::(py, name, values, row_count, nullable), + ChType::Int32 => buf::(py, name, values, row_count, nullable), + ChType::Int64 => buf::(py, name, values, row_count, nullable), + ChType::UInt8 => buf::(py, name, values, row_count, nullable), + ChType::UInt16 => buf::(py, name, values, row_count, nullable), + ChType::UInt32 => buf::(py, name, values, row_count, nullable), + ChType::UInt64 => buf::(py, name, values, row_count, nullable), + ChType::Float32 => buf::(py, name, values, row_count, nullable), + ChType::Float64 => buf::(py, name, values, row_count, nullable), + ChType::BFloat16 => buf::(py, name, values, row_count, nullable), + ChType::Interval(_) => buf::(py, name, values, row_count, nullable), + _ => Ok(None), + } +} + +/// Bool over a numpy bool buffer. Only format '?' with a 1-byte item is +/// accepted (pyo3's `Element for u8` rejects '?'); numpy bool storage is +/// already the 0/1 wire byte, so the buffer copies straight into the column. +fn bool_buffer_column( + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> Option { + let buffer = RawBuffer::matching(values, b'?', 1, row_count)?; + // SAFETY: `matching` validated a C-contiguous buffer of row_count 1-byte + // items, and no Python code runs while the slice is read. + let bytes = unsafe { std::slice::from_raw_parts(buffer.data(), row_count) }; + Some(Column::Bool(if nullable { + // A buffer holds no Python objects, so a nullable column is all-valid. + BoolColumn::from_wire_bytes_nullable(bytes, Bitmap::all_valid(row_count)) + } else { + BoolColumn::from_wire_bytes(bytes) + })) +} + +/// Reinterprets a `Vec` of a `#[repr(transparent)]` wrapper as its inner type, +/// or the reverse, without copying. +/// +/// # Safety +/// Every `T` bit pattern must be a valid `U`. +pub(super) unsafe fn cast_vec(values: Vec) -> Vec { + const { + assert!(std::mem::size_of::() == std::mem::size_of::()); + assert!(std::mem::align_of::() == std::mem::align_of::()); + } + let mut values = std::mem::ManuallyDrop::new(values); + Vec::from_raw_parts( + values.as_mut_ptr().cast::(), + values.len(), + values.capacity(), + ) +} + +/// Multiplicative hasher for pointer-identity keys; object addresses are not +/// attacker-controlled hash-DoS inputs, so a fast mix beats SipHash here. +#[derive(Default)] +pub(super) struct PtrHasher(u64); + +impl std::hash::Hasher for PtrHasher { + fn finish(&self) -> u64 { + self.0 + } + + fn write(&mut self, bytes: &[u8]) { + // The map's usize keys normally arrive via write_usize; fold byte + // input the same way so a std Hash impl change cannot panic. + for &byte in bytes { + self.0 = (self.0 ^ u64::from(byte)).wrapping_mul(0x9E37_79B9_7F4A_7C15); + } + } + + fn write_usize(&mut self, value: usize) { + // Fibonacci hashing: object addresses have aligned low zero bits, so + // mix before the map takes the low bits of the hash. + self.0 = (value as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15); + } +} + +/// Cap on pointer-identity entries so a column of millions of distinct str +/// objects with few distinct contents cannot grow a cache without bound. +pub(super) const PTR_CACHE_CAP: usize = 1 << 16; + +/// Error if a mutable container changed size after a conversion that may have +/// run Python code, before the next borrowed read could go out of bounds. +#[inline] +pub(super) fn check_not_resized(seq: &S, name: &str, row_count: usize) -> PyResult<()> { + if S::MUTABLE && seq.size() != row_count { + return Err(PyValueError::new_err(format!( + "column {name:?} was resized during encoding" + ))); + } + Ok(()) +} diff --git a/rust/ch-core-py/src/insert/json.rs b/rust/ch-core-py/src/insert/json.rs new file mode 100644 index 00000000..62707d9c --- /dev/null +++ b/rust/ch-core-py/src/insert/json.rs @@ -0,0 +1,350 @@ +use super::*; + +/// Build a Dynamic column as its String insert form: each value stringifies +/// with `str()` and None becomes the literal "NULL", matching the python +/// codec's `write_str_values`. +pub(super) fn build_dynamic_string_column( + name: &str, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + let column_values = ColumnValues::new(values, name)?; + check_row_count(name, &column_values, row_count)?; + dynamic_string_column(&column_values, row_count) +} + +pub(super) fn dynamic_string_column<'py, R: RowAccess<'py>>( + rows: &R, + row_count: usize, +) -> PyResult { + let mut offsets = Vec::with_capacity(row_count + 1); + offsets.push(0i32); + let mut data = Vec::new(); + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + data.extend_from_slice(b"NULL"); + } else { + let text = value.str()?; + data.extend_from_slice(text.to_str()?.as_bytes()); + } + let offset = i32::try_from(data.len()) + .map_err(|_| PyValueError::new_err("String column data exceeds i32 offset capacity"))?; + offsets.push(offset); + } + Ok(Column::Utf8(Utf8Column::new(offsets, data))) +} + +/// Build the core's STRING-mode JSON representation. The server performs its +/// normal JSON path/type inference from each document, while the Native header +/// remains the declared JSON type and the core emits structure word 1. +pub(super) fn build_json_text_column( + py: Python<'_>, + name: &str, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult { + let rows = ColumnValues::new(values, name)?; + check_row_count(name, &rows, row_count)?; + if let Ok(list) = values.cast_exact::() { + return json_text_column_from_rows( + py, + name, + &ListRows { + py, + list, + name, + expected: row_count, + }, + row_count, + nullable, + ); + } + if let Ok(tuple) = values.cast_exact::() { + return json_text_column_from_rows(py, name, &TupleRows { py, tuple }, row_count, nullable); + } + json_text_column_from_rows(py, name, &rows, row_count, nullable) +} + +pub(super) fn json_text_column_from_rows<'py, R: RowAccess<'py>>( + py: Python<'py>, + name: &str, + rows: &R, + row_count: usize, + nullable: bool, +) -> PyResult { + // Match the Python codec's `first_value` sniff: nullable scans for the + // first non-null value, non-nullable inspects row 0 alone (None included). + // A str there marks the whole column as pre-serialized JSON text. + let mut direct_text = false; + if nullable { + for row in 0..row_count { + let value = rows.value(row)?; + if !value.is_none() { + direct_text = value.cast::().is_ok(); + break; + } + } + } else if row_count > 0 { + direct_text = rows.value(0)?.cast::().is_ok(); + } + + // Resolved lazily on the first row the native writer cannot serialize. + let mut serializer: Option> = None; + let mut offsets = Vec::with_capacity(row_count + 1); + let mut data = Vec::with_capacity(row_count.saturating_mul(64)); + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + offsets.push(0i32); + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + if let Some(null_map) = &mut null_map { + null_map.push(1); + } + data.extend_from_slice(b"null"); + } else { + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + if direct_text { + append_json_document(&value, name, row, &mut data, false)?; + } else { + let row_start = data.len(); + if write_json_value(value.as_ptr(), &mut data, 0).is_err() { + data.truncate(row_start); + let serializer = match &mut serializer { + Some(serializer) => &*serializer, + slot => { + // Import and attribute lookup run Python, so the + // source list must be revalidated before the next + // unchecked row read. + let resolved = py + .import("clickhouse_connect.datatypes.dynamic")? + .getattr("any_to_json")?; + let serializer = slot.insert(resolved); + rows.validate()?; + &*serializer + } + }; + let encoded = serializer + .call1((&value,)) + .map_err(|err| json_serialize_err(py, name, row, err))?; + rows.validate()?; + append_json_document(&encoded, name, row, &mut data, true)?; + } + } + } + offsets.push(i32::try_from(data.len()).map_err(|_| { + PyValueError::new_err(format!( + "column {name:?} JSON document data exceeds i32 offset capacity" + )) + })?); + } + let validity = null_map.map(|nulls| Bitmap::from_ch_null_map(&nulls)); + Ok(Column::Json( + JsonColumn::text(Utf8Column::new(offsets, data)).with_validity(validity), + )) +} + +fn append_json_document( + value: &Bound<'_, PyAny>, + name: &str, + row: usize, + data: &mut Vec, + serialized: bool, +) -> PyResult<()> { + if let Ok(text) = value.cast::() { + data.extend_from_slice( + text.to_str() + .map_err(|err| json_serialize_err(value.py(), name, row, err))? + .as_bytes(), + ); + return Ok(()); + } + if let Ok(bytes) = value.cast::() { + data.extend_from_slice(bytes.as_bytes()); + return Ok(()); + } + if let Ok(bytes) = value.cast::() { + // SAFETY: the GIL is held and the complete bytearray is copied before + // any Python API can run and resize it. + data.extend_from_slice(unsafe { bytes.as_bytes() }); + return Ok(()); + } + let type_name = python_type_name(value.as_ptr()); + Err(PyValueError::new_err(if serialized { + format!( + "column {name:?} row {row} JSON serializer returned {type_name}, expected str or bytes" + ) + } else { + format!( + "column {name:?} row {row} is {type_name}, expected str because the first value marked \ + this column as pre-serialized JSON strings" + ) + })) +} + +fn json_serialize_err(py: Python<'_>, name: &str, row: usize, err: PyErr) -> PyErr { + let wrapped = PyValueError::new_err(format!( + "column {name:?} row {row} cannot be serialized as JSON" + )); + wrapped.set_cause(py, Some(err)); + wrapped +} + +/// Marker: the native JSON writer cannot serialize this value; the caller +/// rewinds the row and uses the Python serializer. +struct JsonUnsupported; + +/// Nesting depth past which the native writer defers to the Python serializer. +const JSON_NATIVE_MAX_DEPTH: usize = 128; + +/// Write one Python value as JSON text using exact-type C-API fast paths. +/// Exact None/bool/int/float/str/dict/list/tuple traversal runs no user +/// Python, so the source column cannot mutate mid-row; anything else +/// (subclasses, ints past i64, non-finite floats, non-str dict keys, depth +/// past the cap) is `JsonUnsupported` and goes through the Python serializer. +fn write_json_value( + value: *mut ffi::PyObject, + data: &mut Vec, + depth: usize, +) -> Result<(), JsonUnsupported> { + if depth > JSON_NATIVE_MAX_DEPTH { + return Err(JsonUnsupported); + } + // SAFETY: value is live, the GIL is held for the whole traversal, and the + // borrowed container items read below stay valid because no user Python + // runs on this path. + unsafe { + if value == ffi::Py_None() { + data.extend_from_slice(b"null"); + return Ok(()); + } + if value == ffi::Py_True() { + data.extend_from_slice(b"true"); + return Ok(()); + } + if value == ffi::Py_False() { + data.extend_from_slice(b"false"); + return Ok(()); + } + if ffi::PyLong_CheckExact(value) != 0 { + let long = ffi::PyLong_AsLongLong(value); + if long == -1 && !ffi::PyErr_Occurred().is_null() { + ffi::PyErr_Clear(); + return Err(JsonUnsupported); + } + // io::Write into a Vec is infallible. + let _ = write!(data, "{long}"); + return Ok(()); + } + if ffi::PyFloat_CheckExact(value) != 0 { + let float = ffi::PyFloat_AS_DOUBLE(value); + if !float.is_finite() { + return Err(JsonUnsupported); + } + let start = data.len(); + let _ = write!(data, "{float}"); + // Rust Display never uses exponents; a '.'-free rendering is a + // whole float and keeps its float-ness with an explicit ".0". + if !data[start..].contains(&b'.') { + data.extend_from_slice(b".0"); + } + return Ok(()); + } + if ffi::PyUnicode_CheckExact(value) != 0 { + write_json_string(json_utf8(value)?, data); + return Ok(()); + } + if ffi::PyDict_CheckExact(value) != 0 { + data.push(b'{'); + let mut pos: ffi::Py_ssize_t = 0; + let mut key: *mut ffi::PyObject = std::ptr::null_mut(); + let mut item: *mut ffi::PyObject = std::ptr::null_mut(); + let mut first = true; + while ffi::PyDict_Next(value, &mut pos, &mut key, &mut item) != 0 { + if ffi::PyUnicode_CheckExact(key) == 0 { + return Err(JsonUnsupported); + } + if !first { + data.push(b','); + } + first = false; + write_json_string(json_utf8(key)?, data); + data.push(b':'); + write_json_value(item, data, depth + 1)?; + } + data.push(b'}'); + return Ok(()); + } + if ffi::PyList_CheckExact(value) != 0 { + data.push(b'['); + for index in 0..ffi::PyList_GET_SIZE(value) { + if index > 0 { + data.push(b','); + } + write_json_value(ffi::PyList_GET_ITEM(value, index), data, depth + 1)?; + } + data.push(b']'); + return Ok(()); + } + if ffi::PyTuple_CheckExact(value) != 0 { + data.push(b'['); + for index in 0..ffi::PyTuple_GET_SIZE(value) { + if index > 0 { + data.push(b','); + } + write_json_value(ffi::PyTuple_GET_ITEM(value, index), data, depth + 1)?; + } + data.push(b']'); + return Ok(()); + } + } + Err(JsonUnsupported) +} + +/// Borrow the UTF-8 bytes of an exact str; a lone surrogate cannot encode and +/// defers to the Python serializer. +/// +/// # Safety +/// +/// `value` must be a live exact `str` and the GIL must be held. The returned +/// slice borrows the object's cached UTF-8 buffer. +unsafe fn json_utf8<'a>(value: *mut ffi::PyObject) -> Result<&'a [u8], JsonUnsupported> { + let mut size: ffi::Py_ssize_t = 0; + let ptr = ffi::PyUnicode_AsUTF8AndSize(value, &mut size); + if ptr.is_null() { + ffi::PyErr_Clear(); + return Err(JsonUnsupported); + } + Ok(std::slice::from_raw_parts(ptr.cast::(), size as usize)) +} + +/// Emit a JSON string with mandatory-only escapes: `"`, `\`, and control +/// bytes below 0x20 (`\n`/`\r`/`\t` short forms, `\u00XX` otherwise). +/// Non-ASCII bytes pass through as raw UTF-8. +fn write_json_string(bytes: &[u8], data: &mut Vec) { + data.push(b'"'); + let mut start = 0; + for (index, &byte) in bytes.iter().enumerate() { + let escape: &[u8] = match byte { + b'"' => b"\\\"", + b'\\' => b"\\\\", + b'\n' => b"\\n", + b'\r' => b"\\r", + b'\t' => b"\\t", + 0x00..=0x1f => b"", + _ => continue, + }; + data.extend_from_slice(&bytes[start..index]); + if escape.is_empty() { + let _ = write!(data, "\\u{byte:04x}"); + } else { + data.extend_from_slice(escape); + } + start = index + 1; + } + data.extend_from_slice(&bytes[start..]); + data.push(b'"'); +} diff --git a/rust/ch-core-py/src/insert/mod.rs b/rust/ch-core-py/src/insert/mod.rs new file mode 100644 index 00000000..88b70c32 --- /dev/null +++ b/rust/ch-core-py/src/insert/mod.rs @@ -0,0 +1,418 @@ +use std::collections::hash_map::Entry; +use std::collections::HashMap; +use std::ffi::{c_int, c_long}; +use std::io::Write as _; +use std::net::{IpAddr, Ipv4Addr}; + +use pyo3::buffer::{Element, PyBuffer}; +use pyo3::exceptions::{PyMemoryError, PyNotImplementedError, PyRuntimeError, PyValueError}; +use pyo3::ffi; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::{ + PyAnyMethods, PyBool, PyByteArray, PyByteArrayMethods, PyBytes, PyDate, PyDateTime, PyDelta, + PyDeltaAccess, PyDict, PyFloat, PyFrozenSet, PyList, PySet, PyString, PyStringMethods, PyTime, + PyTimeAccess, PyTuple, PyTypeMethods, +}; + +use ch_core_rs::batch::ColBatch as RustColBatch; +use ch_core_rs::bitmap::Bitmap; +use ch_core_rs::column::{ + AggregateStateColumn, ArrayColumn, BoolColumn, Column, DecimalColumn, DictionaryColumn, + FixedBinaryColumn, JsonColumn, MapColumn, NothingColumn, PrimitiveColumn, QBitColumn, + TupleColumn, Utf8Column, VariantColumn, +}; +use ch_core_rs::native::decode::{low_cardinality_dict_value_type, parse_ch_type}; +use ch_core_rs::native::encode::{encode_block, EncodeError, EncodeOptions}; +use ch_core_rs::schema::{ChType, Field, QBitElementType, Schema}; + +use crate::decoder::buffer_to_vec; + +mod containers; +mod fastpath; +mod json; +mod qbit; +mod scalar; +mod special; +mod temporal; +mod variant; + +use containers::*; +use fastpath::*; +use json::*; +use qbit::*; +use scalar::*; +use special::*; +use temporal::*; +use variant::*; + +const EPOCH_DATE_ORDINAL: i64 = 719_163; +const IPV4_V6_PREFIX: [u8; 12] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff]; +const MAX_TIME_SECONDS: i64 = 999 * 3_600 + 59 * 60 + 59; + +/// Match a typed, native-endian PEP 3118 buffer with exactly the requested +/// shape. PyO3 0.23 accepts opposite-endian primitive formats and maps `c` +/// to a one-byte unsigned integer, so both cases need explicit rejection. +fn matching_native_buffer( + value: &Bound<'_, PyAny>, + expected_shape: &[usize], +) -> Option> { + let Ok(buffer) = PyBuffer::::get(value) else { + return None; + }; + let (order, code) = match *buffer.format().to_bytes() { + [code] => (b'@', code), + [order, code] => (order, code), + _ => return None, + }; + let native_order = match order { + b'@' | b'=' => true, + b'<' => cfg!(target_endian = "little"), + b'>' | b'!' => cfg!(target_endian = "big"), + _ => false, + }; + if !native_order || code == b'c' { + return None; + } + if buffer.dimensions() != expected_shape.len() || buffer.shape() != expected_shape { + return None; + } + Some(buffer) +} + +/// Owned PEP 3118 view for formats pyo3's `Element` does not cover ('O' +/// object pointers, '?' bools). `matching` accepts only a 1-D C-contiguous +/// buffer with the exact format code, item size, and row count; any other +/// exporter or view (strided or multi-dimensional) returns `None`. The view +/// is released on drop. +struct RawBuffer<'py> { + py: Python<'py>, + view: Box, +} + +impl<'py> RawBuffer<'py> { + fn matching( + values: &Bound<'py, PyAny>, + code: u8, + itemsize: usize, + row_count: usize, + ) -> Option { + let py = values.py(); + // SAFETY: Py_buffer is a plain C struct; PyObject_GetBuffer fills it + // on success and PyBUF_ND rejects non-C-contiguous exporters. + let mut view = Box::new(unsafe { std::mem::zeroed::() }); + let result = unsafe { + ffi::PyObject_GetBuffer( + values.as_ptr(), + &mut *view, + ffi::PyBUF_FORMAT | ffi::PyBUF_ND, + ) + }; + if result != 0 { + // SAFETY: GIL held; discards the pending TypeError/BufferError. + unsafe { ffi::PyErr_Clear() }; + return None; + } + // Wrapped immediately so every rejection below releases the view. + let buffer = Self { py, view }; + let view = &*buffer.view; + if view.ndim != 1 + || view.buf.is_null() + || view.itemsize != itemsize as ffi::Py_ssize_t + || view.shape.is_null() + { + return None; + } + // SAFETY: PyBUF_ND guarantees a shape array of ndim entries. + if unsafe { *view.shape } != row_count as ffi::Py_ssize_t { + return None; + } + if buffer.format_code() != Some(code) { + return None; + } + if view.buf.align_offset(itemsize) != 0 { + return None; + } + Some(buffer) + } + + /// The single format character, native byte order; a NULL format string + /// means 'B' per the buffer protocol. + fn format_code(&self) -> Option { + if self.view.format.is_null() { + return Some(b'B'); + } + // SAFETY: a non-null Py_buffer format is a NUL-terminated C string. + let bytes = unsafe { std::ffi::CStr::from_ptr(self.view.format) }.to_bytes(); + match *bytes { + [code] => Some(code), + [b'@', code] => Some(code), + _ => None, + } + } + + fn data(&self) -> *const u8 { + self.view.buf.cast() + } +} + +impl Drop for RawBuffer<'_> { + fn drop(&mut self) { + // SAFETY: the held Python token proves the GIL, and the view was + // filled by a successful PyObject_GetBuffer. + let _ = self.py; + unsafe { ffi::PyBuffer_Release(&mut *self.view) }; + } +} + +#[pyfunction] +#[pyo3(signature = (column_names, column_type_names, column_data, row_count, prefix=None))] +pub(crate) fn encode_native_block( + py: Python<'_>, + column_names: Vec, + column_type_names: Vec, + column_data: &Bound<'_, PyAny>, + row_count: usize, + prefix: Option<&Bound<'_, PyAny>>, +) -> PyResult> { + let prefix = match prefix { + Some(obj) if !obj.is_none() => buffer_to_vec(obj).map_err(|_| { + PyValueError::new_err("prefix must be a bytes-like object when supplied") + })?, + _ => Vec::new(), + }; + + if column_names.len() != column_type_names.len() { + return Err(PyValueError::new_err(format!( + "column_names has {} entries but column_type_names has {}", + column_names.len(), + column_type_names.len() + ))); + } + + let data_seq = Indexable::new(column_data, "column_data")?; + if data_seq.len != column_names.len() { + return Err(PyValueError::new_err(format!( + "column_data has {} columns but column_names has {}", + data_seq.len, + column_names.len() + ))); + } + + let mut fields = Vec::with_capacity(column_names.len()); + let mut columns = Vec::with_capacity(column_names.len()); + for (index, (name, type_name)) in column_names.iter().zip(&column_type_names).enumerate() { + let ch_type = parse_ch_type(type_name).ok_or_else(|| { + PyNotImplementedError::new_err(format!( + "unsupported ClickHouse type {type_name:?} for column {name:?}" + )) + })?; + let values = data_seq.get_item(index, "column_data")?; + let column = build_column(py, name, &ch_type, &values, row_count)?; + // Dynamic inserts ship a String column the server casts back with + // type inference, so the Native header carries the substituted type. + let header_type = dynamic_insert_type(&ch_type).unwrap_or(ch_type); + fields.push(Field { + name: name.clone(), + ch_type: header_type, + }); + columns.push(column); + } + + let batch = RustColBatch::new(Schema::new(fields), columns, row_count); + // Defaults: protocol_revision 0 (HTTP INSERT body) and flattened_dynamic + // false (V1/V2 Dynamic layout: pre-25.6 servers reject FLATTENED). + let options = EncodeOptions::default(); + let mut encoded = py + .detach(|| encode_block(&batch, &options)) + .map_err(encode_err)?; + + if !prefix.is_empty() { + let mut out = Vec::with_capacity(prefix.len() + encoded.len()); + out.extend_from_slice(&prefix); + out.append(&mut encoded); + encoded = out; + } + Ok(encoded) +} + +fn encode_err(err: EncodeError) -> PyErr { + match err { + EncodeError::UnsupportedType { column, ch_type } => PyNotImplementedError::new_err( + format!("unsupported ClickHouse type {ch_type} for column {column:?}"), + ), + EncodeError::InconsistentBatch { detail } => PyValueError::new_err(detail), + _ => PyRuntimeError::new_err(format!("Encode error: {err}")), + } +} + +struct Indexable<'py> { + obj: Bound<'py, PyAny>, + len: usize, +} + +impl<'py> Indexable<'py> { + fn new(obj: &Bound<'py, PyAny>, label: &str) -> PyResult { + if is_string_or_bytes_like(obj) { + return Err(PyValueError::new_err(format!( + "{label} must be an indexable collection, not str or bytes" + ))); + } + let len = obj.len().map_err(|_| { + PyValueError::new_err(format!("{label} must be an indexable collection")) + })?; + Ok(Self { + obj: obj.clone(), + len, + }) + } + + fn get_item(&self, index: usize, label: &str) -> PyResult> { + self.obj.get_item(index).map_err(|_| { + PyValueError::new_err(format!("{label} cannot be indexed at position {index}")) + }) + } +} + +fn positional_source<'py>(obj: &Bound<'py, PyAny>) -> Bound<'py, PyAny> { + match obj.getattr("iloc") { + Ok(iloc) if !iloc.is_none() => iloc, + _ => obj.clone(), + } +} + +fn is_string_or_bytes_like(obj: &Bound<'_, PyAny>) -> bool { + obj.cast::().is_ok() + || obj.cast::().is_ok() + || obj.cast::().is_ok() +} + +fn build_column( + py: Python<'_>, + name: &str, + ch_type: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + // Expand SimpleAggregateFunction/geo/Nested aliases to the physical type + // whose column the encoder actually builds; the block header still renders + // the alias spelling from the Field's ChType. + if let Some(delegate) = ch_type.physical_delegate() { + return build_column(py, name, &delegate, values, row_count); + } + match ch_type { + ChType::Nothing => build_nothing_column(name, values, row_count, false), + ChType::AggregateFunction { .. } => { + build_aggregate_state_column(py, name, ch_type, values, row_count) + } + ChType::Nullable(inner) => build_nullable_column(py, name, inner, values, row_count), + ChType::LowCardinality(inner) => { + build_low_cardinality_column(py, name, inner, values, row_count) + } + ChType::Array(inner) => build_array_column(py, name, inner, values, row_count), + ChType::QBit { + element_type, + dimension, + } => build_qbit_column( + py, + name, + *element_type, + *dimension, + values, + row_count, + false, + ), + ChType::Tuple(elements) => build_tuple_column(py, name, elements, values, row_count, false), + ChType::Map(key, value) => build_map_column(py, name, key, value, values, row_count), + ChType::Variant(alternatives) => { + build_variant_column(py, name, alternatives, values, row_count) + } + ChType::Dynamic { .. } => build_dynamic_string_column(name, values, row_count), + ChType::Json { .. } => build_json_text_column(py, name, values, row_count, false), + _ => build_plain_column(py, name, ch_type, values, row_count), + } +} + +/// The insert-header substitution for a type containing Dynamic: Dynamic +/// encodes as a String column the server casts back with +/// `cast_string_to_dynamic_use_inference`, recursing into Array/Tuple/Map +/// exactly like the python codec's `insert_name` chain. `None` when the type +/// contains no Dynamic. +fn dynamic_insert_type(ch_type: &ChType) -> Option { + // Expand name-decoration aliases (Nested, geo) to the physical type first, + // matching `build_column`'s dispatch, so the substituted header describes + // the column actually built. + if let Some(delegate) = ch_type.physical_delegate() { + return dynamic_insert_type(&delegate); + } + match ch_type { + ChType::Dynamic { .. } => Some(ChType::String), + ChType::Array(inner) => { + dynamic_insert_type(inner).map(|inner| ChType::Array(Box::new(inner))) + } + ChType::Map(key, value) => { + let new_key = dynamic_insert_type(key); + let new_value = dynamic_insert_type(value); + if new_key.is_none() && new_value.is_none() { + return None; + } + Some(ChType::Map( + Box::new(new_key.unwrap_or_else(|| (**key).clone())), + Box::new(new_value.unwrap_or_else(|| (**value).clone())), + )) + } + ChType::Tuple(elements) => { + let substituted: Vec> = elements + .iter() + .map(|(_, element)| dynamic_insert_type(element)) + .collect(); + if substituted.iter().all(Option::is_none) { + return None; + } + Some(ChType::Tuple( + elements + .iter() + .zip(substituted) + .map(|((name, element), sub)| { + (name.clone(), sub.unwrap_or_else(|| element.clone())) + }) + .collect(), + )) + } + _ => None, + } +} + +// Insert-supported subset of core's LowCardinality inner-type allowlist, not +// a copy of it; extend only as insert support for a type lands. +fn is_low_cardinality_inner(ch_type: &ChType) -> bool { + matches!( + ch_type, + ChType::Bool + | ChType::Int8 + | ChType::Int16 + | ChType::Int32 + | ChType::Int64 + | ChType::UInt8 + | ChType::UInt16 + | ChType::UInt32 + | ChType::UInt64 + | ChType::Int128 + | ChType::UInt128 + | ChType::Int256 + | ChType::UInt256 + | ChType::Float32 + | ChType::Float64 + | ChType::BFloat16 + | ChType::String + | ChType::FixedString(_) + | ChType::Date + | ChType::Date32 + | ChType::DateTime { .. } + | ChType::Time + | ChType::Interval(_) + | ChType::Uuid + | ChType::Ipv4 + | ChType::Ipv6 + ) +} diff --git a/rust/ch-core-py/src/insert/qbit.rs b/rust/ch-core-py/src/insert/qbit.rs new file mode 100644 index 00000000..cae0cf13 --- /dev/null +++ b/rust/ch-core-py/src/insert/qbit.rs @@ -0,0 +1,553 @@ +use super::*; + +fn qbit_value_count(name: &str, row_count: usize, dimension: usize) -> PyResult { + row_count.checked_mul(dimension).ok_or_else(|| { + PyValueError::new_err(format!( + "column {name:?} QBit value count exceeds usize capacity" + )) + }) +} + +fn reserve_qbit_values(name: &str, value_count: usize) -> PyResult> { + let mut values = Vec::new(); + values.try_reserve_exact(value_count).map_err(|_| { + PyMemoryError::new_err(format!( + "column {name:?} QBit value buffer cannot hold {value_count} elements" + )) + })?; + Ok(values) +} + +fn reserve_qbit_null_map(name: &str, row_count: usize) -> PyResult> { + let mut nulls = Vec::new(); + nulls.try_reserve_exact(row_count).map_err(|_| { + PyMemoryError::new_err(format!( + "column {name:?} QBit null map cannot hold {row_count} rows" + )) + })?; + Ok(nulls) +} + +fn qbit_dimension_error(name: &str, row: usize, dimension: usize, actual: usize) -> PyErr { + PyValueError::new_err(format!( + "column {name:?} row {row} QBit dimension mismatch: expected {dimension}, got {actual}" + )) +} + +fn qbit_vector_error(name: &str, row: usize) -> PyErr { + PyValueError::new_err(format!("column {name:?} row {row} is not a QBit vector")) +} + +fn qbit_element_error(name: &str, row: usize, element: usize, type_name: &str) -> PyErr { + PyValueError::new_err(format!( + "column {name:?} row {row} element {element} cannot be converted to {type_name}" + )) +} + +fn map_buffer(py: Python<'_>, name: &str, buffer: &PyBuffer, map: F) -> PyResult> +where + T: Element + Copy, + F: FnMut(usize, T) -> PyResult, +{ + let mut out = reserve_qbit_values(name, buffer.item_count())?; + extend_buffer(py, buffer, &mut out, map)?; + Ok(out) +} + +fn extend_buffer( + py: Python<'_>, + buffer: &PyBuffer, + out: &mut Vec, + mut map: F, +) -> PyResult<()> +where + T: Element + Copy, + F: FnMut(usize, T) -> PyResult, +{ + if let Some(values) = buffer.as_slice(py) { + for (index, value) in values.iter().enumerate() { + out.push(map(index, value.get())?); + } + return Ok(()); + } + for (index, value) in buffer.to_vec(py)?.into_iter().enumerate() { + out.push(map(index, value)?); + } + Ok(()) +} + +trait QBitValue: Copy { + const DEFAULT: Self; + const TYPE_NAME: &'static str; + + /// Fast exact-float/int conversion that cannot execute Python code. + /// + /// # Safety + /// + /// Requires the GIL and a valid, non-null object pointer. + unsafe fn from_exact(ptr: *mut ffi::PyObject) -> Result; + + fn from_object(value: &Bound<'_, PyAny>) -> Result; + + fn matrix_buffer( + py: Python<'_>, + name: &str, + value: &Bound<'_, PyAny>, + row_count: usize, + dimension: usize, + ) -> PyResult>>; + + fn append_vector_buffer( + py: Python<'_>, + name: &str, + row: usize, + value: &Bound<'_, PyAny>, + dimension: usize, + out: &mut Vec, + ) -> PyResult; + + fn into_child(values: Vec) -> Column; +} + +unsafe fn exact_f64(ptr: *mut ffi::PyObject) -> Result { + ::from_exact(ptr, &ChType::Float64, 0) +} + +impl QBitValue for f32 { + const DEFAULT: Self = 0.0; + const TYPE_NAME: &'static str = "Float32"; + + unsafe fn from_exact(ptr: *mut ffi::PyObject) -> Result { + exact_f64(ptr).map(|value| value as f32) + } + + fn from_object(value: &Bound<'_, PyAny>) -> Result { + value + .extract::() + .map(|value| value as f32) + .map_err(|_| ()) + } + + fn matrix_buffer( + py: Python<'_>, + name: &str, + value: &Bound<'_, PyAny>, + row_count: usize, + dimension: usize, + ) -> PyResult>> { + let shape = [row_count, dimension]; + if let Some(buffer) = matching_native_buffer::(value, &shape) { + // Keep PyO3's single-allocation PyBuffer_ToContiguous path here. + // The shape is already exact, and avoiding a zero-initialization + // pass preserves memcpy-level throughput for large matrices. + return buffer.to_vec(py).map(Some); + } + matching_native_buffer::(value, &shape) + .map(|buffer| map_buffer(py, name, &buffer, |_index, value| Ok(value as f32))) + .transpose() + } + + fn append_vector_buffer( + py: Python<'_>, + _name: &str, + _row: usize, + value: &Bound<'_, PyAny>, + dimension: usize, + out: &mut Vec, + ) -> PyResult { + let shape = [dimension]; + if let Some(buffer) = matching_native_buffer::(value, &shape) { + extend_buffer(py, &buffer, out, |_index, value| Ok(value))?; + return Ok(true); + } + if let Some(buffer) = matching_native_buffer::(value, &shape) { + extend_buffer(py, &buffer, out, |_index, value| Ok(value as f32))?; + return Ok(true); + } + Ok(false) + } + + fn into_child(values: Vec) -> Column { + Column::Float32(PrimitiveColumn::new(values)) + } +} + +impl QBitValue for f64 { + const DEFAULT: Self = 0.0; + const TYPE_NAME: &'static str = "Float64"; + + unsafe fn from_exact(ptr: *mut ffi::PyObject) -> Result { + exact_f64(ptr) + } + + fn from_object(value: &Bound<'_, PyAny>) -> Result { + value.extract::().map_err(|_| ()) + } + + fn matrix_buffer( + py: Python<'_>, + name: &str, + value: &Bound<'_, PyAny>, + row_count: usize, + dimension: usize, + ) -> PyResult>> { + let shape = [row_count, dimension]; + if let Some(buffer) = matching_native_buffer::(value, &shape) { + // See the Float32 path above. PyO3 owns the exact-sized matrix + // allocation so the matching-dtype case stays a contiguous copy. + return buffer.to_vec(py).map(Some); + } + matching_native_buffer::(value, &shape) + .map(|buffer| map_buffer(py, name, &buffer, |_index, value| Ok(value.into()))) + .transpose() + } + + fn append_vector_buffer( + py: Python<'_>, + _name: &str, + _row: usize, + value: &Bound<'_, PyAny>, + dimension: usize, + out: &mut Vec, + ) -> PyResult { + let shape = [dimension]; + if let Some(buffer) = matching_native_buffer::(value, &shape) { + extend_buffer(py, &buffer, out, |_index, value| Ok(value))?; + return Ok(true); + } + if let Some(buffer) = matching_native_buffer::(value, &shape) { + extend_buffer(py, &buffer, out, |_index, value| Ok(value.into()))?; + return Ok(true); + } + Ok(false) + } + + fn into_child(values: Vec) -> Column { + Column::Float64(PrimitiveColumn::new(values)) + } +} + +impl QBitValue for [u8; 2] { + const DEFAULT: Self = [0; 2]; + const TYPE_NAME: &'static str = "BFloat16"; + + unsafe fn from_exact(ptr: *mut ffi::PyObject) -> Result { + exact_f64(ptr).and_then(checked_f64_to_bfloat16) + } + + fn from_object(value: &Bound<'_, PyAny>) -> Result { + value + .extract::() + .map_err(|_| ()) + .and_then(checked_f64_to_bfloat16) + } + + fn matrix_buffer( + py: Python<'_>, + name: &str, + value: &Bound<'_, PyAny>, + row_count: usize, + dimension: usize, + ) -> PyResult>> { + let shape = [row_count, dimension]; + let convert = |index: usize, value: f64| { + checked_f64_to_bfloat16(value).map_err(|_| { + qbit_element_error(name, index / dimension, index % dimension, Self::TYPE_NAME) + }) + }; + if let Some(buffer) = matching_native_buffer::(value, &shape) { + return map_buffer(py, name, &buffer, |index, value| { + convert(index, value.into()) + }) + .map(Some); + } + matching_native_buffer::(value, &shape) + .map(|buffer| map_buffer(py, name, &buffer, convert)) + .transpose() + } + + fn append_vector_buffer( + py: Python<'_>, + name: &str, + row: usize, + value: &Bound<'_, PyAny>, + dimension: usize, + out: &mut Vec, + ) -> PyResult { + let shape = [dimension]; + let convert = |element: usize, value: f64| { + checked_f64_to_bfloat16(value) + .map_err(|_| qbit_element_error(name, row, element, Self::TYPE_NAME)) + }; + if let Some(buffer) = matching_native_buffer::(value, &shape) { + extend_buffer(py, &buffer, out, |element, value| { + convert(element, value.into()) + })?; + return Ok(true); + } + if let Some(buffer) = matching_native_buffer::(value, &shape) { + extend_buffer(py, &buffer, out, convert)?; + return Ok(true); + } + Ok(false) + } + + fn into_child(values: Vec) -> Column { + Column::BFloat16(PrimitiveColumn::new(values)) + } +} + +fn append_qbit_seq( + py: Python<'_>, + name: &str, + row: usize, + seq: &S, + dimension: usize, + out: &mut Vec, +) -> PyResult<()> { + let actual = seq.size(); + if actual != dimension { + return Err(qbit_dimension_error(name, row, dimension, actual)); + } + for element in 0..dimension { + // The sequence length is checked above and after every fallback that + // may execute Python code. The temporary strong reference must drop + // before that check because its finalizer can also resize `seq`. + let ptr = unsafe { seq.get(element) }; + match unsafe { T::from_exact(ptr) } { + Ok(value) => out.push(value), + Err(()) => { + let value = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + let converted = T::from_object(&value) + .map_err(|_| qbit_element_error(name, row, element, T::TYPE_NAME))?; + out.push(converted); + drop(value); + check_not_resized(seq, name, dimension)?; + } + } + } + Ok(()) +} + +fn append_qbit_vector( + py: Python<'_>, + name: &str, + row: usize, + value: &Bound<'_, PyAny>, + dimension: usize, + out: &mut Vec, +) -> PyResult<()> { + if let Ok(list) = value.cast_exact::() { + return append_qbit_seq(py, name, row, &ListSeq(list), dimension, out); + } + if let Ok(tuple) = value.cast_exact::() { + return append_qbit_seq(py, name, row, &TupleSeq(tuple), dimension, out); + } + if T::append_vector_buffer(py, name, row, value, dimension, out)? { + return Ok(()); + } + if is_string_or_bytes_like(value) { + return Err(qbit_vector_error(name, row)); + } + let actual = value.len().map_err(|_| qbit_vector_error(name, row))?; + if actual != dimension { + return Err(qbit_dimension_error(name, row, dimension, actual)); + } + for element in 0..dimension { + let item = value + .get_item(element) + .map_err(|_| qbit_vector_error(name, row))?; + out.push( + T::from_object(&item) + .map_err(|_| qbit_element_error(name, row, element, T::TYPE_NAME))?, + ); + } + Ok(()) +} + +fn finish_qbit(values: Vec, dimension: usize, validity: Option) -> Column { + let values = T::into_child(values); + Column::QBit(match validity { + Some(validity) => QBitColumn::new_nullable(values, dimension, validity), + None => QBitColumn::new(values, dimension), + }) +} + +fn qbit_column_from_seq( + py: Python<'_>, + name: &str, + qbit_type: &ChType, + dimension: usize, + seq: &S, + row_count: usize, + nullable: bool, +) -> PyResult { + if seq.size() != row_count { + return Err(PyValueError::new_err(format!( + "column {name:?} has {} values but row_count is {row_count}", + seq.size() + ))); + } + let value_count = qbit_value_count(name, row_count, dimension)?; + let mut values = reserve_qbit_values(name, value_count)?; + let mut null_map = if nullable { + Some(reserve_qbit_null_map(name, row_count)?) + } else { + None + }; + for row in 0..row_count { + let ptr = unsafe { seq.get(row) }; + if ptr == unsafe { ffi::Py_None() } { + let Some(nulls) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {qbit_type} is not Nullable" + ))); + }; + nulls.push(1); + values.resize(values.len() + dimension, T::DEFAULT); + continue; + } + if let Some(nulls) = &mut null_map { + nulls.push(0); + } + let value = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + append_qbit_vector(py, name, row, &value, dimension, &mut values)?; + // Dropping the row can run a finalizer that resizes `seq`, so it must + // happen before the next unchecked list read is declared safe. + drop(value); + check_not_resized(seq, name, row_count)?; + } + let validity = null_map.map(|nulls| Bitmap::from_ch_null_map(&nulls)); + Ok(finish_qbit(values, dimension, validity)) +} + +fn qbit_column_from_rows<'py, T: QBitValue, R: RowAccess<'py>>( + py: Python<'py>, + name: &str, + qbit_type: &ChType, + dimension: usize, + rows: &R, + row_count: usize, + nullable: bool, +) -> PyResult { + let value_count = qbit_value_count(name, row_count, dimension)?; + let mut values = reserve_qbit_values(name, value_count)?; + let mut null_map = if nullable { + Some(reserve_qbit_null_map(name, row_count)?) + } else { + None + }; + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + let Some(nulls) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {qbit_type} is not Nullable" + ))); + }; + nulls.push(1); + values.resize(values.len() + dimension, T::DEFAULT); + continue; + } + if let Some(nulls) = &mut null_map { + nulls.push(0); + } + append_qbit_vector(py, name, row, &value, dimension, &mut values)?; + drop(value); + rows.validate()?; + } + let validity = null_map.map(|nulls| Bitmap::from_ch_null_map(&nulls)); + Ok(finish_qbit(values, dimension, validity)) +} + +fn build_typed_qbit( + py: Python<'_>, + name: &str, + qbit_type: &ChType, + dimension: usize, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult { + if let Some(matrix) = T::matrix_buffer(py, name, values, row_count, dimension)? { + let validity = nullable.then(|| Bitmap::all_valid(row_count)); + return Ok(finish_qbit(matrix, dimension, validity)); + } + if let Ok(list) = values.cast_exact::() { + return qbit_column_from_seq::( + py, + name, + qbit_type, + dimension, + &ListSeq(list), + row_count, + nullable, + ); + } + if let Ok(tuple) = values.cast_exact::() { + return qbit_column_from_seq::( + py, + name, + qbit_type, + dimension, + &TupleSeq(tuple), + row_count, + nullable, + ); + } + let rows = ColumnValues::new(values, name)?; + check_row_count(name, &rows, row_count)?; + qbit_column_from_rows::(py, name, qbit_type, dimension, &rows, row_count, nullable) +} + +pub(super) fn build_qbit_column( + py: Python<'_>, + name: &str, + element_type: QBitElementType, + dimension: usize, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult { + let qbit_type = ChType::QBit { + element_type, + dimension, + }; + match element_type { + QBitElementType::BFloat16 => build_typed_qbit::<[u8; 2]>( + py, name, &qbit_type, dimension, values, row_count, nullable, + ), + QBitElementType::Float32 => { + build_typed_qbit::(py, name, &qbit_type, dimension, values, row_count, nullable) + } + QBitElementType::Float64 => { + build_typed_qbit::(py, name, &qbit_type, dimension, values, row_count, nullable) + } + } +} + +pub(super) fn build_qbit_column_from_seq( + py: Python<'_>, + name: &str, + element_type: QBitElementType, + dimension: usize, + seq: &S, + row_count: usize, + nullable: bool, +) -> PyResult { + let qbit_type = ChType::QBit { + element_type, + dimension, + }; + match element_type { + QBitElementType::BFloat16 => qbit_column_from_seq::<[u8; 2], _>( + py, name, &qbit_type, dimension, seq, row_count, nullable, + ), + QBitElementType::Float32 => qbit_column_from_seq::( + py, name, &qbit_type, dimension, seq, row_count, nullable, + ), + QBitElementType::Float64 => qbit_column_from_seq::( + py, name, &qbit_type, dimension, seq, row_count, nullable, + ), + } +} diff --git a/rust/ch-core-py/src/insert/scalar.rs b/rust/ch-core-py/src/insert/scalar.rs new file mode 100644 index 00000000..c2df7199 --- /dev/null +++ b/rust/ch-core-py/src/insert/scalar.rs @@ -0,0 +1,1692 @@ +use pyo3::sync::PyOnceLock; +use pyo3::types::PyType; + +use super::*; + +#[derive(Debug)] +pub(super) enum Scalar { + Bool(bool), + Int8(i8), + Int16(i16), + Int32(i32), + Int64(i64), + UInt8(u8), + UInt16(u16), + UInt32(u32), + UInt64(u64), + WideInt(Vec), + Float32(f32), + Float64(f64), + BFloat16([u8; 2]), + Date(u16), + Date32(i32), + DateTime(u32), + DateTime64(i64), + Time(i32), + Time64(i64), + Interval(i64), + Bytes(Vec), + Ipv4(u32), + Enum8(i8), + Enum16(i16), +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(super) enum ScalarKey { + Bool(bool), + Int8(i8), + Int16(i16), + Int32(i32), + Int64(i64), + UInt8(u8), + UInt16(u16), + UInt32(u32), + UInt64(u64), + WideInt(Vec), + Float32(u32), + Float64(u64), + BFloat16([u8; 2]), + Date(u16), + Date32(i32), + DateTime(u32), + DateTime64(i64), + Time(i32), + Time64(i64), + Interval(i64), + Bytes(Vec), + Ipv4(u32), + Enum8(i8), + Enum16(i16), +} + +impl Scalar { + pub(super) fn key(&self) -> ScalarKey { + match self { + Scalar::Bool(v) => ScalarKey::Bool(*v), + Scalar::Int8(v) => ScalarKey::Int8(*v), + Scalar::Int16(v) => ScalarKey::Int16(*v), + Scalar::Int32(v) => ScalarKey::Int32(*v), + Scalar::Int64(v) => ScalarKey::Int64(*v), + Scalar::UInt8(v) => ScalarKey::UInt8(*v), + Scalar::UInt16(v) => ScalarKey::UInt16(*v), + Scalar::UInt32(v) => ScalarKey::UInt32(*v), + Scalar::UInt64(v) => ScalarKey::UInt64(*v), + Scalar::WideInt(v) => ScalarKey::WideInt(v.clone()), + Scalar::Float32(v) => ScalarKey::Float32(v.to_bits()), + Scalar::Float64(v) => ScalarKey::Float64(v.to_bits()), + Scalar::BFloat16(v) => ScalarKey::BFloat16(*v), + Scalar::Date(v) => ScalarKey::Date(*v), + Scalar::Date32(v) => ScalarKey::Date32(*v), + Scalar::DateTime(v) => ScalarKey::DateTime(*v), + Scalar::DateTime64(v) => ScalarKey::DateTime64(*v), + Scalar::Time(v) => ScalarKey::Time(*v), + Scalar::Time64(v) => ScalarKey::Time64(*v), + Scalar::Interval(v) => ScalarKey::Interval(*v), + Scalar::Bytes(v) => ScalarKey::Bytes(v.clone()), + Scalar::Ipv4(v) => ScalarKey::Ipv4(*v), + Scalar::Enum8(v) => ScalarKey::Enum8(*v), + Scalar::Enum16(v) => ScalarKey::Enum16(*v), + } + } +} + +macro_rules! primitive_column { + ($scalars:expr, $validity:expr, $scalar_variant:ident, $column_variant:ident, $ty:ty) => {{ + let mut values = Vec::<$ty>::with_capacity($scalars.len()); + for scalar in $scalars { + match scalar { + Scalar::$scalar_variant(value) => values.push(value), + _ => return Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + Ok(Column::$column_variant(match $validity { + Some(validity) => PrimitiveColumn::new_nullable(values, validity), + None => PrimitiveColumn::new(values), + })) + }}; +} + +pub(super) fn column_from_scalars( + ch_type: &ChType, + scalars: Vec, + validity: Option, +) -> PyResult { + match ch_type { + ChType::Bool => { + let mut bytes = Vec::with_capacity(scalars.len()); + for scalar in scalars { + match scalar { + Scalar::Bool(value) => bytes.push(u8::from(value)), + _ => return Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + Ok(Column::Bool(match validity { + Some(validity) => BoolColumn::from_wire_bytes_nullable(&bytes, validity), + None => BoolColumn::from_wire_bytes(&bytes), + })) + } + ChType::Int8 => primitive_column!(scalars, validity, Int8, Int8, i8), + ChType::Int16 => primitive_column!(scalars, validity, Int16, Int16, i16), + ChType::Int32 => primitive_column!(scalars, validity, Int32, Int32, i32), + ChType::Int64 => primitive_column!(scalars, validity, Int64, Int64, i64), + ChType::UInt8 => primitive_column!(scalars, validity, UInt8, UInt8, u8), + ChType::UInt16 => primitive_column!(scalars, validity, UInt16, UInt16, u16), + ChType::UInt32 => primitive_column!(scalars, validity, UInt32, UInt32, u32), + ChType::UInt64 => primitive_column!(scalars, validity, UInt64, UInt64, u64), + ChType::Int128 | ChType::UInt128 | ChType::Int256 | ChType::UInt256 => { + build_wide_int_column(ch_type, scalars, validity) + } + ChType::Float32 => primitive_column!(scalars, validity, Float32, Float32, f32), + ChType::Float64 => primitive_column!(scalars, validity, Float64, Float64, f64), + ChType::BFloat16 => { + primitive_column!(scalars, validity, BFloat16, BFloat16, [u8; 2]) + } + ChType::QBit { .. } => Err(PyNotImplementedError::new_err(format!( + "unsupported ClickHouse type {ch_type}" + ))), + ChType::AggregateFunction { .. } => Err(PyNotImplementedError::new_err( + "AggregateFunction insert conversion is not implemented", + )), + ChType::Nothing => Err(PyRuntimeError::new_err( + "internal error: Nothing columns use the length-only builder", + )), + ChType::Date => primitive_column!(scalars, validity, Date, Date, u16), + ChType::Date32 => primitive_column!(scalars, validity, Date32, Date32, i32), + ChType::DateTime { .. } => { + primitive_column!(scalars, validity, DateTime, DateTime, u32) + } + ChType::DateTime64 { .. } => { + primitive_column!(scalars, validity, DateTime64, DateTime64, i64) + } + ChType::Time => primitive_column!(scalars, validity, Time, Time, i32), + ChType::Time64 { .. } => primitive_column!(scalars, validity, Time64, Time64, i64), + ChType::Interval(_) => primitive_column!(scalars, validity, Interval, Interval, i64), + ChType::String => build_utf8_column(scalars, validity), + ChType::FixedString(width) => build_fixed_binary_column(scalars, *width, validity), + ChType::Uuid => build_uuid_column(scalars, validity), + ChType::Ipv4 => primitive_column!(scalars, validity, Ipv4, Ipv4, u32), + ChType::Ipv6 => build_ipv6_column(scalars, validity), + ChType::Enum8 { .. } => primitive_column!(scalars, validity, Enum8, Enum8, i8), + ChType::Enum16 { .. } => primitive_column!(scalars, validity, Enum16, Enum16, i16), + ChType::Decimal { + precision, scale, .. + } => { + let width = decimal_width(*precision)?; + let mut data = Vec::with_capacity(width * scalars.len()); + for scalar in scalars { + match scalar { + Scalar::Bytes(value) if value.len() == width => data.extend_from_slice(&value), + _ => return Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + Ok(Column::Decimal(match validity { + Some(validity) => { + DecimalColumn::new_nullable(data, width, *precision, *scale, validity) + } + None => DecimalColumn::new(data, width, *precision, *scale), + })) + } + ChType::Array(_) => Err(PyNotImplementedError::new_err( + "Array columns are built by build_array_column, not the scalar path", + )), + ChType::Tuple(_) | ChType::Map(..) | ChType::Variant(_) => { + Err(PyNotImplementedError::new_err( + "Tuple, Map, and Variant columns are built by their container paths, not the scalar path", + )) + } + ChType::Nullable(_) | ChType::LowCardinality(_) => Err(PyNotImplementedError::new_err( + "nested wrapper conversion is not supported", + )), + ChType::SimpleAggregateFunction { .. } + | ChType::Geo(_) + | ChType::Geometry + | ChType::Nested(_) => Err(PyNotImplementedError::new_err( + "name-decoration aliases are expanded to their physical type before the scalar path", + )), + ChType::Dynamic { .. } => Err(PyNotImplementedError::new_err( + "Dynamic columns are built by the String insert path, not the scalar path", + )), + ChType::Json { .. } => Err(PyNotImplementedError::new_err( + "JSON columns are built by the JSON text insert path, not the scalar path", + )), + } +} + +fn build_wide_int_column( + ch_type: &ChType, + scalars: Vec, + validity: Option, +) -> PyResult { + let (width, _, _) = wide_int_layout(ch_type) + .ok_or_else(|| PyValueError::new_err("internal wide integer type mismatch"))?; + let byte_len = width + .checked_mul(scalars.len()) + .ok_or_else(|| PyValueError::new_err("wide integer column byte size overflow"))?; + let mut data = Vec::with_capacity(byte_len); + for scalar in scalars { + match scalar { + Scalar::WideInt(value) if value.len() == width => data.extend_from_slice(&value), + _ => return Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + finish_wide_int_column(ch_type, data, validity) +} + +pub(super) fn wide_int_layout(ch_type: &ChType) -> Option<(usize, bool, &'static str)> { + match ch_type { + ChType::Int128 => Some((16, true, "Int128")), + ChType::UInt128 => Some((16, false, "UInt128")), + ChType::Int256 => Some((32, true, "Int256")), + ChType::UInt256 => Some((32, false, "UInt256")), + _ => None, + } +} + +pub(super) fn wide_data_buffer(name: &str, width: usize, row_count: usize) -> PyResult> { + let byte_len = width.checked_mul(row_count).ok_or_else(|| { + PyValueError::new_err(format!( + "column {name:?} wide integer byte size exceeds usize capacity" + )) + })?; + Ok(vec![0u8; byte_len]) +} + +pub(super) fn finish_wide_int_column( + ch_type: &ChType, + data: Vec, + validity: Option, +) -> PyResult { + let (width, _, _) = wide_int_layout(ch_type) + .ok_or_else(|| PyValueError::new_err("internal wide integer type mismatch"))?; + let column = match validity { + Some(validity) => FixedBinaryColumn::new_nullable(data, width, validity), + None => FixedBinaryColumn::new(data, width), + }; + Ok(match ch_type { + ChType::Int128 => Column::Int128(column), + ChType::UInt128 => Column::UInt128(column), + ChType::Int256 => Column::Int256(column), + ChType::UInt256 => Column::UInt256(column), + _ => return Err(PyValueError::new_err("internal wide integer type mismatch")), + }) +} + +fn build_utf8_column(scalars: Vec, validity: Option) -> PyResult { + let mut offsets = Vec::with_capacity(scalars.len() + 1); + let mut data = Vec::new(); + offsets.push(0); + for scalar in scalars { + let value = match scalar { + Scalar::Bytes(value) => value, + _ => return Err(PyValueError::new_err("internal scalar type mismatch")), + }; + data.extend_from_slice(&value); + let offset = i32::try_from(data.len()) + .map_err(|_| PyValueError::new_err("String column data exceeds i32 offset capacity"))?; + offsets.push(offset); + } + Ok(Column::Utf8(match validity { + Some(validity) => Utf8Column::new_nullable(offsets, data, validity), + None => Utf8Column::new(offsets, data), + })) +} + +fn build_fixed_binary_column( + scalars: Vec, + width: usize, + validity: Option, +) -> PyResult { + let mut data = Vec::with_capacity(width * scalars.len()); + for scalar in scalars { + match scalar { + Scalar::Bytes(value) if value.len() == width => data.extend_from_slice(&value), + _ => return Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + Ok(Column::FixedBinary(match validity { + Some(validity) => FixedBinaryColumn::new_nullable(data, width, validity), + None => FixedBinaryColumn::new(data, width), + })) +} + +fn build_uuid_column(scalars: Vec, validity: Option) -> PyResult { + let mut data = Vec::with_capacity(16 * scalars.len()); + for scalar in scalars { + match scalar { + Scalar::Bytes(value) if value.len() == 16 => data.extend_from_slice(&value), + _ => return Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + Ok(Column::Uuid(match validity { + Some(validity) => FixedBinaryColumn::new_nullable(data, 16, validity), + None => FixedBinaryColumn::new(data, 16), + })) +} + +fn build_ipv6_column(scalars: Vec, validity: Option) -> PyResult { + let mut data = Vec::with_capacity(16 * scalars.len()); + for scalar in scalars { + match scalar { + Scalar::Bytes(value) if value.len() == 16 => data.extend_from_slice(&value), + _ => return Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + Ok(Column::Ipv6(match validity { + Some(validity) => FixedBinaryColumn::new_nullable(data, 16, validity), + None => FixedBinaryColumn::new(data, 16), + })) +} + +pub(super) fn convert_scalar( + py: Python<'_>, + ch_type: &ChType, + value: &Bound<'_, PyAny>, + column: &str, + row: usize, +) -> PyResult { + macro_rules! integer_scalar { + ($ty:ty, $variant:ident, $type_name:literal) => {{ + let integer = integer_object( + py, + value, + column, + row, + $type_name, + "pass an integer value", + false, + )?; + integer + .extract::<$ty>() + .map(Scalar::$variant) + .map_err(|_| integer_range_error(value, column, row, $type_name)) + }}; + } + + match ch_type { + ChType::Bool => value + .extract::() + .map(Scalar::Bool) + .map_err(|_| conversion_error(column, row, "Bool")), + ChType::Int8 => integer_scalar!(i8, Int8, "Int8"), + ChType::Int16 => integer_scalar!(i16, Int16, "Int16"), + ChType::Int32 => integer_scalar!(i32, Int32, "Int32"), + ChType::Int64 => integer_scalar!(i64, Int64, "Int64"), + ChType::UInt8 => integer_scalar!(u8, UInt8, "UInt8"), + ChType::UInt16 => integer_scalar!(u16, UInt16, "UInt16"), + ChType::UInt32 => integer_scalar!(u32, UInt32, "UInt32"), + ChType::UInt64 => integer_scalar!(u64, UInt64, "UInt64"), + ChType::Int128 => { + wide_int_bytes(py, value, 16, true, column, row, "Int128").map(Scalar::WideInt) + } + ChType::UInt128 => { + wide_int_bytes(py, value, 16, false, column, row, "UInt128").map(Scalar::WideInt) + } + ChType::Int256 => { + wide_int_bytes(py, value, 32, true, column, row, "Int256").map(Scalar::WideInt) + } + ChType::UInt256 => { + wide_int_bytes(py, value, 32, false, column, row, "UInt256").map(Scalar::WideInt) + } + ChType::Float32 => value + .extract::() + .map(Scalar::Float32) + .map_err(|_| float_conversion_error(value, column, row, "Float32")), + ChType::Float64 => value + .extract::() + .map(Scalar::Float64) + .map_err(|_| float_conversion_error(value, column, row, "Float64")), + ChType::BFloat16 => value + .extract::() + .map_err(|_| float_conversion_error(value, column, row, "BFloat16")) + .and_then(|value| { + checked_f64_to_bfloat16(value) + .map(Scalar::BFloat16) + .map_err(|_| conversion_error(column, row, "BFloat16")) + }), + ChType::QBit { .. } => Err(PyNotImplementedError::new_err(format!( + "unsupported ClickHouse type {ch_type} for column {column:?}" + ))), + ChType::AggregateFunction { .. } => Err(PyNotImplementedError::new_err( + "AggregateFunction insert conversion is not implemented", + )), + ChType::Nothing => Err(PyRuntimeError::new_err( + "internal error: Nothing columns use the length-only builder", + )), + ChType::String => Ok(Scalar::Bytes(bytes_value(value, column, row, "String")?)), + ChType::FixedString(width) => Ok(Scalar::Bytes(fixed_string_value( + value, *width, column, row, + )?)), + ChType::Date => Ok(Scalar::Date(date_days(value, column, row).and_then( + |days| { + u16::try_from(days).map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} Date value {days} is outside UInt16 range" + )) + }) + }, + )?)), + ChType::Date32 => Ok(Scalar::Date32(date_days(value, column, row).and_then( + |days| { + i32::try_from(days).map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} Date32 value {days} is outside Int32 range" + )) + }) + }, + )?)), + ChType::DateTime { .. } => { + let secs = datetime_seconds(value, column, row)?; + Ok(Scalar::DateTime(u32::try_from(secs).map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} DateTime value {secs} is outside UInt32 range" + )) + })?)) + } + ChType::DateTime64 { precision, .. } => Ok(Scalar::DateTime64(datetime64_ticks( + py, value, *precision, column, row, + )?)), + ChType::Time => Ok(Scalar::Time(time_ticks(value, column, row)?)), + ChType::Time64 { precision } => Ok(Scalar::Time64(time64_ticks( + value, *precision, column, row, + )?)), + ChType::Interval(_) => value + .extract::() + .map(Scalar::Interval) + .map_err(|_| conversion_error(column, row, &ch_type.to_string())), + ChType::Uuid => Ok(Scalar::Bytes(uuid_bytes(value, column, row)?)), + ChType::Ipv4 => Ok(Scalar::Ipv4(ipv4_value(value, column, row)?)), + ChType::Ipv6 => Ok(Scalar::Bytes(ipv6_bytes(value, column, row)?)), + ChType::Enum8 { variants } => Ok(Scalar::Enum8(enum8_value(value, variants, column, row)?)), + ChType::Enum16 { variants } => { + Ok(Scalar::Enum16(enum16_value(value, variants, column, row)?)) + } + ChType::Decimal { + precision, scale, .. + } => { + let width = decimal_width(*precision)?; + let text = decimal_text(value, column, row)?; + Ok(Scalar::Bytes(decimal_to_le_bytes( + &text, width, *precision, *scale, column, row, + )?)) + } + ChType::Array(_) => Err(PyNotImplementedError::new_err( + "Array columns are built by build_array_column, not the scalar path", + )), + ChType::Tuple(_) | ChType::Map(..) | ChType::Variant(_) => { + Err(PyNotImplementedError::new_err( + "Tuple, Map, and Variant columns are built by their container paths, not the scalar path", + )) + } + ChType::Nullable(_) | ChType::LowCardinality(_) => Err(PyNotImplementedError::new_err( + "nested wrapper conversion is not supported", + )), + ChType::SimpleAggregateFunction { .. } + | ChType::Geo(_) + | ChType::Geometry + | ChType::Nested(_) => { + Err(PyNotImplementedError::new_err( + "name-decoration aliases are expanded to their physical type before the scalar path", + )) + } + ChType::Dynamic { .. } => Err(PyNotImplementedError::new_err( + "Dynamic columns are built by the String insert path, not the scalar path", + )), + ChType::Json { .. } => Err(PyNotImplementedError::new_err( + "JSON columns are built by the JSON text insert path, not the scalar path", + )), + } +} + +pub(super) fn default_scalar(ch_type: &ChType) -> PyResult { + match ch_type { + ChType::Bool => Ok(Scalar::Bool(false)), + ChType::Int8 => Ok(Scalar::Int8(0)), + ChType::Int16 => Ok(Scalar::Int16(0)), + ChType::Int32 => Ok(Scalar::Int32(0)), + ChType::Int64 => Ok(Scalar::Int64(0)), + ChType::UInt8 => Ok(Scalar::UInt8(0)), + ChType::UInt16 => Ok(Scalar::UInt16(0)), + ChType::UInt32 => Ok(Scalar::UInt32(0)), + ChType::UInt64 => Ok(Scalar::UInt64(0)), + ChType::Int128 | ChType::UInt128 => Ok(Scalar::WideInt(vec![0; 16])), + ChType::Int256 | ChType::UInt256 => Ok(Scalar::WideInt(vec![0; 32])), + ChType::Float32 => Ok(Scalar::Float32(0.0)), + ChType::Float64 => Ok(Scalar::Float64(0.0)), + ChType::BFloat16 => Ok(Scalar::BFloat16([0; 2])), + ChType::QBit { .. } => Err(PyNotImplementedError::new_err(format!( + "unsupported ClickHouse type {ch_type}" + ))), + ChType::AggregateFunction { .. } => Err(PyNotImplementedError::new_err( + "AggregateFunction has no generic default state; provide exact serialized state bytes", + )), + ChType::Nothing => Err(PyRuntimeError::new_err( + "internal error: Nothing columns use the length-only builder", + )), + ChType::String => Ok(Scalar::Bytes(Vec::new())), + ChType::FixedString(width) => Ok(Scalar::Bytes(vec![0; *width])), + ChType::Date => Ok(Scalar::Date(0)), + ChType::Date32 => Ok(Scalar::Date32(0)), + ChType::DateTime { .. } => Ok(Scalar::DateTime(0)), + ChType::DateTime64 { .. } => Ok(Scalar::DateTime64(0)), + ChType::Time => Ok(Scalar::Time(0)), + ChType::Time64 { .. } => Ok(Scalar::Time64(0)), + ChType::Interval(_) => Ok(Scalar::Interval(0)), + ChType::Uuid => Ok(Scalar::Bytes(vec![0; 16])), + ChType::Ipv4 => Ok(Scalar::Ipv4(0)), + ChType::Ipv6 => Ok(Scalar::Bytes(vec![0; 16])), + ChType::Enum8 { .. } => Ok(Scalar::Enum8(0)), + ChType::Enum16 { .. } => Ok(Scalar::Enum16(0)), + ChType::Decimal { precision, .. } => Ok(Scalar::Bytes(vec![0; decimal_width(*precision)?])), + ChType::Array(_) => Err(PyNotImplementedError::new_err( + "Array columns are built by build_array_column, not the scalar path", + )), + ChType::Tuple(_) | ChType::Map(..) | ChType::Variant(_) => { + Err(PyNotImplementedError::new_err( + "Tuple, Map, and Variant columns have no scalar placeholder", + )) + } + ChType::Nullable(_) | ChType::LowCardinality(_) => Err(PyNotImplementedError::new_err( + "nested wrapper conversion is not supported", + )), + ChType::SimpleAggregateFunction { .. } + | ChType::Geo(_) + | ChType::Geometry + | ChType::Nested(_) => Err(PyNotImplementedError::new_err( + "name-decoration aliases are expanded to their physical type before the scalar path", + )), + ChType::Dynamic { .. } => Err(PyNotImplementedError::new_err( + "Dynamic columns are built by the String insert path, not the scalar path", + )), + ChType::Json { .. } => Err(PyNotImplementedError::new_err( + "JSON columns are built by the JSON text insert path, not the scalar path", + )), + } +} + +pub(super) fn conversion_error(column: &str, row: usize, type_name: &str) -> PyErr { + PyValueError::new_err(format!( + "column {column:?} row {row} cannot be converted to {type_name}" + )) +} + +const VALUE_REPR_MAX_CHARS: usize = 100; + +fn value_repr(value: &Bound<'_, PyAny>) -> String { + let repr = value + .repr() + .and_then(|repr| repr.to_str().map(str::to_owned)) + .unwrap_or_else(|_| { + let type_name = value + .get_type() + .name() + .ok() + .and_then(|name| name.to_str().ok().map(str::to_owned)) + .unwrap_or_else(|| "value".to_string()); + format!("<{type_name}>") + }); + match repr.char_indices().nth(VALUE_REPR_MAX_CHARS) { + Some((cut, _)) => format!("{}...", &repr[..cut]), + None => repr, + } +} + +fn conversion_error_detail( + value: &Bound<'_, PyAny>, + column: &str, + row: usize, + type_name: &str, + detail: &str, +) -> PyErr { + PyValueError::new_err(format!( + "column {column:?} row {row} cannot be converted to {type_name}: {} {detail}", + value_repr(value) + )) +} + +fn float_conversion_error( + value: &Bound<'_, PyAny>, + column: &str, + row: usize, + type_name: &str, +) -> PyErr { + if value.cast::().is_ok() { + return conversion_error_detail( + value, + column, + row, + type_name, + "strings are not accepted; pass a float instead", + ); + } + conversion_error(column, row, type_name) +} + +fn integer_range_error( + value: &Bound<'_, PyAny>, + column: &str, + row: usize, + type_name: &str, +) -> PyErr { + conversion_error_detail( + value, + column, + row, + type_name, + "is outside the target range; use a value within the column's range", + ) +} + +static DECIMAL_TYPE: PyOnceLock> = PyOnceLock::new(); + +/// isinstance check against decimal.Decimal (imported once), so Decimal +/// subclasses are accepted like the python codec accepts them. +fn is_decimal(value: &Bound<'_, PyAny>) -> bool { + let py = value.py(); + DECIMAL_TYPE + .get_or_try_init(py, || { + py.import("decimal")? + .getattr("Decimal")? + .cast_into::() + .map(Bound::unbind) + .map_err(PyErr::from) + }) + .and_then(|ty| value.is_instance(ty.bind(py))) + .unwrap_or(false) +} + +fn is_numpy_float(value: &Bound<'_, PyAny>) -> bool { + let value_type = value.get_type(); + let Ok(name) = value_type.name() else { + return false; + }; + let Ok(module) = value_type.module() else { + return false; + }; + // Only widths an f64 holds exactly. longdouble/float128 (x86 80/128-bit) + // is excluded: extracting it through f64 can silently round a fraction + // away and accept a value the exactness rules should reject. + name.to_str() + .is_ok_and(|name| matches!(name, "float16" | "float32" | "float64")) + && module + .to_str() + .is_ok_and(|module| module == "numpy" || module.starts_with("numpy.")) +} + +/// Pandas ships missing values in a numeric enum column as float NaN; a +/// Nullable enum maps those to NULL. The non-nullable path keeps its code-0 +/// sentinel via `nan_as_zero`. +pub(super) fn is_enum_nan(ch_type: &ChType, value: &Bound<'_, PyAny>) -> bool { + if !matches!(ch_type, ChType::Enum8 { .. } | ChType::Enum16 { .. }) { + return false; + } + if unsafe { ffi::PyFloat_Check(value.as_ptr()) } != 0 { + // SAFETY: PyFloat_Check above guarantees PyFloat_AsDouble succeeds. + return unsafe { ffi::PyFloat_AsDouble(value.as_ptr()) }.is_nan(); + } + is_numpy_float(value) && value.extract::().is_ok_and(f64::is_nan) +} + +/// Error context for the integer coercion helpers. +struct IntConvCtx<'a> { + column: &'a str, + row: usize, + type_name: &'a str, + guidance: &'a str, +} + +impl IntConvCtx<'_> { + fn detail_err(&self, value: &Bound<'_, PyAny>, detail: &str) -> PyErr { + conversion_error_detail(value, self.column, self.row, self.type_name, detail) + } + + fn fallback_err(&self, value: &Bound<'_, PyAny>) -> PyErr { + self.detail_err(value, &format!("is not an integer; {}", self.guidance)) + } +} + +fn long_from_float<'py>( + py: Python<'py>, + value: &Bound<'py, PyAny>, + number: f64, + ctx: &IntConvCtx<'_>, + nan_as_zero: bool, +) -> PyResult> { + if number.is_nan() && nan_as_zero { + // Pandas represents a numeric enum column with missing values as float64. + // The Python Enum serializer uses zero as its non-nullable missing sentinel. + return unsafe { Bound::from_owned_ptr_or_err(py, ffi::PyLong_FromLong(0)) }; + } + if !number.is_finite() { + return Err(ctx.detail_err(value, &format!("is not finite; {}", ctx.guidance))); + } + if number.fract() != 0.0 { + return Err(ctx.detail_err( + value, + &format!("would lose fractional data; {}", ctx.guidance), + )); + } + unsafe { Bound::from_owned_ptr_or_err(py, ffi::PyNumber_Long(value.as_ptr())) } + .map_err(|_| ctx.fallback_err(value)) +} + +/// Return a Python integer while preserving exactness. Integral floats and +/// decimal.Decimal values are accepted; strings are deliberately rejected. +fn integer_object<'py>( + py: Python<'py>, + value: &Bound<'py, PyAny>, + column: &str, + row: usize, + type_name: &str, + guidance: &str, + nan_as_zero: bool, +) -> PyResult> { + let ctx = IntConvCtx { + column, + row, + type_name, + guidance, + }; + if unsafe { ffi::PyLong_Check(value.as_ptr()) } != 0 { + return Ok(value.clone()); + } + if value.cast::().is_ok() { + return Err(ctx.detail_err(value, "strings are not accepted; pass an int instead")); + } + if unsafe { ffi::PyFloat_Check(value.as_ptr()) } != 0 { + // SAFETY: PyFloat_Check above guarantees PyFloat_AsDouble succeeds. + let number = unsafe { ffi::PyFloat_AsDouble(value.as_ptr()) }; + return long_from_float(py, value, number, &ctx, nan_as_zero); + } + if is_numpy_float(value) { + let number = value + .extract::() + .map_err(|_| ctx.fallback_err(value))?; + return long_from_float(py, value, number, &ctx, nan_as_zero); + } + if is_decimal(value) { + let finite = value + .call_method0(intern!(py, "is_finite")) + .and_then(|result| result.is_truthy()) + .map_err(|_| ctx.fallback_err(value))?; + if !finite { + return Err(ctx.detail_err(value, &format!("is not finite; {guidance}"))); + } + let integral = value + .call_method0(intern!(py, "to_integral_value")) + .map_err(|_| ctx.fallback_err(value))?; + if !integral.eq(value).map_err(|_| ctx.fallback_err(value))? { + return Err(ctx.detail_err(value, &format!("would lose fractional data; {guidance}"))); + } + return unsafe { Bound::from_owned_ptr_or_err(py, ffi::PyNumber_Long(value.as_ptr())) } + .map_err(|_| ctx.fallback_err(value)); + } + unsafe { Bound::from_owned_ptr_or_err(py, ffi::PyNumber_Index(value.as_ptr())) } + .map_err(|_| ctx.fallback_err(value)) +} + +/// Convert a Python integer/index object to an owned +/// fixed-width Native representation. LowCardinality retains these bytes per +/// distinct dictionary value; ordinary columns use `wide_int_into` directly. +fn wide_int_bytes( + py: Python<'_>, + value: &Bound<'_, PyAny>, + width: usize, + signed: bool, + column: &str, + row: usize, + type_name: &str, +) -> PyResult> { + let mut bytes = vec![0u8; width]; + wide_int_into(py, value, &mut bytes, signed, column, row, type_name)?; + Ok(bytes) +} + +/// Outcome of the exact-int i64 probe. +#[derive(Clone, Copy, PartialEq, Eq)] +pub(super) enum WideFast { + /// Bytes written; conversion complete. + Done, + /// Exact int beyond i64; the slow path converts it directly. + WideInt, + /// Not an exact int; the slow path routes through index/`int()`. + NotInt, +} + +/// Try the exact-int fast path: an exact `int` within i64 writes its +/// sign-extended little-endian bytes directly (the target width is always 16 +/// or 32, so any i64 fits). A miss touches nothing and reports which slow +/// path applies; a negative value for an unsigned target is the standard +/// conversion error, matching the full conversion. +/// +/// # Safety +/// +/// Requires the GIL; `ptr` must be a valid, non-null object pointer. Never +/// executes Python code. +#[inline] +pub(super) unsafe fn wide_int_fast_into( + ptr: *mut ffi::PyObject, + bytes: &mut [u8], + signed: bool, + column: &str, + row: usize, + type_name: &str, +) -> PyResult { + debug_assert!(bytes.len() >= 8); + if ffi::PyLong_CheckExact(ptr) == 0 { + return Ok(WideFast::NotInt); + } + let mut overflow: c_int = 0; + let value = ffi::PyLong_AsLongLongAndOverflow(ptr, &mut overflow); + if overflow != 0 { + return Ok(WideFast::WideInt); + } + // An exact int should never error here; the no-error contract is + // implementation behavior, so clear defensively and fall back. + if value == -1 && !ffi::PyErr_Occurred().is_null() { + ffi::PyErr_Clear(); + return Ok(WideFast::WideInt); + } + if !signed && value < 0 { + return Err(conversion_error(column, row, type_name)); + } + bytes[..8].copy_from_slice(&value.to_le_bytes()); + bytes[8..].fill(if value < 0 { 0xff } else { 0 }); + Ok(WideFast::Done) +} + +/// Write one Python integer-compatible object into its final fixed-width +/// little-endian slice, preserving the exactness rules used by narrow integers. +pub(super) fn wide_int_into( + py: Python<'_>, + value: &Bound<'_, PyAny>, + bytes: &mut [u8], + signed: bool, + column: &str, + row: usize, + type_name: &str, +) -> PyResult<()> { + // SAFETY: GIL held via `py`; `value` is a valid object. + match unsafe { wide_int_fast_into(value.as_ptr(), bytes, signed, column, row, type_name)? } { + WideFast::Done => Ok(()), + outcome => wide_int_slow_into( + py, + value, + outcome == WideFast::WideInt, + bytes, + signed, + column, + row, + type_name, + ), + } +} + +/// Slow-path completion after a fast-probe miss; `is_exact_int` carries the +/// probe's type check so it is not repeated. An exact int beyond i64 is +/// already its own index result; anything else routes through the shared +/// exact integer coercion policy. +#[allow(clippy::too_many_arguments)] +pub(super) fn wide_int_slow_into( + py: Python<'_>, + value: &Bound<'_, PyAny>, + is_exact_int: bool, + bytes: &mut [u8], + signed: bool, + column: &str, + row: usize, + type_name: &str, +) -> PyResult<()> { + let owned; + let integer_ptr = if is_exact_int { + value.as_ptr() + } else { + owned = integer_object( + py, + value, + column, + row, + type_name, + "pass an integer value", + false, + )?; + owned.as_ptr() + }; + #[cfg(not(Py_3_13))] + let result = unsafe { + ffi::_PyLong_AsByteArray( + integer_ptr.cast(), + bytes.as_mut_ptr(), + bytes.len(), + 1, + i32::from(signed), + ) + }; + #[cfg(not(Py_3_13))] + if result < 0 { + // SAFETY: GIL held; discards the pending OverflowError. + unsafe { ffi::PyErr_Clear() }; + return Err(conversion_error(column, row, type_name)); + } + #[cfg(Py_3_13)] + { + let mut flags = ffi::Py_ASNATIVEBYTES_LITTLE_ENDIAN; + if !signed { + flags |= ffi::Py_ASNATIVEBYTES_UNSIGNED_BUFFER | ffi::Py_ASNATIVEBYTES_REJECT_NEGATIVE; + } + let required = unsafe { + ffi::PyLong_AsNativeBytes( + integer_ptr, + bytes.as_mut_ptr().cast(), + bytes.len() as ffi::Py_ssize_t, + flags, + ) + }; + if required < 0 { + // SAFETY: GIL held; discards the pending ValueError. + unsafe { ffi::PyErr_Clear() }; + return Err(conversion_error(column, row, type_name)); + } + if required as usize > bytes.len() { + return Err(conversion_error(column, row, type_name)); + } + } + Ok(()) +} + +fn bytes_value( + value: &Bound<'_, PyAny>, + column: &str, + row: usize, + type_name: &str, +) -> PyResult> { + if let Ok(s) = value.cast::() { + return Ok(s.to_str()?.as_bytes().to_vec()); + } + buffer_to_vec(value).map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} cannot be converted to {type_name} bytes" + )) + }) +} + +fn fixed_string_value( + value: &Bound<'_, PyAny>, + width: usize, + column: &str, + row: usize, +) -> PyResult> { + if let Ok(s) = value.cast::() { + let mut bytes = s.to_str()?.as_bytes().to_vec(); + if bytes.len() > width { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} UTF-8 encoded FixedString value is {} bytes, exceeding width {width}", + bytes.len() + ))); + } + bytes.resize(width, 0); + return Ok(bytes); + } + + let bytes = buffer_to_vec(value).map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} cannot be converted to FixedString bytes" + )) + })?; + if bytes.len() != width { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} FixedString binary value is {} bytes, expected {width}", + bytes.len() + ))); + } + Ok(bytes) +} +fn uuid_bytes(value: &Bound<'_, PyAny>, column: &str, row: usize) -> PyResult> { + // Precedence: str -> int (incl. subclasses) -> `int` attribute -> duck + // int (__index__) -> 16 raw bytes. Type checks come before extract so + // UUID inputs do not construct a failed extraction per value. + if let Ok(s) = value.cast::() { + return uuid_int_to_wire(parse_uuid_hex(s.to_str()?, column, row)?); + } + if unsafe { ffi::PyLong_Check(value.as_ptr()) } != 0 { + let x = value + .extract::() + .map_err(|_| conversion_error(column, row, "UUID"))?; + return uuid_int_to_wire(x); + } + if let Ok(x) = value + .getattr(intern!(value.py(), "int")) + .and_then(|i| i.extract::()) + { + return uuid_int_to_wire(x); + } + if let Ok(x) = value.extract::() { + return uuid_int_to_wire(x); + } + let bytes = buffer_to_vec(value).map_err(|_| conversion_error(column, row, "UUID"))?; + if bytes.len() != 16 { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} UUID bytes length is {}, expected 16", + bytes.len() + ))); + } + let mut out = Vec::with_capacity(16); + out.extend(bytes[..8].iter().rev()); + out.extend(bytes[8..].iter().rev()); + Ok(out) +} + +fn parse_uuid_hex(value: &str, column: &str, row: usize) -> PyResult { + let hex: String = value.chars().filter(|&ch| ch != '-').collect(); + if hex.len() != 32 { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} UUID string must contain 32 hex digits" + ))); + } + u128::from_str_radix(&hex, 16).map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} UUID string contains non-hex characters" + )) + }) +} + +fn uuid_int_to_wire(value: u128) -> PyResult> { + let mut out = Vec::with_capacity(16); + out.extend_from_slice(&((value >> 64) as u64).to_le_bytes()); + out.extend_from_slice(&(value as u64).to_le_bytes()); + Ok(out) +} + +fn ipv4_value(value: &Bound<'_, PyAny>, column: &str, row: usize) -> PyResult { + // Precedence: int (incl. subclasses) -> str -> `_ip` attribute -> duck + // int (__index__) -> `packed` attribute. Type checks come before extract + // so IPv4Address inputs do not construct a failed extraction per value. + if unsafe { ffi::PyLong_Check(value.as_ptr()) } != 0 { + return value + .extract::() + .map_err(|_| conversion_error(column, row, "IPv4")); + } + if let Ok(s) = value.cast::() { + return s + .to_str()? + .parse::() + .map(|addr| u32::from_be_bytes(addr.octets())) + .map_err(|_| conversion_error(column, row, "IPv4")); + } + if let Ok(ip) = value + .getattr(intern!(value.py(), "_ip")) + .and_then(|o| o.extract::()) + { + return Ok(ip); + } + if let Ok(v) = value.extract::() { + return Ok(v); + } + if let Ok(packed) = value.getattr("packed").and_then(|o| buffer_to_vec(&o)) { + if packed.len() == 4 { + return Ok(u32::from_be_bytes([ + packed[0], packed[1], packed[2], packed[3], + ])); + } + } + Err(conversion_error(column, row, "IPv4")) +} + +fn ipv6_bytes(value: &Bound<'_, PyAny>, column: &str, row: usize) -> PyResult> { + if let Ok(s) = value.cast::() { + return ip_string_to_ipv6(s.to_str()?, column, row); + } + if let Ok(v) = value.extract::() { + return Ok(v.to_be_bytes().to_vec()); + } + if let Ok(packed) = value.getattr("packed").and_then(|o| buffer_to_vec(&o)) { + return packed_to_ipv6(packed, true, column, row); + } + let bytes = buffer_to_vec(value).map_err(|_| conversion_error(column, row, "IPv6"))?; + packed_to_ipv6(bytes, false, column, row) +} + +fn ip_string_to_ipv6(value: &str, column: &str, row: usize) -> PyResult> { + match value.parse::() { + Ok(IpAddr::V6(addr)) => Ok(addr.octets().to_vec()), + Ok(IpAddr::V4(addr)) => { + let mut out = Vec::with_capacity(16); + out.extend_from_slice(&IPV4_V6_PREFIX); + out.extend_from_slice(&addr.octets()); + Ok(out) + } + Err(_) => Err(conversion_error(column, row, "IPv6")), + } +} + +fn packed_to_ipv6(bytes: Vec, allow_ipv4: bool, column: &str, row: usize) -> PyResult> { + if bytes.len() == 16 { + return Ok(bytes); + } + if allow_ipv4 && bytes.len() == 4 { + let mut out = Vec::with_capacity(16); + out.extend_from_slice(&IPV4_V6_PREFIX); + out.extend_from_slice(&bytes); + return Ok(out); + } + Err(PyValueError::new_err(format!( + "column {column:?} row {row} IPv6 bytes length is {}, expected 16", + bytes.len() + ))) +} + +fn enum8_value( + value: &Bound<'_, PyAny>, + variants: &[(String, i8)], + column: &str, + row: usize, +) -> PyResult { + if let Ok(name) = value.cast::() { + let name = name.to_str()?; + return variants + .iter() + .find_map(|(variant, code)| (variant == name).then_some(*code)) + .ok_or_else(|| { + PyValueError::new_err(format!( + "column {column:?} row {row} Enum8 label {name:?} is not defined" + )) + }); + } + let integer = integer_object( + value.py(), + value, + column, + row, + "Enum8", + "pass a valid enum label or integral code", + true, + )?; + integer + .extract::() + .map_err(|_| integer_range_error(value, column, row, "Enum8")) +} + +fn enum16_value( + value: &Bound<'_, PyAny>, + variants: &[(String, i16)], + column: &str, + row: usize, +) -> PyResult { + if let Ok(name) = value.cast::() { + let name = name.to_str()?; + return variants + .iter() + .find_map(|(variant, code)| (variant == name).then_some(*code)) + .ok_or_else(|| { + PyValueError::new_err(format!( + "column {column:?} row {row} Enum16 label {name:?} is not defined" + )) + }); + } + let integer = integer_object( + value.py(), + value, + column, + row, + "Enum16", + "pass a valid enum label or integral code", + true, + )?; + integer + .extract::() + .map_err(|_| integer_range_error(value, column, row, "Enum16")) +} + +pub(super) fn decimal_width(precision: u8) -> PyResult { + match precision { + 1..=9 => Ok(4), + 10..=18 => Ok(8), + 19..=38 => Ok(16), + 39..=76 => Ok(32), + _ => Err(PyValueError::new_err(format!( + "Decimal precision {precision} is outside 1..=76" + ))), + } +} + +pub(super) fn decimal_text(value: &Bound<'_, PyAny>, column: &str, row: usize) -> PyResult { + value + .str() + .and_then(|s| s.to_str().map(str::to_owned)) + .map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} Decimal value cannot be stringified" + )) + }) +} + +/// Decimal text-to-wire failure classes; `decimal_err` renders the exact +/// user-facing message for each. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum DecimalError { + Invalid, + UnsupportedExponent, + ExceedsPrecision, + OverflowsWidth, + DoesNotFit, +} + +pub(super) fn decimal_err( + err: DecimalError, + text: &str, + width: usize, + precision: u8, + column: &str, + row: usize, +) -> PyErr { + match err { + DecimalError::Invalid => conversion_error(column, row, "Decimal"), + DecimalError::UnsupportedExponent => PyValueError::new_err(format!( + "column {column:?} row {row} Decimal value {text:?} has an unsupported exponent" + )), + DecimalError::ExceedsPrecision => PyValueError::new_err(format!( + "column {column:?} row {row} Decimal value {text:?} exceeds precision {precision}" + )), + DecimalError::OverflowsWidth => PyValueError::new_err(format!( + "column {column:?} row {row} Decimal value overflows target width" + )), + DecimalError::DoesNotFit => PyValueError::new_err(format!( + "column {column:?} row {row} Decimal value {text:?} does not fit in {width} bytes" + )), + } +} + +fn decimal_to_le_bytes( + text: &str, + width: usize, + precision: u8, + scale: u8, + column: &str, + row: usize, +) -> PyResult> { + let mut bytes = vec![0u8; width]; + decimal_wire_into(text, &mut bytes, precision, scale) + .map_err(|err| decimal_err(err, text, width, precision, column, row))?; + Ok(bytes) +} + +const POW10: [u64; 20] = [ + 1, + 10, + 100, + 1_000, + 10_000, + 100_000, + 1_000_000, + 10_000_000, + 100_000_000, + 1_000_000_000, + 10_000_000_000, + 100_000_000_000, + 1_000_000_000_000, + 10_000_000_000_000, + 100_000_000_000_000, + 1_000_000_000_000_000, + 10_000_000_000_000_000, + 100_000_000_000_000_000, + 1_000_000_000_000_000_000, + 10_000_000_000_000_000_000, +]; + +/// limbs = limbs * pow + add over little-endian u64 limbs; carry out of the +/// limb window is an overflow. +fn limbs_mul_add(limbs: &mut [u64], pow: u64, add: u64) -> Result<(), DecimalError> { + let mut carry = u128::from(add); + for limb in limbs.iter_mut() { + let next = u128::from(*limb) * u128::from(pow) + carry; + *limb = next as u64; + carry = next >> 64; + } + if carry != 0 { + return Err(DecimalError::OverflowsWidth); + } + Ok(()) +} + +/// Parse decimal text and write the scaled two's-complement little-endian +/// integer into `out` (4, 8, 16, or 32 bytes). Fractional digits beyond +/// `scale` are truncated toward zero before accumulating. Runs no Python. +pub(super) fn decimal_wire_into( + text: &str, + out: &mut [u8], + precision: u8, + scale: u8, +) -> Result<(), DecimalError> { + let width = out.len(); + let bytes = text.trim().as_bytes(); + if bytes.is_empty() { + return Err(DecimalError::Invalid); + } + + let mut pos = 0; + let negative = match bytes[0] { + b'-' => { + pos += 1; + true + } + b'+' => { + pos += 1; + false + } + _ => false, + }; + + let int_start = pos; + while pos < bytes.len() && bytes[pos].is_ascii_digit() { + pos += 1; + } + let int_run = &bytes[int_start..pos]; + let mut frac_run: &[u8] = &[]; + if pos < bytes.len() && bytes[pos] == b'.' { + pos += 1; + let frac_start = pos; + while pos < bytes.len() && bytes[pos].is_ascii_digit() { + pos += 1; + } + frac_run = &bytes[frac_start..pos]; + } + if int_run.is_empty() && frac_run.is_empty() { + return Err(DecimalError::Invalid); + } + + let mut exp10 = 0i32; + if pos < bytes.len() { + if !matches!(bytes[pos], b'e' | b'E') { + return Err(DecimalError::Invalid); + } + pos += 1; + let mut exp_sign = 1i32; + match bytes.get(pos) { + Some(b'-') => { + exp_sign = -1; + pos += 1; + } + Some(b'+') => { + pos += 1; + } + _ => {} + } + if pos == bytes.len() { + return Err(DecimalError::Invalid); + } + let mut exp = 0i32; + for &byte in &bytes[pos..] { + if !byte.is_ascii_digit() { + return Err(DecimalError::Invalid); + } + exp = exp + .checked_mul(10) + .and_then(|v| v.checked_add(i32::from(byte - b'0'))) + .ok_or(DecimalError::Invalid)?; + } + exp10 = exp_sign * exp; + } + + let frac_count = i32::try_from(frac_run.len()).map_err(|_| DecimalError::Invalid)?; + let exponent = exp10.checked_sub(frac_count).ok_or(DecimalError::Invalid)?; + let shift = exponent + .checked_add(i32::from(scale)) + .ok_or(DecimalError::UnsupportedExponent)?; + + // Significant digits: strip leading zeros across the combined runs. + let mut sig_int = int_run; + while let Some((b'0', rest)) = sig_int.split_first() { + sig_int = rest; + } + let mut sig_frac = frac_run; + if sig_int.is_empty() { + while let Some((b'0', rest)) = sig_frac.split_first() { + sig_frac = rest; + } + } + let sig_len = sig_int.len() + sig_frac.len(); + + // A non-negative shift pads with zeros; a negative shift drops fractional + // digits beyond the scale (truncation toward zero). + let precision = usize::from(precision); + let (keep, zeros) = if shift >= 0 { + let shift = usize::try_from(shift).map_err(|_| DecimalError::Invalid)?; + let new_len = sig_len + .checked_add(shift) + .ok_or(DecimalError::UnsupportedExponent)?; + if sig_len > 0 && new_len > precision { + return Err(DecimalError::ExceedsPrecision); + } + if new_len > width * 4 { + return Err(DecimalError::DoesNotFit); + } + (sig_len, shift) + } else { + let remove = usize::try_from(shift.unsigned_abs()).map_err(|_| DecimalError::Invalid)?; + (sig_len.saturating_sub(remove), 0) + }; + if keep + zeros > precision { + return Err(DecimalError::ExceedsPrecision); + } + + // Accumulate the kept digits in chunks of up to 19 (the largest power of + // ten a u64 holds), then fold the zero padding in as pure multiplies. + let limb_count = width.div_ceil(8).max(1); + let mut limbs = [0u64; 4]; + let limbs = &mut limbs[..limb_count]; + let mut chunk = 0u64; + let mut chunk_len = 0usize; + for &digit in sig_int.iter().chain(sig_frac).take(keep) { + chunk = chunk * 10 + u64::from(digit - b'0'); + chunk_len += 1; + if chunk_len == 19 { + limbs_mul_add(limbs, POW10[19], chunk)?; + chunk = 0; + chunk_len = 0; + } + } + if chunk_len > 0 { + limbs_mul_add(limbs, POW10[chunk_len], chunk)?; + } + let mut zeros = zeros; + while zeros > 0 { + let step = zeros.min(19); + limbs_mul_add(limbs, POW10[step], 0)?; + zeros -= step; + } + if width < 8 && limbs[0] >> (8 * width) != 0 { + return Err(DecimalError::OverflowsWidth); + } + + for (index, byte) in out.iter_mut().enumerate() { + *byte = (limbs[index / 8] >> ((index % 8) * 8)) as u8; + } + if !fits_signed_width(out, negative) { + return Err(DecimalError::DoesNotFit); + } + if negative && out.iter().any(|&b| b != 0) { + twos_complement_in_place(out); + } + Ok(()) +} + +fn fits_signed_width(bytes: &[u8], negative: bool) -> bool { + if bytes.is_empty() { + return true; + } + let last = bytes.len() - 1; + if negative { + bytes[last] < 0x80 || (bytes[last] == 0x80 && bytes[..last].iter().all(|&b| b == 0)) + } else { + bytes[last] < 0x80 + } +} + +fn twos_complement_in_place(bytes: &mut [u8]) { + for byte in bytes.iter_mut() { + *byte = !*byte; + } + let mut carry = 1u16; + for byte in bytes { + let next = u16::from(*byte) + carry; + *byte = next as u8; + carry = next >> 8; + if carry == 0 { + break; + } + } +} + +#[cfg(test)] +mod tests { + use super::{decimal_wire_into, DecimalError}; + + fn wire(text: &str, width: usize, precision: u8, scale: u8) -> Result, DecimalError> { + let mut out = vec![0u8; width]; + decimal_wire_into(text, &mut out, precision, scale)?; + Ok(out) + } + + fn le(value: i128, width: usize) -> Vec { + value.to_le_bytes()[..width].to_vec() + } + + #[test] + fn scaled_values_per_width() { + assert_eq!(wire("1.5", 4, 9, 3), Ok(le(1_500, 4))); + assert_eq!(wire("-1.5", 4, 9, 3), Ok(le(-1_500, 4))); + assert_eq!(wire("+1.5", 4, 9, 3), Ok(le(1_500, 4))); + assert_eq!(wire(" 12.5 ", 4, 9, 3), Ok(le(12_500, 4))); + assert_eq!(wire("7", 8, 18, 6), Ok(le(7_000_000, 8))); + assert_eq!( + wire("123456789.123456", 8, 18, 6), + Ok(le(123_456_789_123_456, 8)) + ); + assert_eq!( + wire("-123456789012345678901234.567890123456", 16, 38, 12), + Ok(le(-123_456_789_012_345_678_901_234_567_890_123_456, 16)) + ); + } + + #[test] + fn boundary_values() { + // +-(10^precision - 1) scaled + assert_eq!(wire("999999.999", 4, 9, 3), Ok(le(999_999_999, 4))); + assert_eq!(wire("-999999.999", 4, 9, 3), Ok(le(-999_999_999, 4))); + assert_eq!( + wire("999999999999.999999", 8, 18, 6), + Ok(le(999_999_999_999_999_999, 8)) + ); + let top38 = 10i128.pow(38) - 1; + assert_eq!( + wire("99999999999999999999999999.999999999999", 16, 38, 12), + Ok(le(top38, 16)) + ); + // the exact -2^(8w-1) two's-complement boundary fits; its positive + // counterpart does not (width deliberately narrower than precision) + assert_eq!( + wire("-2147483648", 4, 10, 0), + Ok(vec![0x00, 0x00, 0x00, 0x80]) + ); + assert_eq!(wire("2147483648", 4, 10, 0), Err(DecimalError::DoesNotFit)); + } + + #[test] + fn zero_forms() { + for text in ["0", "-0", "0.0", "-0.000", "0E-10", "0E-200"] { + assert_eq!(wire(text, 4, 9, 3), Ok(vec![0; 4]), "{text}"); + } + assert_eq!(wire("0E+10", 8, 18, 6), Ok(vec![0; 8])); + // zero padded beyond the precision or width keeps the legacy errors + assert_eq!(wire("0E+10", 4, 9, 3), Err(DecimalError::ExceedsPrecision)); + assert_eq!(wire("0E+200", 4, 9, 3), Err(DecimalError::DoesNotFit)); + } + + #[test] + fn excess_fraction_truncates_toward_zero() { + assert_eq!(wire("1.0005", 4, 9, 3), Ok(le(1_000, 4))); + assert_eq!(wire("-1.0005", 4, 9, 3), Ok(le(-1_000, 4))); + assert_eq!( + wire("123456.789999999999999999999", 4, 9, 3), + Ok(le(123_456_789, 4)) + ); + assert_eq!(wire("0.0001", 4, 9, 3), Ok(vec![0; 4])); + assert_eq!(wire("-0.0001", 4, 9, 3), Ok(vec![0; 4])); + // truncation happens before the precision check + assert_eq!( + wire("1234567890.1234", 4, 9, 3), + Err(DecimalError::ExceedsPrecision) + ); + } + + #[test] + fn scientific_notation() { + assert_eq!(wire("1E+2", 4, 9, 3), Ok(le(100_000, 4))); + assert_eq!(wire("1E-7", 4, 9, 3), Ok(vec![0; 4])); + assert_eq!(wire("-3E+1", 4, 9, 3), Ok(le(-30_000, 4))); + assert_eq!(wire("5e2", 4, 9, 3), Ok(le(500_000, 4))); + assert_eq!(wire("1.25E+3", 4, 9, 3), Ok(le(1_250_000, 4))); + } + + #[test] + fn scale_zero_and_max_scale() { + assert_eq!(wire("123456789", 4, 9, 0), Ok(le(123_456_789, 4))); + assert_eq!(wire("1.9", 4, 9, 0), Ok(le(1, 4))); + assert_eq!(wire("-1.9", 4, 9, 0), Ok(le(-1, 4))); + assert_eq!(wire("0.123456789", 4, 9, 9), Ok(le(123_456_789, 4))); + assert_eq!(wire("1.0", 4, 9, 9), Err(DecimalError::ExceedsPrecision)); + } + + #[test] + fn wide_widths_match_per_digit_reference() { + // Reference accumulator: the per-digit byte-carry algorithm this + // implementation replaced. + fn reference(digits: &str, negative: bool, width: usize) -> Vec { + let mut mag = vec![0u8; width]; + for digit in digits.bytes() { + let mut carry = u16::from(digit - b'0'); + for byte in mag.iter_mut() { + let next = u16::from(*byte) * 10 + carry; + *byte = next as u8; + carry = next >> 8; + } + assert_eq!(carry, 0); + } + if negative && mag.iter().any(|&b| b != 0) { + super::twos_complement_in_place(&mut mag); + } + mag + } + + for digit_count in 1..=76usize { + let digits: String = (0..digit_count) + .map(|i| char::from(b'1' + (i % 9) as u8)) + .collect(); + let (int_part, frac_part) = digits.split_at(digit_count.min(46)); + let text = if frac_part.is_empty() { + format!("{int_part}.0") + } else { + format!("{int_part}.{frac_part}") + }; + let scaled: String = digits + .chars() + .chain(std::iter::repeat('0')) + .take(int_part.len() + 30) + .collect(); + for negative in [false, true] { + let signed = if negative { + format!("-{text}") + } else { + text.clone() + }; + assert_eq!( + wire(&signed, 32, 76, 30), + Ok(reference(&scaled, negative, 32)), + "{signed}" + ); + } + } + } + + #[test] + fn invalid_text() { + for text in [ + "", + " ", + "abc", + "1.2.3", + "1e", + "1e+", + "--5", + "+", + ".", + "1x", + "1e5x", + "nan", + "inf", + "1e999999999999", + ] { + assert_eq!(wire(text, 4, 9, 3), Err(DecimalError::Invalid), "{text:?}"); + } + } + + #[test] + fn exponent_and_precision_errors() { + assert_eq!( + wire("1e2147483647", 4, 9, 3), + Err(DecimalError::UnsupportedExponent) + ); + assert_eq!( + wire("1234567890", 4, 9, 3), + Err(DecimalError::ExceedsPrecision) + ); + assert_eq!( + wire("1E+70", 32, 76, 30), + Err(DecimalError::ExceedsPrecision) + ); + // magnitude overflow of a width narrower than the claimed precision + assert_eq!( + wire("9999999999", 4, 76, 0), + Err(DecimalError::OverflowsWidth) + ); + } +} diff --git a/rust/ch-core-py/src/insert/special.rs b/rust/ch-core-py/src/insert/special.rs new file mode 100644 index 00000000..39341821 --- /dev/null +++ b/rust/ch-core-py/src/insert/special.rs @@ -0,0 +1,932 @@ +use super::*; + +/// Build Nothing from a Python column without constructing per-row scalar +/// values. Plain Nothing ignores every placeholder, matching the Python +/// codec. Nullable(Nothing) additionally retains which placeholders were +/// Python None so the core can emit the structural null map before the +/// canonical Nothing marker bytes. +pub(super) fn build_nothing_column( + name: &str, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult { + let column_values = ColumnValues::new(values, name)?; + check_row_count(name, &column_values, row_count)?; + + if let Ok(list) = values.cast_exact::() { + return Ok(nothing_column_from_seq(&ListSeq(list), row_count, nullable)); + } + if let Ok(tuple) = values.cast_exact::() { + return Ok(nothing_column_from_seq( + &TupleSeq(tuple), + row_count, + nullable, + )); + } + + let validity = if nullable { + let mut null_map = Vec::with_capacity(row_count); + for row in 0..row_count { + null_map.push(u8::from(column_values.get_item(row)?.is_none())); + } + Some(Bitmap::from_ch_null_map(&null_map)) + } else { + None + }; + Ok(nothing_column(row_count, validity)) +} + +pub(super) fn nothing_column_from_seq( + seq: &S, + row_count: usize, + nullable: bool, +) -> Column { + // The unsafe seq reads below are in bounds only under this equality. + assert_eq!(seq.size(), row_count); + let validity = nullable.then(|| { + let mut null_map = Vec::with_capacity(row_count); + for row in 0..row_count { + // SAFETY: row < row_count == seq.size(), the sequence is borrowed + // while the GIL is held, and comparing to Py_None runs no Python. + let value = unsafe { seq.get(row) }; + null_map.push(u8::from(value == unsafe { ffi::Py_None() })); + } + Bitmap::from_ch_null_map(&null_map) + }); + nothing_column(row_count, validity) +} + +fn nothing_column(row_count: usize, validity: Option) -> Column { + Column::Nothing(match validity { + Some(validity) => NothingColumn::new_nullable(row_count, validity), + None => NothingColumn::new(row_count), + }) +} + +/// Build one AggregateFunction state column from exact serialized state bytes. +/// State framing and validation are function-specific, so the binding only +/// recovers Python buffer boundaries; the core encoder validates every slice +/// against the codec registered for the declared ChType. +pub(super) fn build_aggregate_state_column( + py: Python<'_>, + name: &str, + ch_type: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + let column_values = ColumnValues::new(values, name)?; + check_row_count(name, &column_values, row_count)?; + if let Ok(list) = values.cast_exact::() { + return aggregate_state_column_from_seq(py, name, ch_type, &ListSeq(list), row_count); + } + if let Ok(tuple) = values.cast_exact::() { + return aggregate_state_column_from_seq(py, name, ch_type, &TupleSeq(tuple), row_count); + } + aggregate_state_column_from_rows(name, ch_type, &column_values, row_count) +} + +fn aggregate_state_column_from_rows<'py, R: RowAccess<'py>>( + name: &str, + ch_type: &ChType, + rows: &R, + row_count: usize, +) -> PyResult { + let mut offsets = Vec::with_capacity(row_count + 1); + let mut data = Vec::new(); + offsets.push(0); + for row in 0..row_count { + let value = rows.value(row)?; + append_aggregate_state(&value, name, ch_type, row, &mut offsets, &mut data)?; + } + Ok(Column::AggregateState(AggregateStateColumn::new( + offsets, data, + ))) +} + +pub(super) fn aggregate_state_column_from_seq( + py: Python<'_>, + name: &str, + ch_type: &ChType, + seq: &S, + row_count: usize, +) -> PyResult { + let mut offsets = Vec::with_capacity(row_count + 1); + let mut data = Vec::with_capacity(aggregate_state_size_hint(py, seq, row_count)); + offsets.push(0); + for row in 0..row_count { + // SAFETY: row is in bounds; the size is checked against row_count up + // front and revalidated after any row that ran Python. + // `from_borrowed_ptr` takes a strong reference, so container mutation + // during generic buffer conversion cannot invalidate `value`. + let value = unsafe { Bound::from_borrowed_ptr(py, seq.get(row)) }; + let ran_python = + append_aggregate_state(&value, name, ch_type, row, &mut offsets, &mut data)?; + if ran_python { + // Dropping the strong reference can run __del__, which can resize + // the container, so it must precede the size revalidation. + drop(value); + check_not_resized(seq, name, row_count)?; + } + } + Ok(Column::AggregateState(AggregateStateColumn::new( + offsets, data, + ))) +} + +/// Exact data size when every row is exact bytes, read via `Py_SIZE` without +/// running Python; 0 (no reservation) otherwise. +fn aggregate_state_size_hint(_py: Python<'_>, seq: &S, row_count: usize) -> usize { + let mut total = 0usize; + for row in 0..row_count { + // SAFETY: the GIL is held, row is in bounds, and no Python runs in + // this loop, so the borrowed pointer stays valid. + let ptr = unsafe { seq.get(row) }; + if unsafe { ffi::PyBytes_CheckExact(ptr) } == 0 { + return 0; + } + total = total.saturating_add(unsafe { ffi::Py_SIZE(ptr) } as usize); + } + total +} + +fn agg_offset_overflow(name: &str) -> PyErr { + PyValueError::new_err(format!( + "column {name:?} AggregateFunction state data exceeds i64 offset capacity" + )) +} + +fn agg_state_convert_err(py: Python<'_>, name: &str, row: usize, err: PyErr) -> PyErr { + let wrapped = PyValueError::new_err(format!( + "column {name:?} row {row} cannot be converted to AggregateFunction state bytes" + )); + wrapped.set_cause(py, Some(err)); + wrapped +} + +/// Append one Python bytes-like state directly into the column data run. +/// Returns true when generic buffer conversion may have executed Python. +fn append_aggregate_state( + value: &Bound<'_, PyAny>, + name: &str, + ch_type: &ChType, + row: usize, + offsets: &mut Vec, + data: &mut Vec, +) -> PyResult { + if value.is_none() { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + } + let ran_python = if let Ok(bytes) = value.cast::() { + data.extend_from_slice(bytes.as_bytes()); + false + } else if let Ok(bytes) = value.cast::() { + // SAFETY: the GIL is held and extend_from_slice copies the complete + // buffer before any Python or PyO3 API can run and invalidate it. + data.extend_from_slice(unsafe { bytes.as_bytes() }); + false + } else { + let buffer = PyBuffer::::get(value) + .map_err(|err| agg_state_convert_err(value.py(), name, row, err))?; + let start = data.len(); + let end = start + .checked_add(buffer.item_count()) + .ok_or_else(|| agg_offset_overflow(name))?; + data.resize(end, 0); + if let Err(err) = buffer.copy_to_slice(value.py(), &mut data[start..end]) { + data.truncate(start); + return Err(agg_state_convert_err(value.py(), name, row, err)); + } + true + }; + offsets.push(i64::try_from(data.len()).map_err(|_| agg_offset_overflow(name))?); + Ok(ran_python) +} + +pub(super) fn build_low_cardinality_column( + py: Python<'_>, + name: &str, + inner: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + // Resolve the physical dict value type: strip the SAF chain, unwrap an + // optional Nullable, strip again. So LowCardinality(SAF(anyLast, String)) and + // LowCardinality(SAF(anyLast, Nullable(String))) reach the String path. + let (nullable, value_type) = low_cardinality_dict_value_type(inner); + + if !is_low_cardinality_inner(value_type) { + return Err(PyNotImplementedError::new_err(format!( + "unsupported LowCardinality inner type {value_type} for column {name:?}" + ))); + } + + match ColumnSource::resolve(values, name, row_count)? { + ColumnSource::List(list) => lc_column_from_seq( + py, + name, + value_type, + &ListSeq(&list), + values, + row_count, + nullable, + ), + ColumnSource::Tuple(tuple) => lc_column_from_seq( + py, + name, + value_type, + &TupleSeq(&tuple), + values, + row_count, + nullable, + ), + ColumnSource::ObjectArray(seq) => { + if wide_int_layout(value_type).is_some() { + return lc_wide_column(py, name, value_type, &seq, row_count, nullable); + } + if matches!(value_type, ChType::String) { + return lc_string_seq(py, name, value_type, &seq, row_count, nullable); + } + lc_scalar_column(py, name, value_type, &seq, row_count, nullable) + } + ColumnSource::Values(column_values) => { + if wide_int_layout(value_type).is_some() { + return lc_wide_column(py, name, value_type, &column_values, row_count, nullable); + } + lc_scalar_column(py, name, value_type, &column_values, row_count, nullable) + } + } +} + +/// LowCardinality over an exact list or tuple: the String dictionary fast +/// path reads borrowed pointers; other value types keep safe indexed reads. +#[allow(clippy::too_many_arguments)] +fn lc_column_from_seq( + py: Python<'_>, + name: &str, + value_type: &ChType, + seq: &S, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult { + if matches!(value_type, ChType::String) { + return lc_string_seq(py, name, value_type, seq, row_count, nullable); + } + let column_values = ColumnValues::new(values, name)?; + if wide_int_layout(value_type).is_some() { + return lc_wide_column(py, name, value_type, &column_values, row_count, nullable); + } + lc_scalar_column(py, name, value_type, &column_values, row_count, nullable) +} + +pub(super) fn lc_scalar_column<'py, R: RowAccess<'py>>( + py: Python<'py>, + name: &str, + value_type: &ChType, + rows: &R, + row_count: usize, + nullable: bool, +) -> PyResult { + let mut indices = Vec::with_capacity(row_count); + let mut dict_values = Vec::new(); + let mut slots = HashMap::::new(); + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + // Time64 is rejected by is_low_cardinality_inner, so only Time probes. + let mut time_probe = match value_type { + ChType::Time => TimeScalarProbe::new(value_type), + _ => None, + }; + + if nullable && row_count > 0 { + dict_values.push(default_scalar(value_type)?); + } + + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + if !nullable { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but LowCardinality({value_type}) is not nullable" + ))); + } + indices.push(0); + if let Some(nulls) = &mut null_map { + nulls.push(1); + } + continue; + } + + let time_hit = match time_probe.as_mut() { + Some(probe) => probe.probe(&value, name, row)?, + None => None, + }; + if let Some(TimeProbe::Nat) = time_hit { + if !nullable { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is NaT but LowCardinality({value_type}) is not nullable" + ))); + } + indices.push(0); + if let Some(nulls) = &mut null_map { + nulls.push(1); + } + continue; + } + + if let Some(nulls) = &mut null_map { + nulls.push(0); + } + let scalar = match time_hit { + Some(TimeProbe::Ticks(ticks)) => time_ticks_scalar(value_type, ticks, name, row)?, + _ => convert_scalar(py, value_type, &value, name, row)?, + }; + let key = scalar.key(); + let slot = if let Some(slot) = slots.get(&key) { + *slot + } else { + let slot = i32::try_from(dict_values.len()).map_err(|_| { + PyValueError::new_err(format!( + "column {name:?} LowCardinality dictionary exceeds i32 index capacity" + )) + })?; + dict_values.push(scalar); + slots.insert(key, slot); + slot + }; + indices.push(slot); + } + + let dict_column = column_from_scalars(value_type, dict_values, None)?; + match null_map { + Some(nulls) => Ok(Column::Dictionary(DictionaryColumn::new_nullable( + indices, + dict_column, + Bitmap::from_ch_null_map(&nulls), + ))), + None => Ok(Column::Dictionary(DictionaryColumn::new( + indices, + dict_column, + ))), + } +} + +/// Build LowCardinality over a wide integer without a heap allocation per +/// input cell. Conversion writes into a fixed stack buffer; only distinct +/// dictionary entries are retained, and the final dictionary data is copied +/// once into the core's contiguous fixed-binary representation. +pub(super) fn lc_wide_column<'py, R: RowAccess<'py>>( + py: Python<'py>, + name: &str, + value_type: &ChType, + rows: &R, + row_count: usize, + nullable: bool, +) -> PyResult { + let (width, signed, type_name) = wide_int_layout(value_type) + .ok_or_else(|| PyValueError::new_err("internal wide integer type mismatch"))?; + let mut indices = Vec::with_capacity(row_count); + let mut dict_values = Vec::<[u8; 32]>::new(); + let mut slots = HashMap::<[u8; 32], i32>::with_capacity(row_count.min(1024)); + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + + if nullable && row_count > 0 { + dict_values.push([0u8; 32]); + } + + for row in 0..row_count { + let value = rows.value(row)?; + if value.is_none() { + if !nullable { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but LowCardinality({value_type}) is not nullable" + ))); + } + indices.push(0); + if let Some(nulls) = &mut null_map { + nulls.push(1); + } + continue; + } + if let Some(nulls) = &mut null_map { + nulls.push(0); + } + + let mut bytes = [0u8; 32]; + wide_int_into( + py, + &value, + &mut bytes[..width], + signed, + name, + row, + type_name, + )?; + let slot = match slots.entry(bytes) { + Entry::Occupied(slot) => *slot.get(), + Entry::Vacant(vacant) => { + let slot = i32::try_from(dict_values.len()).map_err(|_| { + PyValueError::new_err(format!( + "column {name:?} LowCardinality dictionary exceeds i32 index capacity" + )) + })?; + dict_values.push(bytes); + *vacant.insert(slot) + } + }; + indices.push(slot); + } + + let byte_len = width.checked_mul(dict_values.len()).ok_or_else(|| { + PyValueError::new_err(format!( + "column {name:?} wide integer dictionary byte size exceeds usize capacity" + )) + })?; + let mut data = Vec::with_capacity(byte_len); + for value in dict_values { + data.extend_from_slice(&value[..width]); + } + let dict_column = finish_wide_int_column(value_type, data, None)?; + match null_map { + Some(nulls) => Ok(Column::Dictionary(DictionaryColumn::new_nullable( + indices, + dict_column, + Bitmap::from_ch_null_map(&nulls), + ))), + None => Ok(Column::Dictionary(DictionaryColumn::new( + indices, + dict_column, + ))), + } +} + +/// Build a LowCardinality(String) dictionary column with two cache levels: +/// repeated str objects hit a pointer-identity map (no content read at all), +/// and new exact-str objects read their UTF-8 once for a content-keyed map, +/// allocating only when the content is genuinely new to the dictionary. +/// Non-str values fall back to `convert_scalar` for identical accepted-type +/// and error semantics; that fallback can run arbitrary Python (exotic buffer +/// types), so it holds a strong reference and revalidates the container size. +/// Pointer identity is only meaningful while no Python code has run since the +/// key was cached: a list can drop the last reference to an already-scanned +/// str mid-call, and the allocator can hand its address to a new str, so +/// every fallback return clears the pointer cache. Entries cached after the +/// clear are valid until the next fallback. A tuple keeps every item alive +/// for the whole call, so its cache never needs clearing. +pub(super) fn lc_string_seq( + py: Python<'_>, + name: &str, + value_type: &ChType, + seq: &S, + row_count: usize, + nullable: bool, +) -> PyResult { + let mut indices = Vec::with_capacity(row_count); + let mut dict_values: Vec = Vec::new(); + let mut ptr_slots: HashMap> = + HashMap::default(); + let mut content_slots: HashMap, i32> = HashMap::new(); + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + + if nullable && row_count > 0 { + dict_values.push(default_scalar(value_type)?); + } + + let dict_slot = |dict_values: &Vec| { + i32::try_from(dict_values.len()).map_err(|_| { + PyValueError::new_err(format!( + "column {name:?} LowCardinality dictionary exceeds i32 index capacity" + )) + }) + }; + + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every fallback revalidates; the borrowed pointer is consumed before + // any Python code can run. + let ptr = unsafe { seq.get(row) }; + if ptr == unsafe { ffi::Py_None() } { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but LowCardinality({value_type}) is not nullable" + ))); + }; + indices.push(0); + null_map.push(1); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + if let Some(&slot) = ptr_slots.get(&(ptr as usize)) { + indices.push(slot); + continue; + } + let slot = if unsafe { ffi::PyUnicode_CheckExact(ptr) } != 0 { + // SAFETY: ptr is a valid borrowed reference, verified an exact + // str; reading its UTF-8 runs no Python code. + let obj = + unsafe { Bound::from_borrowed_ptr(py, ptr).cast_into_unchecked::() }; + let bytes = obj.to_str()?.as_bytes(); + let slot = match content_slots.get(bytes) { + Some(&slot) => slot, + None => { + let slot = dict_slot(&dict_values)?; + content_slots.insert(bytes.to_vec(), slot); + dict_values.push(Scalar::Bytes(bytes.to_vec())); + slot + } + }; + // The container still holds this item (no Python ran since the + // borrowed read), so its address is stable and unique among the + // column's live values. + if ptr_slots.len() < PTR_CACHE_CAP { + ptr_slots.insert(ptr as usize, slot); + } + slot + } else { + // SAFETY: ptr is valid here; the strong reference keeps the item + // alive across any Python code the fallback runs. + let obj = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + let scalar = convert_scalar(py, value_type, &obj, name, row)?; + if S::MUTABLE { + // The fallback may have run Python code: a same-size item + // replacement can free a cached str whose address the + // allocator may reuse, so drop all pointer-identity entries. + check_not_resized(seq, name, row_count)?; + ptr_slots.clear(); + } + let Scalar::Bytes(bytes) = scalar else { + return Err(PyValueError::new_err("internal scalar type mismatch")); + }; + match content_slots.get(bytes.as_slice()) { + Some(&slot) => slot, + None => { + let slot = dict_slot(&dict_values)?; + content_slots.insert(bytes.clone(), slot); + dict_values.push(Scalar::Bytes(bytes)); + slot + } + } + }; + indices.push(slot); + } + + let dict_column = column_from_scalars(value_type, dict_values, None)?; + Ok(match null_map { + Some(nulls) => Column::Dictionary(DictionaryColumn::new_nullable( + indices, + dict_column, + Bitmap::from_ch_null_map(&nulls), + )), + None => Column::Dictionary(DictionaryColumn::new(indices, dict_column)), + }) +} + +/// Build a UUID column: exact `uuid.UUID` items read the `int` attribute and +/// write the 16 wire bytes (hi u64 LE, then lo u64 LE) straight into the +/// column buffer; anything else (str, int, bytes) falls back to +/// `convert_scalar`. Both the attribute read (a patched class attribute) and +/// the fallback can run Python code, so the container size is revalidated +/// after each before the next borrowed read. +pub(super) fn uuid_seq( + py: Python<'_>, + seq: &S, + ch_type: &ChType, + name: &str, + row_count: usize, + nullable: bool, +) -> PyResult { + let uuid_type = py + .import(intern!(py, "uuid"))? + .getattr(intern!(py, "UUID"))?; + let uuid_type_ptr = uuid_type.as_ptr(); + let int_attr = intern!(py, "int"); + let mut data = Vec::with_capacity(16 * row_count); + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every conversion revalidates; the borrowed pointer is consumed + // before any Python code can run. + let ptr = unsafe { seq.get(row) }; + if ptr == unsafe { ffi::Py_None() } { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + }; + null_map.push(1); + data.extend_from_slice(&[0u8; 16]); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + // SAFETY: ptr is valid here; the strong reference keeps the item + // alive across any Python code the conversion runs. + let obj = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + let fast = if unsafe { ffi::Py_TYPE(ptr) }.cast::() == uuid_type_ptr { + obj.getattr(int_attr).and_then(|i| i.extract::()).ok() + } else { + None + }; + match fast { + Some(v) => { + data.extend_from_slice(&((v >> 64) as u64).to_le_bytes()); + data.extend_from_slice(&(v as u64).to_le_bytes()); + } + None => { + let Scalar::Bytes(bytes) = convert_scalar(py, ch_type, &obj, name, row)? else { + return Err(PyValueError::new_err("internal scalar type mismatch")); + }; + data.extend_from_slice(&bytes); + } + } + check_not_resized(seq, name, row_count)?; + } + + Ok(Column::Uuid(match null_map { + Some(nulls) => FixedBinaryColumn::new_nullable(data, 16, Bitmap::from_ch_null_map(&nulls)), + None => FixedBinaryColumn::new(data, 16), + })) +} + +/// Build an IPv4 column: exact `ipaddress.IPv4Address` items read the `_ip` +/// slot directly (or the attribute when the slot offset cannot be resolved), +/// exact ints and strs convert without constructing exceptions, and anything +/// else falls back to `convert_scalar`. The attribute read (a patched class +/// attribute) and the fallback can run Python code, so the container size is +/// revalidated after each before the next borrowed read. +pub(super) fn ipv4_seq( + py: Python<'_>, + seq: &S, + ch_type: &ChType, + name: &str, + row_count: usize, + nullable: bool, +) -> PyResult { + let ipv4_type = py + .import(intern!(py, "ipaddress"))? + .getattr(intern!(py, "IPv4Address"))?; + let ipv4_type_ptr = ipv4_type.as_ptr(); + let ip_attr = intern!(py, "_ip"); + // Re-resolved after any edge that runs Python, which can patch the class. + let mut ip_slot = slot_object_offset(&ipv4_type, ip_attr); + let mut values = Vec::::with_capacity(row_count); + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every conversion revalidates; the borrowed pointer is consumed + // before any Python code can run. + let ptr = unsafe { seq.get(row) }; + if ptr == unsafe { ffi::Py_None() } { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + }; + null_map.push(1); + values.push(0); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + if unsafe { ffi::Py_TYPE(ptr) }.cast::() == ipv4_type_ptr { + if let Some(offset) = ip_slot { + // SAFETY: ptr is an exact instance of the class the offset + // was resolved from; the slot holds a strong reference (kept + // alive by the instance) or NULL, and reading it runs no + // Python code. + let slot = unsafe { *ptr.cast::().offset(offset).cast::<*mut ffi::PyObject>() }; + if !slot.is_null() { + // SAFETY: GIL held; slot is a valid object pointer. + if let Ok(v) = unsafe { ::from_exact(slot, ch_type, 0) } { + values.push(v); + continue; + } + } + } + // SAFETY: ptr is valid here; the strong reference keeps the item + // alive across any Python code the conversion runs. + let obj = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + match obj.getattr(ip_attr).and_then(|o| o.extract::()) { + Ok(v) => values.push(v), + Err(_) => values.push(ipv4_fallback(py, ch_type, &obj, name, row)?), + } + check_not_resized(seq, name, row_count)?; + ip_slot = slot_object_offset(&ipv4_type, ip_attr); + continue; + } + // SAFETY: GIL held; ptr is a valid borrowed item pointer. + if let Ok(v) = unsafe { ::from_exact(ptr, ch_type, 0) } { + values.push(v); + continue; + } + if unsafe { ffi::PyUnicode_CheckExact(ptr) } != 0 { + // SAFETY: ptr is a valid borrowed reference, verified an exact + // str; reading its UTF-8 runs no Python code. + let obj = + unsafe { Bound::from_borrowed_ptr(py, ptr).cast_into_unchecked::() }; + let v = obj + .to_str()? + .parse::() + .map(|addr| u32::from_be_bytes(addr.octets())) + .map_err(|_| conversion_error(name, row, "IPv4"))?; + values.push(v); + continue; + } + // SAFETY: ptr is valid here; the strong reference keeps the item + // alive across any Python code the fallback runs. + let obj = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + values.push(ipv4_fallback(py, ch_type, &obj, name, row)?); + check_not_resized(seq, name, row_count)?; + ip_slot = slot_object_offset(&ipv4_type, ip_attr); + } + + Ok(Column::Ipv4(match null_map { + Some(nulls) => PrimitiveColumn::new_nullable(values, Bitmap::from_ch_null_map(&nulls)), + None => PrimitiveColumn::new(values), + })) +} + +/// Instance offset of `attr` when it resolves to a plain object slot +/// (`__slots__` member descriptor, `Py_T_OBJECT_EX`, no flags) defined on +/// exactly `class`. Lets a loop over exact instances read the slot directly, +/// like CPython's LOAD_ATTR_SLOT specialization. `None` (patched or exotic +/// class attribute) means callers must use a normal attribute read. +fn slot_object_offset(class: &Bound<'_, PyAny>, attr: &Bound<'_, PyString>) -> Option { + let descr = class.getattr(attr).ok()?; + if unsafe { ffi::Py_TYPE(descr.as_ptr()) != std::ptr::addr_of_mut!(ffi::PyMemberDescr_Type) } { + return None; + } + // SAFETY: descr is verified an exact member_descriptor, so its layout is + // PyMemberDescrObject; d_member points at the defining PyMemberDef (the + // ffi binding mistypes the field, hence the cast). + unsafe { + let descr = descr.as_ptr().cast::(); + if (*descr).d_common.d_type.cast::() != class.as_ptr() { + return None; + } + let member = (*descr).d_member.cast::(); + if member.is_null() || (*member).name.is_null() { + return None; + } + if (*member).type_code != ffi::Py_T_OBJECT_EX || (*member).flags != 0 { + return None; + } + let offset = (*member).offset; + (offset > 0).then_some(offset) + } +} + +fn ipv4_fallback( + py: Python<'_>, + ch_type: &ChType, + obj: &Bound<'_, PyAny>, + name: &str, + row: usize, +) -> PyResult { + match convert_scalar(py, ch_type, obj, name, row)? { + Scalar::Ipv4(v) => Ok(v), + _ => Err(PyValueError::new_err("internal scalar type mismatch")), + } +} + +/// Enum code width plumbing for the enum sequence fast path. +pub(super) trait EnumCode: Copy + Default { + const TYPE_NAME: &'static str; + + fn from_enum_scalar(scalar: Scalar) -> PyResult; + + fn into_enum_column(values: Vec, validity: Option) -> Column; +} + +macro_rules! impl_enum_code { + ($ty:ty, $variant:ident, $type_name:literal) => { + impl EnumCode for $ty { + const TYPE_NAME: &'static str = $type_name; + + fn from_enum_scalar(scalar: Scalar) -> PyResult { + match scalar { + Scalar::$variant(value) => Ok(value), + _ => Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + + fn into_enum_column(values: Vec, validity: Option) -> Column { + Column::$variant(match validity { + Some(validity) => PrimitiveColumn::new_nullable(values, validity), + None => PrimitiveColumn::new(values), + }) + } + } + }; +} + +impl_enum_code!(i8, Enum8, "Enum8"); +impl_enum_code!(i16, Enum16, "Enum16"); + +/// Build an Enum column with two lookup levels, following `lc_string_seq`: +/// repeated str objects hit a pointer-identity map, and new exact-str objects +/// read their UTF-8 once for a content lookup against the variant labels. An +/// exact str whose label is not defined is an error; anything else (raw int +/// codes, str subclasses) falls back to `convert_scalar`. The fallback can +/// run arbitrary Python, so it revalidates the container size and clears the +/// pointer-identity cache, whose entries are only valid while no Python code +/// has run since they were cached. +pub(super) fn enum_seq( + py: Python<'_>, + seq: &S, + ch_type: &ChType, + variants: &[(String, C)], + name: &str, + row_count: usize, + nullable: bool, +) -> PyResult { + let mut content_codes: HashMap<&[u8], C> = HashMap::with_capacity(variants.len()); + for (label, code) in variants { + // First definition wins, matching enum8_value's linear scan. + content_codes.entry(label.as_bytes()).or_insert(*code); + } + let mut ptr_codes: HashMap> = + HashMap::default(); + let mut codes = Vec::with_capacity(row_count); + let mut null_map = nullable.then(|| Vec::with_capacity(row_count)); + + for row in 0..row_count { + // SAFETY: row < row_count, the container size the caller checked and + // every fallback revalidates; the borrowed pointer is consumed before + // any Python code can run. + let ptr = unsafe { seq.get(row) }; + if ptr == unsafe { ffi::Py_None() } { + let Some(null_map) = &mut null_map else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is None but {ch_type} is not Nullable" + ))); + }; + null_map.push(1); + codes.push(C::default()); + continue; + } + if let Some(null_map) = &mut null_map { + null_map.push(0); + } + if let Some(&code) = ptr_codes.get(&(ptr as usize)) { + codes.push(code); + continue; + } + if unsafe { ffi::PyUnicode_CheckExact(ptr) } != 0 { + // SAFETY: ptr is a valid borrowed reference, verified an exact + // str; reading its UTF-8 runs no Python code. + let obj = + unsafe { Bound::from_borrowed_ptr(py, ptr).cast_into_unchecked::() }; + let label = obj.to_str()?; + let Some(&code) = content_codes.get(label.as_bytes()) else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} {} label {label:?} is not defined", + C::TYPE_NAME + ))); + }; + // The container still holds this item (no Python ran since the + // borrowed read), so its address is stable and unique among the + // column's live values. + if ptr_codes.len() < PTR_CACHE_CAP { + ptr_codes.insert(ptr as usize, code); + } + codes.push(code); + continue; + } + // SAFETY: ptr is valid here; the strong reference keeps the item + // alive across any Python code the fallback runs. + let obj = unsafe { Bound::from_borrowed_ptr(py, ptr) }; + if null_map.is_some() && is_enum_nan(ch_type, &obj) { + if let Some(entry) = null_map.as_mut().and_then(|nulls| nulls.last_mut()) { + *entry = 1; + } + codes.push(C::default()); + if S::MUTABLE { + check_not_resized(seq, name, row_count)?; + ptr_codes.clear(); + } + continue; + } + let scalar = convert_scalar(py, ch_type, &obj, name, row)?; + codes.push(C::from_enum_scalar(scalar)?); + if S::MUTABLE { + check_not_resized(seq, name, row_count)?; + // The fallback may have freed a cached str whose address the + // allocator can reuse, so drop all pointer-identity entries. + ptr_codes.clear(); + } + } + + Ok(C::into_enum_column( + codes, + null_map.map(|nulls| Bitmap::from_ch_null_map(&nulls)), + )) +} diff --git a/rust/ch-core-py/src/insert/temporal.rs b/rust/ch-core-py/src/insert/temporal.rs new file mode 100644 index 00000000..3e4e4157 --- /dev/null +++ b/rust/ch-core-py/src/insert/temporal.rs @@ -0,0 +1,1043 @@ +use super::*; + +#[derive(Clone, Copy)] +enum NumpyByteOrder { + Little, + Big, +} + +#[derive(Clone, Copy)] +struct NumpyTimedeltaMeta { + order: NumpyByteOrder, + /// Duration of one stored unit as numerator/denominator seconds. + numerator: i128, + denominator: i128, +} + +/// Parse a NumPy dtype object's `str` form without importing NumPy. Examples +/// are `m8[ms]`, and `) -> Option<(NumpyTimedeltaMeta, bool)> { + let dtype_str = dtype.getattr(intern!(dtype.py(), "str")).ok()?; + let dtype_str = dtype_str.cast::().ok()?.to_str().ok()?; + let (order, rest) = match dtype_str.as_bytes().split_first()? { + (b'<', rest) => (NumpyByteOrder::Little, rest), + (b'>', rest) => (NumpyByteOrder::Big, rest), + (b'=', rest) | (b'|', rest) => ( + if cfg!(target_endian = "little") { + NumpyByteOrder::Little + } else { + NumpyByteOrder::Big + }, + rest, + ), + _ => return None, + }; + let rest = std::str::from_utf8(rest).ok()?; + if rest == "m8" { + let meta = NumpyTimedeltaMeta { + order, + numerator: 1, + denominator: 1, + }; + return Some((meta, true)); + } + let unit = rest.strip_prefix("m8[")?.strip_suffix(']')?; + let digit_count = unit.bytes().take_while(u8::is_ascii_digit).count(); + let multiplier = if digit_count == 0 { + 1i128 + } else { + unit[..digit_count].parse::().ok()? + }; + let base = &unit[digit_count..]; + let (unit_numerator, denominator) = match base { + "W" => (604_800, 1), + "D" => (86_400, 1), + "h" => (3_600, 1), + "m" => (60, 1), + "s" => (1, 1), + "ms" => (1, 1_000), + "us" => (1, 1_000_000), + "ns" => (1, 1_000_000_000), + "ps" => (1, 1_000_000_000_000), + "fs" => (1, 1_000_000_000_000_000), + "as" => (1, 1_000_000_000_000_000_000), + _ => return None, + }; + Some(( + NumpyTimedeltaMeta { + order, + numerator: multiplier.checked_mul(unit_numerator)?, + denominator, + }, + false, + )) +} + +/// Non-generic timedelta meta read from a value's `dtype` attribute. +fn numpy_timedelta_meta(value: &Bound<'_, PyAny>) -> Option { + let dtype = value.getattr(intern!(value.py(), "dtype")).ok()?; + match parse_timedelta_dtype(&dtype) { + Some((meta, false)) => Some(meta), + _ => None, + } +} + +fn numpy_i64(bytes: &[u8], order: NumpyByteOrder) -> Option { + let raw: [u8; 8] = bytes.try_into().ok()?; + Some(match order { + NumpyByteOrder::Little => i64::from_le_bytes(raw), + NumpyByteOrder::Big => i64::from_be_bytes(raw), + }) +} + +fn numpy_timedelta_scalar_raw( + value: &Bound<'_, PyAny>, +) -> PyResult> { + let Ok(dtype) = value.getattr(intern!(value.py(), "dtype")) else { + return Ok(None); + }; + let Some((meta, generic)) = parse_timedelta_dtype(&dtype) else { + return Ok(None); + }; + let bytes = value.call_method0(intern!(value.py(), "tobytes"))?; + let bytes = bytes + .cast::() + .map_err(|_| PyValueError::new_err("NumPy timedelta tobytes() did not return bytes"))?; + let raw = numpy_i64(bytes.as_bytes(), meta.order) + .ok_or_else(|| PyValueError::new_err("NumPy timedelta scalar is not 8 bytes"))?; + if generic && raw != i64::MIN { + return Ok(None); + } + Ok(Some((raw, meta))) +} + +/// NumPy scalar probing requires Python `dtype`/`tobytes` lookups. Keep those +/// off the established Time input paths, especially the per-row timedelta hot +/// path. Exact timedelta/time objects are excluded; their subclasses +/// (pd.Timedelta, pandas NaT) go through the probe. These checks use only +/// CPython/PyO3 type predicates and run no Python. +#[inline] +fn should_probe_numpy_timedelta(value: &Bound<'_, PyAny>) -> bool { + // SAFETY: value is a valid object pointer; PyLong_Check and + // PyUnicode_Check are C-level type predicates that run no Python. + if unsafe { + ffi::PyLong_Check(value.as_ptr()) != 0 || ffi::PyUnicode_Check(value.as_ptr()) != 0 + } { + return false; + } + if value.cast::().is_ok() { + return false; + } + if value.cast::().is_ok() { + // SAFETY: the successful downcast imported the datetime C-API. + return unsafe { ffi::PyDelta_CheckExact(value.as_ptr()) } == 0; + } + if value.cast::().is_ok() { + // SAFETY: the successful downcast imported the datetime C-API. + return unsafe { ffi::PyTime_CheckExact(value.as_ptr()) } == 0; + } + true +} + +/// pandas NaT detection by C-level type name; runs no Python. +fn is_pandas_nat(value: &Bound<'_, PyAny>) -> bool { + // SAFETY: Py_TYPE of a valid object; tp_name is a NUL-terminated string + // owned by the type. + let name = unsafe { std::ffi::CStr::from_ptr((*ffi::Py_TYPE(value.as_ptr())).tp_name) }; + let name = name.to_bytes(); + name == b"NaTType" || name.ends_with(b".NaTType") +} + +/// Non-null outcome of probing one Time/Time64 cell. +pub(super) enum TimeProbe { + Nat, + Ticks(i64), +} + +/// Per-column state for probing Time/Time64 cells that are not one of the +/// exact fast types. The parsed dtype and unit ratio are cached and reused +/// while subsequent cells carry an identical or equal dtype object. +pub(super) struct TimeScalarProbe { + precision: u8, + fractional: bool, + type_name: &'static str, + cache: Option, +} + +struct TimeDtypeCache { + dtype: Py, + order: NumpyByteOrder, + generic: bool, + ratio: NumpyTimeRatio, +} + +impl TimeScalarProbe { + pub(super) fn new(ch_type: &ChType) -> Option { + match ch_type { + ChType::Time => Some(Self { + precision: 0, + fractional: false, + type_name: "Time", + cache: None, + }), + ChType::Time64 { precision } => Some(Self { + precision: *precision, + fractional: true, + type_name: "Time64", + cache: None, + }), + _ => None, + } + } + + /// Probe one cell. `Ok(None)` means the value is neither a NumPy + /// timedelta scalar nor pandas NaT; the caller converts it normally. + pub(super) fn probe( + &mut self, + value: &Bound<'_, PyAny>, + column: &str, + row: usize, + ) -> PyResult> { + if !should_probe_numpy_timedelta(value) { + return Ok(None); + } + let py = value.py(); + let Ok(dtype) = value.getattr(intern!(py, "dtype")) else { + return Ok(is_pandas_nat(value).then_some(TimeProbe::Nat)); + }; + let cached = self.cache.as_ref().and_then(|cached| { + (cached.dtype.as_ptr() == dtype.as_ptr() + || cached.dtype.bind(py).eq(&dtype).unwrap_or(false)) + .then_some((cached.order, cached.generic, cached.ratio)) + }); + let (order, generic, ratio) = match cached { + Some(hit) => hit, + None => { + let Some((meta, generic)) = parse_timedelta_dtype(&dtype) else { + return Ok(None); + }; + let ratio = numpy_time_ratio(meta, self.precision, column)?; + self.cache = Some(TimeDtypeCache { + dtype: dtype.unbind(), + order: meta.order, + generic, + ratio, + }); + (meta.order, generic, ratio) + } + }; + let bytes = value.call_method0(intern!(py, "tobytes"))?; + let bytes = bytes.cast::().map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} NumPy timedelta tobytes() did not return bytes" + )) + })?; + let raw = numpy_i64(bytes.as_bytes(), order).ok_or_else(|| { + PyValueError::new_err(format!( + "column {column:?} row {row} NumPy timedelta scalar is not 8 bytes" + )) + })?; + if raw == i64::MIN { + return Ok(Some(TimeProbe::Nat)); + } + if generic { + return Ok(None); + } + let ticks = rescale_numpy_timedelta( + raw, + ratio, + self.fractional, + self.precision, + column, + row, + self.type_name, + )?; + Ok(Some(TimeProbe::Ticks(ticks))) + } +} + +/// Build the Time/Time64 scalar for probed ticks. +pub(super) fn time_ticks_scalar( + ch_type: &ChType, + ticks: i64, + column: &str, + row: usize, +) -> PyResult { + match ch_type { + ChType::Time => { + Ok(Scalar::Time(i32::try_from(ticks).map_err(|_| { + time_range_error(column, row, "Time", ticks) + })?)) + } + ChType::Time64 { .. } => Ok(Scalar::Time64(ticks)), + _ => Err(PyValueError::new_err("internal Time scalar type mismatch")), + } +} + +#[derive(Clone, Copy)] +enum NumpyTimeRatio { + Identity, + Multiply(i64), + Divide(i64), + General { numerator: i128, denominator: i128 }, +} + +fn gcd_i128(mut left: i128, mut right: i128) -> i128 { + while right != 0 { + let remainder = left % right; + left = right; + right = remainder; + } + left +} + +/// Resolve source units -> target ticks once per column. Common NumPy units +/// become identity, checked i64 multiply, or signed i64 division; only unusual +/// dtype multipliers retain the general i128 path. +fn numpy_time_ratio( + meta: NumpyTimedeltaMeta, + precision: u8, + column: &str, +) -> PyResult { + let numerator = meta + .numerator + .checked_mul(i128::from(time64_scale(precision))) + .ok_or_else(|| { + PyValueError::new_err(format!( + "column {column:?} NumPy timedelta unit scale overflows" + )) + })?; + let common = gcd_i128(numerator, meta.denominator); + let numerator = numerator / common; + let denominator = meta.denominator / common; + if numerator == 1 && denominator == 1 { + Ok(NumpyTimeRatio::Identity) + } else if denominator == 1 { + Ok(match i64::try_from(numerator) { + Ok(multiplier) => NumpyTimeRatio::Multiply(multiplier), + Err(_) => NumpyTimeRatio::General { + numerator, + denominator, + }, + }) + } else if numerator == 1 { + Ok(match i64::try_from(denominator) { + Ok(divisor) => NumpyTimeRatio::Divide(divisor), + Err(_) => NumpyTimeRatio::General { + numerator, + denominator, + }, + }) + } else { + Ok(NumpyTimeRatio::General { + numerator, + denominator, + }) + } +} + +/// Convert one raw NumPy timedelta count using a precomputed unit ratio. Signed +/// division deliberately truncates toward zero for sub-tick negatives. +fn rescale_numpy_timedelta( + raw: i64, + ratio: NumpyTimeRatio, + fractional: bool, + precision: u8, + column: &str, + row: usize, + type_name: &str, +) -> PyResult { + let ticks = match ratio { + NumpyTimeRatio::Identity => raw, + NumpyTimeRatio::Multiply(multiplier) => raw + .checked_mul(multiplier) + .ok_or_else(|| time_range_error(column, row, type_name, raw))?, + NumpyTimeRatio::Divide(divisor) => raw / divisor, + NumpyTimeRatio::General { + numerator, + denominator, + } => { + let ticks = i128::from(raw) + .checked_mul(numerator) + .ok_or_else(|| time_range_error(column, row, type_name, raw))? + / denominator; + i64::try_from(ticks).map_err(|_| time_range_error(column, row, type_name, ticks))? + } + }; + let max = if fractional { + max_time64_ticks(precision) + } else { + MAX_TIME_SECONDS + }; + if ticks < -max || ticks > max { + return Err(time_range_error(column, row, type_name, ticks)); + } + Ok(ticks) +} + +fn numpy_timedelta_source<'py>( + py: Python<'py>, + values: &Bound<'py, PyAny>, +) -> PyResult, NumpyTimedeltaMeta)>> { + let Some(meta) = numpy_timedelta_meta(values) else { + return Ok(None); + }; + if values.getattr("tobytes").is_ok() { + return Ok(Some((values.clone(), meta))); + } + let Ok(to_numpy) = values.getattr("to_numpy") else { + return Ok(None); + }; + let kwargs = PyDict::new(py); + kwargs.set_item("copy", false)?; + let array = to_numpy.call((), Some(&kwargs))?; + let Some(array_meta) = numpy_timedelta_meta(&array) else { + return Ok(None); + }; + Ok(Some((array, array_meta))) +} + +/// Dependency-free ndarray/pandas fast path. `tobytes` performs one safe bulk +/// copy for arbitrary strides; conversion then walks the contiguous bytes in +/// Rust with no per-cell Python calls or temporary objects. +pub(super) fn try_numpy_timedelta_column( + py: Python<'_>, + name: &str, + ch_type: &ChType, + values: &Bound<'_, PyAny>, + row_count: usize, + nullable: bool, +) -> PyResult> { + let (precision, fractional, type_name) = match ch_type { + ChType::Time => (0, false, "Time"), + ChType::Time64 { precision } => (*precision, true, "Time64"), + _ => return Ok(None), + }; + let Some((source, meta)) = numpy_timedelta_source(py, values)? else { + return Ok(None); + }; + let ndim = source + .getattr("ndim") + .and_then(|value| value.extract::()) + .map_err(|_| { + PyValueError::new_err(format!( + "column {name:?} NumPy timedelta source has no valid ndim" + )) + })?; + if ndim != 1 { + return Err(PyValueError::new_err(format!( + "column {name:?} NumPy timedelta source must be one-dimensional, got {ndim} dimensions" + ))); + } + let bytes = source.call_method0("tobytes")?; + let bytes = bytes.cast::().map_err(|_| { + PyValueError::new_err(format!( + "column {name:?} NumPy timedelta tobytes() did not return bytes" + )) + })?; + let expected_len = row_count + .checked_mul(8) + .ok_or_else(|| PyValueError::new_err(format!("column {name:?} byte size overflow")))?; + let actual_len = bytes.as_bytes().len(); + if actual_len != expected_len { + return Err(PyValueError::new_err(format!( + "column {name:?} NumPy timedelta data is {} bytes, expected {expected_len}", + actual_len + ))); + } + + let ratio = numpy_time_ratio(meta, precision, name)?; + let mut nulls = nullable.then(|| Vec::with_capacity(row_count)); + macro_rules! build_numpy_time_column { + ($rust_type:ty, $column_variant:ident) => {{ + let mut out = Vec::<$rust_type>::with_capacity(row_count); + for (row, raw_bytes) in bytes.as_bytes().chunks_exact(8).enumerate() { + let raw = numpy_i64(raw_bytes, meta.order).ok_or_else(|| { + PyValueError::new_err(format!( + "column {name:?} row {row} has invalid NumPy timedelta bytes" + )) + })?; + if raw == i64::MIN { + let Some(nulls) = &mut nulls else { + return Err(PyValueError::new_err(format!( + "column {name:?} row {row} is NaT but {ch_type} is not Nullable" + ))); + }; + nulls.push(1); + out.push(0); + continue; + } + if let Some(nulls) = &mut nulls { + nulls.push(0); + } + let ticks = rescale_numpy_timedelta( + raw, ratio, fractional, precision, name, row, type_name, + )?; + out.push( + <$rust_type>::try_from(ticks) + .map_err(|_| time_range_error(name, row, type_name, ticks))?, + ); + } + let validity = nulls.map(|map| Bitmap::from_ch_null_map(&map)); + Some(Column::$column_variant(match validity { + Some(validity) => PrimitiveColumn::new_nullable(out, validity), + None => PrimitiveColumn::new(out), + })) + }}; + } + Ok(match ch_type { + ChType::Time => build_numpy_time_column!(i32, Time), + ChType::Time64 { .. } => build_numpy_time_column!(i64, Time64), + _ => None, + }) +} + +/// Narrowing i64 conversion for the fast paths. A generic helper so macro +/// expansions do not trip clippy's fallible-conversion lint when the target +/// is i64 itself. +#[inline] +fn narrow_i64>(value: i64) -> Result { + T::try_from(value).map_err(|_| ()) +} + +#[inline] +pub(super) fn time64_scale(precision: u8) -> i64 { + // parse_ch_type rejects precisions above 9. + debug_assert!(precision <= 9); + match precision { + 0 => 1, + 1 => 10, + 2 => 100, + 3 => 1_000, + 4 => 10_000, + 5 => 100_000, + 6 => 1_000_000, + 7 => 10_000_000, + 8 => 100_000_000, + 9 => 1_000_000_000, + // parse_ch_type rejects this before the binding builds a column. + _ => 0, + } +} + +#[inline] +pub(super) fn max_time64_ticks(precision: u8) -> i64 { + let scale = time64_scale(precision); + MAX_TIME_SECONDS * scale + (scale - 1) +} + +/// Temporal wire values: the fast path accepts exact raw ints only (the same +/// values `convert_scalar` accepts without touching the object protocol); +/// date/datetime/str objects and out-of-range ints go through the fallback, +/// which carries each type's specific conversion and range errors. +macro_rules! impl_fast_temporal { + ($name:ident, $prim:ty, $variant:ident) => { + impl_fast_temporal!($name, $prim, $variant, |_value, _ch_type, _fast_limit| true); + }; + ($name:ident, $prim:ty, $variant:ident, $validate:expr) => { + #[derive(Clone, Copy)] + #[repr(transparent)] + pub(super) struct $name($prim); + + impl FastValue for $name { + const DEFAULT: Self = $name(0); + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + ch_type: &ChType, + fast_limit: i64, + ) -> Result { + let value = narrow_i64::<$prim>(exact_long_as_i64(ptr)?)?; + if ($validate)(value, ch_type, fast_limit) { + Ok(Self(value)) + } else { + Err(()) + } + } + + fn from_scalar(scalar: Scalar) -> PyResult { + match scalar { + Scalar::$variant(value) => Ok(Self(value)), + _ => Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + + fn into_column(values: Vec, validity: Option) -> Column { + // SAFETY: $name is #[repr(transparent)] over $prim. + let values = unsafe { cast_vec::(values) }; + Column::$variant(match validity { + Some(validity) => PrimitiveColumn::new_nullable(values, validity), + None => PrimitiveColumn::new(values), + }) + } + } + }; +} + +impl_fast_temporal!(DateVal, u16, Date); +impl_fast_temporal!(Date32Val, i32, Date32); +impl_fast_temporal!(DateTimeVal, u32, DateTime); +impl_fast_temporal!(DateTime64Val, i64, DateTime64); + +#[derive(Clone, Copy)] +#[repr(transparent)] +pub(super) struct IntervalVal(i64); + +impl FastValue for IntervalVal { + const DEFAULT: Self = Self(0); + + #[inline] + unsafe fn from_exact( + ptr: *mut ffi::PyObject, + _ch_type: &ChType, + _fast_limit: i64, + ) -> Result { + exact_long_as_i64(ptr).map(Self) + } + + fn from_scalar(scalar: Scalar) -> PyResult { + match scalar { + Scalar::Interval(value) => Ok(Self(value)), + _ => Err(PyValueError::new_err("internal scalar type mismatch")), + } + } + + fn from_buffer( + py: Python<'_>, + _name: &str, + values: &Bound<'_, PyAny>, + row_count: usize, + ) -> PyResult>> { + // SAFETY: IntervalVal is #[repr(transparent)] over i64. + Ok(buffer_values::(py, values, row_count)? + .map(|values| unsafe { cast_vec::(values) })) + } + + fn into_column(values: Vec, validity: Option) -> Column { + // SAFETY: IntervalVal is #[repr(transparent)] over i64. + let values = unsafe { cast_vec::(values) }; + Column::Interval(match validity { + Some(validity) => PrimitiveColumn::new_nullable(values, validity), + None => PrimitiveColumn::new(values), + }) + } +} + +impl_fast_temporal!( + TimeVal, + i32, + Time, + |value: i32, _ch_type: &ChType, fast_limit: i64| { + i64::from(value).unsigned_abs() <= fast_limit as u64 + } +); +impl_fast_temporal!( + Time64Val, + i64, + Time64, + |value: i64, _ch_type: &ChType, fast_limit: i64| { value.unsigned_abs() <= fast_limit as u64 } +); + +pub(super) fn date_days(value: &Bound<'_, PyAny>, column: &str, row: usize) -> PyResult { + // Precedence: int (incl. subclasses) -> date/datetime instance -> duck + // tail (__index__, then toordinal). Cheap type checks avoid constructing + // a TypeError per value; an object that is both a date instance and + // int-convertible resolves as a date here. PyDate_Check also matches + // datetime, a date subclass, like the tail's toordinal call. + if unsafe { ffi::PyLong_Check(value.as_ptr()) } != 0 { + // A wider-than-i64 int maps to the tail's conversion error without + // re-running the failing extraction. + return value + .extract::() + .map_err(|_| conversion_error(column, row, "Date")); + } + if value.is_instance_of::() { + let ordinal = value + .call_method0("toordinal") + .and_then(|o| o.extract::()) + .map_err(|_| conversion_error(column, row, "Date"))?; + return Ok(ordinal - EPOCH_DATE_ORDINAL); + } + // Duck-typed tail, e.g. numpy ints via __index__. + if let Ok(days) = value.extract::() { + return Ok(days); + } + let ordinal = value + .call_method0("toordinal") + .and_then(|o| o.extract::()) + .map_err(|_| conversion_error(column, row, "Date"))?; + Ok(ordinal - EPOCH_DATE_ORDINAL) +} + +pub(super) fn datetime_seconds( + value: &Bound<'_, PyAny>, + column: &str, + row: usize, +) -> PyResult { + // Precedence: int -> datetime instance -> duck tail; see date_days. + if unsafe { ffi::PyLong_Check(value.as_ptr()) } != 0 { + return value + .extract::() + .map_err(|_| conversion_error(column, row, "DateTime")); + } + if value.is_instance_of::() { + let ts = value + .call_method0("timestamp") + .and_then(|o| o.extract::()) + .map_err(|_| conversion_error(column, row, "DateTime"))?; + return finite_trunc_to_i64(ts, column, row, "DateTime"); + } + // Duck-typed tail, e.g. numpy ints via __index__. + if let Ok(secs) = value.extract::() { + return Ok(secs); + } + let ts = value + .call_method0("timestamp") + .and_then(|o| o.extract::()) + .map_err(|_| conversion_error(column, row, "DateTime"))?; + finite_trunc_to_i64(ts, column, row, "DateTime") +} + +pub(super) fn datetime64_ticks( + py: Python<'_>, + value: &Bound<'_, PyAny>, + precision: u8, + column: &str, + row: usize, +) -> PyResult { + // Precedence: int -> datetime instance -> duck tail; see date_days. + if unsafe { ffi::PyLong_Check(value.as_ptr()) } != 0 { + return value + .extract::() + .map_err(|_| conversion_error(column, row, "DateTime64")); + } + if value.is_instance_of::() { + let secs = dt64_timestamp_secs(value, column, row)?; + // is_instance_of imported the datetime C API, so the raw exact check + // and struct accessor are safe. A subclass may override the + // microsecond attribute, so only an exact datetime skips the getattr. + let micros = if unsafe { ffi::PyDateTime_CheckExact(value.as_ptr()) } != 0 { + i64::from(unsafe { ffi::PyDateTime_DATE_GET_MICROSECOND(value.as_ptr()) }) + } else { + dt64_microsecond(value, column, row)? + }; + return dt64_ticks_math(secs, micros, precision, column, row); + } + + if let Ok(ticks) = value.extract::() { + return Ok(ticks); + } + + let value = if let Ok(s) = value.cast::() { + py.import("datetime")? + .getattr("datetime")? + .call_method1("fromisoformat", (s.to_str()?,))? + } else { + value.clone() + }; + + let secs = dt64_timestamp_secs(&value, column, row)?; + let micros = dt64_microsecond(&value, column, row)?; + dt64_ticks_math(secs, micros, precision, column, row) +} + +fn dt64_timestamp_secs(value: &Bound<'_, PyAny>, column: &str, row: usize) -> PyResult { + let secs = value + .call_method0("timestamp") + .and_then(|o| o.extract::()) + .map_err(|_| conversion_error(column, row, "DateTime64"))?; + finite_floor_to_i64(secs, column, row, "DateTime64") +} + +fn dt64_microsecond(value: &Bound<'_, PyAny>, column: &str, row: usize) -> PyResult { + value + .getattr("microsecond") + .and_then(|o| o.extract::()) + .map_err(|_| conversion_error(column, row, "DateTime64")) +} + +fn dt64_ticks_math( + secs: i64, + micros: i64, + precision: u8, + column: &str, + row: usize, +) -> PyResult { + if !(0..1_000_000).contains(µs) { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} DateTime64 microsecond {micros} is outside range" + ))); + } + let scale = 10i128.pow(u32::from(precision)); + let total_micros = i128::from(secs) + .checked_mul(1_000_000) + .and_then(|v| v.checked_add(i128::from(micros))) + .ok_or_else(|| { + PyValueError::new_err(format!( + "column {column:?} row {row} DateTime64 value overflows" + )) + })?; + let ticks = total_micros + .checked_mul(scale) + .ok_or_else(|| { + PyValueError::new_err(format!( + "column {column:?} row {row} DateTime64 value overflows" + )) + })? + .div_euclid(1_000_000); + i64::try_from(ticks).map_err(|_| { + PyValueError::new_err(format!( + "column {column:?} row {row} DateTime64 value {ticks} is outside Int64 range" + )) + }) +} + +pub(super) fn time_ticks(value: &Bound<'_, PyAny>, column: &str, row: usize) -> PyResult { + let ticks = time_like_ticks(value, 0, false, column, row, "Time")?; + i32::try_from(ticks).map_err(|_| time_range_error(column, row, "Time", ticks)) +} + +pub(super) fn time64_ticks( + value: &Bound<'_, PyAny>, + precision: u8, + column: &str, + row: usize, +) -> PyResult { + time_like_ticks(value, precision, true, column, row, "Time64") +} + +/// Convert the accepted Time/Time64 Python values to raw signed ticks. Exact +/// timedelta/time objects use PyO3's C-API accessors, so the common object path +/// performs no Python calls or attribute lookups. Integers are raw ticks; +/// floats follow TimeBase's `int(value)` truncation policy. +fn time_like_ticks( + value: &Bound<'_, PyAny>, + precision: u8, + fractional: bool, + column: &str, + row: usize, + type_name: &str, +) -> PyResult { + let scale = time64_scale(precision); + // SAFETY: value is a valid object pointer; PyLong_Check and PyFloat_Check + // are C-level type predicates that run no Python. + let (is_long, is_float) = unsafe { + ( + ffi::PyLong_Check(value.as_ptr()) != 0, + ffi::PyFloat_Check(value.as_ptr()) != 0, + ) + }; + let ticks = if is_long { + value + .extract::() + .map_err(|_| conversion_error(column, row, type_name))? + } else if let Ok(delta) = value.cast::() { + timedelta_ticks(delta, scale, precision, fractional, column, row, type_name)? + } else if let Ok(time) = value.cast::() { + let total_micros = ((i128::from(time.get_hour()) * 3_600 + + i128::from(time.get_minute()) * 60 + + i128::from(time.get_second())) + * 1_000_000) + + i128::from(time.get_microsecond()); + let ticks = total_micros * i128::from(scale) / 1_000_000; + i64::try_from(ticks).map_err(|_| time_range_error(column, row, type_name, ticks))? + } else if let Ok(s) = value.cast::() { + parse_time_literal(s.to_str()?, precision, fractional, column, row, type_name)? + } else if is_float { + // SAFETY: value is a float, so PyFloat_AsDouble cannot fail. + let number = unsafe { ffi::PyFloat_AsDouble(value.as_ptr()) }; + finite_trunc_to_i64(number, column, row, type_name)? + } else if let Some((raw, meta)) = numpy_timedelta_scalar_raw(value)? { + if raw == i64::MIN { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} is NaT and cannot be converted to {type_name}" + ))); + } + let ratio = numpy_time_ratio(meta, precision, column)?; + rescale_numpy_timedelta(raw, ratio, fractional, precision, column, row, type_name)? + } else if is_pandas_nat(value) { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} is NaT and cannot be converted to {type_name}" + ))); + } else { + value + .extract::() + .map_err(|_| conversion_error(column, row, type_name))? + }; + + let max = if fractional { + max_time64_ticks(precision) + } else { + MAX_TIME_SECONDS + }; + if ticks < -max || ticks > max { + return Err(time_range_error(column, row, type_name, ticks)); + } + Ok(ticks) +} + +/// Ticks from a datetime.timedelta. Exact objects read the C-struct fields; +/// subclasses (pd.Timedelta) first try the ns-resolution `asm8` scalar so +/// sub-microsecond values are not truncated. +fn timedelta_ticks( + delta: &Bound<'_, PyDelta>, + scale: i64, + precision: u8, + fractional: bool, + column: &str, + row: usize, + type_name: &str, +) -> PyResult { + // SAFETY: a successful PyDelta downcast produced `delta`, so the datetime + // C-API is imported; the exact check runs no Python. + if unsafe { ffi::PyDelta_CheckExact(delta.as_ptr()) } == 0 { + if let Some(ticks) = + timedelta_subclass_ticks(delta, precision, fractional, column, row, type_name)? + { + return Ok(ticks); + } + } + const MICROS_PER_SECOND: i128 = 1_000_000; + const MICROS_PER_DAY: i128 = 86_400 * MICROS_PER_SECOND; + let total_micros = i128::from(delta.get_days()) * MICROS_PER_DAY + + i128::from(delta.get_seconds()) * MICROS_PER_SECOND + + i128::from(delta.get_microseconds()); + let scaled = total_micros * i128::from(scale); + // Signed division truncates sub-tick values toward zero, consistent + // with Time strings and the Python object decode policy. + let ticks = scaled / MICROS_PER_SECOND; + i64::try_from(ticks).map_err(|_| time_range_error(column, row, type_name, ticks)) +} + +/// ns-resolution ticks for timedelta subclasses exposing `asm8` +/// (pd.Timedelta). `Ok(None)` means the attribute or its dtype probe does not +/// apply and the caller falls back to the struct fields. +fn timedelta_subclass_ticks( + value: &Bound<'_, PyAny>, + precision: u8, + fractional: bool, + column: &str, + row: usize, + type_name: &str, +) -> PyResult> { + let Ok(asm8) = value.getattr(intern!(value.py(), "asm8")) else { + return Ok(None); + }; + let Some((raw, meta)) = numpy_timedelta_scalar_raw(&asm8)? else { + return Ok(None); + }; + if raw == i64::MIN { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} is NaT and cannot be converted to {type_name}" + ))); + } + let ratio = numpy_time_ratio(meta, precision, column)?; + rescale_numpy_timedelta(raw, ratio, fractional, precision, column, row, type_name).map(Some) +} + +fn parse_time_literal( + value: &str, + precision: u8, + fractional: bool, + column: &str, + row: usize, + type_name: &str, +) -> PyResult { + let value = value.trim(); + let (negative, unsigned) = match value.strip_prefix('-') { + Some(rest) => (true, rest), + None => (false, value), + }; + let (hms, fraction) = match unsigned.split_once('.') { + Some((hms, fraction)) + if !fraction.is_empty() + && !fraction.contains('.') + && fraction.bytes().all(|b| b.is_ascii_digit()) => + { + (hms, Some(fraction)) + } + Some(_) => return Err(time_literal_error(column, row, type_name, value)), + None => (unsigned, None), + }; + let mut parts = hms.split(':'); + let hours = parse_time_part(parts.next(), column, row, type_name, value)?; + let minutes = parse_time_part(parts.next(), column, row, type_name, value)?; + let seconds = parse_time_part(parts.next(), column, row, type_name, value)?; + if parts.next().is_some() || hours > 999 || minutes > 59 || seconds > 59 { + return Err(time_literal_error(column, row, type_name, value)); + } + + let scale = time64_scale(precision); + let mut fraction_ticks = 0i64; + if fractional { + let mut digits = 0u8; + for byte in fraction + .unwrap_or_default() + .bytes() + .take(usize::from(precision)) + { + fraction_ticks = fraction_ticks * 10 + i64::from(byte - b'0'); + digits += 1; + } + for _ in digits..precision { + fraction_ticks *= 10; + } + } + let ticks = (hours * 3_600 + minutes * 60 + seconds) * scale + fraction_ticks; + Ok(if negative { -ticks } else { ticks }) +} + +fn parse_time_part( + part: Option<&str>, + column: &str, + row: usize, + type_name: &str, + literal: &str, +) -> PyResult { + let part = part + .filter(|part| !part.is_empty()) + .ok_or_else(|| time_literal_error(column, row, type_name, literal))?; + if !part.bytes().all(|b| b.is_ascii_digit()) { + return Err(time_literal_error(column, row, type_name, literal)); + } + part.parse::() + .map_err(|_| time_literal_error(column, row, type_name, literal)) +} + +fn time_range_error( + column: &str, + row: usize, + type_name: &str, + value: T, +) -> PyErr { + PyValueError::new_err(format!( + "column {column:?} row {row} {type_name} value {value} is outside logical range" + )) +} + +fn time_literal_error(column: &str, row: usize, type_name: &str, value: &str) -> PyErr { + PyValueError::new_err(format!( + "column {column:?} row {row} invalid {type_name} literal {value:?}" + )) +} + +fn finite_trunc_to_i64(value: f64, column: &str, row: usize, type_name: &str) -> PyResult { + if !value.is_finite() || value < i64::MIN as f64 || value > i64::MAX as f64 { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} {type_name} timestamp is outside Int64 range" + ))); + } + Ok(value.trunc() as i64) +} + +fn finite_floor_to_i64(value: f64, column: &str, row: usize, type_name: &str) -> PyResult { + if !value.is_finite() || value < i64::MIN as f64 || value > i64::MAX as f64 { + return Err(PyValueError::new_err(format!( + "column {column:?} row {row} {type_name} timestamp is outside Int64 range" + ))); + } + Ok(value.floor() as i64) +} diff --git a/rust/ch-core-py/src/insert/variant.rs b/rust/ch-core-py/src/insert/variant.rs new file mode 100644 index 00000000..f78b2222 --- /dev/null +++ b/rust/ch-core-py/src/insert/variant.rs @@ -0,0 +1,376 @@ +use super::*; + +/// Build a Variant column in one row scan: one discriminator byte per logical +/// row and one strong-reference run per dense alternative. Each alternative is +/// then built once through the existing column fast paths. +pub(super) fn build_variant_column( + py: Python<'_>, + name: &str, + alternatives: &[ChType], + values: &Bound<'_, PyAny>, + row_count: usize, +) -> PyResult { + let column_values = ColumnValues::new(values, name)?; + check_row_count(name, &column_values, row_count)?; + if let Ok(list) = values.cast_exact::() { + return variant_column_from_seq(py, name, alternatives, &ListSeq(list), row_count); + } + if let Ok(tuple) = values.cast_exact::() { + return variant_column_from_seq(py, name, alternatives, &TupleSeq(tuple), row_count); + } + + let mut builder = VariantBuilder::new(py, name, alternatives, row_count)?; + for row in 0..row_count { + let value = column_values.get_item(row)?; + builder.push_row(&value, row)?; + } + builder.finish() +} + +/// Variant rows over an exact list/tuple or a flattened container element run. +pub(super) fn variant_column_from_seq( + py: Python<'_>, + name: &str, + alternatives: &[ChType], + seq: &S, + row_count: usize, +) -> PyResult { + let mut builder = VariantBuilder::new(py, name, alternatives, row_count)?; + for row in 0..row_count { + // SAFETY: row < row_count, which the caller checked against seq.size(). + // A strong reference protects the value while the builder retains its + // selected payload in an alternative run. + let value = unsafe { Bound::from_borrowed_ptr(py, seq.get(row)) }; + let ran_python = builder.push_row(&value, row)?; + if ran_python { + // Dispatch that ran Python (subclass attribute access or a + // str-subclass name lookup) can resize the source. Drop the + // temporary strong reference first so a finalizer cannot resize + // it after this validation. + drop(value); + check_not_resized(seq, name, row_count)?; + } + } + builder.finish() +} + +/// The driver Variant instance's dispatch tables, resolved once per column. +struct VariantDispatch<'py> { + /// Snapshot of `_python_map`: one strong type-object reference and its + /// discriminator per unambiguous alternative. + types: Vec<(Py, u8)>, + name_index: Bound<'py, PyDict>, + typed_variant_type: Bound<'py, PyAny>, +} + +/// Per-column Variant dispatch state. The Python driver's Variant instance is +/// cached by canonical type name; snapshotting its `_python_map` and borrowing +/// its `_name_index` keeps this binding's exact-type inference and +/// `typed_variant` policy identical to the established Python encoder, +/// including collision removal. +struct VariantBuilder<'a, 'py> { + py: Python<'py>, + name: &'a str, + alternatives: &'a [ChType], + type_name: String, + dispatch: VariantDispatch<'py>, + discriminators: Vec, + values: Vec, + /// Explicit tag names already resolved through `_name_index`, with their + /// discriminators. Rows reuse the same exact-str tag objects, so a pointer + /// compare replaces the dict lookup. The strong references keep each + /// cached name alive, so a cached pointer can never be a freed-and-reused + /// address; exact-str-only entries mean the hit path runs no Python code. + tag_cache: Vec<(Py, usize)>, +} + +impl<'a, 'py> VariantBuilder<'a, 'py> { + fn new( + py: Python<'py>, + name: &'a str, + alternatives: &'a [ChType], + row_count: usize, + ) -> PyResult { + let type_name = ChType::Variant(alternatives.to_vec()).to_string(); + // Any driver-interop failure becomes NotImplementedError so the + // insert probe falls back to the Python codec. + let dispatch = + Self::driver_dispatch(py, &type_name, alternatives.len()).map_err(|err: PyErr| { + PyNotImplementedError::new_err(format!( + "Variant type {type_name} is not supported by the native encoder: {err}" + )) + })?; + Ok(Self { + py, + name, + alternatives, + type_name, + dispatch, + discriminators: Vec::with_capacity(row_count), + values: alternatives.iter().map(|_| FlatRefs::default()).collect(), + tag_cache: Vec::new(), + }) + } + + /// Resolve the driver's cached Variant instance and snapshot its dispatch + /// tables. + fn driver_dispatch( + py: Python<'py>, + type_name: &str, + alternative_count: usize, + ) -> PyResult> { + let py_variant = py + .import("clickhouse_connect.datatypes.registry")? + .getattr("get_from_name")? + .call1((type_name,))?; + let python_map = py_variant.getattr("_python_map")?.cast_into::()?; + let name_index = py_variant.getattr("_name_index")?.cast_into::()?; + let typed_variant_type = py + .import("clickhouse_connect.datatypes.dynamic")? + .getattr("TypedVariant")?; + let mut types = Vec::with_capacity(python_map.len()); + for (key, value) in python_map.iter() { + // SAFETY: value is a live dict value kept alive by its Bound. + // PyLong_AsSize_t raises TypeError on a non-int without running + // Python code. + let index = unsafe { ffi::PyLong_AsSize_t(value.as_ptr()) }; + if index == usize::MAX && unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } + let discriminator = u8::try_from(index) + .ok() + .filter(|&d| usize::from(d) < alternative_count) + .ok_or_else(|| PyValueError::new_err("Variant dispatch index is out of range"))?; + types.push((key.unbind(), discriminator)); + } + Ok(VariantDispatch { + types, + name_index, + typed_variant_type, + }) + } + + /// Select and retain one row. The boolean reports whether successful + /// dispatch ran Python code, so borrowed-sequence callers only pay for a + /// resize check when it did. + fn push_row(&mut self, value: &Bound<'py, PyAny>, row: usize) -> PyResult { + if value.is_none() { + self.discriminators.push(u8::MAX); + return Ok(false); + } + + // SAFETY: value is a live Python object. Reading its type pointer and + // comparing it against snapshot pointers cannot invoke user Python + // code. + let value_type = unsafe { ffi::Py_TYPE(value.as_ptr()) }.cast::(); + if value_type == self.dispatch.typed_variant_type.as_ptr() + && unsafe { ffi::PyTuple_GET_SIZE(value.as_ptr()) } == 2 + { + // Exact two-slot namedtuple layout: slots 0/1 are direct reads. + // An exact instance with any other size falls through to the safe + // attribute path below. + let payload = unsafe { + Bound::from_borrowed_ptr(self.py, ffi::PyTuple_GET_ITEM(value.as_ptr(), 0)) + }; + let explicit_name = unsafe { + Bound::from_borrowed_ptr(self.py, ffi::PyTuple_GET_ITEM(value.as_ptr(), 1)) + }; + return self.push_explicit(&payload, &explicit_name, row); + } + + if let Some(discriminator) = + self.dispatch.types.iter().find_map(|(key, discriminator)| { + (value_type == key.as_ptr()).then_some(*discriminator) + }) + { + self.push_selected(usize::from(discriminator), value)?; + return Ok(false); + } + + // A dispatch miss checks for a TypedVariant subtype before failing. + let is_typed_variant = unsafe { + ffi::PyObject_TypeCheck( + value.as_ptr(), + self.dispatch + .typed_variant_type + .as_ptr() + .cast::(), + ) != 0 + }; + if is_typed_variant { + // A subclass (or a mis-sized exact instance) is not guaranteed to + // have the base tuple's physical two-slot layout. Match Python's + // isinstance policy through safe attribute reads; the outer + // sequence path rechecks resize after this call because these + // lookups may execute Python. + let payload = value.getattr("value")?; + let explicit_name = value.getattr("type_name")?; + self.push_explicit(&payload, &explicit_name, row)?; + return Ok(true); + } + + Err(PyValueError::new_err(format!( + "column {:?} row {row} cannot map Python type {} to any member of {}", + self.name, + python_type_name(value.as_ptr()), + self.type_name + ))) + } + + /// Push an explicitly-tagged row. The boolean reports whether the name + /// lookup may have run Python code (a str-subclass name's `__hash__` and + /// `__eq__` run during the dict lookup). + fn push_explicit( + &mut self, + payload: &Bound<'py, PyAny>, + explicit_name: &Bound<'py, PyAny>, + row: usize, + ) -> PyResult { + if unsafe { ffi::PyUnicode_Check(explicit_name.as_ptr()) } == 0 { + return Err(PyValueError::new_err(format!( + "column {:?} row {row} typed Variant name must be a str", + self.name + ))); + } + let ran_python = unsafe { ffi::PyUnicode_CheckExact(explicit_name.as_ptr()) } == 0; + if !ran_python { + if let Some(&(_, discriminator)) = self + .tag_cache + .iter() + .find(|(cached, _)| cached.as_ptr() == explicit_name.as_ptr()) + { + self.push_selected(discriminator, payload)?; + return Ok(false); + } + } + let discriminator = + unsafe { dict_index(&self.dispatch.name_index, explicit_name.as_ptr())? }.ok_or_else( + || { + let rendered = explicit_name + .extract::() + .unwrap_or_else(|_| "".to_string()); + PyValueError::new_err(format!( + "column {:?} row {row} type {rendered:?} is not a member of {}", + self.name, self.type_name + )) + }, + )?; + // Bound so a column with many distinct exact-str tag objects cannot + // grow an unbounded linear scan; misses fall back to the dict lookup. + if !ran_python && self.tag_cache.len() < 16 { + self.tag_cache + .push((explicit_name.clone().unbind(), discriminator)); + } + self.push_selected(discriminator, payload)?; + Ok(ran_python) + } + + fn push_selected(&mut self, discriminator: usize, payload: &Bound<'py, PyAny>) -> PyResult<()> { + let discriminator = u8::try_from(discriminator).map_err(|_| { + PyValueError::new_err("internal error: Variant discriminator exceeds UInt8 range") + })?; + if usize::from(discriminator) >= self.values.len() { + return Err(PyValueError::new_err( + "internal error: Variant dispatch index is out of range", + )); + } + self.discriminators.push(discriminator); + self.values[usize::from(discriminator)].push_ref(payload); + Ok(()) + } + + fn finish(self) -> PyResult { + let mut children = Vec::with_capacity(self.alternatives.len()); + for (alternative_index, (alternative, values)) in + self.alternatives.iter().zip(&self.values).enumerate() + { + children.push( + build_element_column(self.py, self.name, alternative, &values.ptrs).map_err( + |err| { + remap_variant_child_err( + self.py, + self.name, + &self.discriminators, + alternative_index, + err, + ) + }, + )?, + ); + } + let variant = VariantColumn::try_new(&self.discriminators, children).map_err(|err| { + PyValueError::new_err(format!( + "column {:?} has an invalid Variant layout: {err}", + self.name + )) + })?; + Ok(Column::Variant(variant)) + } +} + +/// Borrow a non-negative Python integer index from an exact dict lookup. +/// +/// # Safety +/// +/// `key` must be a live Python object and the GIL must be held. +unsafe fn dict_index(dict: &Bound<'_, PyDict>, key: *mut ffi::PyObject) -> PyResult> { + let value = ffi::PyDict_GetItemWithError(dict.as_ptr(), key); + if value.is_null() { + if ffi::PyErr_Occurred().is_null() { + return Ok(None); + } + return Err(PyErr::fetch(dict.py())); + } + let index = ffi::PyLong_AsSize_t(value); + if index == usize::MAX && !ffi::PyErr_Occurred().is_null() { + return Err(PyErr::fetch(dict.py())); + } + Ok(Some(index)) +} + +/// Exact CPython type name without allocating or invoking Python code. +pub(super) fn python_type_name(value: *mut ffi::PyObject) -> String { + // SAFETY: value is live and tp_name is a NUL-terminated string owned by its + // type for at least the lifetime of the object. + unsafe { std::ffi::CStr::from_ptr((*ffi::Py_TYPE(value)).tp_name) } + .to_string_lossy() + .into_owned() +} + +/// Default placeholder Python object for a null `Nullable(Tuple)` row's +/// element, matching the wire defaults the server writes for null rows. +pub(super) fn default_pyobject(py: Python<'_>, ch_type: &ChType) -> PyResult> { + if let Some(delegate) = ch_type.physical_delegate() { + return default_pyobject(py, &delegate); + } + Ok(match ch_type { + // A Dynamic placeholder stays None; the String insert path renders it + // as the literal "NULL". + ChType::Nullable(_) | ChType::Variant(_) | ChType::Dynamic { .. } => py.None(), + ChType::LowCardinality(inner) => default_pyobject(py, inner)?, + ChType::Nothing => py.None(), + ChType::Bool => PyBool::new(py, false).to_owned().into_any().unbind(), + ChType::String | ChType::FixedString(_) => PyString::new(py, "").into_any().unbind(), + ChType::QBit { dimension, .. } => PyList::new(py, std::iter::repeat_n(0.0, *dimension))? + .into_any() + .unbind(), + ChType::Array(_) => PyList::empty(py).into_any().unbind(), + ChType::Map(..) => PyDict::new(py).into_any().unbind(), + ChType::Json { .. } => PyDict::new(py).into_any().unbind(), + ChType::Tuple(elements) => { + let defaults = elements + .iter() + .map(|(_, element_type)| default_pyobject(py, element_type)) + .collect::>>()?; + PyTuple::new(py, defaults)?.into_any().unbind() + } + ChType::AggregateFunction { .. } => { + return Err(PyNotImplementedError::new_err( + "null Nullable(Tuple(...)) rows containing AggregateFunction are unsupported until the core provides a canonical placeholder state", + )); + } + // Every remaining scalar accepts a raw 0 (epoch units, enum code, + // integer UUID/IP forms, Decimal("0")). + _ => 0i64.into_pyobject(py)?.into_any().unbind(), + }) +} diff --git a/rust/ch-core-py/src/lib.rs b/rust/ch-core-py/src/lib.rs new file mode 100644 index 00000000..c7ab943b --- /dev/null +++ b/rust/ch-core-py/src/lib.rs @@ -0,0 +1,21 @@ +use pyo3::prelude::*; + +mod batch; +mod decoder; +mod insert; +mod pyval; + +#[pymodule] +fn _ch_core(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + // Binding API contract number checked by clickhouse_connect/driver/rustcodec.py + // at import; bump when the Python-visible binding surface changes incompatibly. + m.add("BINDING_API_VERSION", 2)?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + #[cfg(unix)] + m.add_class::()?; + m.add_function(wrap_pyfunction!(insert::encode_native_block, m)?)?; + Ok(()) +} diff --git a/rust/ch-core-py/src/pyval/containers.rs b/rust/ch-core-py/src/pyval/containers.rs new file mode 100644 index 00000000..9a24ffb7 --- /dev/null +++ b/rust/ch-core-py/src/pyval/containers.rs @@ -0,0 +1,226 @@ +use super::*; + +/// Column-major fill for a Tuple column: allocate every row container up +/// front (tuple, presized dict, or None for a null row), hand each to `sink`, +/// then fill field by field through `fill_column`, one dispatch per field. +/// +/// # Safety +/// +/// Requires the GIL; `fill_column`'s sink contract applies. Containers are +/// filled after they reach the sink, so the sink keeping items alive until +/// this call returns is load-bearing. +pub(super) unsafe fn fill_tuple<'py>( + py: Python<'py>, + c: &TupleColumn, + ctx: &ColumnCtx<'py>, + rows: usize, + sink: DynSink<'_>, +) -> PyResult<()> { + let fctx = ctx.fields.as_deref().ok_or_else(|| ctx_missing("Tuple"))?; + if fctx.len() != c.fields.len() { + return Err(ctx_count_mismatch("Tuple")); + } + if c.fields.iter().any(|f| f.len() < rows) { + return Err(tuple_shape_err()); + } + let validity = c.validity.as_ref(); + let num_fields = c.fields.len() as ffi::Py_ssize_t; + let names = ctx.tuple_names.as_deref(); + + // Borrowed container pointers; the sink owns them and keeps them alive. + let mut containers: Vec<*mut ffi::PyObject> = Vec::with_capacity(rows); + for i in 0..rows { + let ptr = if validity.is_some_and(|bm| !bm.is_valid(i)) { + none_owned_ptr() + } else if names.is_some() { + ptr_to_result(py, dict_new_presized(num_fields))? + } else { + ptr_to_result(py, ffi::PyTuple_New(num_fields))? + }; + containers.push(ptr); + sink(i, ptr); + } + + // Items produced for null rows collect here and drop only after the field + // fill returns, keeping the sink's items-stay-alive contract. + let mut discarded: Vec> = Vec::new(); + for (field_idx, (field_col, field_ctx)) in c.fields.iter().zip(fctx).enumerate() { + match names { + None => { + let mut field_sink = |i: usize, item: *mut ffi::PyObject| { + if validity.is_some_and(|bm| !bm.is_valid(i)) { + // Safety: item is an owned reference the sink takes over. + discarded.push(unsafe { Bound::from_owned_ptr(py, item) }.unbind()); + return; + } + // Safety: containers[i] is a live tuple with num_fields + // slots; the tuple takes over the owned item. + unsafe { + ffi::PyTuple_SET_ITEM(containers[i], field_idx as ffi::Py_ssize_t, item); + } + }; + let mut erased: DynSink<'_> = &mut field_sink; + fill_column(py, field_col, field_ctx, rows, &mut erased)?; + } + Some(names) => { + let name_ptr = names[field_idx].as_ptr(); + let mut err: Option = None; + let mut field_sink = |i: usize, item: *mut ffi::PyObject| { + // Safety: item is an owned reference the sink takes over. + let item = unsafe { Bound::from_owned_ptr(py, item) }.unbind(); + if err.is_some() || validity.is_some_and(|bm| !bm.is_valid(i)) { + discarded.push(item); + return; + } + // Safety: containers[i] is a live dict and name_ptr a live + // str key; SetItem increfs both, our `item` ref drops after. + if unsafe { ffi::PyDict_SetItem(containers[i], name_ptr, item.as_ptr()) } < 0 { + err = Some(PyErr::fetch(py)); + discarded.push(item); + } + }; + let mut erased: DynSink<'_> = &mut field_sink; + fill_column(py, field_col, field_ctx, rows, &mut erased)?; + if let Some(e) = err { + return Err(e); + } + } + } + } + Ok(()) +} + +/// Column-major fill for a Map column: validate the offsets run once, +/// materialize the flat key and value runs through `fill_column`, then zip +/// each row into a presized dict in wire order (last duplicate key wins). +/// +/// # Safety +/// +/// Requires the GIL; `fill_column`'s sink contract applies. +pub(super) unsafe fn fill_map<'py>( + py: Python<'py>, + c: &MapColumn, + ctx: &ColumnCtx<'py>, + rows: usize, + sink: DynSink<'_>, +) -> PyResult<()> { + if rows == 0 { + return Ok(()); + } + let fctx = ctx.fields.as_deref().ok_or_else(|| ctx_missing("Map"))?; + if fctx.len() != 2 { + return Err(ctx_missing("Map")); + } + let entries = match c.entries.as_ref() { + Column::Tuple(t) if t.fields.len() == 2 => t, + _ => return Err(map_entries_err()), + }; + let keys_col = &entries.fields[0]; + let values_col = &entries.fields[1]; + if c.offsets.len() <= rows { + return Err(map_bounds_err()); + } + // One monotonicity pass over the offsets; the per-row casts below are + // then in range. + let offsets = &c.offsets[..=rows]; + let mut prev: i64 = 0; + for &o in offsets { + if o < prev { + return Err(map_bounds_err()); + } + prev = o; + } + let total = offsets[rows] as usize; + if total > keys_col.len().min(values_col.len()) { + return Err(map_bounds_err()); + } + + let keys = materialize_run(py, keys_col, &fctx[0], total)?; + let values = materialize_run(py, values_col, &fctx[1], total)?; + + for (i, pair) in offsets.windows(2).enumerate() { + let (start, end) = (pair[0] as usize, pair[1] as usize); + let dict_ptr = dict_new_presized((end - start) as ffi::Py_ssize_t); + if dict_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: dict_ptr is an owned dict; binding it drops the + // partially-filled dict on the error path. + let dict = Bound::from_owned_ptr(py, dict_ptr); + for slot in start..end { + // SetItem increfs key and value; the run vectors keep our refs. + if ffi::PyDict_SetItem(dict.as_ptr(), keys[slot].as_ptr(), values[slot].as_ptr()) < 0 { + return Err(PyErr::fetch(py)); + } + } + sink(i, dict.into_ptr()); + } + Ok(()) +} + +/// The two flat Float64 runs of an unnamed non-nullable +/// `Tuple(Float64, Float64)` element column (the geo point shape), when the +/// tight point-list loop applies. +pub(super) fn point_pair_slices<'a>( + values: &'a Column, + ctx: &ColumnCtx<'_>, +) -> Option<(&'a [f64], &'a [f64])> { + let Column::Tuple(t) = values else { + return None; + }; + if t.validity.is_some() || ctx.tuple_names.is_some() || t.fields.len() != 2 { + return None; + } + let (Column::Float64(x), Column::Float64(y)) = (&t.fields[0], &t.fields[1]) else { + return None; + }; + if x.validity.is_some() || y.validity.is_some() { + return None; + } + Some((&x.values, &y.values)) +} + +/// Build one Array row's point list from the flat coordinate runs: one +/// presized list, one 2-tuple and two floats per point, with no per-element +/// column dispatch. Allocation stays row-major (tuple, then its floats), the +/// same order as the generic per-cell path; a column-major variant measured +/// faster to fill but slower overall because deallocation and GC traversal of +/// the interleaved result lose locality. +/// +/// # Safety +/// +/// Requires the GIL. Returns an owned reference the caller takes over. +/// `start..end` must be in range for both slices. +pub(super) unsafe fn point_list_owned_ptr( + py: Python<'_>, + xs: &[f64], + ys: &[f64], + start: usize, + end: usize, +) -> PyResult<*mut ffi::PyObject> { + let count = end - start; + let list_ptr = ffi::PyList_New(count as ffi::Py_ssize_t); + if list_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: list_ptr came from PyList_New, so it is a list and this is the + // sole owned reference. Binding it makes the error path drop the + // partially-filled list; list_dealloc tolerates the NULL slots. + let list = Bound::from_owned_ptr(py, list_ptr).cast_into_unchecked::(); + for (slot, k) in (start..end).enumerate() { + let tuple_ptr = ffi::PyTuple_New(2); + if tuple_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: slot < count, the list's allocated length; the list takes + // over the owned tuple, so an error below drops it through the list + // (tuple_dealloc tolerates its NULL slots). + ffi::PyList_SET_ITEM(list.as_ptr(), slot as ffi::Py_ssize_t, tuple_ptr); + let x = ptr_to_result(py, ffi::PyFloat_FromDouble(xs[k]))?; + // Safety: the fresh 2-tuple takes over each owned float. + ffi::PyTuple_SET_ITEM(tuple_ptr, 0, x); + let y = ptr_to_result(py, ffi::PyFloat_FromDouble(ys[k]))?; + ffi::PyTuple_SET_ITEM(tuple_ptr, 1, y); + } + Ok(list.into_ptr()) +} diff --git a/rust/ch-core-py/src/pyval/ctx.rs b/rust/ch-core-py/src/pyval/ctx.rs new file mode 100644 index 00000000..0a8a3425 --- /dev/null +++ b/rust/ch-core-py/src/pyval/ctx.rs @@ -0,0 +1,343 @@ +use super::*; + +// --------------------------------------------------------------------------- +// Temporal value construction +// +// The core decodes temporal columns to their faithful native integer width and +// keeps timezone and precision in the schema's ChType, not in the buffers. The +// binding turns those integers into Python date/datetime objects here, matching +// the value policy of clickhouse-connect's Native reader: +// +// Date, Date32 -> datetime.date (naive) +// DateTime / DateTime64, no tz or -> datetime.datetime, no tzinfo, +// a UTC-equivalent tz decoded as the UTC wall clock +// DateTime / DateTime64, named non-UTC -> datetime.datetime, tz-aware +// tz (zoneinfo.ZoneInfo) +// +// The naive path is pure epoch arithmetic with no per-row Python datetime +// parsing. The tz-aware path defers to datetime.fromtimestamp, which handles +// DST for the named zone. +// --------------------------------------------------------------------------- + +/// Timezone names ClickHouse treats as UTC. A column in one of these renders as +/// a naive datetime, matching clickhouse-connect (tzutil.UTC_EQUIVALENTS). +const UTC_EQUIVALENTS: &[&str] = &[ + "UTC", + "Etc/UTC", + "UCT", + "Etc/UCT", + "GMT", + "Etc/GMT", + "GMT0", + "GMT-0", + "GMT+0", + "Etc/GMT0", + "Etc/GMT-0", + "Etc/GMT+0", + "Universal", + "Etc/Universal", + "Zulu", + "Etc/Zulu", + "Greenwich", + "Etc/Greenwich", +]; + +/// Per-column context for materializing a column's host values, resolved once +/// per column rather than per cell. For a temporal column it carries the +/// timezone policy: a naive column (Date/Date32, or a DateTime/DateTime64 with +/// no timezone or a UTC-equivalent one) has `tz` `None` and is built by epoch +/// arithmetic, while a named non-UTC timezone holds its `zoneinfo.ZoneInfo` in +/// `tz` and the bound `datetime.datetime.fromtimestamp` in `fromtimestamp`. For +/// an Enum8/Enum16 column it carries `enum_names`, the value -> label-string map. +pub(crate) struct ColumnCtx<'py> { + pub(super) tz: Option>, + pub(super) fromtimestamp: Option>, + /// DateTime64/Time64 fractional precision (decimal digits). 0 otherwise. + pub(super) precision: u8, + /// Time64 ticks per second, precomputed once per column. 1 otherwise. + pub(super) time_scale: u64, + /// Materialize Time/Time64 leaves as raw signed ticks instead of timedelta. + pub(super) raw_time_ticks: bool, + /// Enum value -> pre-built label string, for an Enum8/Enum16 column; `None` + /// for any other type. A value missing from the map materializes as None, + /// matching clickhouse-connect's `int_map.get(value, None)`. + pub(super) enum_names: Option>>, + /// `uuid.UUID` construction machinery, for a UUID column. + pub(super) uuid: Option>, + /// `ipaddress` class machinery, for an IPv4/IPv6 column. + pub(super) ip: Option>, + /// The `decimal.Decimal` class, for a Decimal column. + pub(super) decimal_cls: Option>, + /// Recursively-prepared context for the element type of an Array column; + /// `None` for any other type. + pub(super) element: Option>>, + /// Recursively-prepared per-field contexts. For a Tuple column, one per + /// element in declaration order; for a Map column, exactly two (the key + /// context then the value context); for a Variant column, one per dense + /// alternative in the server's canonical discriminator order. JSON uses + /// one entry per declared typed path. Dynamic children are block-local, so + /// their contexts are prepared once per child in `fill_dynamic` instead of + /// being stored here. `None` for any other type. + pub(super) fields: Option>>, + /// Pre-built element-name keys for a NAMED Tuple column, materialized as a + /// `dict` keyed by these (clickhouse-connect's default read format). + /// `None` for an unnamed Tuple (materialized as a `tuple`) and every + /// non-Tuple type. + pub(super) tuple_names: Option>>, + /// Pre-split, percent-decoded keys for each declared JSON typed path. + /// Dynamic and shared paths are block-local and are prepared once per path + /// by `fill_json`. + pub(super) json_paths: Option>>, +} + +/// One ClickHouse JSON path prepared for repeated insertion into Python dicts. +/// Splitting precedes percent decoding because `%2E` represents a literal dot +/// inside one key rather than a nesting separator. `raw` keeps the declared +/// path spelling for the typed-path order guard. +pub(super) struct JsonPath<'py> { + pub(super) raw: String, + pub(super) keys: Vec>, +} + +/// Cached objects to build a `uuid.UUID` the way the Cython codec does: +/// allocate via `UUID.__new__` and set the fields with `object.__setattr__`, +/// bypassing the parsing constructor and the immutability guard. +pub(super) struct UuidCtx<'py> { + pub(super) cls: Bound<'py, PyAny>, + pub(super) new: Bound<'py, PyAny>, + pub(super) object_setattr: Bound<'py, PyAny>, + pub(super) unsafe_marker: Bound<'py, PyAny>, +} + +/// Cached objects to build an `ipaddress.IPv4Address`/`IPv6Address` via +/// `__new__` plus a plain `_ip` setattr (neither class guards setattr). +pub(super) struct IpCtx<'py> { + pub(super) cls: Bound<'py, PyAny>, + pub(super) new: Bound<'py, PyAny>, + /// `IPv6Address.__slots__` includes `_scope_id` (Python 3.9+); each value + /// gets `_scope_id = None`. + pub(super) set_scope_id: bool, +} + +/// Build an enum's value -> label-string map, one Python str per variant created +/// once for the whole column. +fn enum_name_map<'py, V: Copy + Into>( + py: Python<'py>, + variants: &[(String, V)], +) -> HashMap> { + variants + .iter() + .map(|(name, value)| ((*value).into(), PyString::new(py, name))) + .collect() +} + +pub(super) fn prepare_json_path<'py>(py: Python<'py>, path: &str) -> PyResult> { + let unquote = path + .contains('%') + .then(|| py.import("urllib.parse")?.getattr("unquote")) + .transpose()?; + let keys = path + .split('.') + .map(|segment| match &unquote { + Some(unquote) => unquote + .call1((segment,))? + .cast_into::() + .map_err(Into::into), + None => Ok(PyString::new(py, segment)), + }) + .collect::>>()?; + Ok(JsonPath { + raw: path.to_owned(), + keys, + }) +} + +/// Resolve a column's ChType into a ColumnCtx. Cheap for the common naive case +/// (no Python calls); imports zoneinfo only for a named non-UTC zone and builds +/// the enum map only for an enum. Safe to call for any column type; types that +/// need neither yield the naive, non-enum default. +pub(crate) fn prepare_column_ctx<'py>( + py: Python<'py>, + ch_type: &ChType, + raw_time_ticks: bool, +) -> PyResult> { + // Unwrap LowCardinality to reach the value type before inner() strips any + // Nullable, so LowCardinality(DateTime(tz)) still applies timezone policy. + let resolved = match ch_type { + ChType::LowCardinality(inner) => inner.inner(), + other => other.inner(), + }; + + // SimpleAggregateFunction, the geo aliases, and Nested carry only a custom + // name over a physical type; the decoded Column is that physical type's + // column, so build the context from the delegate and recurse. `resolved` + // has already stripped any LowCardinality/Nullable, so Nullable(Point) + // reaches here as Geo(Point). + if let Some(delegate) = resolved.physical_delegate() { + return prepare_column_ctx(py, &delegate, raw_time_ticks); + } + + let enum_names = match resolved { + ChType::Enum8 { variants } => Some(enum_name_map(py, variants)), + ChType::Enum16 { variants } => Some(enum_name_map(py, variants)), + _ => None, + }; + + let (timezone, precision, time_scale) = match resolved { + ChType::DateTime { timezone } => (timezone.as_deref(), 0u8, 1), + ChType::DateTime64 { + precision, + timezone, + } => (timezone.as_deref(), *precision, 1), + ChType::Time64 { precision } => (None, *precision, 10u64.pow(u32::from(*precision))), + _ => (None, 0u8, 1), + }; + + let (tz, fromtimestamp) = match timezone { + Some(tz) if !UTC_EQUIVALENTS.contains(&tz) => { + let zone = py.import("zoneinfo")?.getattr("ZoneInfo")?.call1((tz,))?; + let fromtimestamp = py + .import("datetime")? + .getattr("datetime")? + .getattr("fromtimestamp")?; + (Some(zone), Some(fromtimestamp)) + } + _ => (None, None), + }; + + let (uuid, ip, decimal_cls) = match resolved { + ChType::Uuid => { + let module = py.import("uuid")?; + let cls = module.getattr("UUID")?; + let new = cls.getattr("__new__")?; + let object_setattr = py + .import("builtins")? + .getattr("object")? + .getattr("__setattr__")?; + let unsafe_marker = module.getattr("SafeUUID")?.getattr("unsafe")?; + ( + Some(UuidCtx { + cls, + new, + object_setattr, + unsafe_marker, + }), + None, + None, + ) + } + ChType::Ipv4 => { + let cls = py.import("ipaddress")?.getattr("IPv4Address")?; + let new = cls.getattr("__new__")?; + ( + None, + Some(IpCtx { + cls, + new, + set_scope_id: false, + }), + None, + ) + } + ChType::Ipv6 => { + let cls = py.import("ipaddress")?.getattr("IPv6Address")?; + let new = cls.getattr("__new__")?; + let set_scope_id = cls + .getattr("__slots__")? + .contains(intern!(py, "_scope_id"))?; + ( + None, + Some(IpCtx { + cls, + new, + set_scope_id, + }), + None, + ) + } + ChType::Decimal { .. } => (None, None, Some(py.import("decimal")?.getattr("Decimal")?)), + _ => (None, None, None), + }; + + // `.inner()` only strips Nullable, so an Array column's resolved type is the + // `Array(elem)` itself. Recurse into the element to build its machinery, + // which transparently covers every element shape (Nullable, LowCardinality, + // nested Array, temporal-with-tz, enum, uuid, ip, decimal). + let element = match resolved { + ChType::Array(elem) => Some(Box::new(prepare_column_ctx(py, elem, raw_time_ticks)?)), + _ => None, + }; + + // A Tuple builds one field context per element and, for a named tuple, the + // pre-built name keys. A Map builds exactly two field contexts (key then + // value); its entries live in a nested Tuple column reached directly, so it + // needs no tuple_names. `resolved` is the Nullable-unwrapped type, so a + // `Nullable(Tuple(...))` reaches the Tuple arm here. + let (fields, tuple_names, json_paths) = match resolved { + ChType::Tuple(elements) => { + let fields = elements + .iter() + .map(|(_, t)| prepare_column_ctx(py, t, raw_time_ticks)) + .collect::>>()?; + let named = !elements.is_empty() && elements.iter().all(|(name, _)| name.is_some()); + // `named` guarantees every element has a name; the empty-string + // fallback keeps the map total (one key per field) without an + // unwrap even though it is never taken. + let names = if named { + Some( + elements + .iter() + .map(|(name, _)| PyString::new(py, name.as_deref().unwrap_or_default())) + .collect(), + ) + } else { + None + }; + (Some(fields), names, None) + } + ChType::Map(key, value) => { + let fields = vec![ + prepare_column_ctx(py, key, raw_time_ticks)?, + prepare_column_ctx(py, value, raw_time_ticks)?, + ]; + (Some(fields), None, None) + } + // Variant alternatives always materialize finalized values; the + // driver's raw-ticks materializer does not walk Variant cells. + ChType::Variant(alternatives) => { + let fields = alternatives + .iter() + .map(|alternative| prepare_column_ctx(py, alternative, false)) + .collect::>>()?; + (Some(fields), None, None) + } + ChType::Json { typed_paths, .. } => { + let fields = typed_paths + .iter() + .map(|(_, ch_type)| prepare_column_ctx(py, ch_type, raw_time_ticks)) + .collect::>>()?; + let paths = typed_paths + .iter() + .map(|(path, _)| prepare_json_path(py, path)) + .collect::>>()?; + (Some(fields), None, Some(paths)) + } + _ => (None, None, None), + }; + + Ok(ColumnCtx { + tz, + fromtimestamp, + precision, + time_scale, + raw_time_ticks, + enum_names, + uuid, + ip, + decimal_cls, + element, + fields, + tuple_names, + json_paths, + }) +} diff --git a/rust/ch-core-py/src/pyval/errors.rs b/rust/ch-core-py/src/pyval/errors.rs new file mode 100644 index 00000000..07d58e32 --- /dev/null +++ b/rust/ch-core-py/src/pyval/errors.rs @@ -0,0 +1,69 @@ +use super::*; + +/// Error for a malformed Dynamic SharedVariant cell. +pub(super) fn shared_cell_err(err: &BinaryValueError) -> PyErr { + PyValueError::new_err(format!( + "Malformed payload: invalid Dynamic SharedVariant cell: {err}" + )) +} + +/// Error for a LowCardinality index outside its dictionary. +pub(super) fn lc_index_err() -> PyErr { + PyValueError::new_err("Malformed payload: LowCardinality index out of dictionary range") +} + +/// Error for a column whose ColumnCtx was prepared for a different type; an +/// internal invariant violation, not a payload condition. +pub(super) fn ctx_missing(what: &str) -> PyErr { + PyValueError::new_err(format!("internal error: missing {what} column context")) +} + +/// Error for a container column whose ColumnCtx carries a different number of +/// field contexts than the column has fields; an internal invariant violation. +pub(super) fn ctx_count_mismatch(what: &str) -> PyErr { + PyValueError::new_err(format!( + "internal error: {what} field context count mismatch" + )) +} + +/// Error for an Array column whose offsets are out of range for the element +/// buffer (out of order, negative, or an end past the element count). +pub(super) fn array_bounds_err() -> PyErr { + PyValueError::new_err("Malformed payload: Array offsets are out of range") +} + +/// Error for a Map column whose offsets are out of range for the entries +/// buffer (out of order, negative, or an end past the entry count). +pub(super) fn map_bounds_err() -> PyErr { + PyValueError::new_err("Malformed payload: Map offsets are out of range") +} + +/// Error for a Tuple column whose field columns are shorter than the row count. +pub(super) fn tuple_shape_err() -> PyErr { + PyValueError::new_err("Malformed payload: Tuple field length mismatch") +} + +/// Error for a Map column whose entries are not a two-field key/value tuple; +/// an internal invariant violation (the core always builds this shape). +pub(super) fn map_entries_err() -> PyErr { + PyValueError::new_err("internal error: Map entries are not a two-field tuple") +} + +/// Error for malformed Variant dense-union routing or child lengths. +pub(super) fn variant_shape_err() -> PyErr { + PyValueError::new_err("Malformed payload: Variant routing or child length mismatch") +} + +/// Error for malformed Dynamic dense-union routing or child lengths. +pub(super) fn dynamic_shape_err() -> PyErr { + PyValueError::new_err("Malformed payload: Dynamic routing or child length mismatch") +} + +pub(super) fn json_shape_err() -> PyErr { + PyValueError::new_err("Malformed payload: JSON path or child layout mismatch") +} + +/// Error for a UUID/IPv6 cell whose fixed width is not 16 bytes. +pub(super) fn fixed_width_err(what: &str) -> PyErr { + PyValueError::new_err(format!("Malformed payload: {what} cell is not 16 bytes")) +} diff --git a/rust/ch-core-py/src/pyval/fixed.rs b/rust/ch-core-py/src/pyval/fixed.rs new file mode 100644 index 00000000..7bcf0bc3 --- /dev/null +++ b/rust/ch-core-py/src/pyval/fixed.rs @@ -0,0 +1,472 @@ +use super::*; + +/// The column's validity bitmap, if any. +pub(super) fn column_validity(col: &Column) -> Option<&Bitmap> { + match col { + Column::Bool(c) => c.validity.as_ref(), + Column::Int8(c) => c.validity.as_ref(), + Column::Int16(c) => c.validity.as_ref(), + Column::Int32(c) => c.validity.as_ref(), + Column::Int64(c) => c.validity.as_ref(), + Column::UInt8(c) => c.validity.as_ref(), + Column::UInt16(c) => c.validity.as_ref(), + Column::UInt32(c) => c.validity.as_ref(), + Column::UInt64(c) => c.validity.as_ref(), + Column::Float32(c) => c.validity.as_ref(), + Column::Float64(c) => c.validity.as_ref(), + Column::BFloat16(c) => c.validity.as_ref(), + Column::QBit(c) => c.validity.as_ref(), + Column::Nothing(c) => c.validity.as_ref(), + Column::Date(c) => c.validity.as_ref(), + Column::Date32(c) => c.validity.as_ref(), + Column::DateTime(c) => c.validity.as_ref(), + Column::DateTime64(c) => c.validity.as_ref(), + Column::Time(c) => c.validity.as_ref(), + Column::Time64(c) => c.validity.as_ref(), + Column::Interval(c) => c.validity.as_ref(), + Column::Utf8(c) => c.validity.as_ref(), + Column::AggregateState(_) => None, + Column::FixedBinary(c) => c.validity.as_ref(), + Column::Ipv4(c) => c.validity.as_ref(), + Column::Ipv6(c) => c.validity.as_ref(), + Column::Uuid(c) => c.validity.as_ref(), + Column::Enum8(c) => c.validity.as_ref(), + Column::Enum16(c) => c.validity.as_ref(), + Column::Int128(c) => c.validity.as_ref(), + Column::UInt128(c) => c.validity.as_ref(), + Column::Int256(c) => c.validity.as_ref(), + Column::UInt256(c) => c.validity.as_ref(), + // A LowCardinality column's nulls live in the index validity, the Arrow + // dictionary convention, not as a dictionary entry. + Column::Dictionary(c) => c.validity.as_ref(), + Column::Decimal(c) => c.validity.as_ref(), + // Arrays carry no array-level validity; element nulls live on `values`. + Column::Array(_) => None, + // A Nullable(Tuple) carries tuple-level validity here; element nulls + // live on the field columns. A plain Tuple has `validity == None`. + Column::Tuple(c) => c.validity.as_ref(), + // Dynamic has intrinsic NULL rows in its dense-union routing buffers, + // not a top-level validity bitmap. + Column::Dynamic(_) => None, + // Maps are never nullable at the map level; value nulls live on the + // values column inside `entries`. + Column::Map(_) => None, + // Variant has intrinsic NULL rows in its dense-union routing buffers, + // not a top-level validity bitmap. + Column::Variant(_) => None, + // JSON only carries top-level validity under Nullable(JSON). + Column::Json(c) => c.validity.as_ref(), + } +} + +/// Widen one exact little-endian ClickHouse BFloat16 word to Float32. +#[inline] +pub(super) fn bfloat16_to_f32(word: [u8; 2]) -> f32 { + f32::from_bits(u32::from(u16::from_le_bytes(word)) << 16) +} + +/// Tight per-value loop for one primitive column: `make` is the per-value FFI +/// constructor, resolved once by the caller's variant dispatch rather than per +/// cell. The nullable arm checks the bitmap per cell but keeps the single +/// dispatch. +/// +/// # Safety +/// +/// Requires the GIL. `make` must return an owned reference or null; each +/// pointer passed to `sink` is an owned reference the sink must take over +/// exactly once. +unsafe fn fill_prim( + py: Python<'_>, + values: &[T], + validity: Option<&Bitmap>, + make: F, + sink: &mut S, +) -> PyResult<()> +where + T: Copy, + F: Fn(T) -> *mut ffi::PyObject, + S: FnMut(usize, *mut ffi::PyObject), +{ + match validity { + None => { + for (i, &v) in values.iter().enumerate() { + let item = make(v); + if item.is_null() { + return Err(PyErr::fetch(py)); + } + sink(i, item); + } + } + Some(bm) => { + for (i, &v) in values.iter().enumerate() { + let item = if bm.is_valid(i) { + let made = make(v); + if made.is_null() { + return Err(PyErr::fetch(py)); + } + made + } else { + none_owned_ptr() + }; + sink(i, item); + } + } + } + Ok(()) +} + +/// Tight per-cell loop for a column whose constructor works by row index: +/// `make` is the per-cell builder, with its ctx lookups hoisted by the caller. +/// The nullable arm checks the bitmap per cell but keeps the single dispatch. +/// +/// # Safety +/// +/// Requires the GIL. `make` must return an owned reference on Ok; each pointer +/// passed to `sink` is an owned reference the sink must take over exactly once. +unsafe fn fill_indexed( + rows: usize, + validity: Option<&Bitmap>, + mut make: F, + sink: &mut S, +) -> PyResult<()> +where + F: FnMut(usize) -> PyResult<*mut ffi::PyObject>, + S: FnMut(usize, *mut ffi::PyObject), +{ + match validity { + None => { + for i in 0..rows { + let item = make(i)?; + sink(i, item); + } + } + Some(bm) => { + for i in 0..rows { + let item = if bm.is_valid(i) { + make(i)? + } else { + none_owned_ptr() + }; + sink(i, item); + } + } + } + Ok(()) +} + +/// Materialize the first `rows` cells of a fixed-width column into `sink`, +/// dispatching the Column variant once and iterating the values buffer +/// directly, with any per-column ctx lookups hoisted out of the loop. Returns +/// Ok(false), touching nothing, for a variant with no fast path (strings, +/// temporal, enum, LowCardinality, ...); those stay on the per-cell route. +/// +/// # Safety +/// +/// Requires the GIL. Each pointer passed to `sink` is an owned reference the +/// sink must take over exactly once. +pub(super) unsafe fn fill_fixed_width( + py: Python<'_>, + col: &Column, + ctx: &ColumnCtx<'_>, + rows: usize, + sink: &mut S, +) -> PyResult +where + S: FnMut(usize, *mut ffi::PyObject), +{ + // Safety, for the constructor closures below: pure CPython constructors + // called with the GIL held; fill_prim null-checks every result. + match col { + Column::Int8(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLong(c_long::from(v)) }, + sink, + )?, + Column::Int16(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLong(c_long::from(v)) }, + sink, + )?, + Column::Int32(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLong(c_long::from(v)) }, + sink, + )?, + Column::Int64(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLongLong(v) }, + sink, + )?, + Column::UInt8(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLong(c_long::from(v)) }, + sink, + )?, + Column::UInt16(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLong(c_long::from(v)) }, + sink, + )?, + Column::UInt32(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromUnsignedLongLong(v.into()) }, + sink, + )?, + Column::UInt64(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromUnsignedLongLong(v) }, + sink, + )?, + Column::Int128(c) | Column::Int256(c) => fill_indexed( + rows, + c.validity.as_ref(), + |i| unsafe { wide_int_value_ptr(py, c.value(i), true) }, + sink, + )?, + Column::UInt128(c) | Column::UInt256(c) => fill_indexed( + rows, + c.validity.as_ref(), + |i| unsafe { wide_int_value_ptr(py, c.value(i), false) }, + sink, + )?, + Column::Float32(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyFloat_FromDouble(v.into()) }, + sink, + )?, + Column::Float64(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyFloat_FromDouble(v) }, + sink, + )?, + Column::BFloat16(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyFloat_FromDouble(bfloat16_to_f32(v).into()) }, + sink, + )?, + Column::Interval(c) => fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLongLong(v) }, + sink, + )?, + Column::Time(c) => { + if ctx.raw_time_ticks { + fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLong(c_long::from(v)) }, + sink, + )? + } else { + fill_indexed( + rows, + c.validity.as_ref(), + |i| Ok(PyDelta::new(py, 0, c.values[i], 0, true)?.into_ptr()), + sink, + )? + } + } + Column::Time64(c) => { + if ctx.raw_time_ticks { + fill_prim( + py, + &c.values[..rows], + c.validity.as_ref(), + |v| unsafe { ffi::PyLong_FromLongLong(v) }, + sink, + )? + } else { + let scale = ctx.time_scale; + fill_indexed( + rows, + c.validity.as_ref(), + |i| Ok(make_time64(py, c.values[i], scale)?.into_ptr()), + sink, + )? + } + } + Column::Bool(c) => match &c.validity { + None => { + for i in 0..rows { + sink(i, bool_owned_ptr(c.get(i))); + } + } + Some(bm) => { + for i in 0..rows { + let item = if bm.is_valid(i) { + bool_owned_ptr(c.get(i)) + } else { + none_owned_ptr() + }; + sink(i, item); + } + } + }, + Column::Uuid(c) => { + let uctx = ctx.uuid.as_ref().ok_or_else(|| ctx_missing("UUID"))?; + fill_indexed( + rows, + c.validity.as_ref(), + |i| uuid_value_ptr(py, uctx, c.value(i)), + sink, + )? + } + Column::Ipv4(c) => { + let ictx = ctx.ip.as_ref().ok_or_else(|| ctx_missing("IPv4"))?; + fill_indexed( + rows, + c.validity.as_ref(), + |i| ipv4_value_ptr(py, ictx, c.values[i]), + sink, + )? + } + Column::Ipv6(c) => { + let ictx = ctx.ip.as_ref().ok_or_else(|| ctx_missing("IPv6"))?; + fill_indexed( + rows, + c.validity.as_ref(), + |i| ipv6_value_ptr(py, ictx, c.value(i)), + sink, + )? + } + Column::Decimal(c) => { + let cls = ctx + .decimal_cls + .as_ref() + .ok_or_else(|| ctx_missing("Decimal"))?; + let mut scratch = DecimalScratch::default(); + fill_indexed( + rows, + c.validity.as_ref(), + |i| decimal_value_ptr(cls, &mut scratch, c, i), + sink, + )? + } + _ => return Ok(false), + } + Ok(true) +} + +/// Materialize serialized AggregateFunction states as exact Python bytes. +/// The core recovers state boundaries into one offsets run and one contiguous +/// data buffer; the offsets content is trusted as decoder-constructed, the +/// same model as the Utf8 and FixedBinary arms. Python necessarily copies +/// each state into its bytes object; the Arrow exit remains zero-copy for +/// columnar consumers. +/// +/// # Safety +/// +/// Requires the GIL. Each pointer passed to `sink` is an owned reference the +/// sink must take over exactly once. +pub(super) unsafe fn fill_aggregate_states( + py: Python<'_>, + col: &AggregateStateColumn, + rows: usize, + sink: &mut S, +) -> PyResult<()> +where + S: FnMut(usize, *mut ffi::PyObject), +{ + if rows == 0 { + return Ok(()); + } + if col.offsets.len() <= rows { + return Err(PyValueError::new_err( + "Malformed payload: invalid AggregateFunction state offsets", + )); + } + let mut start = col.offsets[0] as usize; + for (row, &end) in col.offsets[1..=rows].iter().enumerate() { + let end = end as usize; + sink(row, bytes_owned_ptr(py, &col.data[start..end])?); + start = end; + } + Ok(()) +} + +/// Materialize a dictionary (LowCardinality) column into `sink`: build each +/// referenced dictionary value once through `column_value_nonnull_ptr`, then +/// emit every cell as an INCREF of its cached object — the python codec's +/// object-reuse policy. Nulls live in the index validity (Arrow dictionary +/// convention); invalid cells emit None without touching the dictionary. The +/// cache fills lazily on first reference by a valid index, so an all-null +/// column over an inner type with no object-exit support still reads as None +/// and unreferenced slots cost nothing. +/// +/// # Safety +/// +/// Requires the GIL. Each pointer passed to `sink` is an owned reference the +/// sink must take over exactly once. +pub(super) unsafe fn fill_dictionary<'py, S>( + py: Python<'py>, + col: &DictionaryColumn, + ctx: &ColumnCtx<'py>, + rows: usize, + sink: &mut S, +) -> PyResult<()> +where + S: FnMut(usize, *mut ffi::PyObject), +{ + let mut cache: Vec>> = Vec::with_capacity(col.values.len()); + cache.resize_with(col.values.len(), || None); + let cached_item = + |cache: &mut Vec>>, index: i32| -> PyResult<*mut ffi::PyObject> { + let slot = usize::try_from(index).map_err(|_| lc_index_err())?; + let entry = cache.get_mut(slot).ok_or_else(lc_index_err)?; + if entry.is_none() { + // Safety: slot < col.values.len() (checked by get_mut); the + // returned pointer is a valid owned reference; Py takes it + // over and the Vec drops every cached entry on any exit path. + let ptr = unsafe { column_value_nonnull_ptr(py, &col.values, ctx, slot, None)? }; + *entry = Some(unsafe { Bound::from_owned_ptr(py, ptr) }.unbind()); + } + Ok(entry + .as_ref() + .expect("entry filled above") + .clone_ref(py) + .into_ptr()) + }; + match &col.validity { + None => { + for (i, &index) in col.indices[..rows].iter().enumerate() { + sink(i, cached_item(&mut cache, index)?); + } + } + Some(bm) => { + for (i, &index) in col.indices[..rows].iter().enumerate() { + let item = if bm.is_valid(i) { + cached_item(&mut cache, index)? + } else { + none_owned_ptr() + }; + sink(i, item); + } + } + } + Ok(()) +} diff --git a/rust/ch-core-py/src/pyval/json.rs b/rust/ch-core-py/src/pyval/json.rs new file mode 100644 index 00000000..641fbc39 --- /dev/null +++ b/rust/ch-core-py/src/pyval/json.rs @@ -0,0 +1,434 @@ +use super::*; + +/// Insert one materialized value at a prepared dotted JSON path. Intermediate +/// dicts are allocated only on first use and then reused by later paths in the +/// same row. +/// +/// # Safety +/// +/// `root` must be a live exact dict and `value` a live Python object. The GIL +/// must be held. +unsafe fn set_json_path( + py: Python<'_>, + root: *mut ffi::PyObject, + path: &JsonPath<'_>, + value: *mut ffi::PyObject, +) -> PyResult<()> { + let Some((last, parents)) = path.keys.split_last() else { + return Err(json_shape_err()); + }; + let mut current = root; + for key in parents { + let mut child = ffi::PyDict_GetItemWithError(current, key.as_ptr()); + if child.is_null() && !ffi::PyErr_Occurred().is_null() { + return Err(PyErr::fetch(py)); + } + if child.is_null() || child == ffi::Py_None() { + let created = ptr_to_result(py, ffi::PyDict_New())?; + if ffi::PyDict_SetItem(current, key.as_ptr(), created) < 0 { + ffi::Py_DECREF(created); + return Err(PyErr::fetch(py)); + } + // The parent dict owns a reference now; keep only a borrowed + // pointer while descending. + ffi::Py_DECREF(created); + child = created; + } else if ffi::PyDict_Check(child) == 0 { + return Err(PyValueError::new_err( + "Malformed payload: JSON paths collide at a non-object value", + )); + } + current = child; + } + if ffi::PyDict_SetItem(current, last.as_ptr(), value) < 0 { + return Err(PyErr::fetch(py)); + } + Ok(()) +} + +/// Fill one JSON path column directly into the already-allocated row dicts. +/// Dynamic/shared NULL values are absent paths, while typed NULL values remain +/// explicit keys. +unsafe fn fill_json_path_column<'py>( + py: Python<'py>, + column: &Column, + ctx: &ColumnCtx<'py>, + path: &JsonPath<'py>, + rows: usize, + containers: &[*mut ffi::PyObject], + validity: Option<&Bitmap>, +) -> PyResult<()> { + let mut err: Option = None; + // Items rejected after the child sink sees them must stay alive until the + // child fill returns. Some child fills populate containers after sinking. + let mut discarded: Vec> = Vec::new(); + let mut path_sink = |row: usize, item: *mut ffi::PyObject| { + let item = Bound::from_owned_ptr(py, item); + if err.is_some() || validity.is_some_and(|bitmap| !bitmap.is_valid(row)) { + discarded.push(item); + return; + } + if let Err(path_err) = set_json_path(py, containers[row], path, item.as_ptr()) { + err = Some(path_err); + discarded.push(item); + } + }; + let mut erased: DynSink<'_> = &mut path_sink; + fill_column(py, column, ctx, rows, &mut erased)?; + if let Some(err) = err { + return Err(err); + } + Ok(()) +} + +unsafe fn fill_json_dynamic_path<'py>( + py: Python<'py>, + column: &DynamicColumn, + path: &JsonPath<'py>, + rows: usize, + containers: &[*mut ffi::PyObject], + validity: Option<&Bitmap>, +) -> PyResult<()> { + let mut err: Option = None; + let mut discarded: Vec> = Vec::new(); + let mut path_sink = |row: usize, item: *mut ffi::PyObject| { + let item = Bound::from_owned_ptr(py, item); + if err.is_some() + || validity.is_some_and(|bitmap| !bitmap.is_valid(row)) + || item.as_ptr() == ffi::Py_None() + { + discarded.push(item); + return; + } + if let Err(path_err) = set_json_path(py, containers[row], path, item.as_ptr()) { + err = Some(path_err); + discarded.push(item); + } + }; + fill_dynamic(py, column, rows, &mut path_sink)?; + if let Some(err) = err { + return Err(err); + } + Ok(()) +} + +/// Materialize a JSON column as nested Python dicts. Typed and dynamic paths +/// are filled column-major, preserving the binding's one-dispatch-per-column +/// policy. Shared-data descriptors and path keys are cached across rows. +/// +/// # Safety +/// +/// Requires the GIL; `fill_column`'s sink contract applies. +pub(super) unsafe fn fill_json<'py, S>( + py: Python<'py>, + col: &JsonColumn, + ctx: &ColumnCtx<'py>, + rows: usize, + sink: &mut S, +) -> PyResult<()> +where + S: FnMut(usize, *mut ffi::PyObject), +{ + if rows != col.len() { + return Err(json_shape_err()); + } + let validity = col.validity.as_ref(); + let JsonBody::Structured(structured) = &col.body else { + let JsonBody::Text(values) = &col.body else { + unreachable!() + }; + let loads = py.import("json")?.getattr("loads")?; + for row in 0..rows { + if validity.is_some_and(|bitmap| !bitmap.is_valid(row)) { + sink(row, none_owned_ptr()); + continue; + } + let document = pyo3::types::PyBytes::new(py, values.value(row)); + sink(row, loads.call1((document,))?.into_ptr()); + } + return Ok(()); + }; + + let typed_ctxs = ctx.fields.as_deref().ok_or_else(|| ctx_missing("JSON"))?; + let typed_paths = ctx + .json_paths + .as_deref() + .ok_or_else(|| ctx_missing("JSON"))?; + if structured.typed.len() != typed_ctxs.len() + || structured.typed.len() != typed_paths.len() + || structured.len != rows + { + return Err(json_shape_err()); + } + + let estimated_keys = json_estimated_keys(structured); + // Borrowed pointers after `sink` takes ownership; the sink contract keeps + // every valid-row dict alive until this fill returns. + let mut containers: Vec<*mut ffi::PyObject> = Vec::with_capacity(rows); + for row in 0..rows { + let item = if validity.is_some_and(|bitmap| !bitmap.is_valid(row)) { + none_owned_ptr() + } else { + ptr_to_result(py, dict_new_presized(estimated_keys))? + }; + containers.push(item); + sink(row, item); + } + + // The fills below zip the decoded typed columns with the declared-type + // contexts positionally; guard the order once per column per chunk. + check_typed_path_order(&structured.typed, typed_paths)?; + for (((_, column), child_ctx), path) in structured.typed.iter().zip(typed_ctxs).zip(typed_paths) + { + fill_json_path_column(py, column, child_ctx, path, rows, &containers, validity)?; + } + + for (path_name, dynamic) in &structured.dynamic { + let path = prepare_json_path(py, path_name)?; + fill_json_dynamic_path(py, dynamic, &path, rows, &containers, validity)?; + } + + if structured.shared_offsets.len() != rows.saturating_add(1) + || structured.shared_paths.len() != structured.shared_values.len() + { + return Err(json_shape_err()); + } + let mut path_cache: JsonPathCache<'py> = Vec::new(); + let mut value_ctx_cache: SharedCtxCache<'py> = Vec::new(); + for (row, &container) in containers.iter().enumerate() { + let start = + usize::try_from(structured.shared_offsets[row]).map_err(|_| json_shape_err())?; + let end = + usize::try_from(structured.shared_offsets[row + 1]).map_err(|_| json_shape_err())?; + if start > end || end > structured.shared_paths.len() { + return Err(json_shape_err()); + } + if validity.is_some_and(|bitmap| !bitmap.is_valid(row)) { + continue; + } + for index in start..end { + let path = cached_json_path(py, &mut path_cache, structured.shared_paths.value(index))?; + let item = Bound::from_owned_ptr( + py, + shared_cell_owned_ptr( + py, + structured.shared_values.value(index), + Some(&mut value_ctx_cache), + )?, + ); + // JSON nulls are absent paths, matching the Python codec. + if item.as_ptr() != ffi::Py_None() { + set_json_path(py, container, path, item.as_ptr())?; + } + } + } + Ok(()) +} + +/// Row-dict presize estimate: typed/dynamic dotted paths collapse into their +/// distinct top-level keys, shared pairs contribute their per-row mean. +fn json_estimated_keys(structured: &StructuredJson) -> ffi::Py_ssize_t { + let mean_shared = if structured.len == 0 { + 0 + } else { + structured.shared_paths.len().div_ceil(structured.len) + }; + distinct_top_level_keys(structured) + .saturating_add(mean_shared) + .min(ffi::Py_ssize_t::MAX as usize) as ffi::Py_ssize_t +} + +/// Count the distinct top-level key segments across the typed and dynamic +/// paths, for row-dict presizing. +fn distinct_top_level_keys(structured: &StructuredJson) -> usize { + let mut seen: Vec<&str> = Vec::new(); + let paths = structured + .typed + .iter() + .map(|(path, _)| path.as_str()) + .chain(structured.dynamic.iter().map(|(path, _)| path.as_str())); + for path in paths { + let top = path.split('.').next().unwrap_or(path); + if !seen.contains(&top) { + seen.push(top); + } + } + seen.len() +} + +/// Guard for the positional zip of decoded typed columns with declared-type +/// contexts: each column's path must equal the declared path at its index. +/// The caller has already checked the lengths match. +fn check_typed_path_order(typed: &[(String, Column)], paths: &[JsonPath<'_>]) -> PyResult<()> { + for ((declared, _), path) in typed.iter().zip(paths) { + if declared != &path.raw { + return Err(PyValueError::new_err(format!( + "Malformed payload: JSON typed path '{declared}' does not match the declared path '{}'", + path.raw + ))); + } + } + Ok(()) +} + +/// Tiny per-fill cache of prepared JSON paths keyed by the raw path bytes. +/// Unique dynamic/shared paths per block are typically single digits, so a +/// linear scan beats hashing. +pub(super) type JsonPathCache<'py> = Vec<(Vec, JsonPath<'py>)>; + +/// Look up or build the prepared path for `path_bytes`. +fn cached_json_path<'cache, 'py>( + py: Python<'py>, + cache: &'cache mut JsonPathCache<'py>, + path_bytes: &[u8], +) -> PyResult<&'cache JsonPath<'py>> { + match cache.iter().position(|(key, _)| key == path_bytes) { + Some(found) => Ok(&cache[found].1), + None => { + let path_name = std::str::from_utf8(path_bytes).map_err(|_| json_shape_err())?; + let path = prepare_json_path(py, path_name)?; + cache.push((path_bytes.to_vec(), path)); + Ok(&cache.last().expect("pushed above").1) + } + } +} + +/// Materialize one non-null JSON row for recursive container exits. The common +/// top-level path uses `fill_json`'s column-major implementation instead. +/// `cache` is the Array element chain's per-fill `ChainCache::Json`, carrying +/// the resolved `json.loads`, prepared paths, and shared-cell contexts across +/// cells of the same column. +pub(super) unsafe fn json_value_owned_ptr<'py>( + py: Python<'py>, + col: &JsonColumn, + ctx: &ColumnCtx<'py>, + row: usize, + cache: Option<&mut ChainCache<'py>>, +) -> PyResult<*mut ffi::PyObject> { + if row >= col.len() { + return Err(json_shape_err()); + } + let json_cache = match cache { + Some(ChainCache::Json(json_cache)) => Some(json_cache), + _ => None, + }; + match &col.body { + JsonBody::Text(values) => { + let document = pyo3::types::PyBytes::new(py, values.value(row)); + match json_cache { + Some(json_cache) => { + if json_cache.loads.is_none() { + json_cache.loads = Some(py.import("json")?.getattr("loads")?); + } + let loads = json_cache.loads.as_ref().expect("filled above"); + Ok(loads.call1((document,))?.into_ptr()) + } + None => { + let loads = py.import("json")?.getattr("loads")?; + Ok(loads.call1((document,))?.into_ptr()) + } + } + } + JsonBody::Structured(structured) => { + let typed_ctxs = ctx.fields.as_deref().ok_or_else(|| ctx_missing("JSON"))?; + let typed_paths = ctx + .json_paths + .as_deref() + .ok_or_else(|| ctx_missing("JSON"))?; + if structured.typed.len() != typed_ctxs.len() + || structured.typed.len() != typed_paths.len() + { + return Err(json_shape_err()); + } + let (mut cached_paths, cached_shared, order_checked, cached_estimate) = match json_cache + { + Some(json_cache) => ( + Some(&mut json_cache.paths), + Some(&mut json_cache.shared), + Some(&mut json_cache.typed_order_checked), + Some(&mut json_cache.estimated_keys), + ), + None => (None, None, None, None), + }; + // Once per column when cached, once per cell otherwise. + if !order_checked.as_ref().is_some_and(|flag| **flag) { + check_typed_path_order(&structured.typed, typed_paths)?; + if let Some(flag) = order_checked { + *flag = true; + } + } + let estimated = match cached_estimate { + Some(slot) => *slot.get_or_insert_with(|| json_estimated_keys(structured)), + None => json_estimated_keys(structured), + }; + let root = ptr_to_result(py, dict_new_presized(estimated))?; + // Own the dict until the row is complete, so every error path drops + // partially inserted values. + let root = Bound::from_owned_ptr(py, root).cast_into_unchecked::(); + for (((_, column), child_ctx), path) in + structured.typed.iter().zip(typed_ctxs).zip(typed_paths) + { + let item = Bound::from_owned_ptr( + py, + column_value_to_owned_ptr(py, column, child_ctx, row, None)?, + ); + set_json_path(py, root.as_ptr(), path, item.as_ptr())?; + } + for (path_name, dynamic) in &structured.dynamic { + let item = + Bound::from_owned_ptr(py, dynamic_value_owned_ptr(py, dynamic, row, None)?); + if item.as_ptr() != ffi::Py_None() { + let prepared; + let path = match cached_paths.as_mut() { + Some(cache) => cached_json_path(py, cache, path_name.as_bytes())?, + None => { + prepared = prepare_json_path(py, path_name)?; + &prepared + } + }; + set_json_path(py, root.as_ptr(), path, item.as_ptr())?; + } + } + if structured.shared_offsets.len() <= row + 1 { + return Err(json_shape_err()); + } + let start = + usize::try_from(structured.shared_offsets[row]).map_err(|_| json_shape_err())?; + let end = usize::try_from(structured.shared_offsets[row + 1]) + .map_err(|_| json_shape_err())?; + if start > end + || end > structured.shared_paths.len() + || end > structured.shared_values.len() + { + return Err(json_shape_err()); + } + let mut local_shared: SharedCtxCache<'py> = Vec::new(); + let shared_cache = cached_shared.unwrap_or(&mut local_shared); + for index in start..end { + let path_bytes = structured.shared_paths.value(index); + let prepared; + let path = match cached_paths.as_mut() { + Some(cache) => cached_json_path(py, cache, path_bytes)?, + None => { + let path_name = + std::str::from_utf8(path_bytes).map_err(|_| json_shape_err())?; + prepared = prepare_json_path(py, path_name)?; + &prepared + } + }; + let item = Bound::from_owned_ptr( + py, + shared_cell_owned_ptr( + py, + structured.shared_values.value(index), + Some(&mut *shared_cache), + )?, + ); + if item.as_ptr() != ffi::Py_None() { + set_json_path(py, root.as_ptr(), path, item.as_ptr())?; + } + } + Ok(root.into_ptr()) + } + } +} diff --git a/rust/ch-core-py/src/pyval/mod.rs b/rust/ch-core-py/src/pyval/mod.rs new file mode 100644 index 00000000..0470aac1 --- /dev/null +++ b/rust/ch-core-py/src/pyval/mod.rs @@ -0,0 +1,284 @@ +use std::collections::HashMap; +#[cfg(not(Py_3_13))] +use std::ffi::c_int; +use std::ffi::{c_char, c_long}; + +use pyo3::exceptions::{PyUnicodeDecodeError, PyValueError}; +use pyo3::ffi; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::{PyDate, PyDateTime, PyDelta, PyDict, PyList, PyString, PyTuple}; + +use ch_core_rs::bitmap::Bitmap; +use ch_core_rs::column::{ + AggregateStateColumn, Column, DecimalColumn, DictionaryColumn, DynamicChild, DynamicColumn, + JsonBody, JsonColumn, MapColumn, QBitColumn, StructuredJson, TupleColumn, VariantColumn, +}; +use ch_core_rs::native::binary_value::{ + decode_binary_value, read_binary_type_prefix, BinaryValueError, +}; +use ch_core_rs::native::varint::ByteReader; +use ch_core_rs::schema::ChType; + +mod containers; +mod ctx; +mod errors; +mod fixed; +mod json; +mod ptr; +mod qbit; +mod scalar; +mod temporal; +mod variant_dynamic; + +use containers::{fill_map, fill_tuple, point_list_owned_ptr, point_pair_slices}; +pub(crate) use ctx::{prepare_column_ctx, ColumnCtx}; +use ctx::{prepare_json_path, IpCtx, JsonPath, UuidCtx}; +use errors::*; +use fixed::{ + bfloat16_to_f32, column_validity, fill_aggregate_states, fill_dictionary, fill_fixed_width, +}; +use json::{fill_json, json_value_owned_ptr, JsonPathCache}; +use ptr::*; +use qbit::{fill_qbit, qbit_value_owned_ptr}; +use scalar::{ + column_value_nonnull_ptr, decimal_value_ptr, ipv4_value_ptr, ipv6_value_ptr, uuid_value_ptr, + wide_int_value_ptr, DecimalScratch, +}; +use temporal::{dt64_secs_micros, make_date, make_datetime, make_time64}; +use variant_dynamic::{dynamic_value_owned_ptr, fill_dynamic, fill_variant, shared_cell_owned_ptr}; + +/// Allocate a dict sized for `len` entries. `_PyDict_NewPresized` is gone +/// from pyo3-ffi 0.28+, so this is a plain `PyDict_New`; `len` is kept so +/// callers still document the expected entry count. +/// +/// # Safety +/// +/// Requires the GIL. Returns an owned reference, or null with an error set. +unsafe fn dict_new_presized(len: ffi::Py_ssize_t) -> *mut ffi::PyObject { + let _ = len; + ffi::PyDict_New() +} + +/// Type-erased sink. The Tuple/Map fills recurse through `fill_column` with +/// fresh closure types per nesting level; erasing at the container boundary +/// keeps monomorphization finite while top-level fills stay static. +type DynSink<'a> = &'a mut dyn FnMut(usize, *mut ffi::PyObject); + +/// Materialize `rows` cells of `col` into `sink`: bulk fixed-width and +/// dictionary fills first, hoisted Tuple/Map fills next, per-cell fallback +/// last with the validity branch hoisted. +/// +/// # Safety +/// +/// Requires the GIL. `sink` is called exactly once per row with the cell's +/// positional row index and an owned reference it must take over. Variant and +/// Dynamic fills call it in child-major order; every other fill calls it in +/// ascending row order, which `materialize_run`'s push sink enforces. The sink must keep +/// every item alive until this call returns because Tuple fills write into +/// containers after sinking them. +pub(crate) unsafe fn fill_column<'py, S>( + py: Python<'py>, + col: &Column, + ctx: &ColumnCtx<'py>, + rows: usize, + sink: &mut S, +) -> PyResult<()> +where + S: FnMut(usize, *mut ffi::PyObject), +{ + if matches!(col, Column::Nothing(_)) { + // Nothing has no host value. Its optional bitmap only preserves the + // structural null map of Nullable(Nothing) for Native re-encoding; + // both valid and invalid bits materialize as Python None. Keep this + // as a column-wide fill so flat results do not check that bitmap or + // redispatch the Column enum for every row. + for i in 0..rows { + sink(i, none_owned_ptr()); + } + return Ok(()); + } + if let Column::AggregateState(states) = col { + return fill_aggregate_states(py, states, rows, sink); + } + if let Column::QBit(qbit) = col { + return fill_qbit(py, qbit, rows, sink); + } + if fill_fixed_width(py, col, ctx, rows, sink)? { + return Ok(()); + } + match col { + Column::Dictionary(dict) => fill_dictionary(py, dict, ctx, rows, sink), + Column::Variant(c) => fill_variant(py, c, ctx, rows, sink), + Column::Dynamic(c) => fill_dynamic(py, c, rows, sink), + Column::Json(c) => fill_json(py, c, ctx, rows, sink), + Column::Tuple(c) => fill_tuple(py, c, ctx, rows, sink), + Column::Map(c) => fill_map(py, c, ctx, rows, sink), + _ => { + let mut chain_cache = new_array_chain_cache(col); + match column_validity(col) { + None => { + for i in 0..rows { + let item = column_value_nonnull_ptr(py, col, ctx, i, chain_cache.as_mut())?; + sink(i, item); + } + } + Some(bm) => { + for i in 0..rows { + let item = if bm.is_valid(i) { + column_value_nonnull_ptr(py, col, ctx, i, chain_cache.as_mut())? + } else { + none_owned_ptr() + }; + sink(i, item); + } + } + } + Ok(()) + } + } +} + +/// Materialize the first `rows` cells of `col` into owned objects through the +/// bulk fill machinery. The indexed slots accept Variant and Dynamic's +/// child-major scatter while preserving logical row order for Map key/value +/// runs. Error paths drop whatever was already produced. +unsafe fn materialize_run<'py>( + py: Python<'py>, + col: &Column, + ctx: &ColumnCtx<'py>, + rows: usize, +) -> PyResult>> { + if !matches!(col, Column::Variant(_) | Column::Dynamic(_)) { + let mut out: Vec> = Vec::with_capacity(rows); + // Items sunk out of ascending order are an internal error; they stay + // alive here until the fill returns, per the sink contract. + let mut misordered: Vec> = Vec::new(); + { + let mut sink = |i: usize, item: *mut ffi::PyObject| { + // Safety: item is an owned reference the sink takes over. + let item = unsafe { Bound::from_owned_ptr(py, item) }.unbind(); + if i == out.len() { + out.push(item); + } else { + misordered.push(item); + } + }; + let mut erased: DynSink<'_> = &mut sink; + fill_column(py, col, ctx, rows, &mut erased)?; + } + if !misordered.is_empty() { + return Err(PyValueError::new_err( + "internal error: column fill ran out of row order", + )); + } + if out.len() != rows { + return Err(PyValueError::new_err( + "internal error: column fill produced the wrong row count", + )); + } + return Ok(out); + } + + let mut out: Vec>> = Vec::with_capacity(rows); + out.resize_with(rows, || None); + { + let mut sink = |i: usize, item: *mut ffi::PyObject| { + // Safety: item is an owned reference the sink takes over. + out[i] = Some(unsafe { Bound::from_owned_ptr(py, item) }.unbind()); + }; + let mut erased: DynSink<'_> = &mut sink; + fill_column(py, col, ctx, rows, &mut erased)?; + } + out.into_iter() + .map(|item| { + item.ok_or_else(|| PyValueError::new_err("internal error: column fill omitted a row")) + }) + .collect() +} + +/// Lazy cache of materialized dictionary slot objects, one entry per slot. +type DictSlotCache = Vec>>; + +/// Descriptor bytes -> parsed type plus prepared context for SharedVariant +/// and JSON shared-data cells. Distinct descriptors per block are typically +/// single digits, so a linear scan beats hashing. +type SharedCtxCache<'py> = Vec<(Vec, ChType, ColumnCtx<'py>)>; + +/// Per-fill cache for a JSON column inside an Array element chain: the +/// resolved `json.loads`, prepared dynamic/shared paths, shared-cell +/// descriptor contexts, and the once-per-column typed-order check. All fill +/// lazily on first reference. +#[derive(Default)] +struct JsonChainCache<'py> { + loads: Option>, + paths: JsonPathCache<'py>, + shared: SharedCtxCache<'py>, + typed_order_checked: bool, + estimated_keys: Option, +} + +/// A per-fill cache for the terminal of an Array column's element chain. One +/// cache per array column per chunk, threaded through the per-cell path. +/// `Dict` caches materialized LowCardinality dictionary slots so repeated +/// labels share one object, matching `fill_dictionary`'s reuse policy. +/// `Dynamic` caches one prepared `ColumnCtx` per block-local Dynamic child, +/// plus one per distinct SharedVariant descriptor, so +/// `column_value_nonnull_ptr` does not rebuild contexts per cell. `Json` +/// carries the per-column JSON state. All fill lazily on first reference. +enum ChainCache<'py> { + Dict(DictSlotCache), + Dynamic { + contexts: Vec>>, + shared: SharedCtxCache<'py>, + }, + Json(JsonChainCache<'py>), +} + +/// Build the cache for the Dictionary or Dynamic column terminating an Array +/// column's element chain, if any. +fn new_array_chain_cache<'py>(col: &Column) -> Option> { + fn chain_terminal(col: &Column) -> Option<&Column> { + match col { + Column::Array(c) => chain_terminal(&c.values), + Column::Dictionary(_) | Column::Dynamic(_) | Column::Json(_) => Some(col), + _ => None, + } + } + fn empty_slots(len: usize) -> Vec> { + let mut slots = Vec::with_capacity(len); + slots.resize_with(len, || None); + slots + } + let Column::Array(c) = col else { return None }; + match chain_terminal(&c.values)? { + Column::Dictionary(dict) => Some(ChainCache::Dict(empty_slots(dict.values.len()))), + Column::Dynamic(dynamic) => Some(ChainCache::Dynamic { + contexts: empty_slots(dynamic.children.len()), + shared: Vec::new(), + }), + Column::Json(_) => Some(ChainCache::Json(JsonChainCache::default())), + _ => None, + } +} + +/// Build the cell at `index` as an owned pointer, None for a null cell. Used +/// by the row path, where columns interleave; the column paths hoist the +/// validity check instead. +/// +/// # Safety +/// +/// Returns an owned reference; the caller must take over the reference count. +unsafe fn column_value_to_owned_ptr<'py>( + py: Python<'py>, + col: &Column, + ctx: &ColumnCtx<'py>, + index: usize, + cache: Option<&mut ChainCache<'py>>, +) -> PyResult<*mut ffi::PyObject> { + if column_validity(col).is_some_and(|v| !v.is_valid(index)) { + Ok(none_owned_ptr()) + } else { + column_value_nonnull_ptr(py, col, ctx, index, cache) + } +} diff --git a/rust/ch-core-py/src/pyval/ptr.rs b/rust/ch-core-py/src/pyval/ptr.rs new file mode 100644 index 00000000..edbb1ab3 --- /dev/null +++ b/rust/ch-core-py/src/pyval/ptr.rs @@ -0,0 +1,93 @@ +use super::*; + +/// Build a Python str from raw String column bytes. Invalid UTF-8 renders as +/// the lowercase hex of the raw bytes, matching clickhouse-connect's String +/// read fallback. Single scan in the valid case: CPython's decode is the +/// validation, and the hex path runs only after a UnicodeDecodeError. +pub(super) fn utf8_or_hex_owned_ptr(py: Python<'_>, bytes: &[u8]) -> PyResult<*mut ffi::PyObject> { + // Safety: the pointer/length pair is a live borrowed slice and CPython + // copies the bytes before returning. A zero-length slice's dangling + // pointer is never read. + let ptr = unsafe { + ffi::PyUnicode_FromStringAndSize( + bytes.as_ptr() as *const c_char, + bytes.len() as ffi::Py_ssize_t, + ) + }; + if !ptr.is_null() { + return Ok(ptr); + } + let err = PyErr::fetch(py); + if !err.is_instance_of::(py) { + return Err(err); + } + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut hex = Vec::with_capacity(bytes.len() * 2); + for &b in bytes { + hex.push(HEX[(b >> 4) as usize]); + hex.push(HEX[(b & 0x0f) as usize]); + } + // Safety: hex is a live ASCII buffer, valid UTF-8, copied by CPython. + unsafe { + ptr_to_result( + py, + ffi::PyUnicode_FromStringAndSize( + hex.as_ptr() as *const c_char, + hex.len() as ffi::Py_ssize_t, + ), + ) + } +} + +/// # Safety +/// +/// Returns an owned reference; the caller must take over the reference count. +pub(super) unsafe fn none_owned_ptr() -> *mut ffi::PyObject { + let none = ffi::Py_None(); + ffi::Py_INCREF(none); + none +} + +/// # Safety +/// +/// Returns an owned reference; the caller must take over the reference count. +pub(super) unsafe fn bool_owned_ptr(value: bool) -> *mut ffi::PyObject { + let ptr = if value { + ffi::Py_True() + } else { + ffi::Py_False() + }; + ffi::Py_INCREF(ptr); + ptr +} + +/// # Safety +/// +/// `ptr` must be an owned reference or null; on `Ok` the caller takes over +/// the reference count. +pub(super) unsafe fn ptr_to_result( + py: Python<'_>, + ptr: *mut ffi::PyObject, +) -> PyResult<*mut ffi::PyObject> { + if ptr.is_null() { + Err(PyErr::fetch(py)) + } else { + Ok(ptr) + } +} + +/// Copy one binary value into a Python bytes object. +/// +/// # Safety +/// +/// Requires the GIL. Returns an owned reference; the caller must take over the +/// reference count. +pub(super) unsafe fn bytes_owned_ptr(py: Python<'_>, bytes: &[u8]) -> PyResult<*mut ffi::PyObject> { + ptr_to_result( + py, + ffi::PyBytes_FromStringAndSize( + bytes.as_ptr() as *const c_char, + bytes.len() as ffi::Py_ssize_t, + ), + ) +} diff --git a/rust/ch-core-py/src/pyval/qbit.rs b/rust/ch-core-py/src/pyval/qbit.rs new file mode 100644 index 00000000..ce499ae2 --- /dev/null +++ b/rust/ch-core-py/src/pyval/qbit.rs @@ -0,0 +1,120 @@ +use super::*; + +fn qbit_shape_error() -> PyErr { + PyValueError::new_err("Malformed payload: invalid QBit child buffer") +} + +/// Build one logical QBit vector as a Python list from its row-major child +/// slice. The core already performs the Native bit-plane transpose once; this +/// exit only allocates the Python objects requested by a row-oriented consumer. +/// +/// # Safety +/// +/// Requires the GIL. `make` must return a new reference or null. On success the +/// returned pointer is an owned reference that the caller must take over. +unsafe fn float_list_owned_ptr( + py: Python<'_>, + values: &[T], + make: F, +) -> PyResult<*mut ffi::PyObject> +where + T: Copy, + F: Fn(T) -> *mut ffi::PyObject, +{ + let list_ptr = ffi::PyList_New(values.len() as ffi::Py_ssize_t); + if list_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Binding the fresh list makes every error path drop it. CPython list + // deallocation tolerates the still-null slots after a partial fill. + let list = Bound::from_owned_ptr(py, list_ptr).cast_into_unchecked::(); + for (index, &value) in values.iter().enumerate() { + let item = make(value); + if item.is_null() { + return Err(PyErr::fetch(py)); + } + // The fresh list takes ownership of each new float reference. + ffi::PyList_SET_ITEM(list.as_ptr(), index as ffi::Py_ssize_t, item); + } + Ok(list.into_ptr()) +} + +/// Materialize one non-null QBit row as an owned Python list. +/// +/// # Safety +/// +/// Requires the GIL. The returned pointer is an owned reference that the +/// caller must take over. +pub(super) unsafe fn qbit_value_owned_ptr( + py: Python<'_>, + col: &QBitColumn, + index: usize, +) -> PyResult<*mut ffi::PyObject> { + if col.dimension == 0 { + return Err(qbit_shape_error()); + } + let start = index + .checked_mul(col.dimension) + .ok_or_else(qbit_shape_error)?; + let end = start + .checked_add(col.dimension) + .ok_or_else(qbit_shape_error)?; + match col.values.as_ref() { + Column::BFloat16(values) if values.validity.is_none() && end <= values.values.len() => { + float_list_owned_ptr(py, &values.values[start..end], |value| { + ffi::PyFloat_FromDouble(bfloat16_to_f32(value).into()) + }) + } + Column::Float32(values) if values.validity.is_none() && end <= values.values.len() => { + float_list_owned_ptr(py, &values.values[start..end], |value| { + ffi::PyFloat_FromDouble(value.into()) + }) + } + Column::Float64(values) if values.validity.is_none() && end <= values.values.len() => { + float_list_owned_ptr(py, &values.values[start..end], |value| { + ffi::PyFloat_FromDouble(value) + }) + } + _ => Err(qbit_shape_error()), + } +} + +/// Bulk QBit fill for top-level columns and Tuple/Map field runs. Dispatch the +/// column once, then allocate one Python list per valid row. Nullable QBit has +/// list-level validity, so null rows avoid all child float allocation. +/// +/// # Safety +/// +/// Requires the GIL. Each pointer passed to `sink` is an owned reference the +/// sink must take over exactly once. +pub(super) unsafe fn fill_qbit( + py: Python<'_>, + col: &QBitColumn, + rows: usize, + sink: &mut S, +) -> PyResult<()> +where + S: FnMut(usize, *mut ffi::PyObject), +{ + if col.len() < rows { + return Err(qbit_shape_error()); + } + match &col.validity { + None => { + for row in 0..rows { + sink(row, qbit_value_owned_ptr(py, col, row)?); + } + } + Some(validity) => { + for row in 0..rows { + let item = if validity.is_valid(row) { + qbit_value_owned_ptr(py, col, row)? + } else { + none_owned_ptr() + }; + sink(row, item); + } + } + } + Ok(()) +} diff --git a/rust/ch-core-py/src/pyval/scalar.rs b/rust/ch-core-py/src/pyval/scalar.rs new file mode 100644 index 00000000..32c30ea5 --- /dev/null +++ b/rust/ch-core-py/src/pyval/scalar.rs @@ -0,0 +1,592 @@ +use super::*; + +/// Build the cell at `index` as an owned pointer, assuming the cell is not +/// null; callers check validity first. `cache` is the Array element chain's +/// per-fill cache, if the caller materializes one. +/// +/// # Safety +/// +/// Returns an owned reference; the caller must take over the reference count. +pub(super) unsafe fn column_value_nonnull_ptr<'py>( + py: Python<'py>, + col: &Column, + ctx: &ColumnCtx<'py>, + index: usize, + mut cache: Option<&mut ChainCache<'py>>, +) -> PyResult<*mut ffi::PyObject> { + match col { + Column::Bool(c) => ptr_to_result(py, ffi::PyBool_FromLong(c.get(index).into())), + Column::Int8(c) => ptr_to_result(py, ffi::PyLong_FromLongLong(c.values[index].into())), + Column::Int16(c) => ptr_to_result(py, ffi::PyLong_FromLongLong(c.values[index].into())), + Column::Int32(c) => ptr_to_result(py, ffi::PyLong_FromLongLong(c.values[index].into())), + Column::Int64(c) => ptr_to_result(py, ffi::PyLong_FromLongLong(c.values[index])), + Column::UInt8(c) => { + ptr_to_result(py, ffi::PyLong_FromUnsignedLongLong(c.values[index].into())) + } + Column::UInt16(c) => { + ptr_to_result(py, ffi::PyLong_FromUnsignedLongLong(c.values[index].into())) + } + Column::UInt32(c) => { + ptr_to_result(py, ffi::PyLong_FromUnsignedLongLong(c.values[index].into())) + } + Column::UInt64(c) => ptr_to_result(py, ffi::PyLong_FromUnsignedLongLong(c.values[index])), + Column::Int128(c) | Column::Int256(c) => wide_int_value_ptr(py, c.value(index), true), + Column::UInt128(c) | Column::UInt256(c) => wide_int_value_ptr(py, c.value(index), false), + Column::Float32(c) => ptr_to_result(py, ffi::PyFloat_FromDouble(c.values[index].into())), + Column::Float64(c) => ptr_to_result(py, ffi::PyFloat_FromDouble(c.values[index])), + Column::BFloat16(c) => ptr_to_result( + py, + ffi::PyFloat_FromDouble(bfloat16_to_f32(c.values[index]).into()), + ), + Column::QBit(c) => qbit_value_owned_ptr(py, c, index), + Column::AggregateState(c) => bytes_owned_ptr(py, c.value(index)), + // The bulk fill above serves top-level and Tuple/Map column runs. + // This per-cell arm is needed for Array(Nothing) and other recursive + // container paths. + Column::Nothing(_) => Ok(none_owned_ptr()), + Column::Date(c) => Ok(make_date(py, c.values[index] as i64)?.into_ptr()), + Column::Date32(c) => Ok(make_date(py, c.values[index] as i64)?.into_ptr()), + Column::DateTime(c) => Ok(make_datetime(py, c.values[index] as i64, 0, ctx)?.into_ptr()), + Column::DateTime64(c) => { + let (secs, micros) = dt64_secs_micros(c.values[index], ctx.precision); + Ok(make_datetime(py, secs, micros, ctx)?.into_ptr()) + } + Column::Time(c) => { + if ctx.raw_time_ticks { + ptr_to_result(py, ffi::PyLong_FromLong(c_long::from(c.values[index]))) + } else { + Ok(PyDelta::new(py, 0, c.values[index], 0, true)?.into_ptr()) + } + } + Column::Time64(c) => { + if ctx.raw_time_ticks { + ptr_to_result(py, ffi::PyLong_FromLongLong(c.values[index])) + } else { + Ok(make_time64(py, c.values[index], ctx.time_scale)?.into_ptr()) + } + } + Column::Interval(c) => ptr_to_result(py, ffi::PyLong_FromLongLong(c.values[index])), + Column::Utf8(c) => utf8_or_hex_owned_ptr(py, c.value(index)), + Column::FixedBinary(c) => bytes_owned_ptr(py, c.value(index)), + // LowCardinality(T): resolve the row's dictionary index, then build the + // inner value through this same constructor. The cell is known non-null + // here (callers check the index validity first), so the resolved slot is + // a real dictionary entry. The ctx already reflects the inner type, so a + // LowCardinality temporal column gets the right timezone and precision. + Column::Dictionary(c) => { + let slot = c + .indices + .get(index) + .copied() + .and_then(|i| usize::try_from(i).ok()) + .filter(|&slot| slot < c.values.len()) + .ok_or_else(lc_index_err)?; + match cache { + // Array element path: build each referenced slot once per + // chunk and emit clone_ref of the cached object. + Some(ChainCache::Dict(slots)) => { + let entry = slots.get_mut(slot).ok_or_else(lc_index_err)?; + if entry.is_none() { + let ptr = column_value_nonnull_ptr(py, &c.values, ctx, slot, None)?; + *entry = Some(Bound::from_owned_ptr(py, ptr).unbind()); + } + Ok(entry + .as_ref() + .expect("entry filled above") + .clone_ref(py) + .into_ptr()) + } + _ => column_value_nonnull_ptr(py, &c.values, ctx, slot, None), + } + } + // Enum8/Enum16 carry only the physical signed int; map it to its label + // string through the per-column value->name map. A value with no defined + // label becomes None, matching clickhouse-connect's int_map.get default. + Column::Enum8(c) => enum_value_ptr(ctx, c.values[index] as i64), + Column::Enum16(c) => enum_value_ptr(ctx, c.values[index] as i64), + Column::Uuid(c) => { + let uctx = ctx.uuid.as_ref().ok_or_else(|| ctx_missing("UUID"))?; + uuid_value_ptr(py, uctx, c.value(index)) + } + Column::Ipv4(c) => { + let ictx = ctx.ip.as_ref().ok_or_else(|| ctx_missing("IPv4"))?; + ipv4_value_ptr(py, ictx, c.values[index]) + } + Column::Ipv6(c) => { + let ictx = ctx.ip.as_ref().ok_or_else(|| ctx_missing("IPv6"))?; + ipv6_value_ptr(py, ictx, c.value(index)) + } + Column::Decimal(c) => { + let cls = ctx + .decimal_cls + .as_ref() + .ok_or_else(|| ctx_missing("Decimal"))?; + let mut scratch = DecimalScratch::default(); + decimal_value_ptr(cls, &mut scratch, c, index) + } + // Array(T): materialize row `index` as a Python list of its elements. + // The offsets buffer is public and could be hand-built, so guard every + // access: reject a negative offset, out-of-order pair, or an end past + // the element buffer rather than index out of bounds or panic. Element + // nulls are handled by column_value_to_owned_ptr, so Array(Nullable(T)) + // yields None elements correctly. + Column::Array(c) => { + let ectx = ctx.element.as_deref().ok_or_else(|| ctx_missing("Array"))?; + let start = c + .offsets + .get(index) + .copied() + .and_then(|o| usize::try_from(o).ok()) + .ok_or_else(array_bounds_err)?; + let end = c + .offsets + .get(index + 1) + .copied() + .and_then(|o| usize::try_from(o).ok()) + .ok_or_else(array_bounds_err)?; + if start > end || end > c.values.len() { + return Err(array_bounds_err()); + } + // Geo point rows: build the list straight from the two flat + // Float64 runs, skipping the per-element dispatch below. + if let Some((xs, ys)) = point_pair_slices(&c.values, ectx) { + if end <= xs.len().min(ys.len()) { + return point_list_owned_ptr(py, xs, ys, start, end); + } + } + let count = end - start; + let list_ptr = ffi::PyList_New(count as ffi::Py_ssize_t); + if list_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: list_ptr came from PyList_New, so it is a list and this is + // the sole owned reference. Binding it makes the error and panic + // paths drop the partially-filled list; list_dealloc tolerates the + // NULL slots not yet filled. + let list = Bound::from_owned_ptr(py, list_ptr).cast_into_unchecked::(); + for slot in 0..count { + let item = column_value_to_owned_ptr( + py, + &c.values, + ectx, + start + slot, + cache.as_deref_mut(), + )?; + // Safety: slot < count, the list's allocated length, and the + // list takes over the owned item. + ffi::PyList_SET_ITEM(list.as_ptr(), slot as ffi::Py_ssize_t, item); + } + Ok(list.into_ptr()) + } + // Tuple(T1, ...): an unnamed tuple materializes as a Python `tuple`, a + // named tuple as a `dict` keyed by the element names, matching + // clickhouse-connect's default read format. Field values recurse + // through column_value_to_owned_ptr, so a Nullable/LowCardinality/nested + // container element composes and a Nullable element yields None. + Column::Tuple(c) => { + let fctx = ctx.fields.as_deref().ok_or_else(|| ctx_missing("Tuple"))?; + if fctx.len() != c.fields.len() { + return Err(ctx_count_mismatch("Tuple")); + } + match &ctx.tuple_names { + Some(names) => { + let dict_ptr = ffi::PyDict_New(); + if dict_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: dict_ptr came from PyDict_New; binding it drops the + // partially-filled dict on the error path. + let dict = Bound::from_owned_ptr(py, dict_ptr).cast_into_unchecked::(); + for (field_idx, field_col) in c.fields.iter().enumerate() { + let item = column_value_to_owned_ptr( + py, + field_col, + &fctx[field_idx], + index, + None, + )?; + // Take ownership so an error before/at insertion drops it; + // PyDict_SetItem does not steal, it increfs the value. + let value = Bound::from_owned_ptr(py, item); + if ffi::PyDict_SetItem( + dict.as_ptr(), + names[field_idx].as_ptr(), + value.as_ptr(), + ) < 0 + { + return Err(PyErr::fetch(py)); + } + } + Ok(dict.into_ptr()) + } + None => { + let tuple_ptr = ffi::PyTuple_New(c.fields.len() as ffi::Py_ssize_t); + if tuple_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: tuple_ptr came from PyTuple_New; binding it drops the + // partially-filled tuple on the error path (tuple_dealloc + // Py_XDECREFs each slot, tolerating the NULL slots). + let tuple = + Bound::from_owned_ptr(py, tuple_ptr).cast_into_unchecked::(); + for (field_idx, field_col) in c.fields.iter().enumerate() { + let item = column_value_to_owned_ptr( + py, + field_col, + &fctx[field_idx], + index, + None, + )?; + // Safety: field_idx < tuple len, and the tuple takes over + // the owned item. + ffi::PyTuple_SET_ITEM(tuple.as_ptr(), field_idx as ffi::Py_ssize_t, item); + } + Ok(tuple.into_ptr()) + } + } + } + // Map(K, V): materialize row `index` as a Python `dict`. The entries are + // the flattened Tuple(keys, values) column sliced by the Array-shaped + // offsets; guard every offset access like the Array arm. Keys are + // inserted in wire order, so a duplicate key keeps its first position and + // last value, matching clickhouse-connect's dict(zip(keys, values)). + Column::Map(c) => { + let fctx = ctx.fields.as_deref().ok_or_else(|| ctx_missing("Map"))?; + if fctx.len() != 2 { + return Err(ctx_missing("Map")); + } + let entries = match c.entries.as_ref() { + Column::Tuple(t) if t.fields.len() == 2 => t, + _ => return Err(map_entries_err()), + }; + let keys_col = &entries.fields[0]; + let values_col = &entries.fields[1]; + let start = c + .offsets + .get(index) + .copied() + .and_then(|o| usize::try_from(o).ok()) + .ok_or_else(map_bounds_err)?; + let end = c + .offsets + .get(index + 1) + .copied() + .and_then(|o| usize::try_from(o).ok()) + .ok_or_else(map_bounds_err)?; + // Bound against the buffers actually indexed below (the decoder + // guarantees both equal entries.len(), but guard the exact slices + // like the Array arm rather than the declared tuple length). + if start > end || end > keys_col.len().min(values_col.len()) { + return Err(map_bounds_err()); + } + let dict_ptr = ffi::PyDict_New(); + if dict_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + // Safety: dict_ptr came from PyDict_New; binding it drops the + // partially-filled dict on the error path. + let dict = Bound::from_owned_ptr(py, dict_ptr).cast_into_unchecked::(); + for slot in start..end { + let key = column_value_to_owned_ptr(py, keys_col, &fctx[0], slot, None)?; + let key = Bound::from_owned_ptr(py, key); + let value = column_value_to_owned_ptr(py, values_col, &fctx[1], slot, None)?; + let value = Bound::from_owned_ptr(py, value); + // PyDict_SetItem increfs both; the Bounds drop our refs after. + if ffi::PyDict_SetItem(dict.as_ptr(), key.as_ptr(), value.as_ptr()) < 0 { + return Err(PyErr::fetch(py)); + } + } + Ok(dict.into_ptr()) + } + // Variant's row routing selects one value from a dense alternative + // column. This per-cell path is used when Variant is itself inside an + // Array or another container whose elements are materialized by index; + // top-level and Tuple/Map Variant runs use the bulk scatter above. + Column::Variant(c) => { + let contexts = ctx + .fields + .as_deref() + .ok_or_else(|| ctx_missing("Variant"))?; + if contexts.len() != c.variants.len() { + return Err(ctx_count_mismatch("Variant")); + } + let (discriminator, offset) = c.value_position(index).ok_or_else(variant_shape_err)?; + if discriminator == u8::MAX { + return Ok(none_owned_ptr()); + } + let alternative = usize::from(discriminator); + let child = c.variants.get(alternative).ok_or_else(variant_shape_err)?; + let child_ctx = contexts.get(alternative).ok_or_else(variant_shape_err)?; + column_value_to_owned_ptr( + py, + child, + child_ctx, + usize::try_from(offset).map_err(|_| variant_shape_err())?, + None, + ) + } + // Dynamic's per-cell path is used inside Array and other containers + // whose elements are already being materialized by index. Top-level + // and Tuple/Map Dynamic runs use the child-major bulk scatter above. + Column::Dynamic(c) => dynamic_value_owned_ptr(py, c, index, cache), + Column::Json(c) => json_value_owned_ptr(py, c, ctx, index, cache), + } +} + +/// Convert one little-endian fixed-width integer into an exact Python int. +/// The core intentionally keeps wide integers in their Native/Arrow byte +/// representation, so signedness is supplied by the distinct Column variant. +/// +/// # Safety +/// +/// Requires the GIL. Returns an owned reference; the caller must take over +/// the reference count. +pub(super) unsafe fn wide_int_value_ptr( + py: Python<'_>, + bytes: &[u8], + signed: bool, +) -> PyResult<*mut ffi::PyObject> { + // Word-scan fast path: high words all zero (or all ones for a negative + // signed value) means the value fits one C long long constructor. + if bytes.len().is_multiple_of(8) { + if let Some((lo_bytes, high)) = bytes.split_first_chunk::<8>() { + let lo = u64::from_le_bytes(*lo_bytes); + let mut hi_or = 0u64; + let mut hi_and = u64::MAX; + for word in high.chunks_exact(8) { + let word = u64::from_le_bytes(word.try_into().expect("chunks_exact(8)")); + hi_or |= word; + hi_and &= word; + } + if signed { + let lo = lo as i64; + if (hi_or == 0 && lo >= 0) || (hi_and == u64::MAX && lo < 0) { + return ptr_to_result(py, ffi::PyLong_FromLongLong(lo)); + } + } else if hi_or == 0 { + return ptr_to_result(py, ffi::PyLong_FromUnsignedLongLong(lo)); + } + } + } + #[cfg(not(Py_3_13))] + { + ptr_to_result( + py, + ffi::_PyLong_FromByteArray(bytes.as_ptr(), bytes.len(), 1, c_int::from(signed)), + ) + } + #[cfg(Py_3_13)] + { + let ptr = if signed { + ffi::PyLong_FromNativeBytes( + bytes.as_ptr().cast(), + bytes.len(), + ffi::Py_ASNATIVEBYTES_LITTLE_ENDIAN, + ) + } else { + ffi::PyLong_FromUnsignedNativeBytes( + bytes.as_ptr().cast(), + bytes.len(), + ffi::Py_ASNATIVEBYTES_LITTLE_ENDIAN, + ) + }; + ptr_to_result(py, ptr) + } +} + +/// Build a `uuid.UUID` from the 16 raw wire bytes: `UUID.__new__(UUID)`, then +/// `object.__setattr__` of `int` and `is_safe` (SafeUUID.unsafe), matching the +/// Cython codec's read_uuid_col. The wire int is le(b[0..8]) << 64 | le(b[8..16]), +/// which is `from_le_bytes` with the halves swapped. +pub(super) fn uuid_value_ptr( + py: Python<'_>, + ctx: &UuidCtx<'_>, + bytes: &[u8], +) -> PyResult<*mut ffi::PyObject> { + let b: &[u8; 16] = bytes.try_into().map_err(|_| fixed_width_err("UUID"))?; + let int_val = u128::from_le_bytes(*b).rotate_left(64); + let value = ctx.new.call1((&ctx.cls,))?; + ctx.object_setattr + .call1((&value, intern!(py, "int"), int_val))?; + ctx.object_setattr + .call1((&value, intern!(py, "is_safe"), &ctx.unsafe_marker))?; + Ok(value.into_ptr()) +} + +/// Build an `ipaddress.IPv4Address` from the numeric address value. +pub(super) fn ipv4_value_ptr( + py: Python<'_>, + ctx: &IpCtx<'_>, + value: u32, +) -> PyResult<*mut ffi::PyObject> { + let addr = ctx.new.call1((&ctx.cls,))?; + addr.setattr(intern!(py, "_ip"), value)?; + Ok(addr.into_ptr()) +} + +/// Build an `ipaddress.IPv6Address` from the 16 network-order wire bytes, +/// always IPv6Address even for a v4-mapped value, matching _read_binary_ip. +pub(super) fn ipv6_value_ptr( + py: Python<'_>, + ctx: &IpCtx<'_>, + bytes: &[u8], +) -> PyResult<*mut ffi::PyObject> { + let b: &[u8; 16] = bytes.try_into().map_err(|_| fixed_width_err("IPv6"))?; + let int_val = u128::from_be_bytes(*b); + let addr = ctx.new.call1((&ctx.cls,))?; + addr.setattr(intern!(py, "_ip"), int_val)?; + if ctx.set_scope_id { + addr.setattr(intern!(py, "_scope_id"), py.None())?; + } + Ok(addr.into_ptr()) +} + +/// Reusable buffers for Decimal text rendering: the magnitude digits and the +/// composed constructor argument. +#[derive(Default)] +pub(super) struct DecimalScratch { + digits: String, + text: String, +} + +/// Build a `decimal.Decimal` for the cell: render the unscaled value as exact +/// decimal text (sign, integer digits, exactly `scale` fractional digits) and +/// call the class once. The text form yields the same value and exponent as +/// the python codec's `Decimal(unscaled).scaleb(-scale)`. +pub(super) fn decimal_value_ptr( + cls: &Bound<'_, PyAny>, + scratch: &mut DecimalScratch, + col: &DecimalColumn, + index: usize, +) -> PyResult<*mut ffi::PyObject> { + scratch.digits.clear(); + let negative = write_decimal_magnitude(col.value(index), &mut scratch.digits)?; + compose_decimal_text( + &mut scratch.text, + negative, + &scratch.digits, + col.scale as usize, + ); + Ok(cls.call1((scratch.text.as_str(),))?.into_ptr()) +} + +/// Write the magnitude digits (no sign, no leading zeros, "0" for zero) of a +/// little-endian two's-complement integer of width 4/8/16/32 bytes into `out`; +/// returns whether the value is negative. +fn write_decimal_magnitude(bytes: &[u8], out: &mut String) -> PyResult { + use std::fmt::Write as _; + match bytes.len() { + 4 => { + let v = i32::from_le_bytes(bytes.try_into().expect("width checked")); + let _ = write!(out, "{}", v.unsigned_abs()); + Ok(v < 0) + } + 8 => { + let v = i64::from_le_bytes(bytes.try_into().expect("width checked")); + let _ = write!(out, "{}", v.unsigned_abs()); + Ok(v < 0) + } + 16 => { + let v = i128::from_le_bytes(bytes.try_into().expect("width checked")); + let _ = write!(out, "{}", v.unsigned_abs()); + Ok(v < 0) + } + 32 => { + let mut limbs = [0u64; 4]; + for (limb, chunk) in limbs.iter_mut().zip(bytes.chunks_exact(8)) { + *limb = u64::from_le_bytes(chunk.try_into().expect("chunks_exact(8)")); + } + let negative = limbs[3] >> 63 == 1; + if negative { + negate_limbs(&mut limbs); + } + write_u256_digits(limbs, out); + Ok(negative) + } + w => Err(PyValueError::new_err(format!( + "Malformed payload: unsupported Decimal width {w}" + ))), + } +} + +/// Two's-complement negate a 256-bit little-endian limb array in place. +fn negate_limbs(limbs: &mut [u64; 4]) { + let mut carry = 1u64; + for limb in limbs.iter_mut() { + let (v, overflowed) = (!*limb).overflowing_add(carry); + *limb = v; + carry = u64::from(overflowed); + } +} + +/// Divide a 256-bit little-endian limb magnitude in place by `divisor`, +/// returning the remainder. Standard long division, most-significant limb first. +fn div_rem_limbs(limbs: &mut [u64; 4], divisor: u64) -> u64 { + let mut rem: u128 = 0; + for limb in limbs.iter_mut().rev() { + let cur = (rem << 64) | u128::from(*limb); + *limb = (cur / u128::from(divisor)) as u64; + rem = cur % u128::from(divisor); + } + rem as u64 +} + +/// Write the decimal digits of a 256-bit little-endian limb magnitude: repeated +/// divmod by 1e19 yields base-1e19 chunks, most significant unpadded, the rest +/// zero-padded to 19 digits. At most 5 chunks (2^255 has 77 digits). +fn write_u256_digits(mut limbs: [u64; 4], out: &mut String) { + use std::fmt::Write as _; + const CHUNK: u64 = 10_000_000_000_000_000_000; // 1e19 + let mut chunks = [0u64; 5]; + let mut count = 0; + loop { + chunks[count] = div_rem_limbs(&mut limbs, CHUNK); + count += 1; + if limbs == [0u64; 4] { + break; + } + } + let _ = write!(out, "{}", chunks[count - 1]); + for &chunk in chunks[..count - 1].iter().rev() { + let _ = write!(out, "{chunk:019}"); + } +} + +/// Compose the Decimal constructor text: optional '-', integer digits, and for +/// scale > 0 a '.' with exactly `scale` fractional digits. `digits` is the +/// magnitude with no sign or leading zeros ("0" only for zero). +fn compose_decimal_text(out: &mut String, negative: bool, digits: &str, scale: usize) { + out.clear(); + if negative { + out.push('-'); + } + if scale == 0 { + out.push_str(digits); + return; + } + if digits.len() > scale { + let split = digits.len() - scale; + out.push_str(&digits[..split]); + out.push('.'); + out.push_str(&digits[split..]); + } else { + out.push_str("0."); + for _ in 0..(scale - digits.len()) { + out.push('0'); + } + out.push_str(digits); + } +} + +/// Map an enum's physical integer to its label string, or None for a value with +/// no defined label (matching clickhouse-connect's `int_map.get(value, None)`). +/// +/// # Safety +/// +/// Returns an owned reference; the caller must take over the reference count. +unsafe fn enum_value_ptr(ctx: &ColumnCtx<'_>, value: i64) -> PyResult<*mut ffi::PyObject> { + match ctx.enum_names.as_ref().and_then(|m| m.get(&value)) { + Some(name) => Ok(name.clone().into_ptr()), + None => Ok(none_owned_ptr()), + } +} diff --git a/rust/ch-core-py/src/pyval/temporal.rs b/rust/ch-core-py/src/pyval/temporal.rs new file mode 100644 index 00000000..7148fece --- /dev/null +++ b/rust/ch-core-py/src/pyval/temporal.rs @@ -0,0 +1,143 @@ +use super::*; + +/// Split DateTime64 ticks at `precision` into whole seconds and microseconds. +/// Euclidean division so pre-epoch (negative) ticks floor correctly and the +/// microsecond remainder stays in `0..1_000_000`. Sub-microsecond digits +/// (precision 7..9) are truncated, since Python datetime resolves to +/// microseconds, matching clickhouse-connect. +pub(super) fn dt64_secs_micros(ticks: i64, precision: u8) -> (i64, u32) { + let scale = 10i64.pow(precision as u32); + let secs = ticks.div_euclid(scale); + let frac = ticks.rem_euclid(scale); + let micros = if precision <= 6 { + frac * 10i64.pow(6 - precision as u32) + } else { + frac / 10i64.pow(precision as u32 - 6) + }; + (secs, micros as u32) +} + +/// Civil date (year, month, day) from a day count since 1970-01-01. Howard +/// Hinnant's `civil_from_days`; valid across the full Date/Date32 range. Month +/// and day are 1-based. +fn civil_from_days(days: i64) -> (i32, u8, u8) { + let z = days + 719_468; + let era = if z >= 0 { z } else { z - 146_096 } / 146_097; + let doe = z - era * 146_097; // [0, 146096] + let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365; // [0, 399] + let y = yoe + era * 400; + let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); // [0, 365] + let mp = (5 * doy + 2) / 153; // [0, 11] + let d = (doy - (153 * mp + 2) / 5 + 1) as u8; // [1, 31] + let m = if mp < 10 { mp + 3 } else { mp - 9 } as u8; // [1, 12] + let year = (y + if m <= 2 { 1 } else { 0 }) as i32; + (year, m, d) +} + +/// Calendar components (year, month, day, hour, minute, second) from seconds +/// since 1970-01-01 00:00:00 UTC. Euclidean split keeps the time of day in +/// `0..86_400` for negative (pre-epoch) inputs. +fn civil_from_secs(secs: i64) -> (i32, u8, u8, u8, u8, u8) { + let days = secs.div_euclid(86_400); + let tod = secs.rem_euclid(86_400); + let (year, month, day) = civil_from_days(days); + let hour = (tod / 3600) as u8; + let minute = ((tod % 3600) / 60) as u8; + let second = (tod % 60) as u8; + (year, month, day, hour, minute, second) +} + +/// Build a `datetime.date` from a day count since the epoch. +pub(super) fn make_date(py: Python<'_>, days: i64) -> PyResult> { + let (year, month, day) = civil_from_days(days); + Ok(PyDate::new(py, year, month, day)?.into_any()) +} + +/// Build a `datetime.datetime` from epoch seconds plus microseconds, honoring +/// the column's ColumnCtx (naive UTC arithmetic, or tz-aware fromtimestamp). +pub(super) fn make_datetime<'py>( + py: Python<'py>, + secs: i64, + micros: u32, + ctx: &ColumnCtx<'py>, +) -> PyResult> { + match (&ctx.tz, &ctx.fromtimestamp) { + (Some(tz), Some(fromtimestamp)) => { + if micros == 0 { + fromtimestamp.call1((secs, tz)) + } else if secs.unsigned_abs() < (1 << 32) { + // Single call on a float timestamp. Exact: |secs| < 2^32 keeps + // the f64 ulp of the sum under one microsecond, so CPython's + // round-to-nearest-microsecond recovers `micros` exactly. + fromtimestamp.call1((secs as f64 + micros as f64 / 1e6, tz)) + } else { + // Distant timestamps lose sub-microsecond float precision; + // set the exact microsecond on the aware datetime instead. + let dt = fromtimestamp.call1((secs, tz))?; + let kwargs = PyDict::new(py); + kwargs.set_item("microsecond", micros)?; + dt.call_method("replace", (), Some(&kwargs)) + } + } + _ => { + let (year, month, day, hour, minute, second) = civil_from_secs(secs); + Ok( + PyDateTime::new(py, year, month, day, hour, minute, second, micros, None)? + .into_any(), + ) + } + } +} + +/// Build a Python timedelta from a signed second and microsecond total. The +/// components may be negative; normalize them into Python's canonical +/// day/second/microsecond representation without a Python arithmetic call. +fn make_timedelta<'py>( + py: Python<'py>, + seconds: i128, + microseconds: i128, +) -> PyResult> { + const MICROS_PER_SECOND: i128 = 1_000_000; + const MICROS_PER_DAY: i128 = 86_400 * MICROS_PER_SECOND; + + let total_micros = seconds + .checked_mul(MICROS_PER_SECOND) + .and_then(|v| v.checked_add(microseconds)) + .ok_or_else(|| PyValueError::new_err("Time value overflows datetime.timedelta"))?; + let days = total_micros.div_euclid(MICROS_PER_DAY); + let day_micros = total_micros.rem_euclid(MICROS_PER_DAY); + let day_seconds = day_micros / MICROS_PER_SECOND; + let micros = day_micros % MICROS_PER_SECOND; + let days = i32::try_from(days) + .map_err(|_| PyValueError::new_err("Time value is outside datetime.timedelta range"))?; + // day_seconds and micros are normalized to Python's documented ranges. + PyDelta::new(py, days, day_seconds as i32, micros as i32, false) +} + +/// Materialize Time64 ticks as timedelta with microsecond precision. Fractional +/// ticks below one microsecond are truncated toward zero, including negatives, +/// matching clickhouse_connect.datatypes.temporal.Time64._ticks_to_timedelta. +pub(super) fn make_time64<'py>( + py: Python<'py>, + ticks: i64, + scale: u64, +) -> PyResult> { + let negative = ticks < 0; + let magnitude = ticks.unsigned_abs(); + let seconds = magnitude / scale; + let frac = magnitude % scale; + // scale <= 1e9, so the numerator fits u64 and avoids software u128 division + // in this per-cell materialization path. + let micros = frac * 1_000_000 / scale; + let sign = if negative { -1i128 } else { 1i128 }; + if seconds <= i32::MAX as u64 { + return PyDelta::new( + py, + 0, + (sign * i128::from(seconds)) as i32, + (sign * i128::from(micros)) as i32, + true, + ); + } + make_timedelta(py, sign * i128::from(seconds), sign * i128::from(micros)) +} diff --git a/rust/ch-core-py/src/pyval/variant_dynamic.rs b/rust/ch-core-py/src/pyval/variant_dynamic.rs new file mode 100644 index 00000000..74dab614 --- /dev/null +++ b/rust/ch-core-py/src/pyval/variant_dynamic.rs @@ -0,0 +1,417 @@ +use super::*; + +/// Materialize a Variant column child by child: route each logical row to its +/// dense child slot from the discriminator run, then scatter each child's bulk +/// fill into those row positions. +/// +/// # Safety +/// +/// Requires the GIL. `fill_column`'s sink contract applies. +pub(super) unsafe fn fill_variant<'py, S>( + py: Python<'py>, + col: &VariantColumn, + ctx: &ColumnCtx<'py>, + rows: usize, + sink: &mut S, +) -> PyResult<()> +where + S: FnMut(usize, *mut ffi::PyObject), +{ + let contexts = ctx + .fields + .as_deref() + .ok_or_else(|| ctx_missing("Variant"))?; + if contexts.len() != col.variants.len() { + return Err(ctx_count_mismatch("Variant")); + } + let discriminators = col.discriminators(); + if rows != col.len() || discriminators.len() != rows { + return Err(variant_shape_err()); + } + + let mut starts = Vec::with_capacity(col.variants.len() + 1); + starts.push(0usize); + for child in &col.variants { + let next = starts + .last() + .copied() + .unwrap_or_default() + .checked_add(child.len()) + .ok_or_else(variant_shape_err)?; + starts.push(next); + } + // Occurrence ordinals are the dense child offsets, so per-alternative + // running counters assign each non-NULL row a distinct slot; the written + // counter is the completeness check. + let mut destinations = vec![0usize; starts.last().copied().unwrap_or_default()]; + let mut counters = vec![0usize; col.variants.len()]; + let mut written = 0usize; + let mut nulls = 0usize; + for (row, &discriminator) in discriminators.iter().enumerate() { + if discriminator == u8::MAX { + nulls += 1; + sink(row, none_owned_ptr()); + continue; + } + let alternative = usize::from(discriminator); + if alternative >= col.variants.len() { + return Err(variant_shape_err()); + } + let slot = starts[alternative] + counters[alternative]; + if slot >= starts[alternative + 1] { + return Err(variant_shape_err()); + } + counters[alternative] += 1; + destinations[slot] = row; + written += 1; + } + if nulls != col.nulls.len || written != destinations.len() { + return Err(variant_shape_err()); + } + + for (alternative, (child, child_ctx)) in col.variants.iter().zip(contexts).enumerate() { + let positions = &destinations[starts[alternative]..starts[alternative + 1]]; + let mut scatter = |child_row: usize, item: *mut ffi::PyObject| { + sink(positions[child_row], item); + }; + let mut erased: DynSink<'_> = &mut scatter; + fill_column(py, child, child_ctx, child.len(), &mut erased)?; + } + Ok(()) +} + +/// Materialize a Dynamic column child by child. Dynamic carries block-local +/// child ids and occurrence-ordinal offsets rather than Variant's raw UInt8 +/// discriminator run, but both use the same dense child-major layout. +/// SharedVariant cells decode to typed Python values through the core's +/// single-value binary decoder and the same per-type conversion machinery as +/// typed children; AggregateFunction cells and unsupported descriptors keep +/// the raw cell bytes. The Arrow C Stream export keeps every SharedVariant +/// cell as bytes for schema stability; only these object exits decode them. +/// +/// # Safety +/// +/// Requires the GIL. `fill_column`'s sink contract applies. +pub(super) unsafe fn fill_dynamic<'py, S>( + py: Python<'py>, + col: &DynamicColumn, + rows: usize, + sink: &mut S, +) -> PyResult<()> +where + S: FnMut(usize, *mut ffi::PyObject), +{ + if rows != col.len() || col.offsets.len() != rows { + return Err(dynamic_shape_err()); + } + + let mut starts = Vec::with_capacity(col.children.len() + 1); + starts.push(0usize); + for child in &col.children { + let next = starts + .last() + .copied() + .unwrap_or_default() + .checked_add(child.len()) + .ok_or_else(dynamic_shape_err)?; + starts.push(next); + } + let mut destinations = vec![0usize; starts.last().copied().unwrap_or_default()]; + let mut counters = vec![0usize; col.children.len()]; + let mut nulls = 0usize; + for (row, (&type_id, &offset)) in col.type_ids.iter().zip(&col.offsets).enumerate() { + let offset = usize::try_from(offset).map_err(|_| dynamic_shape_err())?; + if type_id == u32::MAX { + if offset != nulls { + return Err(dynamic_shape_err()); + } + nulls += 1; + sink(row, none_owned_ptr()); + continue; + } + let child = usize::try_from(type_id).map_err(|_| dynamic_shape_err())?; + let counter = counters.get_mut(child).ok_or_else(dynamic_shape_err)?; + if offset != *counter || offset >= col.children[child].len() { + return Err(dynamic_shape_err()); + } + destinations[starts[child] + offset] = row; + *counter += 1; + } + if nulls != col.nulls.len + || counters + .iter() + .zip(&col.children) + .any(|(&count, child)| count != child.len()) + { + return Err(dynamic_shape_err()); + } + + for (child_index, child) in col.children.iter().enumerate() { + let positions = &destinations[starts[child_index]..starts[child_index + 1]]; + match child { + DynamicChild::Typed { ch_type, values } => { + // Dynamic cells are type-erased at the driver boundary. As for + // Variant alternatives, finalize temporal leaves here instead + // of exposing raw ticks to a walker that cannot see the child + // type. + let child_ctx = prepare_column_ctx(py, ch_type, false)?; + let mut scatter = |child_row: usize, item: *mut ffi::PyObject| { + sink(positions[child_row], item); + }; + let mut erased: DynSink<'_> = &mut scatter; + fill_column(py, values, &child_ctx, child.len(), &mut erased)?; + } + DynamicChild::Shared(values) => { + // Cells with the same descriptor share one prepared context + // per fill. + let mut ctx_cache: SharedCtxCache<'py> = Vec::new(); + for (child_row, &position) in positions.iter().enumerate() { + let item = + shared_cell_owned_ptr(py, values.value(child_row), Some(&mut ctx_cache))?; + sink(position, item); + } + } + } + } + Ok(()) +} + +/// Materialize one logical Dynamic cell. Used by recursive container/JSON +/// exits; top-level Dynamic columns keep the child-major bulk fill above. +pub(super) unsafe fn dynamic_value_owned_ptr<'py>( + py: Python<'py>, + col: &DynamicColumn, + index: usize, + mut cache: Option<&mut ChainCache<'py>>, +) -> PyResult<*mut ffi::PyObject> { + let (&type_id, &offset) = col + .type_ids + .get(index) + .zip(col.offsets.get(index)) + .ok_or_else(dynamic_shape_err)?; + let offset = usize::try_from(offset).map_err(|_| dynamic_shape_err())?; + if type_id == u32::MAX { + if offset >= col.nulls.len { + return Err(dynamic_shape_err()); + } + return Ok(none_owned_ptr()); + } + let child_index = usize::try_from(type_id).map_err(|_| dynamic_shape_err())?; + let child = col + .children + .get(child_index) + .ok_or_else(dynamic_shape_err)?; + if offset >= child.len() { + return Err(dynamic_shape_err()); + } + match child { + DynamicChild::Typed { ch_type, values } => { + if let Some(ChainCache::Dynamic { contexts, .. }) = cache.as_deref_mut() { + if let Some(entry) = contexts.get_mut(child_index) { + if entry.is_none() { + *entry = Some(prepare_column_ctx(py, ch_type, false)?); + } + let child_ctx = entry.as_ref().expect("entry filled above"); + return column_value_to_owned_ptr(py, values, child_ctx, offset, None); + } + } + let child_ctx = prepare_column_ctx(py, ch_type, false)?; + column_value_to_owned_ptr(py, values, &child_ctx, offset, None) + } + DynamicChild::Shared(values) => { + let shared_cache = match cache { + Some(ChainCache::Dynamic { shared, .. }) => Some(shared), + _ => None, + }; + shared_cell_owned_ptr(py, values.value(offset), shared_cache) + } + } +} + +/// Materialize one SharedVariant cell (binary type descriptor + one +/// `serializeBinary` value) as a typed Python object: decode the value to a +/// one-row Column and route it through the same conversion machinery as a +/// typed column. AggregateFunction cells (opaque state payloads) and cells +/// whose descriptor does not parse or whose value encoding is unsupported +/// stay raw Python bytes. A cell whose descriptor parsed but whose payload +/// fails to decode is a ValueError. +/// +/// # Safety +/// +/// Requires the GIL. Returns an owned reference; the caller must take over +/// the reference count. +pub(super) unsafe fn shared_cell_owned_ptr<'py>( + py: Python<'py>, + cell: &[u8], + mut ctx_cache: Option<&mut SharedCtxCache<'py>>, +) -> PyResult<*mut ffi::PyObject> { + // Cells are varint-length framed, so the raw bytes are always + // recoverable; an unknown descriptor tag is a future server type, not + // corruption. Every descriptor-parse failure falls back to bytes. + if let Some(cache) = ctx_cache.as_deref_mut() { + // Binary type descriptors are a prefix code (the parser consumes + // exactly the descriptor from the front), so a cached descriptor + // prefix identifies the type and the cursor advance without a parse. + if let Some((key, ch_type, ctx)) = cache.iter().find(|(key, ..)| cell.starts_with(key)) { + return shared_payload_owned_ptr(py, ch_type, ctx, cell, key.len()); + } + } + let Ok((ch_type, consumed)) = read_binary_type_prefix(cell) else { + return bytes_owned_ptr(py, cell); + }; + if matches!(ch_type, ChType::AggregateFunction { .. }) { + return bytes_owned_ptr(py, cell); + } + let payload = &cell[consumed..]; + if let Some(fast) = shared_scalar_owned_ptr(py, &ch_type, payload) { + // Scalar contexts need no Python machinery; cache for later cells. + if let Some(cache) = ctx_cache { + let ctx = prepare_column_ctx(py, &ch_type, false)?; + cache.push((cell[..consumed].to_vec(), ch_type, ctx)); + } + return fast; + } + // Decode before preparing a context, so a bytes-fallback descriptor + // (Variant, Dynamic, JSON, ...) never runs the Python ctx machinery, + // which can raise (for one, a nested unknown timezone). + let column = match decode_binary_value(&ch_type, payload) { + Ok(column) => column, + Err(BinaryValueError::Unsupported(_)) => return bytes_owned_ptr(py, cell), + Err(err) => return Err(shared_cell_err(&err)), + }; + let ctx = prepare_column_ctx(py, &ch_type, false)?; + if let Some(cache) = ctx_cache { + cache.push((cell[..consumed].to_vec(), ch_type, ctx)); + let (_, _, ctx) = cache.last().expect("pushed above"); + column_value_to_owned_ptr(py, &column, ctx, 0, None) + } else { + column_value_to_owned_ptr(py, &column, &ctx, 0, None) + } +} + +/// Materialize the single value after a cache-hit cell's `consumed`-byte +/// descriptor. Cached descriptors already passed the AggregateFunction gate; +/// an Unsupported payload (an `Array(Nothing)` value, say) still falls back +/// to bytes. +/// +/// # Safety +/// +/// Requires the GIL. Returns an owned reference; the caller must take over +/// the reference count. +unsafe fn shared_payload_owned_ptr<'py>( + py: Python<'py>, + ch_type: &ChType, + ctx: &ColumnCtx<'py>, + cell: &[u8], + consumed: usize, +) -> PyResult<*mut ffi::PyObject> { + let payload = &cell[consumed..]; + if let Some(fast) = shared_scalar_owned_ptr(py, ch_type, payload) { + return fast; + } + let column = match decode_binary_value(ch_type, payload) { + Ok(column) => column, + Err(BinaryValueError::Unsupported(_)) => return bytes_owned_ptr(py, cell), + Err(err) => return Err(shared_cell_err(&err)), + }; + column_value_to_owned_ptr(py, &column, ctx, 0, None) +} + +/// Direct materialization for the hot scalar shared-cell types, skipping the +/// one-row Column build. Matches `decode_binary_value` semantics: truncation +/// and trailing bytes are errors. `None` falls through to the generic route. +/// +/// # Safety +/// +/// Requires the GIL. Returns an owned reference; the caller must take over +/// the reference count. +unsafe fn shared_scalar_owned_ptr( + py: Python<'_>, + ch_type: &ChType, + payload: &[u8], +) -> Option> { + unsafe fn fixed( + py: Python<'_>, + payload: &[u8], + build: impl FnOnce([u8; N]) -> *mut ffi::PyObject, + ) -> PyResult<*mut ffi::PyObject> { + match payload.try_into() { + Ok(bytes) => ptr_to_result(py, build(bytes)), + // Match decode_binary_value's wording exactly. + Err(_) => Err(shared_cell_err(&BinaryValueError::Invalid( + if payload.len() < N { + "truncated value payload".to_string() + } else { + format!("{} trailing bytes after the value", payload.len() - N) + }, + ))), + } + } + Some(match ch_type { + ChType::String => shared_string_owned_ptr(py, payload), + ChType::Bool => fixed(py, payload, |[b]: [u8; 1]| { + ffi::PyBool_FromLong(c_long::from(b != 0)) + }), + ChType::Int8 => fixed(py, payload, |b: [u8; 1]| { + ffi::PyLong_FromLongLong(i8::from_le_bytes(b).into()) + }), + ChType::Int16 => fixed(py, payload, |b: [u8; 2]| { + ffi::PyLong_FromLongLong(i16::from_le_bytes(b).into()) + }), + ChType::Int32 => fixed(py, payload, |b: [u8; 4]| { + ffi::PyLong_FromLongLong(i32::from_le_bytes(b).into()) + }), + ChType::Int64 => fixed(py, payload, |b: [u8; 8]| { + ffi::PyLong_FromLongLong(i64::from_le_bytes(b)) + }), + ChType::UInt8 => fixed(py, payload, |b: [u8; 1]| { + ffi::PyLong_FromUnsignedLongLong(u8::from_le_bytes(b).into()) + }), + ChType::UInt16 => fixed(py, payload, |b: [u8; 2]| { + ffi::PyLong_FromUnsignedLongLong(u16::from_le_bytes(b).into()) + }), + ChType::UInt32 => fixed(py, payload, |b: [u8; 4]| { + ffi::PyLong_FromUnsignedLongLong(u32::from_le_bytes(b).into()) + }), + ChType::UInt64 => fixed(py, payload, |b: [u8; 8]| { + ffi::PyLong_FromUnsignedLongLong(u64::from_le_bytes(b)) + }), + ChType::Float32 => fixed(py, payload, |b: [u8; 4]| { + ffi::PyFloat_FromDouble(f32::from_le_bytes(b).into()) + }), + ChType::Float64 => fixed(py, payload, |b: [u8; 8]| { + ffi::PyFloat_FromDouble(f64::from_le_bytes(b)) + }), + _ => return None, + }) +} + +/// String shared-cell fast path: varint length plus raw bytes, materialized +/// with the same invalid-UTF-8 hex fallback as the bulk String fill. +/// +/// # Safety +/// +/// Requires the GIL. Returns an owned reference; the caller must take over +/// the reference count. +unsafe fn shared_string_owned_ptr(py: Python<'_>, payload: &[u8]) -> PyResult<*mut ffi::PyObject> { + let mut reader = ByteReader::new(payload); + let bytes = (|| -> Result<&[u8], BinaryValueError> { + let len = usize::try_from(reader.read_varint()?) + .map_err(|_| BinaryValueError::Invalid("String value length overflows usize".into()))?; + if i32::try_from(len).is_err() { + return Err(BinaryValueError::Invalid( + "String value exceeds i32 offset range".into(), + )); + } + let bytes = reader.read_slice(len)?; + if reader.remaining() != 0 { + return Err(BinaryValueError::Invalid(format!( + "{} trailing bytes after the value", + reader.remaining() + ))); + } + Ok(bytes) + })() + .map_err(|err| shared_cell_err(&err))?; + utf8_or_hex_owned_ptr(py, bytes) +} diff --git a/rust/ch-core-py/tests/helpers.py b/rust/ch-core-py/tests/helpers.py new file mode 100644 index 00000000..dd7f46a3 --- /dev/null +++ b/rust/ch-core-py/tests/helpers.py @@ -0,0 +1,608 @@ +"""Tests for _ch_core Python bindings - Phase 1 types.""" + +import array +import datetime as dt +import decimal +import ipaddress +import math +import os +import struct +import subprocess +import sys +import textwrap +import uuid +from zoneinfo import ZoneInfo + +import pytest + +__all__ = [ + "ZoneInfo", + "array", + "decimal", + "math", + "os", + "subprocess", + "sys", + "textwrap", +] + +_ch_core = pytest.importorskip("_ch_core") + +_EPOCH_DATE = dt.date(1970, 1, 1) +_EPOCH_NAIVE = dt.datetime(1970, 1, 1) + + +class _NdarrayLikeColumn: + def __init__(self, values): + self._values = values + + def __len__(self): + return len(self._values) + + def __getitem__(self, index): + return self._values[index] + + +class _SeriesLikeColumn: + def __init__(self, values): + self._values = values + self.iloc = _NdarrayLikeColumn(values) + + def __len__(self): + return len(self._values) + + def __getitem__(self, index): + raise KeyError(index) + + +def _encode_varint(value: int) -> bytes: + result = bytearray() + while True: + byte = value & 0x7F + value >>= 7 + if value != 0: + byte |= 0x80 + result.append(byte) + if value == 0: + break + return bytes(result) + + +def _encode_varint_string(s: str) -> bytes: + encoded = s.encode("utf-8") + return _encode_varint(len(encoded)) + encoded + + +# Type -> (struct format, byte width) for fixed-width types +_FIXED_TYPES = { + "Bool": ("B", 1), + "Int8": ("b", 1), + "Int16": (" (wire byte width, signedness). Wide integers stay raw little-endian +# fixed-width values in the Native and Arrow representations. +_WIDE_TYPES = { + "Int128": (16, True), + "UInt128": (16, False), + "Int256": (32, True), + "UInt256": (32, False), +} + +_INTERVAL_TYPES = ( + "IntervalYear", + "IntervalQuarter", + "IntervalMonth", + "IntervalWeek", + "IntervalDay", + "IntervalHour", + "IntervalMinute", + "IntervalSecond", + "IntervalMillisecond", + "IntervalMicrosecond", + "IntervalNanosecond", +) + + +def _temporal_struct_fmt(inner_type: str): + """Wire struct format for a temporal type, or None if not temporal. + + Temporal columns are plain bulk integers on the wire; timezone and precision + are type-name metadata only. Values are passed as raw units (Date/Date32 days, + DateTime seconds, DateTime64 ticks, Time seconds, and Time64 ticks). + """ + if inner_type == "Date": + return " int: + """Wire byte width of Decimal(P, S), derived from the precision.""" + if precision <= 9: + return 4 + if precision <= 18: + return 8 + if precision <= 38: + return 16 + return 32 + + +def _uuid_wire_bytes(value: uuid.UUID) -> bytes: + """UUID wire form: little-endian high half then little-endian low half.""" + return (value.int >> 64).to_bytes(8, "little") + ( + value.int & 0xFFFFFFFFFFFFFFFF + ).to_bytes(8, "little") + + +def _bfloat16_bytes(value) -> bytes: + """Truncate a Float32 to its upper 16 bits in Native little-endian order.""" + bits32 = struct.unpack("> 16) + + +def _bfloat16_value(value) -> float: + """Return the Python float represented by ClickHouse's truncated BF16 word.""" + return struct.unpack(" bytes: + """LowCardinality per-column state prefix: the u64 key version.""" + return struct.pack(" 256 or any(i > 255 for i in indices): + raise ValueError("_build_low_cardinality_body: test helper only emits UInt8 indices") + + buf = bytearray() + # HasAdditionalKeysBit | NeedUpdateDictionary | UInt8 index width (tag 0). + buf.extend(struct.pack(" physical nesting (unnamed Point tuple, one Array level per depth). +_GEO_PHYSICAL = { + "Point": _POINT_PHYSICAL, + "Ring": f"Array({_POINT_PHYSICAL})", + "LineString": f"Array({_POINT_PHYSICAL})", + "Polygon": f"Array(Array({_POINT_PHYSICAL}))", + "MultiLineString": f"Array(Array({_POINT_PHYSICAL}))", + "MultiPolygon": f"Array(Array(Array({_POINT_PHYSICAL})))", +} + + +def _expand_alias(type_name): + """Physical type string of a name-decoration alias, else the input unchanged. + + SimpleAggregateFunction, the six geo aliases, and Nested carry a custom name + over a physical type whose wire body is byte-identical, so the helper emits + the alias header but the physical body. Recurses through Nullable + (Nullable(Point) is legal); container recursion is left to the body helpers, + which expand each inner type in turn. + """ + if type_name.startswith("Nullable("): + return f"Nullable({_expand_alias(type_name[len('Nullable('):-1])})" + if type_name in _GEO_PHYSICAL: + return _GEO_PHYSICAL[type_name] + if type_name.startswith("SimpleAggregateFunction("): + return _expand_alias( + _split_top_level_commas(type_name[len("SimpleAggregateFunction("):-1])[1] + ) + if type_name.startswith("Nested("): + return f"Array(Tuple({type_name[len('Nested('):-1]}))" + return type_name + + +def _strip_saf(type_name): + """Strip a SimpleAggregateFunction name-decoration chain to its inner type.""" + while type_name.startswith("SimpleAggregateFunction("): + type_name = _split_top_level_commas(type_name[len("SimpleAggregateFunction("):-1])[1] + return type_name + + +def _lc_dict_value_type(inner): + """(nullable, physical_value_type) for a LowCardinality inner, mirroring the + core: strip the SAF chain, unwrap an optional Nullable, strip the SAF chain + again. LowCardinality(SAF(anyLast, Nullable(String))) is index-level nullable.""" + inner = _strip_saf(inner) + if inner.startswith("Nullable("): + return True, _strip_saf(inner[len("Nullable("):-1]) + return False, inner + + +def _element_state_prefix(type_name): + """State prefix hoisted to the front of a column, recursing into containers. + + Only LowCardinality contributes bytes (its u64 key version). Array, Tuple, + Map, and Nullable recurse into their inner types in serialization order; + every leaf writes nothing. Matches the core's read_state_prefix chain. + """ + type_name = _expand_alias(type_name) + if type_name.startswith("Array("): + return _element_state_prefix(type_name[len("Array("):-1]) + if type_name.startswith("Nullable("): + return _element_state_prefix(type_name[len("Nullable("):-1]) + if type_name.startswith("Tuple("): + return b"".join(_element_state_prefix(t) for _, t in _parse_tuple_elements(type_name)) + if type_name.startswith("Map("): + key, value = _parse_map_types(type_name) + return _element_state_prefix(key) + _element_state_prefix(value) + if type_name.startswith("LowCardinality("): + return _lc_key_version() + return b"" + + +def _build_body_no_prefix(type_name, values): + """Column body with its state prefix omitted (already hoisted by the caller).""" + type_name = _expand_alias(type_name) + if type_name.startswith("Array("): + return _build_array_body(type_name, values) + if type_name.startswith("Map("): + return _build_map_body(type_name, values) + if type_name.startswith("Tuple("): + return _build_tuple_body(type_name, values, nullable=False) + if type_name.startswith("LowCardinality("): + return _build_low_cardinality_body_no_prefix(type_name, values) + if type_name.startswith("Nullable("): + inner = type_name[len("Nullable("):-1] + # Nullable(Tuple) is the only nullable container: the tuple-level null + # map precedes the tuple body (a null row still carries placeholder + # element values on the wire). + if inner.startswith("Tuple("): + return _build_tuple_body(inner, values, nullable=True) + buf = bytearray() + for v in values: + buf.append(0x01 if v is None else 0x00) + buf.extend(_encode_plain_body(inner, values)) + return bytes(buf) + return _encode_plain_body(type_name, values) + + +def _build_array_body(type_name, rows): + """Array(T) body: absolute end-offsets (no leading zero) then the element body. + + The element body is the flattened element column written WITHOUT its state + prefix, which is hoisted to the front of the whole column. + """ + inner = type_name[len("Array("):-1] + flat = [] + offsets = bytearray() + for row in rows: + flat.extend(row) + offsets.extend(struct.pack("0 gate). + if num_rows == 0: + continue + + # Emit the alias header but the physical type's body: a name-decoration + # alias (geo, Nested, SimpleAggregateFunction) is byte-identical on the + # wire to its physical type. + physical = _expand_alias(type_name) + + if physical.startswith("LowCardinality("): + buf.extend(_build_low_cardinality_body(physical, values)) + continue + + if physical.startswith("Array("): + buf.extend(_element_state_prefix(physical)) + buf.extend(_build_array_body(physical, values)) + continue + + nullable_tuple = ( + physical.startswith("Nullable(") + and physical[len("Nullable("):-1].startswith("Tuple(") + ) + if physical.startswith("Tuple(") or physical.startswith("Map(") or nullable_tuple: + buf.extend(_element_state_prefix(physical)) + buf.extend(_build_body_no_prefix(physical, values)) + continue + + is_nullable = physical.startswith("Nullable(") + inner_type = physical[len("Nullable("):-1] if is_nullable else physical + if is_nullable: + for v in values: + buf.append(0x01 if v is None else 0x00) + buf.extend(_encode_plain_body(inner_type, values)) + return bytes(buf) + + +def build_native_block_from_bodies(columns, row_count): + """Build a Native block from already-encoded column bodies.""" + buf = bytearray() + buf.extend(_encode_varint(len(columns))) + buf.extend(_encode_varint(row_count)) + for name, type_name, body in columns: + buf.extend(_encode_varint_string(name)) + buf.extend(_encode_varint_string(type_name)) + buf.extend(body) + return bytes(buf) + + +# --------------------------------------------------------------------------- +# Native insert encoding +# --------------------------------------------------------------------------- diff --git a/rust/ch-core-py/tests/test_containers.py b/rust/ch-core-py/tests/test_containers.py new file mode 100644 index 00000000..a6510005 --- /dev/null +++ b/rust/ch-core-py/tests/test_containers.py @@ -0,0 +1,1427 @@ +from helpers import ( + ZoneInfo, + _ch_core, + _encode_varint, + _encode_varint_string, + _NdarrayLikeColumn, + build_native_block, + build_native_block_from_bodies, + decimal, + dt, + ipaddress, + pytest, + struct, + sys, + uuid, +) + +# --------------------------------------------------------------------------- +# Array +# --------------------------------------------------------------------------- + +_ARR_KNOWN_UUID = uuid.UUID("00112233-4455-6677-8899-aabbccddeeff") +_ARR_NY = ZoneInfo("America/New_York") + +# DateTime element: tz-aware New York input round-trips to tz-aware New York. +_ARR_DT_TZ_ROWS = [ + [dt.datetime(2024, 1, 15, 7, 34, 56, tzinfo=_ARR_NY)], + [], + [ + dt.datetime(2000, 6, 15, 12, 0, 0, tzinfo=_ARR_NY), + dt.datetime(1970, 1, 1, 0, 0, 0, tzinfo=_ARR_NY), + ], +] + +# DateTime64(3) has no timezone, so tz-aware UTC input decodes to the naive UTC +# wall clock; passing UTC keeps the expectation independent of the host locale. +_ARR_DT64_UTC_ROWS = [ + [dt.datetime(2024, 1, 15, 12, 34, 56, 789000, tzinfo=dt.timezone.utc)], + [], + [ + dt.datetime(1970, 1, 1, 0, 0, 0, 1000, tzinfo=dt.timezone.utc), + dt.datetime(2000, 6, 15, 12, 0, 0, 500000, tzinfo=dt.timezone.utc), + ], +] +_ARR_DT64_EXPECTED = [ + [dt.datetime(2024, 1, 15, 12, 34, 56, 789000)], + [], + [ + dt.datetime(1970, 1, 1, 0, 0, 0, 1000), + dt.datetime(2000, 6, 15, 12, 0, 0, 500000), + ], +] + +# (type_name, py_rows, expected) with expected=None meaning it equals py_rows. +# Every case has multiple rows, at least one empty array, ragged lengths, and +# nullable element types mix None with values. +_ARR_ROUND_TRIP = [ + ("Array(Int32)", [[13, 79], [], [-1, 0, 2147483647], [7]], None), + ("Array(Int64)", [[13], [], [9223372036854775807, -9223372036854775808], [5, 5]], None), + ("Array(UInt64)", [[0, 18446744073709551615], [], [13, 79, 5]], None), + ("Array(Float64)", [[1.5, -2.5], [], [0.0, 1e300]], None), + ("Array(Bool)", [[True, False, True], [], [False]], None), + ("Array(String)", [["user_1", "user_2"], [], ["", "sventon"]], None), + ("Array(FixedString(3))", [[b"abc", b"xyz"], [], [b"a\x00\x00"]], None), + ("Array(Date)", [[dt.date(2024, 1, 2), dt.date(1970, 1, 1)], [], [dt.date(2000, 6, 15)]], None), + ("Array(DateTime('America/New_York'))", _ARR_DT_TZ_ROWS, None), + ("Array(DateTime64(3))", _ARR_DT64_UTC_ROWS, _ARR_DT64_EXPECTED), + ("Array(UUID)", [[uuid.UUID(int=0), _ARR_KNOWN_UUID], [], [uuid.UUID(int=79)]], None), + ( + "Array(IPv4)", + [ + [ipaddress.IPv4Address("0.0.0.0"), ipaddress.IPv4Address("1.2.3.4")], + [], + [ipaddress.IPv4Address("255.255.255.255")], + ], + None, + ), + ( + "Array(IPv6)", + [ + [ipaddress.IPv6Address("::1")], + [], + [ipaddress.IPv6Address("2001:db8::1"), ipaddress.IPv6Address("::ffff:1.2.3.4")], + ], + None, + ), + ( + "Array(Decimal(9, 2))", + [[decimal.Decimal("1.00"), decimal.Decimal("-3.50")], [], [decimal.Decimal("9999999.99")]], + None, + ), + ( + "Array(Decimal(38, 10))", + [[decimal.Decimal("1.5"), decimal.Decimal("-2.25")], [], [decimal.Decimal("0")]], + None, + ), + ("Array(Enum8('a' = 1, 'b' = 2))", [["a", "b"], [], ["b", "a", "b"]], None), + ("Array(Nullable(Int32))", [[13, None, 79], [], [None], [7]], None), + ("Array(Nullable(String))", [["user_1", None], [], [None, "x"]], None), + ("Array(LowCardinality(String))", [["red", "green", "red"], [], ["blue"]], None), + ("Array(LowCardinality(Nullable(String)))", [["x", None, "x"], [], [None, "y"]], None), + ("Array(Array(Int32))", [[[13, 79], [5]], [], [[7]], [[], [1, 2, 3]]], None), +] + +# (type_name, py_rows, wire_rows, expected): wire_rows is the raw wire form the +# helper serializes, py_rows is what the encoder is given, expected is the decode. +# Excludes LowCardinality-in-array because the encoder sets the index word's +# NeedUpdateDictionary bit that the helper does not, so their bytes differ. +_ARR_GOLDEN = [ + ("Array(Int32)", [[13, 79], [], [7]], [[13, 79], [], [7]], [[13, 79], [], [7]]), + ("Array(String)", [["user_1", "user_2"], [], ["x"]], [["user_1", "user_2"], [], ["x"]], [["user_1", "user_2"], [], ["x"]]), + ("Array(FixedString(3))", [[b"abc"], [], [b"xyz", b"a\x00\x00"]], [[b"abc"], [], [b"xyz", b"a\x00\x00"]], [[b"abc"], [], [b"xyz", b"a\x00\x00"]]), + ("Array(UInt64)", [[0, 18446744073709551615], [], [13]], [[0, 18446744073709551615], [], [13]], [[0, 18446744073709551615], [], [13]]), + ("Array(Nullable(Int32))", [[13, None], [], [7]], [[13, None], [], [7]], [[13, None], [], [7]]), + ("Array(Array(Int32))", [[[13, 79], [5]], [], [[7]]], [[[13, 79], [5]], [], [[7]]], [[[13, 79], [5]], [], [[7]]]), + ( + "Array(UUID)", + [[uuid.UUID(int=0), _ARR_KNOWN_UUID], [], [uuid.UUID(int=79)]], + [[uuid.UUID(int=0), _ARR_KNOWN_UUID], [], [uuid.UUID(int=79)]], + [[uuid.UUID(int=0), _ARR_KNOWN_UUID], [], [uuid.UUID(int=79)]], + ), + ( + "Array(IPv4)", + [[ipaddress.IPv4Address("1.2.3.4")], [], [ipaddress.IPv4Address("0.0.0.0"), ipaddress.IPv4Address("255.255.255.255")]], + [[ipaddress.IPv4Address("1.2.3.4")], [], [ipaddress.IPv4Address("0.0.0.0"), ipaddress.IPv4Address("255.255.255.255")]], + [[ipaddress.IPv4Address("1.2.3.4")], [], [ipaddress.IPv4Address("0.0.0.0"), ipaddress.IPv4Address("255.255.255.255")]], + ), + ( + "Array(IPv6)", + [[ipaddress.IPv6Address("::1")], [], [ipaddress.IPv6Address("2001:db8::1")]], + [[ipaddress.IPv6Address("::1")], [], [ipaddress.IPv6Address("2001:db8::1")]], + [[ipaddress.IPv6Address("::1")], [], [ipaddress.IPv6Address("2001:db8::1")]], + ), + ( + "Array(Date)", + [[dt.date(2024, 1, 2)], [], [dt.date(1970, 1, 1), dt.date(2000, 6, 15)]], + [ + [dt.date(2024, 1, 2).toordinal() - 719163], + [], + [0, dt.date(2000, 6, 15).toordinal() - 719163], + ], + [[dt.date(2024, 1, 2)], [], [dt.date(1970, 1, 1), dt.date(2000, 6, 15)]], + ), + ( + "Array(DateTime64(3))", + [ + [dt.datetime(2024, 1, 15, 12, 34, 56, 789000, tzinfo=dt.timezone.utc)], + [], + [dt.datetime(1970, 1, 1, 0, 0, 0, 1000, tzinfo=dt.timezone.utc)], + ], + [[1705322096789], [], [1]], + [ + [dt.datetime(2024, 1, 15, 12, 34, 56, 789000)], + [], + [dt.datetime(1970, 1, 1, 0, 0, 0, 1000)], + ], + ), + ( + "Array(Decimal(9, 2))", + [[decimal.Decimal("1.00"), decimal.Decimal("-3.50")], [], [decimal.Decimal("0.00")]], + [[100, -350], [], [0]], + [[decimal.Decimal("1.00"), decimal.Decimal("-3.50")], [], [decimal.Decimal("0.00")]], + ), + ( + "Array(Enum8('a' = 1, 'b' = 2))", + [["a", "b"], [], ["a"]], + [[1, 2], [], [1]], + [["a", "b"], [], ["a"]], + ), +] + + +class TestArray: + @pytest.mark.parametrize("type_name,py_rows,expected", _ARR_ROUND_TRIP) + def test_round_trip(self, type_name, py_rows, expected): + if expected is None: + expected = py_rows + encoded = _ch_core.encode_native_block(["a"], [type_name], [py_rows], len(py_rows)) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == expected + + @pytest.mark.parametrize("type_name,py_rows,wire_rows,expected", _ARR_GOLDEN) + def test_golden_bytes(self, type_name, py_rows, wire_rows, expected): + encoded = _ch_core.encode_native_block(["a"], [type_name], [py_rows], len(py_rows)) + built = build_native_block([("a", type_name, wire_rows)]) + assert encoded == built + assert list(_ch_core.ColBatch.decode_native(built).column_data(0)) == expected + + def test_golden_decode_low_cardinality_in_array(self): + # encode == build is not asserted: the encoder sets the index word's + # NeedUpdateDictionary bit while the helper does not, so their bytes + # differ. Cover it by decode-of-helper-bytes plus an encode round-trip. + rows = [["red", "green", "red"], [], ["blue"]] + built = build_native_block([("c", "Array(LowCardinality(String))", rows)]) + assert list(_ch_core.ColBatch.decode_native(built).column_data(0)) == rows + encoded = _ch_core.encode_native_block( + ["c"], ["Array(LowCardinality(String))"], [rows], len(rows) + ) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == rows + + def test_low_cardinality_in_array_all_empty(self): + # Every array empty (total_elements == 0): the LowCardinality element + # body is absent entirely, so only the hoisted key version and the zero + # offsets remain. With no index word to differ, encode == build holds. + rows = [[], [], []] + built = build_native_block([("c", "Array(LowCardinality(String))", rows)]) + encoded = _ch_core.encode_native_block( + ["c"], ["Array(LowCardinality(String))"], [rows], len(rows) + ) + assert encoded == built + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == rows + + @pytest.mark.parametrize( + "type_name,rows", + [ + ("Array(Int32)", [[13, 79], [], [7, 5, 5]]), + ("Array(Nullable(String))", [["user_1", None], [], [None, "x"]]), + ("Array(Enum8('a' = 1, 'b' = 2))", [["a", "b"], [], ["b", "a", "b"]]), + ("Array(DateTime('America/New_York'))", _ARR_DT_TZ_ROWS), + ("Array(UUID)", [[uuid.UUID(int=0), _ARR_KNOWN_UUID], [], [uuid.UUID(int=79)]]), + ("Array(LowCardinality(Nullable(String)))", [["x", None, "x"], [], [None, "y"]]), + ("Array(Array(Int64))", [[[13, 79], [5]], [], [[7]], [[], [1, 2, 3]]]), + ], + ) + def test_all_exit_paths_agree(self, type_name, rows): + encoded = _ch_core.encode_native_block(["a"], [type_name], [rows], len(rows)) + batch = _ch_core.ColBatch.decode_native(encoded) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == rows + + def test_multi_block_concatenation(self): + rows_a = [[13, 79], [], [7]] + rows_b = [[5], [1, 2], []] + block_a = _ch_core.encode_native_block(["a"], ["Array(Int32)"], [rows_a], len(rows_a)) + block_b = _ch_core.encode_native_block(["a"], ["Array(Int32)"], [rows_b], len(rows_b)) + merged = _ch_core.ColBatch.from_batches( + [ + _ch_core.ColBatch.decode_native(block_a), + _ch_core.ColBatch.decode_native(block_b), + ] + ) + assert merged.num_chunks == 2 + assert list(merged.column_data(0)) == rows_a + rows_b + concat = _ch_core.ColBatch.decode_native(block_a + block_b) + assert list(concat.column_data(0)) == rows_a + rows_b + + def test_zero_rows(self): + encoded = _ch_core.encode_native_block(["a"], ["Array(Int32)"], [[]], 0) + assert encoded == build_native_block([("a", "Array(Int32)", [])]) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.num_rows == 0 + assert batch.column_type_names == ["Array(Int32)"] + assert list(batch.column_data(0)) == [] + + def test_nullable_array_type_rejected(self): + with pytest.raises(NotImplementedError, match="unsupported ClickHouse type"): + _ch_core.encode_native_block(["a"], ["Nullable(Array(Int32))"], [[[13]]], 1) + + def test_none_row_rejected(self): + with pytest.raises(ValueError, match="is None but Array"): + _ch_core.encode_native_block(["a"], ["Array(Int32)"], [[None]], 1) + + def test_bare_str_row_rejected(self): + with pytest.raises(ValueError, match="is a str, not an Array sequence"): + _ch_core.encode_native_block(["a"], ["Array(Int32)"], [["abc"]], 1) + + def test_bytes_like_rows_flatten_as_int_elements(self): + rows = [b"\x01\x02", bytearray(b"\x03"), memoryview(b"\x04\x05"), []] + expected = _ch_core.encode_native_block( + ["a"], ["Array(UInt8)"], [[[1, 2], [3], [4, 5], []]], 4 + ) + assert _ch_core.encode_native_block(["a"], ["Array(UInt8)"], [rows], 4) == expected + + def test_bytes_row_for_string_elements_raises_conversion_error(self): + with pytest.raises(ValueError, match="row 0 element 0 cannot be converted to String"): + _ch_core.encode_native_block(["a"], ["Array(String)"], [[b"ab"]], 1) + + def test_element_error_reports_outer_row_and_element(self): + for outer in ([[1, 2], ["x"]], _NdarrayLikeColumn([[1, 2], ["x"]])): + with pytest.raises( + ValueError, match='column "v" row 1 element 0 cannot be converted to Int32' + ): + _ch_core.encode_native_block(["v"], ["Array(Int32)"], [outer], 2) + with pytest.raises( + ValueError, match='column "v" row 0 element 1 is None but Int32 is not Nullable' + ): + _ch_core.encode_native_block(["v"], ["Array(Int32)"], [[[1, None]]], 1) + + def test_nested_element_error_reports_outer_path(self): + rows = [[[1]], [[2], ["x"]]] + with pytest.raises( + ValueError, + match='column "v" row 1 element 1 element 0 cannot be converted to Int32', + ): + _ch_core.encode_native_block(["v"], ["Array(Array(Int32))"], [rows], 2) + + def test_array_lc_reuses_dictionary_objects(self): + rows = [["red", "red", "green"], [], ["red", "green"]] + encoded = _ch_core.encode_native_block( + ["c"], ["Array(LowCardinality(String))"], [rows], len(rows) + ) + batch = _ch_core.ColBatch.decode_native(encoded) + col = batch.column_data(0) + assert list(col) == rows + assert col[0][0] is col[0][1] + assert col[0][0] is col[2][0] + cols = batch.to_python_columns()[0] + assert cols[0][0] is cols[0][1] and cols[0][0] is cols[2][0] + out_rows = batch.to_python_rows() + assert out_rows[0][0][0] is out_rows[0][0][1] + assert out_rows[0][0][0] is out_rows[2][0][0] + + def test_nested_array_lc_reuses_dictionary_objects(self): + rows = [[["a", "a"], ["a"]], [], [["a"]]] + encoded = _ch_core.encode_native_block( + ["c"], ["Array(Array(LowCardinality(String)))"], [rows], len(rows) + ) + batch = _ch_core.ColBatch.decode_native(encoded) + col = batch.column_data(0) + assert list(col) == rows + assert col[0][0][0] is col[0][1][0] + assert col[0][0][0] is col[2][0][0] + + def test_array_lc_nullable_reuse_and_values(self): + rows = [["x", None, "x"], [], [None, "x"]] + encoded = _ch_core.encode_native_block( + ["c"], ["Array(LowCardinality(Nullable(String)))"], [rows], len(rows) + ) + col = _ch_core.ColBatch.decode_native(encoded).column_data(0) + assert list(col) == rows + assert col[0][0] is col[0][2] + assert col[0][0] is col[2][1] + + def test_json_element_round_trip(self): + rows = [[{"a": 13}, {"b": "user_1"}], [], [{"c": [1, 2]}]] + encoded = _ch_core.encode_native_block(["a"], ["Array(JSON)"], [rows], len(rows)) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == rows + + def test_deeply_nested_type_rejected_not_crash(self): + # Past the parser depth cap the type is rejected without a stack overflow. + deep = "Array(" * 200 + "Int32" + ")" * 200 + with pytest.raises(NotImplementedError): + _ch_core.encode_native_block(["c"], [deep], [[]], 0) + + def test_unordered_row_rejected(self): + for bad in ({3, 1, 2}, {13: 0, 79: 0}): + with pytest.raises(ValueError): + _ch_core.encode_native_block(["v"], ["Array(Int32)"], [[bad]], 1) + + def test_malformed_offsets_decode_error(self): + buf = bytearray() + buf.extend(_encode_varint(1)) + buf.extend(_encode_varint(2)) + buf.extend(_encode_varint_string("a")) + buf.extend(_encode_varint_string("Array(Int32)")) + buf.extend(struct.pack(" invalid + with pytest.raises(ValueError, match="Invalid Array layout"): + _ch_core.ColBatch.decode_native(bytes(buf)) + + +class TestArrayInsertFastPath: + """Exact list/tuple row flattening into one flat element run.""" + + def _encode(self, type_name, rows, row_count=None): + n = len(rows) if row_count is None else row_count + return _ch_core.encode_native_block(["v"], [type_name], [rows], n) + + def test_row_and_outer_container_kinds_agree(self): + rows = [[13, 79], [], [7, 5, 5]] + from_lists = self._encode("Array(Int64)", rows) + assert from_lists == build_native_block([("v", "Array(Int64)", rows)]) + assert self._encode("Array(Int64)", tuple(rows), 3) == from_lists + assert self._encode("Array(Int64)", [tuple(r) for r in rows]) == from_lists + assert self._encode("Array(Int64)", _NdarrayLikeColumn(rows), 3) == from_lists + + def test_exotic_row_containers_match_list_rows(self): + expected = self._encode("Array(Int64)", [[0, 1, 2], [7], [], [3, 4]]) + rows = [range(3), _NdarrayLikeColumn([7]), (x for x in []), (3, 4)] + assert self._encode("Array(Int64)", rows, 4) == expected + + def test_mixed_rows_agree_across_outer_containers(self): + rows = [[1, 2], (3,), range(2)] + expected = self._encode("Array(Int64)", [[1, 2], [3], [0, 1]]) + assert self._encode("Array(Int64)", rows, 3) == expected + assert self._encode("Array(Int64)", _NdarrayLikeColumn(rows), 3) == expected + + def test_nested_tuple_rows(self): + rows = [((1, 2), [3]), [], [(7,)], [[], (1, 2, 3)]] + expected_rows = [[[1, 2], [3]], [], [[7]], [[], [1, 2, 3]]] + encoded = self._encode("Array(Array(Int64))", rows, 4) + assert encoded == self._encode("Array(Array(Int64))", expected_rows) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == expected_rows + + def test_nullable_element_tuple_rows(self): + rows = [(13, None), (), (None, 7)] + encoded = self._encode("Array(Nullable(Int64))", rows, 3) + assert encoded == self._encode("Array(Nullable(Int64))", [list(r) for r in rows]) + expected = [[13, None], [], [None, 7]] + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == expected + + def test_string_and_lc_tuple_rows(self): + rows = [("a", "bb"), (), ("a",)] + for type_name in ("Array(String)", "Array(LowCardinality(String))"): + fast = self._encode(type_name, rows, 3) + assert fast == self._encode(type_name, [list(r) for r in rows]) + + def test_lc_non_string_element_round_trip(self): + rows = [[13, 79, 13], [], [79]] + encoded = self._encode("Array(LowCardinality(UInt32))", rows) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == rows + + def test_empty_rows_tuple_outer(self): + assert self._encode("Array(Int64)", (), 0) == build_native_block( + [("v", "Array(Int64)", [])] + ) + assert self._encode("Array(Int64)", ((), (), ()), 3) == build_native_block( + [("v", "Array(Int64)", [[], [], []])] + ) + + def test_invalid_row_in_fallback_reports_row(self): + with pytest.raises(ValueError, match="row 1 is not a valid Array value"): + self._encode("Array(Int64)", [[1], 5, [3]], 3) + + def test_row_fallback_preserves_cause(self): + class BoomError(Exception): + pass + + class EvilRow: + def __iter__(self): + raise BoomError("iteration exploded") + + with pytest.raises(ValueError, match="row 1 is not a valid Array value") as excinfo: + self._encode("Array(Int64)", [[1], EvilRow()], 2) + assert isinstance(excinfo.value.__cause__, BoomError) + + def test_flat_refs_refcount_balance(self): + shared = 1 << 40 + + rows = [[shared, shared], [], [shared]] + before = sys.getrefcount(shared) + self._encode("Array(Int64)", rows, 3) + assert sys.getrefcount(shared) == before + + rows = [[shared], ["x"]] + before = sys.getrefcount(shared) + with pytest.raises(ValueError, match="element 0 cannot be converted"): + self._encode("Array(Int64)", rows, 2) + assert sys.getrefcount(shared) == before + + rows = [[shared], 5] + before = sys.getrefcount(shared) + with pytest.raises(ValueError, match="not a valid Array value"): + self._encode("Array(Int64)", rows, 2) + assert sys.getrefcount(shared) == before + + def test_outer_list_resized_during_row_fallback_raises(self): + rows = [[1], None, [3], [4]] + + class Evil: + def __iter__(self): + del rows[2:] + return iter([7]) + + rows[1] = Evil() + with pytest.raises(ValueError, match="resized during encoding"): + self._encode("Array(Int64)", rows, 4) + + def test_outer_list_resized_during_element_fallback_encodes_snapshot(self): + # Element conversion runs Python that clears the outer list after the + # flatten pass; the flat run holds strong references, so the original + # rows still encode. + rows = [[1], [], [3]] + + class Evil: + def __index__(self): + rows.clear() + return 9 + + rows[1] = [Evil()] + expected = self._encode("Array(Int64)", [[1], [9], [3]]) + assert self._encode("Array(Int64)", rows, 3) == expected + + +class TestTupleInsert: + """Tuple(T1, ...) encode: positional and named-dict rows.""" + + def _encode(self, type_name, rows, row_count=None): + n = len(rows) if row_count is None else row_count + return _ch_core.encode_native_block(["t"], [type_name], [rows], n) + + def _decode(self, encoded): + return list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) + + def test_unnamed_rows_match_golden_and_round_trip(self): + rows = [(13, "a"), (-1, ""), (79, "sventon")] + encoded = self._encode("Tuple(Int32, String)", rows) + assert encoded == build_native_block([("t", "Tuple(Int32, String)", rows)]) + assert self._decode(encoded) == rows + + def test_row_and_outer_container_kinds_agree(self): + rows = [(13, "a"), (79, "b")] + expected = self._encode("Tuple(Int32, String)", rows) + assert self._encode("Tuple(Int32, String)", [list(r) for r in rows]) == expected + assert self._encode("Tuple(Int32, String)", tuple(rows), 2) == expected + assert self._encode("Tuple(Int32, String)", _NdarrayLikeColumn(rows), 2) == expected + + def test_generic_iterable_rows(self): + expected = self._encode("Tuple(Int64, Int64)", [(0, 1), (5, 6)]) + assert self._encode("Tuple(Int64, Int64)", [range(2), (5, 6)], 2) == expected + + def test_named_tuple_positional_rows(self): + rows = [(13, "x"), (79, "y")] + encoded = self._encode("Tuple(a Int32, b String)", rows) + assert encoded == build_native_block([("t", "Tuple(a Int32, b String)", rows)]) + assert self._decode(encoded) == [{"a": 13, "b": "x"}, {"a": 79, "b": "y"}] + + def test_named_tuple_dict_rows(self): + dict_rows = [{"a": 13, "b": "x"}, {"a": 79, "b": "y"}] + expected = self._encode("Tuple(a Int32, b String)", [(13, "x"), (79, "y")]) + assert self._encode("Tuple(a Int32, b String)", dict_rows) == expected + # Extra keys are ignored. + dict_rows = [{"a": 13, "b": "x", "zz": 5}, {"b": "y", "a": 79}] + assert self._encode("Tuple(a Int32, b String)", dict_rows) == expected + + def test_backtick_named_tuple_dict_rows(self): + type_name = "Tuple(`a b` Int32, c String)" + expected = self._encode(type_name, [(13, "x")]) + assert expected == build_native_block([("t", type_name, [(13, "x")])]) + assert self._encode(type_name, [{"a b": 13, "c": "x"}]) == expected + + def test_dict_row_missing_key_nullable_becomes_none(self): + type_name = "Tuple(a Int32, b Nullable(String))" + assert self._encode(type_name, [{"a": 13}]) == self._encode(type_name, [(13, None)]) + + def test_dict_row_missing_key_non_nullable_raises(self): + with pytest.raises(ValueError, match='column "t" row 0 element "b" is None'): + self._encode("Tuple(a Int32, b String)", [{"a": 13}]) + + def test_dict_subclass_rows_read_via_get(self): + import collections + + rows = [collections.OrderedDict([("a", 13), ("b", "x")])] + assert self._encode("Tuple(a Int32, b String)", rows) == self._encode( + "Tuple(a Int32, b String)", [(13, "x")] + ) + + def test_non_dict_row_in_dict_mode_raises(self): + rows = [{"a": 13, "b": "x"}, (79, "y")] + with pytest.raises(ValueError, match="row 1 cannot be read as a dict"): + self._encode("Tuple(a Int32, b String)", rows, 2) + + def test_dict_row_in_positional_mode_raises(self): + rows = [(13, "x"), {"a": 79, "b": "y"}] + with pytest.raises(ValueError, match="row 1 is a dict but Tuple rows are read"): + self._encode("Tuple(a Int32, b String)", rows, 2) + with pytest.raises(ValueError, match="row 0 is a dict but Tuple rows are read"): + self._encode("Tuple(Int32, String)", [{"a": 1}], 1) + + def test_arity_mismatch_raises(self): + with pytest.raises( + ValueError, match='column "t" row 1 has 3 elements but the Tuple declares 2' + ): + self._encode("Tuple(Int32, String)", [(1, "a"), (1, "a", "extra")], 2) + with pytest.raises(ValueError, match="row 0 has 1 elements but the Tuple declares 2"): + self._encode("Tuple(Int32, String)", [[1]], 1) + + def test_str_row_rejected(self): + with pytest.raises(ValueError, match="row 0 is a str, not a Tuple row"): + self._encode("Tuple(String, String)", ["ab"], 1) + + def test_none_row_non_nullable_raises(self): + with pytest.raises(ValueError, match="row 1 is None but Tuple"): + self._encode("Tuple(Int32, String)", [(1, "a"), None], 2) + + def test_element_error_names_column_row_and_element(self): + with pytest.raises(ValueError, match='column "t" row 1 element 0 cannot be converted'): + self._encode("Tuple(Int64, String)", [(1, "a"), ("x", "b")], 2) + with pytest.raises(ValueError, match='column "t" row 0 element "b" cannot be converted'): + self._encode("Tuple(a Int32, b Int32)", [{"a": 1, "b": "x"}]) + + def test_lc_string_element(self): + # The core encoder's LC flags word differs from the test helper's + # (both valid), so LC types round-trip instead of golden-comparing. + type_name = "Tuple(LowCardinality(String), Int64)" + rows = [("red", 1), ("red", 2), ("blue", 3)] + encoded = self._encode(type_name, rows) + assert self._decode(encoded) == rows + + def test_array_and_nullable_elements(self): + type_name = "Tuple(Array(Int64), Nullable(String))" + rows = [([13, 79], "a"), ([], None)] + encoded = self._encode(type_name, rows) + assert self._decode(encoded) == rows + + def test_nested_tuple_and_map_elements(self): + rows = [(1, (2, "x")), (3, (4, "y"))] + encoded = self._encode("Tuple(UInt8, Tuple(UInt8, String))", rows) + assert self._decode(encoded) == rows + rows = [({"k1": 1, "k2": 2}, 10), ({}, 20)] + encoded = self._encode("Tuple(Map(String, UInt8), Int32)", rows) + assert self._decode(encoded) == rows + + def test_nullable_tuple_none_rows_match_golden_and_round_trip(self): + type_name = "Nullable(Tuple(Nullable(Int32), String))" + rows = [(13, "a"), None, (None, "b"), None] + encoded = self._encode(type_name, rows) + assert encoded == build_native_block([("t", type_name, rows)]) + assert self._decode(encoded) == rows + + def test_nullable_tuple_default_placeholders_cover_scalar_types(self): + type_name = "Nullable(Tuple(UUID, IPv4, Date, Decimal(9, 2), Bool, Float64))" + rows = [None, (uuid.UUID(int=13), ipaddress.IPv4Address("1.2.3.4"), 79, 5, True, 1.5)] + encoded = self._encode(type_name, rows, 2) + decoded = self._decode(encoded) + assert decoded[0] is None + assert decoded[1] == ( + uuid.UUID(int=13), + ipaddress.IPv4Address("1.2.3.4"), + dt.date(1970, 3, 21), + decimal.Decimal("5.00"), + True, + 1.5, + ) + + def test_nullable_named_tuple_dict_rows_with_none(self): + type_name = "Nullable(Tuple(a UInt64, b String))" + rows = [{"a": 13, "b": "x"}, None, {"a": 5, "b": "y"}] + encoded = self._encode(type_name, rows, 3) + assert self._decode(encoded) == [{"a": 13, "b": "x"}, None, {"a": 5, "b": "y"}] + + def test_empty_tuple(self): + rows = [(), (), ()] + encoded = self._encode("Tuple()", rows) + assert encoded == build_native_block([("t", "Tuple()", rows)]) + assert self._decode(encoded) == rows + with pytest.raises(ValueError, match="has 1 elements but the Tuple declares 0"): + self._encode("Tuple()", [(1,)], 1) + + def test_array_of_tuple(self): + type_name = "Array(Tuple(Int64, String))" + rows = [[(1, "a"), (2, "b")], [], [(3, "c")]] + encoded = self._encode(type_name, rows) + assert self._decode(encoded) == rows + + def test_zero_row_probe(self): + encoded = _ch_core.encode_native_block( + ["t", "m"], ["Tuple(a Int32, b String)", "Map(String, UInt8)"], [[], []], 0 + ) + assert encoded == build_native_block( + [("t", "Tuple(a Int32, b String)", []), ("m", "Map(String, UInt8)", [])] + ) + + def test_lc_tuple_and_nullable_map_still_rejected(self): + with pytest.raises(NotImplementedError, match="unsupported"): + _ch_core.encode_native_block( + ["v"], ["LowCardinality(Tuple(Int32, String))"], [[(1, "a")]], 1 + ) + with pytest.raises(NotImplementedError, match="unsupported"): + _ch_core.encode_native_block(["v"], ["Nullable(Map(String, UInt8))"], [[{}]], 1) + + def test_flat_refs_refcount_balance(self): + shared = 1 << 40 + + rows = [(shared, "a"), (shared, "b")] + before = sys.getrefcount(shared) + self._encode("Tuple(Int64, String)", rows) + assert sys.getrefcount(shared) == before + + rows = [(shared, "a"), (shared, 5)] + before = sys.getrefcount(shared) + with pytest.raises(ValueError, match="element 1"): + self._encode("Tuple(Int64, String)", rows, 2) + assert sys.getrefcount(shared) == before + + rows = [(shared, "a"), (shared,)] + before = sys.getrefcount(shared) + with pytest.raises(ValueError, match="declares 2"): + self._encode("Tuple(Int64, String)", rows, 2) + assert sys.getrefcount(shared) == before + + rows = [{"a": shared, "b": "x"}, {"a": shared}] + before = sys.getrefcount(shared) + with pytest.raises(ValueError, match='element "b" is None'): + self._encode("Tuple(a Int64, b String)", rows, 2) + assert sys.getrefcount(shared) == before + + +class TestMapInsert: + """Map(K, V) encode: dict rows flattened into key and value runs.""" + + def _encode(self, type_name, rows, row_count=None): + n = len(rows) if row_count is None else row_count + return _ch_core.encode_native_block(["m"], [type_name], [rows], n) + + def _decode(self, encoded): + return list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) + + def test_dict_rows_match_golden_and_round_trip(self): + rows = [{"k1": 1, "k2": 2}, {}, {"x": 255}] + encoded = self._encode("Map(String, UInt8)", rows) + assert encoded == build_native_block([("m", "Map(String, UInt8)", rows)]) + assert self._decode(encoded) == rows + + def test_outer_container_kinds_agree(self): + rows = [{13: "a", 79: "b"}, {}, {5: "c"}] + expected = self._encode("Map(UInt64, String)", rows) + assert self._encode("Map(UInt64, String)", tuple(rows), 3) == expected + assert self._encode("Map(UInt64, String)", _NdarrayLikeColumn(rows), 3) == expected + + def test_all_empty_dict_rows(self): + rows = [{}, {}, {}] + encoded = self._encode("Map(String, UInt8)", rows) + assert encoded == build_native_block([("m", "Map(String, UInt8)", rows)]) + assert self._decode(encoded) == rows + + def test_dict_subclass_rows_via_items(self): + import collections + + rows = [collections.OrderedDict([("a", 13), ("b", 79)]), {}] + assert self._encode("Map(String, UInt8)", rows) == self._encode( + "Map(String, UInt8)", [{"a": 13, "b": 79}, {}] + ) + + def test_non_dict_row_rejected(self): + with pytest.raises(ValueError, match="row 1 is not a dict for Map"): + self._encode("Map(String, UInt8)", [{"a": 1}, [("b", 2)]], 2) + with pytest.raises(ValueError, match="row 0 is not a dict for Map"): + self._encode("Map(String, UInt8)", ["ab"], 1) + + def test_none_row_rejected(self): + with pytest.raises(ValueError, match="row 1 is None but Map"): + self._encode("Map(String, UInt8)", [{"a": 1}, None], 2) + + def test_nullable_value_type(self): + rows = [{"a": 13, "b": None}, {}, {"c": 79}] + encoded = self._encode("Map(String, Nullable(Int32))", rows) + assert encoded == build_native_block([("m", "Map(String, Nullable(Int32))", rows)]) + assert self._decode(encoded) == rows + + def test_array_value_type(self): + rows = [{"k": [13, 79], "e": []}, {}, {"s": [5]}] + encoded = self._encode("Map(String, Array(Int64))", rows) + assert self._decode(encoded) == rows + + def test_low_cardinality_string_key(self): + rows = [{"red": 1, "blue": 2}, {}, {"red": 3}] + encoded = self._encode("Map(LowCardinality(String), UInt64)", rows) + assert self._decode(encoded) == rows + + def test_map_of_map_and_tuple_values(self): + rows = [{1: {2: "a"}}, {}, {3: {}}] + encoded = self._encode("Map(UInt64, Map(UInt64, String))", rows) + assert self._decode(encoded) == rows + rows = [{"k": (1, "x")}, {}] + encoded = self._encode("Map(String, Tuple(UInt8, String))", rows) + assert self._decode(encoded) == rows + + def test_key_and_value_errors_name_row_and_entry(self): + with pytest.raises(ValueError, match='column "m" row 1 key 0 cannot be converted'): + self._encode("Map(UInt8, UInt8)", [{1: 1}, {"x": 2}], 2) + with pytest.raises(ValueError, match='column "m" row 0 value 1 cannot be converted'): + self._encode("Map(String, UInt8)", [{"a": 1, "b": "x"}], 1) + + def test_flat_refs_refcount_balance(self): + shared = 1 << 40 + + rows = [{shared: shared}, {}, {1: shared}] + before = sys.getrefcount(shared) + self._encode("Map(UInt64, Int64)", rows) + assert sys.getrefcount(shared) == before + + rows = [{shared: shared}, {shared: "x"}] + before = sys.getrefcount(shared) + with pytest.raises(ValueError, match="value 0 cannot be converted"): + self._encode("Map(UInt64, Int64)", rows, 2) + assert sys.getrefcount(shared) == before + + rows = [{shared: shared}, "boom"] + before = sys.getrefcount(shared) + with pytest.raises(ValueError, match="is not a dict for Map"): + self._encode("Map(UInt64, Int64)", rows, 2) + assert sys.getrefcount(shared) == before + + +# --------------------------------------------------------------------------- +# Tuple +# --------------------------------------------------------------------------- + +_TUP_KNOWN_UUID = uuid.UUID("00112233-4455-6677-8899-aabbccddeeff") + +# (type_name, wire_rows, expected). Each wire row is a sequence of element +# values; expected None means the decoded column equals wire_rows (as tuples). +_TUPLE_UNNAMED = [ + ("Tuple(Int32, String)", [(13, "user_1"), (-1, ""), (79, "sventon")], None), + ("Tuple(UInt64, Float64, Bool)", + [(0, 1.5, True), (18446744073709551615, -2.5, False)], None), + ("Tuple(Int32)", [(13,), (79,), (-2147483648,)], None), + ("Tuple(Nullable(Int32), String)", [(13, "a"), (None, "b"), (79, "c")], None), + ("Tuple(UUID, IPv4)", + [(uuid.UUID(int=0), ipaddress.IPv4Address("1.2.3.4")), + (_TUP_KNOWN_UUID, ipaddress.IPv4Address("255.255.255.255"))], None), + ("Tuple(LowCardinality(String), Int32)", + [("red", 1), ("red", 2), ("blue", 3)], None), + ("Tuple(Array(Int64), String)", + [([13, 79], "a"), ([], "b"), ([5, 5, 5], "c")], None), + ("Tuple(UInt8, Tuple(UInt8, String))", + [(1, (2, "x")), (3, (4, "y"))], None), + ("Tuple(Map(String, UInt8), Int32)", + [({"k1": 1, "k2": 2}, 10), ({}, 20)], None), +] + +_TUPLE_NAMED = [ + ("Tuple(a Int32, b String)", + [(13, "x"), (79, "y")], [{"a": 13, "b": "x"}, {"a": 79, "b": "y"}]), + ("Tuple(id UInt64, val Nullable(Int32))", + [(1, 13), (2, None)], [{"id": 1, "val": 13}, {"id": 2, "val": None}]), + ("Tuple(`a b` Int32, c String)", + [(13, "x")], [{"a b": 13, "c": "x"}]), +] + + +class TestTuple: + @pytest.mark.parametrize("type_name,wire_rows,expected", _TUPLE_UNNAMED) + def test_unnamed_to_tuple(self, type_name, wire_rows, expected): + if expected is None: + expected = [tuple(r) for r in wire_rows] + built = build_native_block([("t", type_name, wire_rows)]) + decoded = list(_ch_core.ColBatch.decode_native(built).column_data(0)) + assert decoded == expected + assert all(isinstance(v, tuple) for v in decoded) + + @pytest.mark.parametrize("type_name,wire_rows,expected", _TUPLE_NAMED) + def test_named_to_dict(self, type_name, wire_rows, expected): + built = build_native_block([("t", type_name, wire_rows)]) + decoded = list(_ch_core.ColBatch.decode_native(built).column_data(0)) + assert decoded == expected + assert all(isinstance(v, dict) for v in decoded) + + @pytest.mark.parametrize( + "type_name,wire_rows,expected", + [ + ("Tuple(Int32, String)", [(13, "a"), (79, "b")], [(13, "a"), (79, "b")]), + ("Tuple(x Int32, y String)", [(13, "a")], [{"x": 13, "y": "a"}]), + ("Tuple(Array(Int64), Nullable(String))", + [([13], "a"), ([], None)], [([13], "a"), ([], None)]), + ], + ) + def test_all_exit_paths_agree(self, type_name, wire_rows, expected): + built = build_native_block([("t", type_name, wire_rows)]) + batch = _ch_core.ColBatch.decode_native(built) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == expected + + def test_nullable_tuple_null_rows_are_none(self): + # Nullable(Tuple) decodes null rows to None and non-null rows to their + # value, across all exit paths. A nullable field inside a non-null tuple + # keeps its own None. (clickhouse-connect mis-decodes this rare type, so + # the Rust path is the correct reference, not a parity target.) + type_name = "Nullable(Tuple(Nullable(Int32), String))" + wire_rows = [(13, "a"), None, (None, "b"), None] + expected = [(13, "a"), None, (None, "b"), None] + batch = _ch_core.ColBatch.decode_native(build_native_block([("t", type_name, wire_rows)])) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == expected + + def test_nullable_named_tuple_null_rows_are_none(self): + type_name = "Nullable(Tuple(a UInt64, b String))" + wire_rows = [(13, "x"), None, (5, "y")] + expected = [{"a": 13, "b": "x"}, None, {"a": 5, "b": "y"}] + batch = _ch_core.ColBatch.decode_native(build_native_block([("t", type_name, wire_rows)])) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == expected + + def test_empty_tuple_is_empty_tuple_per_row(self): + # Tuple() has one placeholder byte per row and decodes to an empty tuple. + built = build_native_block([("t", "Tuple()", [(), (), ()])]) + decoded = list(_ch_core.ColBatch.decode_native(built).column_data(0)) + assert decoded == [(), (), ()] + + def test_decimal_and_date_element_machinery(self): + # Decimal/Date elements carry raw wire values in the helper (unscaled + # integer, epoch days) and decode through the recursive per-field + # context: the Decimal field builds its scaled decimal.Decimal. + day = (dt.date(2024, 1, 2) - dt.date(1970, 1, 1)).days + wire_rows = [(100, day), (-350, 0)] + built = build_native_block([("t", "Tuple(Decimal(9, 2), Date)", wire_rows)]) + decoded = list(_ch_core.ColBatch.decode_native(built).column_data(0)) + assert decoded == [ + (decimal.Decimal("1.00"), dt.date(2024, 1, 2)), + (decimal.Decimal("-3.50"), dt.date(1970, 1, 1)), + ] + + def test_multi_block_concatenation(self): + type_name = "Tuple(Int32, String)" + rows_a = [(13, "a"), (79, "b")] + rows_b = [(5, "c")] + block_a = build_native_block([("t", type_name, rows_a)]) + block_b = build_native_block([("t", type_name, rows_b)]) + decoded = list(_ch_core.ColBatch.decode_native(block_a + block_b).column_data(0)) + assert decoded == [(13, "a"), (79, "b"), (5, "c")] + + def test_truncated_tuple_raises_eof(self): + built = build_native_block([("t", "Tuple(Int32, Int32)", [(13, 79)])]) + with pytest.raises(EOFError): + _ch_core.ColBatch.decode_native(built[:-2]) + + +# --------------------------------------------------------------------------- +# Map +# --------------------------------------------------------------------------- + +_MAP_DECODE = [ + ("Map(String, UInt8)", [{"k1": 1, "k2": 2}, {}, {"x": 255}], None), + ("Map(UInt64, String)", [{13: "a", 79: "b"}, {}, {5: "c"}], None), + ("Map(String, Nullable(Int32))", [{"a": 13, "b": None}, {}, {"c": 79}], None), + ("Map(LowCardinality(String), UInt64)", + [{"red": 1, "blue": 2}, {}, {"red": 3}], None), + ("Map(UInt8, Array(Int32))", [{1: [13, 79], 2: []}, {}, {3: [5]}], None), + ("Map(String, Tuple(UInt8, String))", [{"k": (1, "x")}, {}], None), + ("Map(UInt64, Map(UInt64, String))", [{1: {2: "a"}}, {}, {3: {}}], None), +] + + +class TestMap: + @pytest.mark.parametrize("type_name,wire_rows,expected", _MAP_DECODE) + def test_round_trip(self, type_name, wire_rows, expected): + if expected is None: + expected = wire_rows + built = build_native_block([("m", type_name, wire_rows)]) + decoded = list(_ch_core.ColBatch.decode_native(built).column_data(0)) + assert decoded == expected + assert all(isinstance(v, dict) for v in decoded) + + @pytest.mark.parametrize( + "type_name,wire_rows", + [ + ("Map(String, UInt8)", [{"a": 13, "b": 79}, {}, {"c": 5}]), + ("Map(UInt64, Nullable(String))", [{13: "x", 79: None}, {}]), + ("Map(String, Array(Int32))", [{"k": [13, 79]}, {}]), + ], + ) + def test_all_exit_paths_agree(self, type_name, wire_rows): + built = build_native_block([("m", type_name, wire_rows)]) + batch = _ch_core.ColBatch.decode_native(built) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == wire_rows + + def test_duplicate_keys_last_value_wins(self): + # dict input cannot express duplicate keys, so build the entries by hand: + # one row with keys [1, 1] and values [10, 20]. dict(zip(...)) keeps the + # first position and the last value. + body = struct.pack(" + # tuple, array-level -> list nesting exactly. + assert list(batch.column_data(0)) == rows + + def test_point_leaves_are_tuples(self): + decoded = list(_ch_core.ColBatch.decode_native(build_native_block([("g", "Point", [(1.5, 2.5), (-3.25, 4.0)])])).column_data(0)) + assert all(isinstance(v, tuple) and len(v) == 2 for v in decoded) + + @pytest.mark.parametrize("type_name,rows", _GEO_DECODE) + def test_encode_round_trip_and_golden(self, type_name, rows): + encoded = _ch_core.encode_native_block(["g"], [type_name], [rows], len(rows)) + assert encoded == build_native_block([("g", type_name, rows)]) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == rows + + @pytest.mark.parametrize("type_name,rows", _GEO_DECODE) + def test_all_exit_paths_agree(self, type_name, rows): + batch = _ch_core.ColBatch.decode_native(build_native_block([("g", type_name, rows)])) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == rows + + def test_nullable_point_null_rows_are_none(self): + # Nullable(Point) parses as Nullable(Geo(Point)); the delegate is a Tuple, + # so the decoded column carries tuple-level validity and null rows read + # as None across all exit paths. + rows = [(1.5, 2.5), None, (-3.25, 4.0), None] + batch = _ch_core.ColBatch.decode_native( + build_native_block([("g", "Nullable(Point)", rows)]) + ) + assert batch.column_type_names == ["Nullable(Point)"] + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == rows + + def test_nullable_point_encode_round_trip(self): + rows = [(1.5, 2.5), None, (0.0, -1.0)] + encoded = _ch_core.encode_native_block(["g"], ["Nullable(Point)"], [rows], len(rows)) + assert encoded == build_native_block([("g", "Nullable(Point)", rows)]) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == rows + + +# --------------------------------------------------------------------------- +# Nested (physical Array(Tuple(named ...))) +# --------------------------------------------------------------------------- + +class TestNested: + def test_decode_to_list_of_dicts(self): + type_name = "Nested(a UInt32, b String)" + wire_rows = [[(13, "north"), (79, "south")], [], [(5, "east")]] + expected = [ + [{"a": 13, "b": "north"}, {"a": 79, "b": "south"}], + [], + [{"a": 5, "b": "east"}], + ] + batch = _ch_core.ColBatch.decode_native(build_native_block([("n", type_name, wire_rows)])) + assert batch.column_type_names == [type_name] + decoded = list(batch.column_data(0)) + assert decoded == expected + assert all(isinstance(item, dict) for row in decoded for item in row) + + def test_all_exit_paths_agree(self): + type_name = "Nested(a UInt32, b String)" + wire_rows = [[(13, "north")], [], [(5, "east"), (7, "west")]] + expected = [[{"a": 13, "b": "north"}], [], [{"a": 5, "b": "east"}, {"a": 7, "b": "west"}]] + batch = _ch_core.ColBatch.decode_native(build_native_block([("n", type_name, wire_rows)])) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == expected + + def test_encode_positional_golden_and_round_trip(self): + type_name = "Nested(a UInt32, b String)" + wire_rows = [[(13, "north"), (79, "south")], [], [(5, "east")]] + encoded = _ch_core.encode_native_block(["n"], [type_name], [wire_rows], len(wire_rows)) + assert encoded == build_native_block([("n", type_name, wire_rows)]) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == [ + [{"a": 13, "b": "north"}, {"a": 79, "b": "south"}], + [], + [{"a": 5, "b": "east"}], + ] + + def test_encode_dict_rows_match_positional(self): + type_name = "Nested(a UInt32, b String)" + positional = [[(13, "north"), (79, "south")], [], [(5, "east")]] + dict_rows = [ + [{"a": 13, "b": "north"}, {"a": 79, "b": "south"}], + [], + [{"a": 5, "b": "east"}], + ] + expected = _ch_core.encode_native_block(["n"], [type_name], [positional], 3) + assert _ch_core.encode_native_block(["n"], [type_name], [dict_rows], 3) == expected + + def test_nested_low_cardinality_field_round_trip(self): + # A LowCardinality Nested field resolves through the LC path; the core's + # LC index word differs from the helper's, so this round-trips instead of + # golden-comparing. + type_name = "Nested(city LowCardinality(String), pop UInt32)" + dict_rows = [ + [{"city": "harbor", "pop": 13}, {"city": "harbor", "pop": 79}], + [], + [{"city": "ridge", "pop": 5}], + ] + encoded = _ch_core.encode_native_block(["n"], [type_name], [dict_rows], 3) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == dict_rows + + def test_nested_point_field_golden_and_round_trip(self): + # A geo alias (Point) nested inside a non-alias container: Nested expands + # to Array(Tuple(p Point, q UInt32)) and the Point field expands again. + type_name = "Nested(p Point, q UInt32)" + wire_rows = [[((1.5, 2.5), 13), ((3.0, 4.0), 79)], [], [((5.5, 6.5), 5)]] + encoded = _ch_core.encode_native_block(["n"], [type_name], [wire_rows], len(wire_rows)) + assert encoded == build_native_block([("n", type_name, wire_rows)]) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == [ + [{"p": (1.5, 2.5), "q": 13}, {"p": (3.0, 4.0), "q": 79}], + [], + [{"p": (5.5, 6.5), "q": 5}], + ] + + +# --------------------------------------------------------------------------- +# Tuple/Map column-major fill +# --------------------------------------------------------------------------- + +class TestTupleMapFill: + """The buffered exits decode top-level Tuple/Map columns with a hoisted + column-major fill; Tuple/Map nested inside Array stay on the per-cell path.""" + + @staticmethod + def _all_exits(batch): + return ( + list(batch.column_data(0)), + list(batch.to_python_columns()[0]), + [r[0] for r in batch.to_python_rows()], + ) + + @pytest.mark.parametrize( + "type_name,wire_rows,expected", + [ + ("Tuple(LowCardinality(String), Int64)", + [("x", 1), ("y", 2), ("x", 3)], [("x", 1), ("y", 2), ("x", 3)]), + ("Tuple(a LowCardinality(String), b Int64)", + [("x", 1), ("y", 2)], [{"a": "x", "b": 1}, {"a": "y", "b": 2}]), + ("Map(String, Int64)", [{"a": 1, "b": 2}, {}, {"c": 3}], None), + ("Map(String, Array(Int64))", [{"k": [1, 2], "e": []}, {}, {"z": [3]}], None), + ("Nullable(Tuple(Int64, String))", [(1, "a"), None, (2, "b")], None), + ("Tuple()", [(), (), ()], None), + ("Map(String, UInt8)", [{}, {}, {}], None), + ], + ) + def test_all_exit_paths_agree(self, type_name, wire_rows, expected): + if expected is None: + expected = wire_rows + built = build_native_block([("c", type_name, wire_rows)]) + batch = _ch_core.ColBatch.decode_native(built) + via_column_data, via_columns, via_rows = self._all_exits(batch) + assert via_column_data == via_columns == via_rows == expected + + def test_tuple_lc_field_identity_within_chunk(self): + rows = [("red", 1), ("blue", 2), ("red", 3), ("red", 4)] + built = build_native_block([("t", "Tuple(LowCardinality(String), Int64)", rows)]) + batch = _ch_core.ColBatch.decode_native(built) + col = batch.column_data(0) + assert list(col) == rows + assert col[0][0] is col[2][0] + assert col[0][0] is col[3][0] + cols = batch.to_python_columns()[0] + assert cols[0][0] is cols[2][0] + out_rows = batch.to_python_rows() + assert out_rows[0][0][0] is out_rows[2][0][0] + + def test_map_lc_value_identity_within_chunk(self): + rows = [{"a": "red", "b": "blue"}, {}, {"c": "red"}] + built = build_native_block([("m", "Map(String, LowCardinality(String))", rows)]) + batch = _ch_core.ColBatch.decode_native(built) + col = batch.column_data(0) + assert list(col) == rows + assert col[0]["a"] is col[2]["c"] + cols = batch.to_python_columns()[0] + assert cols[0]["a"] is cols[2]["c"] + out_rows = batch.to_python_rows() + assert out_rows[0][0]["a"] is out_rows[2][0]["c"] + + def test_map_lc_key_identity_within_chunk(self): + rows = [{"red": 1, "blue": 2}, {}, {"red": 3}] + built = build_native_block([("m", "Map(LowCardinality(String), UInt64)", rows)]) + col = _ch_core.ColBatch.decode_native(built).column_data(0) + assert list(col) == rows + key0 = next(k for k in col[0] if k == "red") + assert key0 is next(iter(col[2])) + + def test_nullable_tuple_lc_field_null_rows_and_identity(self): + # Null rows discard their field values; valid rows still share the + # dictionary object. + rows = [("x", 1), None, ("x", 2), ("y", 3)] + enc = _ch_core.encode_native_block( + ["c"], ["Nullable(Tuple(LowCardinality(String), Int64))"], [rows], len(rows) + ) + batch = _ch_core.ColBatch.decode_native(enc) + via_column_data, via_columns, via_rows = self._all_exits(batch) + assert via_column_data == via_columns == via_rows == rows + col = batch.column_data(0) + assert col[0][0] is col[2][0] + + def test_fill_path_matches_per_cell_array_element_path(self): + # The same values decode through the column-major fill at top level and + # through the per-cell path one Array level down; they must agree. + tup_rows = [(1, "a"), (2, "b"), (3, "c")] + map_rows = [{"a": 1, "b": 2}, {}, {"c": 3}] + top = _ch_core.ColBatch.decode_native(build_native_block([ + ("t", "Tuple(Int64, String)", tup_rows), + ("m", "Map(String, Int64)", map_rows), + ])) + wrapped = _ch_core.ColBatch.decode_native(build_native_block([ + ("t", "Array(Tuple(Int64, String))", [[r] for r in tup_rows]), + ("m", "Array(Map(String, Int64))", [[r] for r in map_rows]), + ])) + assert [[v] for v in top.column_data(0)] == list(wrapped.column_data(0)) + assert [[v] for v in top.column_data(1)] == list(wrapped.column_data(1)) + + def test_multi_chunk_named_tuple_and_map(self): + tn = "Tuple(a Int64, b String)" + mp = "Map(String, Int64)" + block_a = build_native_block([("t", tn, [(1, "x")]), ("m", mp, [{"a": 1}])]) + block_b = build_native_block( + [("t", tn, [(2, "y"), (3, "z")]), ("m", mp, [{}, {"b": 2}])] + ) + batch = _ch_core.ColBatch.decode_native(block_a + block_b) + assert list(batch.column_data(0)) == [ + {"a": 1, "b": "x"}, {"a": 2, "b": "y"}, {"a": 3, "b": "z"} + ] + assert list(batch.column_data(1)) == [{"a": 1}, {}, {"b": 2}] + assert list(batch.to_python_rows()) == [ + ({"a": 1, "b": "x"}, {"a": 1}), + ({"a": 2, "b": "y"}, {}), + ({"a": 3, "b": "z"}, {"b": 2}), + ] diff --git a/rust/ch-core-py/tests/test_insert_fastpaths.py b/rust/ch-core-py/tests/test_insert_fastpaths.py new file mode 100644 index 00000000..18dd3e70 --- /dev/null +++ b/rust/ch-core-py/tests/test_insert_fastpaths.py @@ -0,0 +1,409 @@ +from helpers import ( + _ch_core, + _NdarrayLikeColumn, + build_native_block, + dt, + ipaddress, + pytest, + uuid, +) + + +class TestEncodeFastPaths: + """Exact list/tuple and buffer-protocol fast paths for primitive columns.""" + + def _encode(self, type_name, values, row_count=None): + n = len(values) if row_count is None else row_count + return _ch_core.encode_native_block(["v"], [type_name], [values], n) + + @pytest.mark.parametrize( + "type_name,values", + [ + ("Int8", [-128, -1, 0, 127]), + ("Int16", [-32768, 0, 32767]), + ("Int32", [-(2**31), 0, 2**31 - 1]), + ("Int64", [-(2**63), -1, 0, 2**63 - 1]), + ("UInt8", [0, 255]), + ("UInt16", [0, 65535]), + ("UInt32", [0, 2**32 - 1]), + ("UInt64", [0, 1, 2**64 - 1]), + ("Float32", [0.0, -1.5, 3.25]), + ("Float64", [0.0, -1.5, 1e300]), + ("Bool", [True, False, True]), + ], + ) + def test_edge_values_match_helper_and_container_kinds_agree(self, type_name, values): + from_list = self._encode(type_name, list(values)) + assert from_list == build_native_block([("v", type_name, values)]) + assert self._encode(type_name, tuple(values), len(values)) == from_list + # Non-list, non-buffer container takes the generic path; same bytes. + assert self._encode(type_name, _NdarrayLikeColumn(values), len(values)) == from_list + + @pytest.mark.parametrize( + "type_name,values", + [ + ("Int8", [0, 128]), + ("Int8", [-129]), + ("Int64", [2**63]), + ("Int64", [-(2**63) - 1]), + ("UInt8", [256]), + ("UInt64", [-1]), + ("UInt64", [2**64]), + ("Float64", [10**400]), + ("IntervalDay", [2**63]), + ("IntervalDay", [-(2**63) - 1]), + ], + ) + def test_out_of_range_raises_conversion_error(self, type_name, values): + with pytest.raises(ValueError, match=f"row {len(values) - 1} cannot be converted to {type_name}"): + self._encode(type_name, values) + + @pytest.mark.parametrize("value", [1.5, "5", dt.timedelta(days=5)]) + def test_interval_rejects_non_int_values(self, value): + with pytest.raises(ValueError, match="row 0 cannot be converted to IntervalDay"): + self._encode("IntervalDay", [value]) + + def test_none_in_non_nullable_raises(self): + with pytest.raises(ValueError, match='column "v" row 1 is None but Int64 is not Nullable'): + self._encode("Int64", [3, None, 5]) + + @pytest.mark.parametrize("make", [list, tuple]) + def test_nullable_list_and_tuple(self, make): + values = [3, None, -(2**63), None, 2**63 - 1] + encoded = self._encode("Nullable(Int64)", make(values), len(values)) + assert encoded == build_native_block([("v", "Nullable(Int64)", values)]) + + def test_fallback_types_still_accepted(self): + import enum + from decimal import Decimal + + class IntLike(enum.IntEnum): + SEVEN = 7 + + values = [True, IntLike.SEVEN, 13.0, Decimal("79.000")] + assert self._encode("Int64", values) == build_native_block([("v", "Int64", [1, 7, 13, 79])]) + # Exact ints take the fast path into floats; bool goes through the fallback. + assert self._encode("Float64", [1, True, 2.5]) == build_native_block( + [("v", "Float64", [1.0, 1.0, 2.5])] + ) + + @pytest.mark.parametrize( + "value,detail", + [ + (13.5, "would lose fractional data; pass an integer value"), + (float("nan"), "is not finite; pass an integer value"), + (float("inf"), "is not finite; pass an integer value"), + ], + ) + def test_float_to_integer_rejection_is_actionable(self, value, detail): + with pytest.raises(ValueError, match=detail): + self._encode("Int32", [value]) + + @pytest.mark.parametrize("value", ["13", "-79"]) + def test_numeric_string_to_integer_rejection_is_actionable(self, value): + with pytest.raises(ValueError, match="strings are not accepted; pass an int instead"): + self._encode("Int32", [value]) + + def test_decimal_to_integer_rejection_is_actionable(self): + from decimal import Decimal + + with pytest.raises(ValueError, match="would lose fractional data; pass an integer value"): + self._encode("Int32", [Decimal("13.5")]) + with pytest.raises(ValueError, match="is not finite; pass an integer value"): + self._encode("Int32", [Decimal("NaN")]) + + def test_decimal_subclass_accepted_like_decimal(self): + from decimal import Decimal + + class D(Decimal): + pass + + assert self._encode("Int64", [D("79"), Decimal("7")]) == build_native_block( + [("v", "Int64", [79, 7])] + ) + with pytest.raises(ValueError, match="would lose fractional data; pass an integer value"): + self._encode("Int64", [D("13.5")]) + + @pytest.mark.parametrize("type_name", ["Float32", "Float64", "BFloat16"]) + def test_string_to_float_rejection_is_actionable(self, type_name): + with pytest.raises(ValueError, match="strings are not accepted; pass a float instead"): + self._encode(type_name, ["1.5"]) + + def test_non_integer_object_rejection_has_connector(self): + with pytest.raises(ValueError, match="is not an integer; pass an integer value"): + self._encode("Int32", [object()]) + + def test_error_value_repr_is_truncated(self): + with pytest.raises(ValueError) as exc: + self._encode("Int32", ["x" * 500]) + message = str(exc.value) + assert "..." in message + assert len(message) < 250 + + def test_numpy_longdouble_rejected_for_integer_targets(self): + np = pytest.importorskip("numpy") + # On arm64 macOS longdouble is f64, so the fraction in + # 2**64 + 0.5 vanishes at construction; on x86 it survives and must + # not be silently truncated through the f64 extraction. + fractional = np.longdouble(2**64) + np.longdouble(0.5) + with pytest.raises(ValueError, match="pass an integer value"): + self._encode("UInt64", [fractional]) + with pytest.raises(ValueError, match="pass an integer value"): + self._encode("UInt64", [np.longdouble(7.0)]) + + def test_float32_overflow_becomes_inf(self): + encoded = self._encode("Float32", [1e300]) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_data(0)[0] == float("inf") + + def test_list_resized_during_fallback_raises(self): + values = [1, None, 3, 4] + + class Evil: + def __index__(self): + del values[2:] + return 7 + + values[1] = Evil() + with pytest.raises(ValueError, match="resized during encoding"): + self._encode("Int64", values, 4) + + @pytest.mark.parametrize( + "type_name,dtype,values", + [ + ("Int8", "int8", [-128, 0, 127]), + ("Int64", "int64", [-(2**63), 0, 2**63 - 1]), + ("IntervalDay", "int64", [-(2**63), 0, 2**63 - 1]), + ("UInt64", "uint64", [0, 2**64 - 1]), + ("Float32", "float32", [0.0, -1.5, 3.25]), + ("Float64", "float64", [0.0, -1.5, 1e300]), + ], + ) + def test_numpy_buffer_matches_list_path(self, type_name, dtype, values): + np = pytest.importorskip("numpy") + arr = np.array(values, dtype=dtype) + assert self._encode(type_name, arr, len(values)) == self._encode(type_name, list(values)) + + def test_numpy_strided_view_matches_list_path(self): + np = pytest.importorskip("numpy") + arr = np.arange(10, dtype="int64")[::2] + assert self._encode("Int64", arr, 5) == self._encode("Int64", list(arr)) + + def test_numpy_mismatched_dtype_falls_back(self): + np = pytest.importorskip("numpy") + arr = np.array([1, 2, 3], dtype="int32") + assert self._encode("Int64", arr, 3) == self._encode("Int64", [1, 2, 3]) + + def test_numpy_nullable_is_all_valid(self): + np = pytest.importorskip("numpy") + arr = np.array([1.5, 2.5], dtype="float64") + expected = build_native_block([("v", "Nullable(Float64)", [1.5, 2.5])]) + assert self._encode("Nullable(Float64)", arr, 2) == expected + + @pytest.mark.parametrize( + "type_name,dtype,values", + [ + ("Int64", ">i8", [1, 2, 3]), + ("Float64", ">f8", [1.5, -2.5, 1e300]), + ], + ) + def test_numpy_non_native_byte_order_matches_list_path(self, type_name, dtype, values): + np = pytest.importorskip("numpy") + arr = np.array(values, dtype=dtype) + assert self._encode(type_name, arr, len(values)) == self._encode(type_name, list(values)) + + def test_char_format_buffer_still_raises_for_uint8(self): + view = memoryview(b"AB").cast("c") + with pytest.raises(ValueError, match="cannot be converted to UInt8"): + self._encode("UInt8", view, 2) + + def test_tuple_fallback_types_still_accepted(self): + import enum + + class IntLike(enum.IntEnum): + SEVEN = 7 + + values = (True, IntLike.SEVEN, 3) + assert self._encode("Int64", values, 3) == build_native_block([("v", "Int64", [1, 7, 3])]) + with pytest.raises(ValueError, match="row 1 cannot be converted to Int64"): + self._encode("Int64", (1, "x", 3), 3) + + +class TestScalarObjectInsertFastPath: + """UUID, IPv4, and Enum8/16 fast paths over exact lists and tuples.""" + + _E8 = "Enum8('alpha' = 1, 'beta' = 2, 'gamma' = 3)" + _E16 = "Enum16('alpha' = -5, 'beta' = 0, 'gamma' = 1000)" + + def _encode(self, type_name, vals, n=None): + n = len(vals) if n is None else n + return _ch_core.encode_native_block(["v"], [type_name], [vals], n) + + def _decode(self, encoded): + return list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) + + def test_uuid_containers_agree_and_round_trip(self): + vals = [ + uuid.UUID(int=0), + uuid.UUID("00112233-4455-6677-8899-aabbccddeeff"), + uuid.UUID(int=(1 << 128) - 1), + ] + fast = self._encode("UUID", vals) + assert fast == self._encode("UUID", tuple(vals), len(vals)) + assert fast == self._encode("UUID", _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == vals + + def test_uuid_mixed_types_use_fallback(self): + class SubUUID(uuid.UUID): + pass + + known = uuid.UUID("00112233-4455-6677-8899-aabbccddeeff") + vals = [known, str(known), known.int, known.bytes, SubUUID(int=known.int)] + fast = self._encode("UUID", vals) + assert fast == self._encode("UUID", _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == [known] * 5 + + def test_nullable_uuid(self): + vals = [uuid.UUID(int=9), None, uuid.UUID(int=0), None] + fast = self._encode("Nullable(UUID)", vals) + assert fast == self._encode("Nullable(UUID)", tuple(vals), len(vals)) + assert fast == self._encode("Nullable(UUID)", _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == vals + + def test_uuid_errors_unchanged(self): + with pytest.raises(ValueError, match="row 1 cannot be converted to UUID"): + self._encode("UUID", [uuid.UUID(int=1), 1.5]) + with pytest.raises(ValueError, match="row 1 is None but UUID is not Nullable"): + self._encode("UUID", [uuid.UUID(int=1), None]) + + def test_ipv4_containers_agree_and_round_trip(self): + vals = [ + ipaddress.IPv4Address("0.0.0.0"), + ipaddress.IPv4Address("255.255.255.255"), + ipaddress.IPv4Address("192.0.2.1"), + ] + fast = self._encode("IPv4", vals) + assert fast == self._encode("IPv4", tuple(vals), len(vals)) + assert fast == self._encode("IPv4", _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == vals + + def test_ipv4_mixed_types_without_object_wrappers(self): + import enum + + class IntLike(enum.IntEnum): + ADDR = 16909060 + + vals = [ipaddress.IPv4Address("1.2.3.4"), "1.2.3.4", 16909060, IntLike.ADDR] + fast = self._encode("IPv4", vals) + assert fast == self._encode("IPv4", _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == [ipaddress.IPv4Address("1.2.3.4")] * 4 + + def test_nullable_ipv4(self): + vals = [ipaddress.IPv4Address("1.2.3.4"), None, ipaddress.IPv4Address("0.0.0.0")] + fast = self._encode("Nullable(IPv4)", vals) + assert fast == self._encode("Nullable(IPv4)", tuple(vals), len(vals)) + assert fast == self._encode("Nullable(IPv4)", _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == vals + + def test_ipv4_errors_unchanged(self): + with pytest.raises(ValueError, match="row 0 cannot be converted to IPv4"): + self._encode("IPv4", ["1.2.3.999"]) + with pytest.raises(ValueError, match="row 0 cannot be converted to IPv4"): + self._encode("IPv4", [_NdarrayLikeColumn([1])]) + with pytest.raises(ValueError, match="row 1 cannot be converted to IPv4"): + self._encode("IPv4", [1, 2**32]) + + def test_enum8_labels_and_codes_round_trip(self): + vals = ["alpha", "beta", 3, "alpha", 99] + fast = self._encode(self._E8, vals) + assert fast == self._encode(self._E8, tuple(vals), len(vals)) + assert fast == self._encode(self._E8, _NdarrayLikeColumn(vals), len(vals)) + # Unknown raw codes pass through and read back as None. + assert self._decode(fast) == ["alpha", "beta", "gamma", "alpha", None] + + def test_enum16_labels_and_codes_round_trip(self): + vals = ["gamma", -5, "beta", "gamma"] + fast = self._encode(self._E16, vals) + assert fast == self._encode(self._E16, _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == ["gamma", "alpha", "beta", "gamma"] + + def test_nullable_enum8(self): + vals = ["alpha", None, "gamma", None] + fast = self._encode(f"Nullable({self._E8})", vals) + assert fast == self._encode(f"Nullable({self._E8})", tuple(vals), len(vals)) + assert fast == self._encode(f"Nullable({self._E8})", _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == vals + + @pytest.mark.parametrize("type_name", [_E8, _E16]) + def test_enum_accepts_integral_float_and_pandas_nan(self, type_name): + vals = [1.0, float("nan"), 2.0] + encoded = self._encode(type_name, vals) + assert encoded == build_native_block([("v", type_name, [1, 0, 2])]) + + @pytest.mark.parametrize("type_name", [_E8, _E16]) + def test_nullable_enum_nan_becomes_null(self, type_name): + vals = [1.0, float("nan"), 2.0] + expected = build_native_block([("v", f"Nullable({type_name})", [1, None, 2])]) + assert self._encode(f"Nullable({type_name})", vals) == expected + assert self._encode(f"Nullable({type_name})", tuple(vals), len(vals)) == expected + assert ( + self._encode(f"Nullable({type_name})", _NdarrayLikeColumn(vals), len(vals)) + == expected + ) + + def test_nullable_enum_numpy_nan_becomes_null(self): + np = pytest.importorskip("numpy") + vals = [np.float32(1.0), np.float32("nan"), np.float64("nan")] + expected = build_native_block([("v", f"Nullable({self._E8})", [1, None, None])]) + assert self._encode(f"Nullable({self._E8})", vals) == expected + + def test_enum_rejects_lossy_float_with_actionable_error(self): + with pytest.raises(ValueError, match="would lose fractional data; pass a valid enum label or integral code"): + self._encode(self._E8, [1.5]) + + @pytest.mark.parametrize("type_name,enum_name", [(_E8, "Enum8"), (_E16, "Enum16")]) + def test_enum_unknown_label_raises(self, type_name, enum_name): + expected = f'row 1 {enum_name} label "delta" is not defined' + with pytest.raises(ValueError, match=expected): + self._encode(type_name, ["alpha", "delta"]) + with pytest.raises(ValueError, match=expected): + self._encode(type_name, _NdarrayLikeColumn(["alpha", "delta"]), 2) + + def test_enum_str_subclass_uses_fallback(self): + class S(str): + pass + + vals = [S("alpha"), "beta", S("gamma")] + fast = self._encode(self._E8, vals) + assert fast == self._encode(self._E8, _NdarrayLikeColumn(vals), len(vals)) + assert self._decode(fast) == ["alpha", "beta", "gamma"] + + @pytest.mark.parametrize("size", [4, 8, 16, 32, 64]) + def test_enum_item_replacement_during_fallback_invalidates_ptr_cache(self, size): + # A fallback __index__ drops the last ref to an already-scanned label; + # the allocator can hand its address to a new same-size str, which + # must not false-hit the pointer-identity cache. + a, b = "A" * size, "B" * size + tname = f"Enum8('{a}' = 1, '{b}' = 2, 'EV' = 3)" + vals = ["A" * size, None, "C" * size] + + class Evil: + def __index__(self): + vals[0] = "x" # drop the sole ref to the scanned label + vals[2] = "B" * size # same size class, may reuse its address + return 3 + + vals[1] = Evil() + assert self._decode(self._encode(tname, vals, 3)) == [a, "EV", b] + + def test_enum_list_resized_during_fallback_raises(self): + vals = ["alpha", None, "beta", "gamma"] + + class Evil: + def __index__(self): + del vals[2:] + return 2 + + vals[1] = Evil() + with pytest.raises(ValueError, match="resized during encoding"): + self._encode(self._E8, vals, 4) diff --git a/rust/ch-core-py/tests/test_intake_exits.py b/rust/ch-core-py/tests/test_intake_exits.py new file mode 100644 index 00000000..d01da97d --- /dev/null +++ b/rust/ch-core-py/tests/test_intake_exits.py @@ -0,0 +1,768 @@ +from helpers import ( + _INTERVAL_TYPES, + _ch_core, + _encode_varint, + _encode_varint_string, + _NdarrayLikeColumn, + _SeriesLikeColumn, + build_native_block, + build_native_block_from_bodies, + decimal, + dt, + ipaddress, + os, + pytest, + subprocess, + sys, + textwrap, + uuid, +) + + +class TestEncodeNativeBlock: + def test_indexable_non_sequence_columns_match_native_helper(self): + names = ["i", "s"] + type_names = ["Int32", "String"] + columns = [ + _NdarrayLikeColumn([13, 79]), + _SeriesLikeColumn(["user_1", "user_2"]), + ] + + encoded = _ch_core.encode_native_block(names, type_names, columns, 2) + expected = build_native_block( + [ + ("i", "Int32", [13, 79]), + ("s", "String", ["user_1", "user_2"]), + ] + ) + assert encoded == expected + + @pytest.mark.parametrize("bad_values", ["xy", b"xy", bytearray(b"xy")]) + def test_bare_string_or_bytes_column_rejected(self, bad_values): + with pytest.raises(ValueError, match="bare str or bytes"): + _ch_core.encode_native_block(["s"], ["String"], [bad_values], len(bad_values)) + + def test_common_types_match_native_helper(self): + enum_type = "Enum8('red' = 1, 'green' = 2)" + names = ["i", "s", "n", "fs", "d", "ts", "e"] + type_names = [ + "Int32", + "String", + "Nullable(UInt16)", + "FixedString(4)", + "Date", + "DateTime64(3)", + enum_type, + ] + aware_ts = dt.datetime(2024, 1, 15, 12, 34, 56, 789000, tzinfo=dt.timezone.utc) + columns = [ + [13, 79], + ["user_1", b"\xff"], + [13, None], + ["ab", b"xy\x00\x00"], + [dt.date(1970, 1, 2), 19737], + [aware_ts, 0], + ["red", 2], + ] + encoded = _ch_core.encode_native_block(names, type_names, columns, 2) + expected = build_native_block( + [ + ("i", "Int32", [13, 79]), + ("s", "String", ["user_1", b"\xff"]), + ("n", "Nullable(UInt16)", [13, None]), + ("fs", "FixedString(4)", [b"ab", b"xy"]), + ("d", "Date", [1, 19737]), + ("ts", "DateTime64(3)", [1705322096789, 0]), + ("e", enum_type, [1, 2]), + ] + ) + assert encoded == expected + + def test_datetime64_pre_epoch_fractional_encode(self): + values = [ + dt.datetime(1969, 12, 31, 23, 59, 59, 500000, tzinfo=dt.timezone.utc), + dt.datetime(1969, 12, 31, 23, 59, 59, 999999, tzinfo=dt.timezone.utc), + dt.datetime(1970, 1, 1, 0, 0, 0, 999000, tzinfo=dt.timezone.utc), + ] + encoded = _ch_core.encode_native_block(["ts"], ["DateTime64(3)"], [values], len(values)) + expected = build_native_block([("ts", "DateTime64(3)", [-500, -1, 999])]) + assert encoded == expected + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == [ + dt.datetime(1969, 12, 31, 23, 59, 59, 500000), + dt.datetime(1969, 12, 31, 23, 59, 59, 999000), + dt.datetime(1970, 1, 1, 0, 0, 0, 999000), + ] + + @pytest.mark.parametrize("type_name", _INTERVAL_TYPES) + def test_interval_round_trip_preserves_kind_and_signed_i64(self, type_name): + values = [-(2**63), -79, 0, 13, 2**63 - 1] + encoded = _ch_core.encode_native_block(["v"], [type_name], [values], len(values)) + + assert encoded == build_native_block([("v", type_name, values)]) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == values + assert list(batch.to_python_columns()[0]) == values + assert batch.to_python_rows() == [(value,) for value in values] + + def test_interval_wrappers_and_containers_round_trip(self): + columns = [ + ("scalar", "IntervalDay", [-13, 0, 79]), + ("nullable", "Nullable(IntervalHour)", [-13, None, 79]), + ("array", "Array(IntervalMinute)", [[-13, 79], [], [0]]), + ("tuple", "Tuple(IntervalSecond, String)", [(-13, "x"), (0, "y"), (79, "z")]), + ( + "array_tuple", + "Array(Tuple(IntervalMillisecond, IntervalMonth))", + [[(-13, 1), (79, -2)], [], [(0, 3)]], + ), + ("map", "Map(IntervalDay, String)", [{-13: "x"}, {}, {79: "z"}]), + ("low_cardinality", "LowCardinality(IntervalHour)", [13, 79, 13]), + ] + encoded = _ch_core.encode_native_block( + [name for name, _, _ in columns], + [type_name for _, type_name, _ in columns], + [values for _, _, values in columns], + 3, + ) + + assert encoded == build_native_block(columns) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name for _, type_name, _ in columns] + assert batch.to_python_columns() == [values for _, _, values in columns] + + @pytest.mark.parametrize( + "type_name", + ["intervalDay", "Intervalday", "INTERVALDAY", "IntervalDays", "IntervalDay()"], + ) + def test_interval_type_names_are_exact(self, type_name): + with pytest.raises(NotImplementedError, match="unsupported ClickHouse type"): + _ch_core.encode_native_block(["v"], [type_name], [[13]], 1) + + def test_prefix_is_prepended(self): + payload = _ch_core.encode_native_block(["v"], ["Int8"], [[13]], 1) + prefix = b"INSERT INTO t FORMAT Native\n" + encoded = _ch_core.encode_native_block(["v"], ["Int8"], [[13]], 1, prefix=prefix) + assert encoded == prefix + payload + + def test_low_cardinality_round_trip(self): + vals = ["x", None, "y", "x", None, ""] + encoded = _ch_core.encode_native_block(["c"], ["LowCardinality(Nullable(String))"], [vals], len(vals)) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == ["LowCardinality(Nullable(String))"] + assert list(batch.column_data(0)) == vals + + def test_uuid_bytes_match_python_serializer_order(self): + value_uuid = uuid.UUID("00112233-4455-6677-8899-aabbccddeeff") + encoded = _ch_core.encode_native_block(["u"], ["UUID"], [[value_uuid.bytes]], 1) + expected_body = bytearray() + expected_body.extend(bytes(reversed(value_uuid.bytes[:8]))) + expected_body.extend(bytes(reversed(value_uuid.bytes[8:]))) + expected = build_native_block_from_bodies([("u", "UUID", expected_body)], 1) + assert encoded == expected + + def test_special_binary_types_exact_bytes(self): + value_uuid = uuid.UUID("00112233-4455-6677-8899-aabbccddeeff") + ipv4_values = ["192.0.2.1", ipaddress.IPv4Address("198.51.100.7")] + ipv6_values = [ipaddress.IPv6Address("2001:db8::1"), "192.0.2.9"] + decimals = [decimal.Decimal("123.4567"), "-1.5"] + + encoded = _ch_core.encode_native_block( + ["u", "v4", "v6", "dec"], + ["UUID", "IPv4", "IPv6", "Decimal(20, 4)"], + [[value_uuid, 0], ipv4_values, ipv6_values, decimals], + 2, + ) + + uuid_body = bytearray() + uuid_int = value_uuid.int + uuid_body.extend((uuid_int >> 64).to_bytes(8, "little")) + uuid_body.extend((uuid_int & 0xFFFFFFFFFFFFFFFF).to_bytes(8, "little")) + uuid_body.extend(bytes(16)) + + ipv4_body = bytearray() + for value in ipv4_values: + ipv4_body.extend(int(ipaddress.IPv4Address(value)).to_bytes(4, "little")) + + ipv6_body = bytearray() + ipv6_body.extend(ipv6_values[0].packed) + ipv6_body.extend(b"\x00" * 10 + b"\xff\xff" + ipaddress.IPv4Address(ipv6_values[1]).packed) + + decimal_body = bytearray() + decimal_body.extend((1234567).to_bytes(16, "little", signed=True)) + decimal_body.extend((-15000).to_bytes(16, "little", signed=True)) + + expected = build_native_block_from_bodies( + [ + ("u", "UUID", uuid_body), + ("v4", "IPv4", ipv4_body), + ("v6", "IPv6", ipv6_body), + ("dec", "Decimal(20, 4)", decimal_body), + ], + 2, + ) + assert encoded == expected + + def test_decimal_precision_boundary_and_overflow(self): + encoded = _ch_core.encode_native_block( + ["d"], + ["Decimal(3, 1)"], + [[decimal.Decimal("99.9"), decimal.Decimal("-99.9")]], + 2, + ) + body = bytearray() + body.extend((999).to_bytes(4, "little", signed=True)) + body.extend((-999).to_bytes(4, "little", signed=True)) + assert encoded == build_native_block_from_bodies([("d", "Decimal(3, 1)", body)], 2) + + for value in (decimal.Decimal("100.0"), decimal.Decimal("-100.0"), "999"): + with pytest.raises(ValueError, match="exceeds precision 3"): + _ch_core.encode_native_block(["d"], ["Decimal(3, 1)"], [[value]], 1) + + def test_decimal_str_failure_is_value_error_with_context(self): + class BadDecimalStr: + def __str__(self): + raise RuntimeError("cannot render") + + with pytest.raises(ValueError, match='column "d" row 0 Decimal value cannot be stringified'): + _ch_core.encode_native_block(["d"], ["Decimal(9, 2)"], [[BadDecimalStr()]], 1) + + def test_encode_errors_are_specific(self): + with pytest.raises(ValueError, match="row_count"): + _ch_core.encode_native_block(["v"], ["Int8"], [[13, 79]], 1) + with pytest.raises(ValueError, match="not Nullable"): + _ch_core.encode_native_block(["v"], ["Int8"], [[None]], 1) + with pytest.raises(ValueError, match="FixedString binary value"): + _ch_core.encode_native_block(["fs"], ["FixedString(4)"], [[b"xy"]], 1) + with pytest.raises(NotImplementedError, match="unsupported ClickHouse type"): + _ch_core.encode_native_block(["v"], ["Object('json')"], [[{"a": 1}]], 1) + with pytest.raises(ValueError, match="label"): + _ch_core.encode_native_block(["e"], ["Enum8('ok' = 1)"], [["missing"]], 1) + + +# --------------------------------------------------------------------------- +# Multi-column +# --------------------------------------------------------------------------- + +class TestMultiColumn: + def test_mixed_types(self): + data = build_native_block([ + ("i", "Int32", [1, 2]), + ("f", "Float64", [1.1, 2.2]), + ("s", "String", ["a", "bb"]), + ("b", "Bool", [1, 0]), + ]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.num_rows == 2 + assert batch.num_columns == 4 + assert batch.column_names == ["i", "f", "s", "b"] + assert list(batch.column_data(0)) == [1, 2] + assert list(batch.column_data(2)) == ["a", "bb"] + assert list(batch.column_data(3)) == [True, False] + + +# --------------------------------------------------------------------------- +# Python access +# --------------------------------------------------------------------------- + +class TestPythonAccess: + def test_to_python_rows(self): + data = build_native_block([ + ("a", "Int32", [10, 20]), + ("b", "String", ["x", "y"]), + ]) + batch = _ch_core.ColBatch.decode_native(data) + rows = list(batch.to_python_rows()) + assert rows == [(10, "x"), (20, "y")] + + def test_to_python_columns(self): + data = build_native_block([ + ("a", "Int64", [1, 2, 3]), + ("b", "Float64", [1.0, 2.0, 3.0]), + ]) + batch = _ch_core.ColBatch.decode_native(data) + cols = list(batch.to_python_columns()) + assert list(cols[0]) == [1, 2, 3] + + def test_column_data_out_of_range(self): + data = build_native_block([("a", "Int64", [1])]) + batch = _ch_core.ColBatch.decode_native(data) + with pytest.raises(ValueError, match="out of range"): + batch.column_data(5) + + +# --------------------------------------------------------------------------- +# Path equivalence +# --------------------------------------------------------------------------- + +class TestPathEquivalence: + _COLUMNS = [ + ("b", "Bool", [1, 0, 1]), + ("i8", "Int8", [-5, 0, 127]), + ("i16", "Int16", [-300, 0, 32767]), + ("i32", "Int32", [-70000, 0, 70000]), + ("i64", "Int64", [-1, 0, 2**62]), + ("u8", "UInt8", [0, 128, 255]), + ("u16", "UInt16", [0, 13, 65535]), + ("u32", "UInt32", [0, 79, 4_000_000_000]), + ("u64", "UInt64", [0, 1, 2**64 - 1]), + ("f32", "Float32", [1.5, -2.25, 0.0]), + ("f64", "Float64", [1.5, -0.1, 1e300]), + ("s", "String", ["u1", "", b"\xff"]), + ("fs", "FixedString(2)", [b"ab", b"cd", b"ef"]), + ("d", "Date", [0, 100, 19737]), + ("d32", "Date32", [-25567, 0, 19737]), + ("dt", "DateTime", [0, 961056000, 1705322096]), + ("dtz", "DateTime('America/New_York')", [0, 961056000, 1705322096]), + ("ts", "DateTime64(3)", [0, 13, 1705322096789]), + ("tsz", "DateTime64(6, 'Asia/Istanbul')", [0, 79, 1705322096789012]), + ("nb", "Nullable(Bool)", [1, None, 0]), + ("ni", "Nullable(Int64)", [13, None, 79]), + ("nf", "Nullable(Float64)", [1.5, None, -2.5]), + ("ns", "Nullable(String)", ["x", None, b"\x80"]), + ("nts", "Nullable(DateTime64(3))", [1705322096789, None, 0]), + ] + + def _assert_paths_agree(self, batch): + via_columns = [list(c) for c in batch.to_python_columns()] + via_column_data = [list(batch.column_data(i)) for i in range(batch.num_columns)] + via_rows = [list(col) for col in zip(*batch.to_python_rows())] + assert via_columns == via_column_data + assert via_columns == via_rows + + def test_all_types_agree_across_paths(self): + data = build_native_block(self._COLUMNS) + self._assert_paths_agree(_ch_core.ColBatch.decode_native(data)) + + def test_mid_column_error_raises_on_all_paths(self): + # A timestamp beyond datetime.MAXYEAR fails partway through the + # column; every path must surface a clean ValueError, not crash or + # return a partial result. + data = build_native_block([("ts", "DateTime64(0)", [0, 300_000_000_000])]) + batch = _ch_core.ColBatch.decode_native(data) + with pytest.raises(ValueError): + batch.column_data(0) + with pytest.raises(ValueError): + batch.to_python_columns() + with pytest.raises(ValueError): + batch.to_python_rows() + + def test_paths_agree_across_chunks(self): + # Two Native blocks in one buffer decode as two chunks; every path + # must concatenate them identically. + first = build_native_block(self._COLUMNS) + second = build_native_block( + [(name, type_name, values[::-1]) for name, type_name, values in self._COLUMNS] + ) + batch = _ch_core.ColBatch.decode_native(first + second) + assert batch.num_chunks == 2 + assert batch.num_rows == 6 + self._assert_paths_agree(batch) + + +# --------------------------------------------------------------------------- +# Arrow export +# --------------------------------------------------------------------------- + +class TestArrowExport: + def test_arrow_c_stream(self): + pa = pytest.importorskip("pyarrow") + data = build_native_block([ + ("i", "Int32", [10, 20, 30]), + ("f", "Float64", [1.5, 2.5, 3.5]), + ("s", "String", ["a", "b", "c"]), + ]) + batch = _ch_core.ColBatch.decode_native(data) + reader = pa.RecordBatchReader.from_stream(batch) + result = reader.read_all() + + assert result.num_rows == 3 + assert result.schema.field("i").type == pa.int32() + assert result.schema.field("f").type == pa.float64() + assert result.schema.field("s").type == pa.utf8() + assert result.column("i").to_pylist() == [10, 20, 30] + + def test_arrow_bool(self): + pa = pytest.importorskip("pyarrow") + data = build_native_block([("b", "Bool", [1, 0, 1])]) + batch = _ch_core.ColBatch.decode_native(data) + reader = pa.RecordBatchReader.from_stream(batch) + result = reader.read_all() + assert result.schema.field("b").type == pa.bool_() + assert result.column("b").to_pylist() == [True, False, True] + + def test_arrow_int_widths(self): + pa = pytest.importorskip("pyarrow") + data = build_native_block([ + ("a", "Int8", [1]), + ("b", "Int16", [2]), + ("c", "UInt32", [3]), + ("d", "UInt64", [4]), + ]) + batch = _ch_core.ColBatch.decode_native(data) + reader = pa.RecordBatchReader.from_stream(batch) + result = reader.read_all() + assert result.schema.field("a").type == pa.int8() + assert result.schema.field("b").type == pa.int16() + assert result.schema.field("c").type == pa.uint32() + assert result.schema.field("d").type == pa.uint64() + + def test_arrow_float32(self): + pa = pytest.importorskip("pyarrow") + data = build_native_block([("f", "Float32", [1.5])]) + batch = _ch_core.ColBatch.decode_native(data) + reader = pa.RecordBatchReader.from_stream(batch) + result = reader.read_all() + assert result.schema.field("f").type == pa.float32() + + def test_arrow_fixed_binary(self): + pa = pytest.importorskip("pyarrow") + data = build_native_block([("fs", "FixedString(3)", [b"abc", b"xyz"])]) + batch = _ch_core.ColBatch.decode_native(data) + reader = pa.RecordBatchReader.from_stream(batch) + result = reader.read_all() + assert result.schema.field("fs").type == pa.binary(3) + assert result.column("fs").to_pylist() == [b"abc", b"xyz"] + + def test_arrow_nullable(self): + pa = pytest.importorskip("pyarrow") + data = build_native_block([("n", "Nullable(Int32)", [1, None, 3])]) + batch = _ch_core.ColBatch.decode_native(data) + reader = pa.RecordBatchReader.from_stream(batch) + result = reader.read_all() + col = result.column("n") + assert col.null_count == 1 + assert col.to_pylist() == [1, None, 3] + + def test_arrow_tuple(self): + # Tuple exports as an Arrow struct; a named tuple keeps its element names, + # an unnamed one uses the core's positional field names. + pa = pytest.importorskip("pyarrow") + data = build_native_block([("t", "Tuple(a Int32, b String)", [(13, "x"), (79, "y")])]) + batch = _ch_core.ColBatch.decode_native(data) + result = pa.RecordBatchReader.from_stream(batch).read_all() + assert pa.types.is_struct(result.schema.field("t").type) + assert result.column("t").to_pylist() == [{"a": 13, "b": "x"}, {"a": 79, "b": "y"}] + + def test_arrow_map(self): + # Map exports as an Arrow large_list of a key/value struct. + pa = pytest.importorskip("pyarrow") + data = build_native_block([("m", "Map(String, UInt8)", [{"k1": 1, "k2": 2}, {}])]) + batch = _ch_core.ColBatch.decode_native(data) + result = pa.RecordBatchReader.from_stream(batch).read_all() + assert result.column("m").to_pylist() == [ + [{"key": "k1", "value": 1}, {"key": "k2", "value": 2}], + [], + ] + + +# --------------------------------------------------------------------------- +# PipeDecoder fd ownership +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(sys.platform == "win32", reason="PipeDecoder is unix-only") +class TestPipeDecoderFd: + def test_invalid_fd_raises(self): + with pytest.raises(OSError): + _ch_core.PipeDecoder(999999) + + def test_negative_fd_raises(self): + with pytest.raises(OSError): + _ch_core.PipeDecoder(-1) + + def test_caller_closing_fd_is_safe(self): + # The decoder reads from its own duplicate, so the caller closing both + # pipe ends must not abort the process when the decoder is dropped. + read_fd, write_fd = os.pipe() + decoder = _ch_core.PipeDecoder(read_fd) + os.close(write_fd) + os.close(read_fd) + assert list(decoder) == [] + del decoder + + def test_streams_blocks_from_pipe(self): + read_fd, write_fd = os.pipe() + os.write(write_fd, build_native_block([("v", "Int64", [13, 79])])) + os.close(write_fd) + decoder = _ch_core.PipeDecoder(read_fd) + batches = [list(b.column_data(0)) for b in decoder] + os.close(read_fd) + assert batches == [[13, 79]] + + +# --------------------------------------------------------------------------- +# Arrow capsule lifecycle +# --------------------------------------------------------------------------- + +class TestArrowCapsuleLifecycle: + def test_unconsumed_capsule_does_not_leak(self): + # Run in a subprocess so ru_maxrss reflects only this workload and not + # the pytest process high-water mark. + script = textwrap.dedent(""" + import resource + import sys + + import _ch_core + + data = sys.stdin.buffer.read() + + def rss_kb(): + usage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # ru_maxrss is bytes on macOS, kilobytes on Linux. + return usage // 1024 if sys.platform == "darwin" else usage + + for _ in range(1000): + _ch_core.ColBatch.decode_native(data).__arrow_c_stream__() + base = rss_kb() + for _ in range(50000): + _ch_core.ColBatch.decode_native(data).__arrow_c_stream__() + print(rss_kb() - base) + """) + data = build_native_block([("v", "Int64", [13, 79, 5])]) + result = subprocess.run( + [sys.executable, "-c", script], + input=data, + capture_output=True, + check=True, + ) + delta_kb = int(result.stdout) + # The leak this guards against grew tens of MB over this loop. + assert delta_kb < 16 * 1024 + + def test_consumed_capsule_no_double_free(self): + pa = pytest.importorskip("pyarrow") + data = build_native_block([("v", "Int64", [13, 79])]) + for _ in range(10000): + table = pa.RecordBatchReader.from_stream( + _ch_core.ColBatch.decode_native(data) + ).read_all() + assert table.num_rows == 2 + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- + +class TestUnsupportedType: + def test_raises_value_error(self): + # A server type the core does not implement surfaces as a clean + # UnsupportedType error before body decoding. + buf = bytearray() + buf.extend(_encode_varint(1)) + buf.extend(_encode_varint(1)) + buf.extend(_encode_varint_string("id")) + buf.extend(_encode_varint_string("QBit(UInt8, 8)")) + with pytest.raises(NotImplementedError, match="Unsupported ClickHouse type 'QBit"): + _ch_core.ColBatch.decode_native(bytes(buf)) + + +class TestBufferInputs: + def test_decode_native_accepts_buffer_types(self): + data = build_native_block([("v", "Int64", [13, 79])]) + for view in (data, bytearray(data), memoryview(data)): + batch = _ch_core.ColBatch.decode_native(view) + assert list(batch.column_data(0)) == [13, 79] + + def test_feed_accepts_buffer_types(self): + data = build_native_block([("v", "Int64", [13, 79])]) + for view in (data, bytearray(data), memoryview(data)): + decoder = _ch_core.StreamDecoder() + batches = list(decoder.feed(view)) + list(decoder.finish()) + assert [list(b.column_data(0)) for b in batches] == [[13, 79]] + + +class TestMismatchedBlocks: + def test_mismatched_second_block_rejected_at_decode(self): + # Every block of a result shares one schema; the core rejects a + # mismatched later block at decode time. The binding also guards its + # raw list and tuple fills as defense in depth. + wide = build_native_block([("a", "Int64", [13]), ("b", "Int64", [79])]) + narrow = build_native_block([("a", "Int64", [5])]) + with pytest.raises(ValueError, match="schema differs"): + _ch_core.ColBatch.decode_native(wide + narrow) + + def test_mismatched_type_rejected_at_decode(self): + first = build_native_block([("a", "Int64", [13])]) + second = build_native_block([("a", "String", ["u1"])]) + with pytest.raises(ValueError, match="schema differs"): + _ch_core.ColBatch.decode_native(first + second) + + def test_stream_decoder_rejects_mismatch(self): + first = build_native_block([("a", "Int64", [13])]) + second = build_native_block([("a", "String", ["u1"])]) + decoder = _ch_core.StreamDecoder() + assert len(decoder.feed(first)) == 1 + with pytest.raises(ValueError, match="schema differs"): + decoder.feed(second) + + def test_block_decoder_rejects_mismatch(self): + first = build_native_block([("a", "Int64", [13])]) + second = build_native_block([("a", "String", ["u1"])]) + decoder = _ch_core.BlockDecoder(first + second) + batch = next(decoder) + assert list(batch.column_data(0)) == [13] + with pytest.raises(ValueError, match="schema differs"): + next(decoder) + + @pytest.mark.skipif(sys.platform == "win32", reason="PipeDecoder is unix-only") + def test_pipe_decoder_rejects_mismatch(self): + read_fd, write_fd = os.pipe() + os.write(write_fd, build_native_block([("a", "Int64", [13])])) + os.write(write_fd, build_native_block([("a", "String", ["u1"])])) + os.close(write_fd) + decoder = _ch_core.PipeDecoder(read_fd) + batches = [] + with pytest.raises(ValueError, match="schema differs"): + for batch in decoder: + batches.append(list(batch.column_data(0))) + os.close(read_fd) + # Whether the first good block is yielded depends on how the pipe + # reads chunk the feed; only the rejection is guaranteed. + assert batches in ([], [[13]]) + + +class TestTruncatedData: + def test_decode_native_truncated_raises_eof(self): + data = build_native_block([("v", "Int64", [13, 79])]) + with pytest.raises(EOFError): + _ch_core.ColBatch.decode_native(data[:-3]) + + def test_stream_decoder_truncated_finish_raises_eof(self): + data = build_native_block([("v", "Int64", [13, 79])]) + decoder = _ch_core.StreamDecoder() + assert list(decoder.feed(data[:-3])) == [] + with pytest.raises(EOFError): + decoder.finish() + + +class TestEmptyBatch: + def test_zero_rows(self): + data = build_native_block([("a", "Int32", []), ("b", "String", [])]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.num_rows == 0 + assert batch.num_columns == 2 + assert list(batch.to_python_rows()) == [] + + +class TestFromBatches: + _COLUMNS_A = [("id", "Int64", [13, 79]), ("name", "String", ["user_1", "user_2"])] + _COLUMNS_B = [("id", "Int64", [5]), ("name", "String", ["user_3"])] + + def test_merge_matches_decode_native(self): + first = build_native_block(self._COLUMNS_A) + second = build_native_block(self._COLUMNS_B) + parts = list(_ch_core.BlockDecoder(first)) + list(_ch_core.BlockDecoder(second)) + merged = _ch_core.ColBatch.from_batches(parts) + reference = _ch_core.ColBatch.decode_native(first + second) + assert merged.num_rows == 3 + assert merged.num_chunks == 2 + assert list(merged.to_python_rows()) == list(reference.to_python_rows()) + pa = pytest.importorskip("pyarrow") + merged_table = pa.RecordBatchReader.from_stream(merged).read_all() + reference_table = pa.RecordBatchReader.from_stream(reference).read_all() + assert merged_table == reference_table + + def test_mismatched_schema_raises(self): + first = _ch_core.ColBatch.decode_native( + build_native_block([("id", "Int64", [13])]) + ) + wrong_type = _ch_core.ColBatch.decode_native( + build_native_block([("id", "String", ["user_1"])]) + ) + wrong_name = _ch_core.ColBatch.decode_native( + build_native_block([("other", "Int64", [79])]) + ) + with pytest.raises(ValueError, match="Batch 1 schema differs"): + _ch_core.ColBatch.from_batches([first, wrong_type]) + with pytest.raises(ValueError, match="Batch 2 schema differs"): + _ch_core.ColBatch.from_batches([first, first, wrong_name]) + + def test_empty_list_raises(self): + with pytest.raises(ValueError, match="at least one batch"): + _ch_core.ColBatch.from_batches([]) + + def test_zero_row_trailer_dropped(self): + rows = _ch_core.ColBatch.decode_native(build_native_block(self._COLUMNS_A)) + # BlockDecoder keeps a zero-row block as a chunk, so the trailer batch + # carries one zero-row chunk into the merge. + (trailer,) = _ch_core.BlockDecoder( + build_native_block([("id", "Int64", []), ("name", "String", [])]) + ) + assert trailer.num_chunks == 1 + merged = _ch_core.ColBatch.from_batches([rows, trailer]) + assert merged.num_rows == 2 + assert merged.num_chunks == 1 + assert list(merged.to_python_rows()) == [(13, "user_1"), (79, "user_2")] + + def test_all_zero_rows(self): + empty_block = build_native_block([("id", "Int64", []), ("name", "String", [])]) + reference = _ch_core.ColBatch.decode_native(empty_block) + # Each part holds one zero-row chunk; the merge must drop them all. + (part,) = _ch_core.BlockDecoder(empty_block) + merged = _ch_core.ColBatch.from_batches([part, part]) + assert merged.num_rows == 0 + assert merged.num_chunks == 0 + assert merged.num_chunks == reference.num_chunks + assert list(merged.to_python_rows()) == list(reference.to_python_rows()) == [] + merged_cols = [list(c) for c in merged.to_python_columns()] + reference_cols = [list(c) for c in reference.to_python_columns()] + assert merged_cols == reference_cols == [[], []] + pa = pytest.importorskip("pyarrow") + table = pa.RecordBatchReader.from_stream(merged).read_all() + assert table.schema.names == ["id", "name"] + assert table.schema == pa.RecordBatchReader.from_stream(reference).read_all().schema + assert table.num_rows == 0 + + +class TestMaterializeFastPath: + """Edge values and multi-chunk slot accounting for the typed fill loops.""" + + _EDGE_COLUMNS = [ + ("i8", "Int8", [-128, -1, 127]), + ("i64", "Int64", [-(2**63), 0, 2**63 - 1]), + ("u32", "UInt32", [0, 2**31, 2**32 - 1]), + ("u64", "UInt64", [0, 2**63, 2**64 - 1]), + ("f64", "Float64", [-0.5, 0.0, 1.5]), + ("b", "Bool", [True, False, True]), + ("n", "Nullable(Int64)", [-(2**63), None, 2**63 - 1]), + ] + + def test_edge_values_rows_columns_and_column_data(self): + batch = _ch_core.ColBatch.decode_native(build_native_block(self._EDGE_COLUMNS)) + expected_cols = [values for _, _, values in self._EDGE_COLUMNS] + assert [list(c) for c in batch.to_python_columns()] == expected_cols + assert list(batch.to_python_rows()) == list(zip(*expected_cols)) + for idx, col in enumerate(expected_cols): + assert list(batch.column_data(idx)) == col + + def test_multi_chunk_mixed_fast_and_fallback_columns(self): + block = build_native_block( + [ + ("v", "Int64", [-(2**63), 2**63 - 1]), + ("s", "String", ["user_1", "user_2"]), + ("u", "UInt32", [2**32 - 1, 7]), + ] + ) + # The same batch three times stresses the per-chunk row-offset + # bookkeeping with duplicate Arc chunks. + (part,) = _ch_core.BlockDecoder(block) + merged = _ch_core.ColBatch.from_batches([part, part, part]) + assert merged.num_chunks == 3 + expected_rows = [(-(2**63), "user_1", 2**32 - 1), (2**63 - 1, "user_2", 7)] * 3 + assert list(merged.to_python_rows()) == expected_rows + assert [list(c) for c in merged.to_python_columns()] == [ + [-(2**63), 2**63 - 1] * 3, + ["user_1", "user_2"] * 3, + [2**32 - 1, 7] * 3, + ] + + +class TestBlockInfo: + def test_with_block_info(self): + data = build_native_block( + [("v", "Int64", [77, 88])], + block_info=True, + ) + batch = _ch_core.ColBatch.decode_native(data, has_block_info=True) + assert batch.num_rows == 2 + assert list(batch.column_data(0)) == [77, 88] diff --git a/rust/ch-core-py/tests/test_json.py b/rust/ch-core-py/tests/test_json.py new file mode 100644 index 00000000..75b6856e --- /dev/null +++ b/rust/ch-core-py/tests/test_json.py @@ -0,0 +1,473 @@ +from helpers import ( + _ch_core, + _encode_plain_body, + _encode_varint, + _encode_varint_string, + build_native_block_from_bodies, + decimal, + dt, + pytest, + struct, + sys, +) + + +class _DictSubclass(dict): + pass + + +class _StrSubclass(str): + pass + + +def _deep_list(depth): + doc = "x" + for _ in range(depth): + doc = [doc] + return doc + + +def _json_v2_prefix(): + """JSON V2 state prefix: one dynamic path nested.value whose Dynamic has + one String alternative plus the implicit SharedVariant, in BASIC mode.""" + prefix = bytearray(struct.pack(" primitive recursion. + vals = [100, 200, 100, 4_000_000_000, 200] + data = build_native_block([("c", "LowCardinality(UInt32)", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == vals + + def test_datetime_named_zone(self): + # LowCardinality(DateTime(tz)) must still apply the timezone policy, which + # means prepare_temporal has to see through the LowCardinality wrapper. + secs = 1705322096 # 2024-01-15 12:34:56 UTC + data = build_native_block( + [("c", "LowCardinality(DateTime('America/New_York'))", [secs, secs])] + ) + batch = _ch_core.ColBatch.decode_native(data) + v = list(batch.column_data(0))[0] + assert v.tzinfo == ZoneInfo("America/New_York") + assert v == dt.datetime(2024, 1, 15, 12, 34, 56, tzinfo=dt.timezone.utc) + + def test_invalid_utf8_hex_fallback(self): + vals = ["ok", b"\xff\xfe", "ok"] + data = build_native_block([("c", "LowCardinality(String)", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == ["ok", "fffe", "ok"] + + def test_paths_agree(self): + vals = ["a", "a", "b", "c", "b", "a", None] + data = build_native_block([("c", "LowCardinality(Nullable(String))", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [row[0] for row in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == vals + + def test_arrow_dictionary(self): + pa = pytest.importorskip("pyarrow") + vals = ["red", "green", "red", "blue"] + data = build_native_block([("c", "LowCardinality(String)", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + result = pa.RecordBatchReader.from_stream(batch).read_all() + assert pa.types.is_dictionary(result.schema.field("c").type) + assert result.column("c").to_pylist() == vals + + def test_arrow_nullable_dictionary(self): + pa = pytest.importorskip("pyarrow") + vals = ["x", None, "y", "x", None] + data = build_native_block([("c", "LowCardinality(Nullable(String))", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + col = pa.RecordBatchReader.from_stream(batch).read_all().column("c") + assert pa.types.is_dictionary(col.type) + assert col.to_pylist() == vals + assert col.null_count == 2 + + def test_across_chunks(self): + # Each Native block carries its own dictionary; the two chunks must + # concatenate correctly even with different per-block dictionaries. + first = build_native_block([("c", "LowCardinality(String)", ["a", "b", "a"])]) + second = build_native_block([("c", "LowCardinality(String)", ["c", "c", "d"])]) + batch = _ch_core.ColBatch.decode_native(first + second) + assert batch.num_chunks == 2 + assert list(batch.column_data(0)) == ["a", "b", "a", "c", "c", "d"] + + def test_with_block_info(self): + # The shape clickhouse-connect actually receives: client_protocol_version + # 54405 emits a BlockInfo preamble and, being below 54454, no per-column + # custom-serialization marker. The LowCardinality state prefix follows. + vals = ["red", "green", "red", None, "green"] + data = build_native_block( + [("c", "LowCardinality(Nullable(String))", vals)], block_info=True + ) + batch = _ch_core.ColBatch.decode_native(data, has_block_info=True) + assert list(batch.column_data(0)) == vals diff --git a/rust/ch-core-py/tests/test_qbit.py b/rust/ch-core-py/tests/test_qbit.py new file mode 100644 index 00000000..0c60ec5f --- /dev/null +++ b/rust/ch-core-py/tests/test_qbit.py @@ -0,0 +1,264 @@ +from helpers import ( + _bfloat16_value, + _ch_core, + _NdarrayLikeColumn, + build_native_block_from_bodies, + math, + pytest, +) + + +def _round_trip(type_name, values): + encoded = _ch_core.encode_native_block(["v"], [type_name], [values], len(values)) + return _ch_core.ColBatch.decode_native(encoded) + + +@pytest.mark.parametrize( + ("type_name", "values", "expected"), + [ + ( + "QBit(BFloat16, 3)", + [[1.0, -2.5, 3.25], [0.1, float("inf"), float("-inf")]], + [ + [_bfloat16_value(v) for v in [1.0, -2.5, 3.25]], + [_bfloat16_value(v) for v in [0.1, float("inf"), float("-inf")]], + ], + ), + ( + "QBit(Float32, 9)", + [ + [float(v) / 3 for v in range(9)], + [-0.0, 1.0, -2.0, 3.5, float("inf"), float("-inf"), 7.25, 8.5, 9.75], + ], + [ + [float(v) / 3 for v in range(9)], + [-0.0, 1.0, -2.0, 3.5, float("inf"), float("-inf"), 7.25, 8.5, 9.75], + ], + ), + ( + "QBit(Float64, 5)", + [[-0.99105519, 1.28887844, -0.43526649, -0.98520696, 0.66154391]], + [[-0.99105519, 1.28887844, -0.43526649, -0.98520696, 0.66154391]], + ), + ], +) +def test_python_exits_round_trip_all_element_widths(type_name, values, expected): + batch = _round_trip(type_name, values) + columns = list(batch.column_data(0)) + assert len(columns) == len(expected) + for actual, wanted in zip(columns, expected): + assert actual == pytest.approx(wanted, rel=1e-6, abs=0.0, nan_ok=True) + assert list(batch.to_python_columns()[0]) == columns + assert [row[0] for row in batch.to_python_rows()] == columns + + +def test_nullable_and_composite_shapes_round_trip(): + names = ["n", "a", "t", "at", "m", "v", "nt"] + types = [ + "Nullable(QBit(Float32, 3))", + "Array(QBit(Float64, 2))", + "Tuple(QBit(Float32, 3), Nullable(QBit(BFloat16, 2)))", + "Array(Tuple(QBit(Float32, 1), UInt8))", + "Map(String, QBit(Float32, 2))", + "Variant(QBit(Float32, 2), String)", + "Nullable(Tuple(QBit(Float32, 2), UInt8))", + ] + columns = [ + [[1.0, 2.0, 3.0], None, [-1.0, -2.0, -3.0]], + [[[1.0, 2.0]], [], [[3.0, 4.0], [5.0, 6.0]]], + [ + ([1.0, 2.0, 3.0], [4.0, 5.0]), + ([6.0, 7.0, 8.0], None), + ([9.0, 10.0, 11.0], [12.0, 13.0]), + ], + [[([1.0], 13)], [], [([2.0], 79), ([3.0], 5)]], + [{"a": [1.0, 2.0]}, {}, {"b": [3.0, 4.0]}], + [[1.0, 2.0], "value", None], + [([1.0, 2.0], 13), None, ([3.0, 4.0], 79)], + ] + encoded = _ch_core.encode_native_block(names, types, columns, 3) + batch = _ch_core.ColBatch.decode_native(encoded) + assert list(batch.to_python_rows()) == list(zip(*columns)) + + +def test_exact_and_generic_outer_containers_encode_identically(): + values = [[1.0, 2.0, 3.0], [-4.0, 5.5, 6.25]] + expected = _ch_core.encode_native_block(["v"], ["QBit(Float64, 3)"], [values], 2) + assert _ch_core.encode_native_block(["v"], ["QBit(Float64, 3)"], [tuple(map(tuple, values))], 2) == expected + assert _ch_core.encode_native_block(["v"], ["QBit(Float64, 3)"], [_NdarrayLikeColumn(values)], 2) == expected + + +@pytest.mark.parametrize("dimension", [1, 7, 8, 9, 15, 16, 17]) +def test_float32_byte_boundary_dimensions(dimension): + values = [[float(index) - 8.5 for index in range(dimension)]] + result = list(_round_trip(f"QBit(Float32, {dimension})", values).column_data(0)) + assert result[0] == pytest.approx(values[0], rel=1e-6) + + +def test_dimension_nine_exact_native_byte_group_order(): + row_1 = [-0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.0, -0.0] + row_2 = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.0] + # Only plane zero, the sign bit, is nonzero. Native is plane-major and + # row-major within the plane. The byte groups inside each row are reversed. + body = b"\x01\x81\x01\x00" + bytes(31 * 2 * 2) + native = build_native_block_from_bodies([("v", "QBit(Float32, 9)", body)], 2) + + assert _ch_core.encode_native_block(["v"], ["QBit(Float32, 9)"], [[row_1, row_2]], 2) == native + decoded = list(_ch_core.ColBatch.decode_native(native).column_data(0)) + assert [math.copysign(1.0, value) for value in decoded[0]] == [ + -1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + -1.0, + -1.0, + ] + assert [math.copysign(1.0, value) for value in decoded[1]] == [1.0] * 8 + [-1.0] + + +@pytest.mark.parametrize( + ("type_name", "plane_count"), + [("QBit(BFloat16, 9)", 16), ("QBit(Float64, 9)", 64)], +) +def test_dimension_nine_exact_native_bytes_other_widths(type_name, plane_count): + row = [-0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.0] + body = b"\x01\x01" + bytes((plane_count - 1) * 2) + native = build_native_block_from_bodies([("v", type_name, body)], 1) + + encoded = _ch_core.encode_native_block(["v"], [type_name], [[row]], 1) + assert encoded == native + decoded = list(_ch_core.ColBatch.decode_native(native).column_data(0))[0] + assert [math.copysign(1.0, value) for value in decoded] == [-1.0] + [1.0] * 7 + [-1.0] + + +@pytest.mark.parametrize( + ("alias", "canonical"), + [ + ("FLOAT", "Float32"), + ("real", "Float32"), + ("SINGLE", "Float32"), + ("DOUBLE", "Float64"), + ("double precision", "Float64"), + ], +) +def test_float_aliases_resolve_to_canonical_qbit_type(alias, canonical): + encoded = _ch_core.encode_native_block(["v"], [f"QBit({alias}, 2)"], [[[1.0, 2.0]]], 1) + assert _ch_core.ColBatch.decode_native(encoded).column_type_names == [f"QBit({canonical}, 2)"] + + +@pytest.mark.parametrize( + ("type_name", "dtype"), + [ + ("QBit(BFloat16, 3)", "float32"), + ("QBit(BFloat16, 3)", "float64"), + ("QBit(Float32, 3)", "float32"), + ("QBit(Float32, 3)", "float64"), + ("QBit(Float64, 3)", "float32"), + ("QBit(Float64, 3)", "float64"), + ], +) +def test_numpy_matrix_fast_path(type_name, dtype): + np = pytest.importorskip("numpy") + values = np.array([[1.0, 2.0, 3.0], [-4.0, 5.5, 6.25]], dtype=dtype) + expected_rows = list(_round_trip(type_name, values.tolist()).column_data(0)) + batch = _round_trip(type_name, values) + for actual, expected in zip(batch.column_data(0), expected_rows): + assert actual == pytest.approx(expected) + + +def test_nullable_numpy_matrix_fast_path_is_all_valid(): + np = pytest.importorskip("numpy") + values = np.array([[1.0, 2.0, 3.0], [-4.0, 5.5, 6.25]], dtype="float32") + batch = _round_trip("Nullable(QBit(Float32, 3))", values) + assert list(batch.column_data(0)) == values.tolist() + + +def test_bfloat16_numpy_matrix_rejects_out_of_range_value(): + np = pytest.importorskip("numpy") + values = np.array([[1e300]], dtype="float64") + with pytest.raises(ValueError, match="row 0 element 0 cannot be converted to BFloat16"): + _ch_core.encode_native_block(["v"], ["QBit(BFloat16, 1)"], [values], 1) + + +def test_zero_row_encode(): + batch = _round_trip("QBit(Float32, 9)", []) + assert batch.num_rows == 0 + assert list(batch.column_data(0)) == [] + + +def test_arrow_fixed_size_list_export(): + pa = pytest.importorskip("pyarrow") + values = [[1.0, 2.0, 3.0], None, [-4.0, 5.5, 6.25]] + batch = _round_trip("Nullable(QBit(Float32, 3))", values) + table = pa.RecordBatchReader.from_stream(batch).read_all() + assert pa.types.is_fixed_size_list(table.schema.field("v").type) + assert table.schema.field("v").type.list_size == 3 + assert table.schema.field("v").type.value_type == pa.float32() + assert table.column("v").to_pylist() == values + + +@pytest.mark.parametrize( + ("type_name", "values", "message"), + [ + ("QBit(Float32, 3)", [[1.0, 2.0]], "dimension mismatch"), + ("QBit(Float32, 3)", [[1.0, 2.0, 3.0, 4.0]], "dimension mismatch"), + ("QBit(Float32, 3)", [None], r"is None but QBit\(Float32, 3\) is not Nullable"), + ("QBit(Float32, 3)", [[1.0, "bad", 3.0]], "element 1 cannot be converted"), + ("QBit(Float32, 3)", [13], "is not a QBit vector"), + ("QBit(BFloat16, 1)", [[1e300]], "element 0 cannot be converted to BFloat16"), + ], +) +def test_insert_errors_identify_row_and_element(type_name, values, message): + with pytest.raises(ValueError, match=message): + _ch_core.encode_native_block(["v"], [type_name], [values], len(values)) + + +def test_special_float_bits_survive_round_trip(): + values = [[float("nan"), float("inf"), float("-inf"), -0.0]] + result = list(_round_trip("QBit(Float64, 4)", values).column_data(0))[0] + assert math.isnan(result[0]) + assert result[1:] == [float("inf"), float("-inf"), -0.0] + + +def test_inner_list_finalizer_resize_is_rejected(): + row = [None, 13.0] + + class MutatingFloat: + def __init__(self, container): + self._container = container + + def __float__(self): + self._container[0] = 5.0 + return 7.0 + + def __del__(self): + self._container.clear() + + row[0] = MutatingFloat(row) + with pytest.raises(ValueError, match="resized during encoding"): + _ch_core.encode_native_block(["v"], ["QBit(Float32, 2)"], [[row]], 1) + + +def test_outer_list_finalizer_resize_is_rejected(): + rows = [None, [13.0]] + + class MutatingVector: + def __init__(self, container): + self._container = container + + def __len__(self): + return 1 + + def __getitem__(self, index): + self._container[0] = [5.0] + return 7.0 + + def __del__(self): + self._container.clear() + + rows[0] = MutatingVector(rows) + with pytest.raises(ValueError, match="resized during encoding"): + _ch_core.encode_native_block(["v"], ["QBit(Float32, 1)"], [rows], 2) diff --git a/rust/ch-core-py/tests/test_scalars.py b/rust/ch-core-py/tests/test_scalars.py new file mode 100644 index 00000000..56aa0530 --- /dev/null +++ b/rust/ch-core-py/tests/test_scalars.py @@ -0,0 +1,447 @@ +import array + +from helpers import ( + _bfloat16_bytes, + _bfloat16_value, + _ch_core, + _NdarrayLikeColumn, + build_native_block, + decimal, + math, + pytest, + struct, +) + +# --------------------------------------------------------------------------- +# Integer types +# --------------------------------------------------------------------------- + + +class TestDecodeInt8: + def test_basic(self): + data = build_native_block([("v", "Int8", [1, -1, 127])]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.num_rows == 3 + assert batch.column_type_names == ["Int8"] + assert list(batch.column_data(0)) == [1, -1, 127] + + +class TestDecodeInt16: + def test_basic(self): + data = build_native_block([("v", "Int16", [256, -256, 32767])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [256, -256, 32767] + + +class TestDecodeInt32: + def test_basic(self): + data = build_native_block([("v", "Int32", [70000, -70000])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [70000, -70000] + + +class TestDecodeInt64: + def test_basic(self): + data = build_native_block([("id", "Int64", [10, 20, 30])]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.num_rows == 3 + assert batch.column_names == ["id"] + assert batch.column_type_names == ["Int64"] + assert list(batch.column_data(0)) == [10, 20, 30] + + def test_negative_values(self): + data = build_native_block([("n", "Int64", [-1, 0, 9223372036854775807])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [-1, 0, 9223372036854775807] + + +class TestDecodeUInt8: + def test_basic(self): + data = build_native_block([("v", "UInt8", [0, 128, 255])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [0, 128, 255] + + +class TestDecodeUInt16: + def test_basic(self): + data = build_native_block([("v", "UInt16", [0, 65535])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [0, 65535] + + +class TestDecodeUInt32: + def test_basic(self): + data = build_native_block([("v", "UInt32", [0, 4_000_000_000])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [0, 4_000_000_000] + + +class TestDecodeUInt64: + def test_basic(self): + data = build_native_block([("v", "UInt64", [0, 2**64 - 1])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [0, 2**64 - 1] + + +# --------------------------------------------------------------------------- +# Float types +# --------------------------------------------------------------------------- + + +class TestDecodeFloat32: + def test_basic(self): + data = build_native_block([("v", "Float32", [1.5, -2.25])]) + batch = _ch_core.ColBatch.decode_native(data) + vals = list(batch.column_data(0)) + assert vals[0] == pytest.approx(1.5) + assert vals[1] == pytest.approx(-2.25) + + +class TestDecodeFloat64: + def test_basic(self): + data = build_native_block([("val", "Float64", [1.5, 2.7, -0.1])]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.num_rows == 3 + vals = list(batch.column_data(0)) + assert vals[0] == pytest.approx(1.5) + assert vals[1] == pytest.approx(2.7) + assert vals[2] == pytest.approx(-0.1) + + +def test_typed_numeric_columns_match_python_array_contract_across_chunks(): + columns = [ + ("i8", "Int8", "b", [-13, 0, 79]), + ("i16", "Int16", "h", [-1300, 0, 7900]), + ("i32", "Int32", "i", [-130_000, 0, 790_000]), + ("i64", "Int64", "q", [-13_000_000_000, 0, 79_000_000_000]), + ("u8", "UInt8", "B", [0, 13, 79]), + ("u16", "UInt16", "H", [0, 1300, 7900]), + ("u32", "UInt32", "I", [0, 130_000, 790_000]), + ("u64", "UInt64", "Q", [0, 13_000_000_000, 79_000_000_000]), + ("f32", "Float32", "f", [-1.25, 0.0, 79.5]), + ("f64", "Float64", "d", [-1.25, 0.0, 79.5]), + ] + batches = [] + for indexes in ([0, 1], [2]): + payload = build_native_block( + [(name, type_name, [values[index] for index in indexes]) for name, type_name, _, values in columns] + ) + batches.append(_ch_core.ColBatch.decode_native(payload)) + + merged = _ch_core.ColBatch.from_batches(batches) + typed = merged.to_python_columns(typed_numeric=True) + for output, (_, _, typecode, values) in zip(typed, columns): + assert isinstance(output, array.array) + assert output.typecode == typecode + assert list(output) == pytest.approx(values) + + assert all(isinstance(output, list) for output in merged.to_python_columns()) + + +def test_typed_numeric_columns_leave_nullable_and_non_numeric_columns_as_lists(): + data = build_native_block( + [ + ("nullable", "Nullable(Int32)", [13, None, 79]), + ("enum", "Enum8('user_1'=1, 'user_2'=2)", [1, 2, 1]), + ("date", "Date", [13, 79, 101]), + ] + ) + + columns = _ch_core.ColBatch.decode_native(data).to_python_columns(typed_numeric=True) + + assert all(isinstance(output, list) for output in columns) + + +def test_typed_numeric_zero_row_columns_are_empty_typed_arrays(): + data = build_native_block( + [ + ("i64", "Int64", []), + ("f32", "Float32", []), + ("nullable", "Nullable(Int32)", []), + ("s", "String", []), + ] + ) + + typed, f32, nullable, s = _ch_core.ColBatch.decode_native(data).to_python_columns( + typed_numeric=True + ) + + assert isinstance(typed, array.array) and typed.typecode == "q" and len(typed) == 0 + assert isinstance(f32, array.array) and f32.typecode == "f" and len(f32) == 0 + assert isinstance(nullable, list) and nullable == [] + assert isinstance(s, list) and s == [] + + +class TestBFloat16: + def test_encode_decode_type_matrix_and_all_object_exits(self): + columns = [ + ("scalar", "BFloat16", [1.1, -1.1, 13.0]), + ("nullable", "Nullable(BFloat16)", [1.1, None, -1.1]), + ("array", "Array(BFloat16)", [[1.1, -1.1], [], [13.0]]), + ( + "tuple", + "Tuple(BFloat16, String)", + [(1.1, "user_1"), (-1.1, "user_2"), (13.0, "user_3")], + ), + ( + "array_tuple", + "Array(Tuple(BFloat16, UInt8))", + [[(1.1, 13), (-1.1, 79)], [], [(13.0, 5)]], + ), + ( + "map", + "Map(BFloat16, String)", + [{1.1: "user_1"}, {}, {-1.1: "user_2"}], + ), + ( + "low_cardinality", + "LowCardinality(BFloat16)", + [1.1, -1.1, 1.1], + ), + ( + "low_cardinality_nullable", + "LowCardinality(Nullable(BFloat16))", + [1.1, None, 1.1], + ), + ] + expected = [ + [_bfloat16_value(1.1), _bfloat16_value(-1.1), 13.0], + [_bfloat16_value(1.1), None, _bfloat16_value(-1.1)], + [[_bfloat16_value(1.1), _bfloat16_value(-1.1)], [], [13.0]], + [ + (_bfloat16_value(1.1), "user_1"), + (_bfloat16_value(-1.1), "user_2"), + (13.0, "user_3"), + ], + [ + [(_bfloat16_value(1.1), 13), (_bfloat16_value(-1.1), 79)], + [], + [(13.0, 5)], + ], + [{_bfloat16_value(1.1): "user_1"}, {}, {_bfloat16_value(-1.1): "user_2"}], + [_bfloat16_value(1.1), _bfloat16_value(-1.1), _bfloat16_value(1.1)], + [_bfloat16_value(1.1), None, _bfloat16_value(1.1)], + ] + + encoded = _ch_core.encode_native_block( + [name for name, _, _ in columns], + [type_name for _, type_name, _ in columns], + [values for _, _, values in columns], + 3, + ) + + assert encoded == build_native_block(columns) + batch = _ch_core.ColBatch.decode_native(encoded) + assert list(batch.to_python_columns()) == expected + assert [list(column) for column in zip(*batch.to_python_rows())] == expected + assert [list(batch.column_data(index)) for index in range(len(columns))] == expected + + @pytest.mark.parametrize("make", [list, tuple, _NdarrayLikeColumn]) + def test_scalar_container_paths_match_golden_bytes(self, make): + values = [3.141592, -2.71828, 13] + encoded = _ch_core.encode_native_block( + ["v"], + ["BFloat16"], + [make(values)], + len(values), + ) + assert encoded == build_native_block([("v", "BFloat16", values)]) + + @pytest.mark.parametrize("dtype", ["float32", "float64"]) + def test_numpy_buffer_matches_list_path(self, dtype): + np = pytest.importorskip("numpy") + values = [3.141592, -2.71828, 13.0] + array = np.array(values, dtype=dtype) + from_buffer = _ch_core.encode_native_block(["v"], ["BFloat16"], [array], len(array)) + from_list = _ch_core.encode_native_block(["v"], ["BFloat16"], [list(array)], len(array)) + assert from_buffer == from_list == build_native_block([("v", "BFloat16", values)]) + + def test_numpy_strided_buffer_matches_list_path(self): + np = pytest.importorskip("numpy") + array = np.array([1.1, 0.0, -1.1, 0.0, 13.0], dtype="float32")[::2] + from_buffer = _ch_core.encode_native_block(["v"], ["BFloat16"], [array], len(array)) + from_list = _ch_core.encode_native_block(["v"], ["BFloat16"], [list(array)], len(array)) + assert from_buffer == from_list + + def test_float32_buffer_signaling_nan_encodes_nan_word(self): + np = pytest.importorskip("numpy") + # sNaN whose payload lives only in the low 16 mantissa bits; plain + # truncation would produce 0x7F80 (+inf). + array = np.array([0x7F800001], dtype=np.uint32).view(np.float32) + encoded = _ch_core.encode_native_block(["v"], ["BFloat16"], [array], 1) + word = struct.unpack(" 32-bit, <=18 -> 64, <=38 -> 128, <=76 -> 256). + _CASES = [ + (9, 4, [0, 5, -5, -12000, 999_999_999, -999_999_999]), + (18, 6, [0, 7, -123456, 123456, 10**18 - 1, -(10**18 - 1)]), + (38, 10, [0, -3, 10**9 + 1, 10**38 - 1, -(10**38 - 1)]), + (76, 20, [0, 42, -42, 10**19, -(10**41 + 7), 10**76 - 1, -(10**76 - 1)]), + (9, 0, [0, -13, 999_999_999]), + (76, 0, [0, 10**76 - 1, -(10**76 - 1)]), + ] + + @staticmethod + def _reference(unscaled, precision, scale): + with decimal.localcontext() as ctx: + ctx.prec = precision + return decimal.Decimal(unscaled).scaleb(-scale) + + @pytest.mark.parametrize("precision,scale,unscaled", _CASES) + def test_matches_python_reference(self, precision, scale, unscaled): + type_name = f"Decimal({precision}, {scale})" + data = build_native_block([("d", type_name, unscaled)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == [type_name] + got = list(batch.column_data(0)) + expected = [self._reference(u, precision, scale) for u in unscaled] + assert got == expected + assert [v.as_tuple() for v in got] == [e.as_tuple() for e in expected] + + def test_paths_agree(self): + unscaled = [0, -3, 10**38 - 1] + data = build_native_block([("d", "Decimal(38, 10)", unscaled)]) + batch = _ch_core.ColBatch.decode_native(data) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [row[0] for row in batch.to_python_rows()] + expected = [self._reference(u, 38, 10) for u in unscaled] + assert via_column_data == via_columns == via_rows == expected + + def test_nullable(self): + vals = [12345, None, -12345, None] + data = build_native_block([("d", "Nullable(Decimal(18, 4))", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + expected = [None if v is None else self._reference(v, 18, 4) for v in vals] + assert list(batch.column_data(0)) == expected + assert [row[0] for row in batch.to_python_rows()] == expected + + def test_round_trip(self): + vals = [decimal.Decimal("123.4567"), decimal.Decimal("-1.5")] + encoded = _ch_core.encode_native_block(["d"], ["Decimal(20, 4)"], [vals], len(vals)) + got = list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) + expected = [decimal.Decimal("123.4567"), decimal.Decimal("-1.5000")] + assert got == expected + assert [v.as_tuple() for v in got] == [e.as_tuple() for e in expected] diff --git a/rust/ch-core-py/tests/test_strings_special.py b/rust/ch-core-py/tests/test_strings_special.py new file mode 100644 index 00000000..2b1fd04b --- /dev/null +++ b/rust/ch-core-py/tests/test_strings_special.py @@ -0,0 +1,843 @@ +from helpers import ( + _ch_core, + _NdarrayLikeColumn, + array, + build_native_block, + build_native_block_from_bodies, + ipaddress, + pytest, + sys, + uuid, +) + + +class TestNothing: + @pytest.mark.parametrize( + ("type_name", "values", "expected"), + [ + ("Nothing", [13, None, "ignored"], [None, None, None]), + ("Nullable(Nothing)", [13, None, "ignored"], [None, None, None]), + ( + "Array(Nothing)", + [[13, None], [], ["ignored"]], + [[None, None], [], [None]], + ), + ( + "Array(Nullable(Nothing))", + [[None, None], [], [None]], + [[None, None], [], [None]], + ), + ( + "Tuple(Nothing, UInt8)", + [(13, 13), (None, 79), ("ignored", 5)], + [(None, 13), (None, 79), (None, 5)], + ), + ( + "Tuple(Nullable(Nothing), UInt8)", + [(None, 13), (None, 79), (None, 5)], + [(None, 13), (None, 79), (None, 5)], + ), + ( + "Array(Tuple(Nothing, UInt8))", + [[(13, 13), (None, 79)], [], [("ignored", 5)]], + [[(None, 13), (None, 79)], [], [(None, 5)]], + ), + ( + "Map(Nothing, UInt8)", + [{13: 13}, {}, {"ignored": 79}], + [{None: 13}, {}, {None: 79}], + ), + ( + "Map(UInt8, Nothing)", + [{13: 13}, {}, {79: "ignored"}], + [{13: None}, {}, {79: None}], + ), + ( + "Map(UInt8, Nullable(Nothing))", + [{13: None}, {}, {79: None}], + [{13: None}, {}, {79: None}], + ), + ], + ) + def test_encode_decode_type_matrix_and_all_object_exits(self, type_name, values, expected): + encoded = _ch_core.encode_native_block(["v"], [type_name], [values], len(values)) + assert encoded == build_native_block([("v", type_name, values)]) + + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == expected + assert list(batch.to_python_columns()[0]) == expected + assert [row[0] for row in batch.to_python_rows()] == expected + + def test_noncanonical_wire_placeholder_bytes_are_ignored(self): + data = build_native_block_from_bodies( + [ + ("v", "Nothing", b"\x00\x7f\xff"), + ("sentinel", "UInt8", bytes([13, 79, 5])), + ], + 3, + ) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [None, None, None] + assert list(batch.column_data(1)) == [13, 79, 5] + + def test_nullable_nothing_with_structurally_valid_row_is_still_none(self): + data = build_native_block_from_bodies([("v", "Nullable(Nothing)", b"\x00\x01" + b"01")], 2) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [None, None] + + @pytest.mark.parametrize("type_name", ["Nothing", "Nullable(Nothing)", "Array(Nothing)"]) + def test_zero_rows(self, type_name): + encoded = _ch_core.encode_native_block(["v"], [type_name], [[]], 0) + assert encoded == build_native_block([("v", type_name, [])]) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == [] + + @pytest.mark.parametrize("container", [list, tuple, _NdarrayLikeColumn]) + def test_python_values_only_affect_nullable_structural_mask(self, container): + values = container([13, None, "ignored"]) + plain = _ch_core.encode_native_block(["v"], ["Nothing"], [values], 3) + nullable = _ch_core.encode_native_block(["v"], ["Nullable(Nothing)"], [values], 3) + assert plain == build_native_block([("v", "Nothing", [13, None, "ignored"])]) + assert nullable == build_native_block([("v", "Nullable(Nothing)", [13, None, "ignored"])]) + + def test_exact_bytes(self): + # Pin the canonical 0x30 marker run independently of build_native_block. + plain = _ch_core.encode_native_block(["v"], ["Nothing"], [[None, None, None]], 3) + assert plain == build_native_block_from_bodies([("v", "Nothing", b"\x30\x30\x30")], 3) + nullable = _ch_core.encode_native_block(["v"], ["Nullable(Nothing)"], [[None, 13, None]], 3) + assert nullable == build_native_block_from_bodies( + [("v", "Nullable(Nothing)", b"\x01\x00\x01" + b"\x30\x30\x30")], 3 + ) + + def test_multi_entry_map_nothing_keys_collapse(self): + # Dict representation: all-None keys collide, keeping the last value. + data = build_native_block_from_bodies( + [("v", "Map(Nothing, UInt8)", b"\x02\x00\x00\x00\x00\x00\x00\x00" + b"00" + bytes([13, 79]))], + 1, + ) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == [{None: 79}] + + @pytest.mark.parametrize("type_name", ["nothing", "NOTHING", "Nothing "]) + def test_type_name_is_case_sensitive(self, type_name): + with pytest.raises(NotImplementedError, match="unsupported ClickHouse type"): + _ch_core.encode_native_block(["v"], [type_name], [[]], 0) + + @pytest.mark.parametrize( + "type_name", + ["LowCardinality(Nothing)", "LowCardinality(Nullable(Nothing))"], + ) + def test_low_cardinality_nothing_rejected(self, type_name): + with pytest.raises(NotImplementedError, match="unsupported LowCardinality"): + _ch_core.encode_native_block(["v"], [type_name], [[]], 0) + + def test_truncated_marker_run_rejected(self): + data = build_native_block_from_bodies([("v", "Nothing", b"00")], 3) + with pytest.raises(EOFError): + _ch_core.ColBatch.decode_native(data) + + def test_arrow_null_type(self): + pa = pytest.importorskip("pyarrow") + encoded = _ch_core.encode_native_block( + ["v", "n", "a"], + ["Nothing", "Nullable(Nothing)", "Array(Nothing)"], + [[None, None], [None, None], [[None], []]], + 2, + ) + table = pa.RecordBatchReader.from_stream(_ch_core.ColBatch.decode_native(encoded)).read_all() + assert table.schema.types == [pa.null(), pa.null(), pa.large_list(pa.null())] + assert table.schema.field("v").nullable + assert table.schema.field("n").nullable + assert table.column("v").to_pylist() == [None, None] + assert table.column("n").to_pylist() == [None, None] + assert table.column("a").to_pylist() == [[None], []] + + +# --------------------------------------------------------------------------- +# String types +# --------------------------------------------------------------------------- + + +class TestDecodeString: + def test_basic(self): + data = build_native_block([("s", "String", ["hello", "", "world!"])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == ["hello", "", "world!"] + + def test_unicode(self): + data = build_native_block([("s", "String", ["héllo", "日本語"])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == ["héllo", "日本語"] + + +class TestStringInvalidUtf8: + def test_hex_fallback_on_all_paths(self): + # Invalid UTF-8 renders as the lowercase hex of the raw bytes on every + # materialization path, matching clickhouse-connect's String fallback. + data = build_native_block([("s", "String", ["ok", b"\xff\xfe", b"\x80abc"])]) + batch = _ch_core.ColBatch.decode_native(data) + expected = ["ok", "fffe", "80616263"] + assert list(batch.column_data(0)) == expected + assert list(batch.to_python_columns()[0]) == expected + assert [row[0] for row in batch.to_python_rows()] == expected + + def test_hex_fallback_nullable(self): + data = build_native_block([("s", "Nullable(String)", [b"\xc3\x28", None, "u1"])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == ["c328", None, "u1"] + + +class TestDecodeFixedString: + def test_basic(self): + data = build_native_block([("fs", "FixedString(3)", [b"abc", b"xyz"])]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.num_rows == 2 + assert batch.column_type_names == ["FixedString(3)"] + result = list(batch.column_data(0)) + assert result[0] == b"abc" + assert result[1] == b"xyz" + + def test_null_padded(self): + data = build_native_block([("fs", "FixedString(4)", [b"ab\x00\x00"])]) + batch = _ch_core.ColBatch.decode_native(data) + result = list(batch.column_data(0)) + assert result[0] == b"ab\x00\x00" + + +# --------------------------------------------------------------------------- +# Enum +# --------------------------------------------------------------------------- + + +class TestDecodeEnum: + _E8 = "Enum8('red' = 1, 'green' = 2, 'blue' = 3)" + _E16 = "Enum16('alpha' = -5, 'beta' = 0, 'gamma' = 1000)" + + def test_enum8_basic(self): + # Wire carries the integer codes; the Python exit yields the labels. + data = build_native_block([("c", self._E8, [1, 2, 3, 1, 2])]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == [self._E8] + assert list(batch.column_data(0)) == ["red", "green", "blue", "red", "green"] + + def test_enum16_negative_and_large(self): + data = build_native_block([("c", self._E16, [-5, 0, 1000, 0, -5])]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == [self._E16] + assert list(batch.column_data(0)) == ["alpha", "beta", "gamma", "beta", "alpha"] + + def test_nullable_enum8(self): + data = build_native_block([("c", f"Nullable({self._E8})", [1, None, 3, None])]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == [f"Nullable({self._E8})"] + assert list(batch.column_data(0)) == ["red", None, "blue", None] + + def test_unknown_code_is_none(self): + # A code with no defined label materializes as None, matching + # clickhouse-connect's int_map.get(code, None). + data = build_native_block([("c", self._E8, [1, 99, 2])]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == ["red", None, "green"] + + def test_paths_agree(self): + data = build_native_block([("c", self._E16, [1000, -5, 0, 1000, -5])]) + batch = _ch_core.ColBatch.decode_native(data) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [row[0] for row in batch.to_python_rows()] + expected = ["gamma", "alpha", "beta", "gamma", "alpha"] + assert via_column_data == via_columns == via_rows == expected + + def test_arrow_exports_int_codes(self): + # The core exports Enum as the raw signed-int codes (no per-cell label + # remapping); the label policy lives only on the Python object exit. + pa = pytest.importorskip("pyarrow") + data = build_native_block([("c", self._E8, [1, 2, 3]), ("d", self._E16, [-5, 0, 1000])]) + batch = _ch_core.ColBatch.decode_native(data) + result = pa.RecordBatchReader.from_stream(batch).read_all() + assert result.schema.field("c").type == pa.int8() + assert result.schema.field("d").type == pa.int16() + assert result.column("c").to_pylist() == [1, 2, 3] + assert result.column("d").to_pylist() == [-5, 0, 1000] + + +# --------------------------------------------------------------------------- +# UUID +# --------------------------------------------------------------------------- + +class TestDecodeUUID: + _KNOWN = uuid.UUID("00112233-4455-6677-8899-aabbccddeeff") + _VALUES = [uuid.UUID(int=0), _KNOWN, uuid.UUID(int=(1 << 128) - 1)] + + def test_values(self): + data = build_native_block([("u", "UUID", self._VALUES)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["UUID"] + got = list(batch.column_data(0)) + assert got == self._VALUES + assert all(type(v) is uuid.UUID for v in got) + assert got[1] == uuid.UUID("00112233-4455-6677-8899-aabbccddeeff") + assert all(v.is_safe is uuid.SafeUUID.unsafe for v in got) + + def test_paths_agree(self): + data = build_native_block([("u", "UUID", self._VALUES)]) + batch = _ch_core.ColBatch.decode_native(data) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [row[0] for row in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == self._VALUES + + def test_nullable(self): + vals = [self._KNOWN, None, uuid.UUID(int=0), None] + data = build_native_block([("u", "Nullable(UUID)", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == vals + assert [row[0] for row in batch.to_python_rows()] == vals + + def test_low_cardinality(self): + vals = [self._KNOWN, uuid.UUID(int=0), self._KNOWN, self._KNOWN] + data = build_native_block([("u", "LowCardinality(UUID)", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + got = list(batch.column_data(0)) + assert got == vals + assert all(v.is_safe is uuid.SafeUUID.unsafe for v in got) + + def test_low_cardinality_nullable(self): + vals = [self._KNOWN, None, self._KNOWN, None] + data = build_native_block([("u", "LowCardinality(Nullable(UUID))", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == vals + assert [row[0] for row in batch.to_python_rows()] == vals + + def test_low_cardinality_all_null(self): + vals = [None, None, None] + data = build_native_block([("u", "LowCardinality(Nullable(UUID))", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == vals + + def test_round_trip(self): + vals = [self._KNOWN, uuid.UUID(int=79)] + encoded = _ch_core.encode_native_block(["u"], ["UUID"], [vals], len(vals)) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == vals + + +# --------------------------------------------------------------------------- +# IPv4 +# --------------------------------------------------------------------------- + +class TestDecodeIPv4: + _VALUES = [ + ipaddress.IPv4Address("0.0.0.0"), + ipaddress.IPv4Address("255.255.255.255"), + ipaddress.IPv4Address("1.2.3.4"), + ] + + def test_values(self): + data = build_native_block([("v", "IPv4", self._VALUES)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["IPv4"] + got = list(batch.column_data(0)) + assert got == self._VALUES + assert all(type(v) is ipaddress.IPv4Address for v in got) + assert got[2] == ipaddress.ip_address("1.2.3.4") + + def test_paths_agree(self): + data = build_native_block([("v", "IPv4", self._VALUES)]) + batch = _ch_core.ColBatch.decode_native(data) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [row[0] for row in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == self._VALUES + + def test_nullable(self): + vals = [ipaddress.IPv4Address("1.2.3.4"), None, ipaddress.IPv4Address("0.0.0.0")] + data = build_native_block([("v", "Nullable(IPv4)", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == vals + assert [row[0] for row in batch.to_python_rows()] == vals + + def test_round_trip(self): + vals = [ipaddress.IPv4Address("192.0.2.1"), ipaddress.IPv4Address("198.51.100.7")] + encoded = _ch_core.encode_native_block(["v"], ["IPv4"], [vals], len(vals)) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == vals + + +# --------------------------------------------------------------------------- +# IPv6 +# --------------------------------------------------------------------------- + +class TestDecodeIPv6: + _VALUES = [ + ipaddress.IPv6Address("::"), + ipaddress.IPv6Address("::1"), + ipaddress.IPv6Address("2001:db8:85a3:8d3:1319:8a2e:370:7348"), + ipaddress.IPv6Address("::ffff:1.2.3.4"), + ] + + def test_values(self): + data = build_native_block([("v", "IPv6", self._VALUES)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["IPv6"] + got = list(batch.column_data(0)) + assert got == self._VALUES + # Always IPv6Address, even for a v4-mapped value. + assert all(type(v) is ipaddress.IPv6Address for v in got) + # str() reads _scope_id, so it verifies the attribute was set. + assert [str(v) for v in got] == [str(v) for v in self._VALUES] + + def test_paths_agree(self): + data = build_native_block([("v", "IPv6", self._VALUES)]) + batch = _ch_core.ColBatch.decode_native(data) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [row[0] for row in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == self._VALUES + + def test_nullable(self): + vals = [ipaddress.IPv6Address("::1"), None, ipaddress.IPv6Address("::ffff:1.2.3.4")] + data = build_native_block([("v", "Nullable(IPv6)", vals)]) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0)) == vals + assert [row[0] for row in batch.to_python_rows()] == vals + + def test_round_trip(self): + vals = [ipaddress.IPv6Address("2001:db8::1"), ipaddress.IPv6Address("::ffff:192.0.2.9")] + encoded = _ch_core.encode_native_block(["v"], ["IPv6"], [vals], len(vals)) + got = list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) + assert got == vals + assert all(type(v) is ipaddress.IPv6Address for v in got) + + +# --------------------------------------------------------------------------- +# SimpleAggregateFunction (physical == the inner type) +# --------------------------------------------------------------------------- + + +class TestSimpleAggregateFunction: + def test_sum_uint64_golden_and_round_trip(self): + type_name = "SimpleAggregateFunction(sum, UInt64)" + values = [13, 79, 0, 18446744073709551615] + encoded = _ch_core.encode_native_block(["s"], [type_name], [values], len(values)) + assert encoded == build_native_block([("s", type_name, values)]) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + decoded = list(batch.column_data(0)) + assert decoded == values + assert all(isinstance(v, int) for v in decoded) + + def test_any_string_golden_and_round_trip(self): + type_name = "SimpleAggregateFunction(any, String)" + values = ["red", "", "green"] + encoded = _ch_core.encode_native_block(["s"], [type_name], [values], len(values)) + assert encoded == build_native_block([("s", type_name, values)]) + decoded = list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) + assert decoded == values + assert all(isinstance(v, str) for v in decoded) + + def test_anylast_low_cardinality_string_round_trip(self): + # The inner LowCardinality resolves through the LC path; the core's LC + # index word differs from the helper's, so this round-trips. + type_name = "SimpleAggregateFunction(anyLast, LowCardinality(String))" + values = ["red", "blue", "red", "red", "green"] + encoded = _ch_core.encode_native_block(["s"], [type_name], [values], len(values)) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == values + + def test_all_exit_paths_agree(self): + type_name = "SimpleAggregateFunction(sum, UInt64)" + values = [13, 79, 5] + batch = _ch_core.ColBatch.decode_native(build_native_block([("s", type_name, values)])) + via_column_data = list(batch.column_data(0)) + via_columns = list(batch.to_python_columns()[0]) + via_rows = [r[0] for r in batch.to_python_rows()] + assert via_column_data == via_columns == via_rows == values + + def test_array_of_saf_golden_and_round_trip(self): + # A SAF alias nested inside a non-alias Array container. + type_name = "Array(SimpleAggregateFunction(sum, UInt64))" + rows = [[13, 79], [], [5, 5, 18446744073709551615]] + encoded = _ch_core.encode_native_block(["a"], [type_name], [rows], len(rows)) + assert encoded == build_native_block([("a", type_name, rows)]) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == rows + + def test_low_cardinality_string_round_trip(self): + # LowCardinality(SAF(anyLast, String)): the SAF chain resolves through the + # LC path. Server round-trips this live; the python codec cannot. The + # core's LC index word differs from the helper's, so no golden compare. + type_name = "LowCardinality(SimpleAggregateFunction(anyLast, String))" + values = ["red", "blue", "red", "red", "green"] + encoded = _ch_core.encode_native_block(["v"], [type_name], [values], len(values)) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == values + built = build_native_block([("v", type_name, values)]) + assert list(_ch_core.ColBatch.decode_native(built).column_data(0)) == values + + def test_low_cardinality_nullable_string_round_trip(self): + # LowCardinality(SAF(anyLast, Nullable(String))) is index-level nullable: + # the SAF chain is stripped before the Nullable is detected. + type_name = "LowCardinality(SimpleAggregateFunction(anyLast, Nullable(String)))" + values = ["red", None, "red", None, "blue"] + encoded = _ch_core.encode_native_block(["v"], [type_name], [values], len(values)) + batch = _ch_core.ColBatch.decode_native(encoded) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == values + built = build_native_block([("v", type_name, values)]) + assert list(_ch_core.ColBatch.decode_native(built).column_data(0)) == values + + +# --------------------------------------------------------------------------- +# AggregateFunction (opaque serialized state bytes) +# --------------------------------------------------------------------------- + + +class TestAggregateFunction: + def test_count_decode_all_python_exits(self): + type_name = "AggregateFunction(count)" + states = [b"\x00", b"\x0d", b"\x80\x01"] + native = build_native_block_from_bodies( + [("c", type_name, b"".join(states))], + len(states), + ) + + batch = _ch_core.ColBatch.decode_native(native) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == states + assert list(batch.to_python_columns()[0]) == states + assert batch.to_python_rows() == [(state,) for state in states] + assert all(isinstance(state, bytes) for state in batch.column_data(0)) + + def test_count_encode_accepts_bytes_like_and_round_trips(self): + type_name = "AggregateFunction(count)" + values = [b"\x00", bytearray(b"\x0d"), memoryview(b"\x80\xff\x01")[::2]] + encoded = _ch_core.encode_native_block(["c"], [type_name], [values], len(values)) + expected = build_native_block_from_bodies( + [("c", type_name, b"\x00\x0d\x80\x01")], + len(values), + ) + + assert encoded == expected + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == [bytes(value) for value in values] + + generic = _NdarrayLikeColumn([b"\x05", memoryview(b"\x0b")]) + generic_encoded = _ch_core.encode_native_block( + ["c"], + [type_name], + [generic], + len(generic), + ) + assert list(_ch_core.ColBatch.decode_native(generic_encoded).column_data(0)) == [ + b"\x05", + b"\x0b", + ] + + @pytest.mark.parametrize( + ("type_name", "width", "states"), + [ + ("AggregateFunction(sum, UInt64)", 8, [13, 79]), + ("AggregateFunction(sum, Int32)", 8, [-13, 79]), + ("AggregateFunction(sum, UInt128)", 16, [13, 79]), + ("AggregateFunction(sum, UInt256)", 32, [13, 79]), + ("AggregateFunction(sum, Decimal(9, 2))", 16, [1300, 7900]), + ], + ) + def test_sum_fixed_width_states_round_trip(self, type_name, width, states): + values = [value.to_bytes(width, "little", signed=value < 0) for value in states] + encoded = _ch_core.encode_native_block(["s"], [type_name], [values], len(values)) + expected = build_native_block_from_bodies( + [("s", type_name, b"".join(values))], + len(values), + ) + + assert encoded == expected + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == values + + @pytest.mark.parametrize( + ("type_name", "width", "value"), + [ + ("AggregateFunction(sum, Nullable(UInt8))", 8, 13), + ("AggregateFunction(sum, Nullable(Int32))", 8, -13), + ("AggregateFunction(sum, Nullable(UInt128))", 16, 79), + ("AggregateFunction(sum, Nullable(Decimal(9, 2)))", 16, 1300), + ("AggregateFunction(sum, Nullable(UInt256))", 32, 79), + ], + ) + def test_nullable_sum_variable_width_states_round_trip(self, type_name, width, value): + accumulator = value.to_bytes(width, "little", signed=value < 0) + states = [b"\x00", b"\x01" + accumulator, b"\x80" + accumulator] + encoded = _ch_core.encode_native_block(["s"], [type_name], [states], len(states)) + + assert encoded == build_native_block_from_bodies( + [("s", type_name, b"".join(states))], + len(states), + ) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == states + + def test_nothing_uint64_state_round_trip(self): + type_name = "AggregateFunction(nothingUInt64, Nullable(Nothing))" + states = [b"\x00", b"\x00", b"\x00"] + encoded = _ch_core.encode_native_block(["s"], [type_name], [states], len(states)) + + assert encoded == build_native_block_from_bodies( + [("s", type_name, b"\x00\x00\x00")], + len(states), + ) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == states + + @pytest.mark.parametrize( + ("type_name", "values"), + [ + ( + "Array(AggregateFunction(count))", + [[b"\x0d", b"\x4f"], [], [b"\x80\x01"]], + ), + ( + "Array(AggregateFunction(sum, Nullable(Int32)))", + [ + [b"\x00", b"\x01" + (-13).to_bytes(8, "little", signed=True)], + [], + [b"\x80" + (79).to_bytes(8, "little")], + ], + ), + ( + "Tuple(AggregateFunction(sum, UInt64), UInt8)", + [((13).to_bytes(8, "little"), 5), ((79).to_bytes(8, "little"), 11)], + ), + ( + "Tuple(state AggregateFunction(count), code UInt8)", + [ + {"state": b"\x0d", "code": 5}, + {"state": b"\x4f", "code": 11}, + ], + ), + ( + "Array(Tuple(AggregateFunction(count), UInt8))", + [[(b"\x0d", 5), (b"\x4f", 11)], [], [(b"\x80\x01", 17)]], + ), + ( + "Array(Tuple(AggregateFunction(sum, Nullable(UInt64)), UInt8))", + [ + [ + (b"\x00", 5), + (b"\x01" + (13).to_bytes(8, "little"), 11), + ], + [], + [(b"\x80" + (79).to_bytes(8, "little"), 17)], + ], + ), + ( + "Nullable(Tuple(AggregateFunction(count)))", + [(b"\x0d",), (b"\x4f",)], + ), + ( + "Map(String, AggregateFunction(count))", + [{"first": b"\x0d", "second": b"\x4f"}, {}, {"third": b"\x80\x01"}], + ), + ], + ) + def test_container_shapes_round_trip(self, type_name, values): + encoded = _ch_core.encode_native_block(["v"], [type_name], [values], len(values)) + assert list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) == values + + @pytest.mark.parametrize( + ("type_name", "states"), + [ + ("AggregateFunction(count)", [b"\x0d", b"\x80\x01", b"\x05"]), + ( + "AggregateFunction(sum, Nullable(UInt64))", + [b"\x00", b"\x01" + (13).to_bytes(8, "little"), b"\x00"], + ), + ], + ) + def test_arrow_export_is_large_binary(self, type_name, states): + pa = pytest.importorskip("pyarrow") + native = build_native_block_from_bodies( + [("s", type_name, b"".join(states))], + len(states), + ) + batch = _ch_core.ColBatch.decode_native(native) + + table = pa.RecordBatchReader.from_stream(batch).read_all() + assert table.schema.field("s").type == pa.large_binary() + assert not table.schema.field("s").nullable + assert table.column("s").to_pylist() == states + + def test_stream_decoder_preserves_variable_state_boundaries_across_chunks(self): + type_name = "AggregateFunction(count)" + first_states = [b"\x0d", b"\x80\x01"] + second_states = [b"\x05", b"\xff\x01"] + native = b"".join( + [ + build_native_block_from_bodies( + [("c", type_name, b"".join(first_states))], + len(first_states), + ), + build_native_block_from_bodies( + [("c", type_name, b"".join(second_states))], + len(second_states), + ), + ] + ) + decoder = _ch_core.StreamDecoder() + batches = [] + for byte in native: + batches.extend(decoder.feed(bytes([byte]))) + batches.extend(decoder.finish()) + + assert len(batches) == 2 + combined = _ch_core.ColBatch.from_batches(batches) + assert combined.num_chunks == 2 + assert list(combined.column_data(0)) == first_states + second_states + + def test_nullable_sum_stream_decoder_preserves_conditional_boundaries(self): + type_name = "AggregateFunction(sum, Nullable(UInt64))" + first_states = [b"\x00", b"\x01" + (13).to_bytes(8, "little")] + second_states = [b"\x80" + (79).to_bytes(8, "little"), b"\x00"] + native = b"".join( + build_native_block_from_bodies( + [("s", type_name, b"".join(states))], + len(states), + ) + for states in (first_states, second_states) + ) + decoder = _ch_core.StreamDecoder() + batches = [] + for byte in native: + batches.extend(decoder.feed(bytes([byte]))) + batches.extend(decoder.finish()) + + combined = _ch_core.ColBatch.from_batches(batches) + expected = first_states + second_states + assert combined.num_chunks == 2 + assert list(combined.column_data(0)) == expected + assert list(combined.to_python_columns()[0]) == expected + assert combined.to_python_rows() == [(state,) for state in expected] + + @pytest.mark.parametrize("state", [b"", b"\x01", b"\x01" + b"\x0d" * 7]) + def test_nullable_sum_truncation_is_eof(self, state): + native = build_native_block_from_bodies( + [("s", "AggregateFunction(sum, Nullable(UInt64))", state)], + 1, + ) + + with pytest.raises(EOFError, match="Truncated Native data"): + _ch_core.ColBatch.decode_native(native) + + @pytest.mark.parametrize( + "type_name", + [ + "AggregateFunction(sum, UInt64)", + "AggregateFunction(sum, Nullable(UInt64))", + ], + ) + def test_zero_rows_preserves_schema(self, type_name): + native = build_native_block_from_bodies([("s", type_name, b"")], 0) + batch = _ch_core.ColBatch.decode_native(native) + + assert batch.num_rows == 0 + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == [] + assert _ch_core.encode_native_block(["s"], [type_name], [[]], 0) == native + + @pytest.mark.parametrize( + ("type_name", "value", "error"), + [ + ("AggregateFunction(count)", "not-bytes", "AggregateFunction state bytes"), + ( + "AggregateFunction(count)", + array.array("b", [13]), + "AggregateFunction state bytes", + ), + ("AggregateFunction(count)", None, "is None but AggregateFunction"), + ("AggregateFunction(count)", b"", "not exactly one valid serialized"), + ("AggregateFunction(count)", b"\x80", "not exactly one valid serialized"), + ( + "AggregateFunction(nothingUInt64, Nullable(Nothing))", + b"\x01", + "not exactly one valid serialized", + ), + ( + "AggregateFunction(sum, UInt64)", + b"\x00" * 7, + "not exactly one valid serialized", + ), + ( + "AggregateFunction(sum, Nullable(UInt64))", + b"", + "not exactly one valid serialized", + ), + ( + "AggregateFunction(sum, Nullable(UInt64))", + b"\x00\x0d", + "not exactly one valid serialized", + ), + ( + "AggregateFunction(sum, Nullable(UInt64))", + b"\x01" + b"\x0d" * 7, + "not exactly one valid serialized", + ), + ( + "AggregateFunction(sum, Nullable(UInt64))", + b"\x01" + b"\x0d" * 9, + "not exactly one valid serialized", + ), + ], + ) + def test_insert_rejects_invalid_state(self, type_name, value, error): + with pytest.raises(ValueError, match=error): + _ch_core.encode_native_block(["s"], [type_name], [[value]], 1) + + @pytest.mark.skipif(sys.version_info < (3, 12), reason="pure-python __buffer__") + def test_resize_during_buffer_conversion_is_rejected(self): + class SlotSwappingState: + def __init__(self, container): + self._container = container + + def __buffer__(self, flags): + self._container[0] = b"\x0d" + return memoryview(b"\x0d") + + def __del__(self): + self._container.pop() + + values = [None, b"\x4f", b"\x05"] + values[0] = SlotSwappingState(values) + with pytest.raises(ValueError, match="resized during encoding"): + _ch_core.encode_native_block( + ["s"], ["AggregateFunction(count)"], [values], 3 + ) + + @pytest.mark.parametrize( + "type_name", + [ + "Nullable(AggregateFunction(count))", + "LowCardinality(AggregateFunction(count))", + "Aggregatefunction(count)", + "AggregateFunction(sum, Nullable(Nothing))", + "AggregateFunction(sum, Nullable(String))", + ], + ) + def test_illegal_or_misspelled_type_is_rejected(self, type_name): + with pytest.raises(NotImplementedError, match="unsupported"): + _ch_core.encode_native_block(["s"], [type_name], [[b"\x00"]], 1) + + def test_null_nullable_tuple_requires_core_placeholder_state(self): + type_name = "Nullable(Tuple(AggregateFunction(count)))" + + with pytest.raises(NotImplementedError, match="canonical placeholder state"): + _ch_core.encode_native_block(["s"], [type_name], [[None]], 1) + + +class TestAliasRejections: + @pytest.mark.parametrize( + "type_name", + ["Nullable(Ring)", "Nullable(Nested(a UInt32, b String))"], + ) + def test_server_illegal_nullable_alias_rejected(self, type_name): + # Ring and Nested expand to an Array, which is not nullable-able, so the + # core parser rejects the type and encode raises at the header stage. + with pytest.raises(NotImplementedError, match="unsupported ClickHouse type"): + _ch_core.encode_native_block(["x"], [type_name], [[None]], 1) diff --git a/rust/ch-core-py/tests/test_temporal.py b/rust/ch-core-py/tests/test_temporal.py new file mode 100644 index 00000000..09d53e0c --- /dev/null +++ b/rust/ch-core-py/tests/test_temporal.py @@ -0,0 +1,759 @@ +from helpers import ( + _EPOCH_DATE, + _EPOCH_NAIVE, + ZoneInfo, + _ch_core, + _encode_plain_body, + _NdarrayLikeColumn, + build_native_block, + build_native_block_from_bodies, + dt, + pytest, + struct, +) + + +class TestTemporalInsertFastPath: + def _encode(self, type_name, vals, n=None): + n = len(vals) if n is None else n + return _ch_core.encode_native_block(["t"], [type_name], [vals], n) + + @pytest.mark.parametrize("type_name", ["DateTime", "DateTime64(3)", "DateTime64(6)"]) + def test_datetime_inputs_match_generic_container(self, type_name): + class SubDT(dt.datetime): + pass + + vals = [ + dt.datetime(2024, 5, 4, 3, 2, 1), + dt.datetime(2024, 1, 15, 12, 34, 56, 789000, tzinfo=dt.timezone.utc), + SubDT(2024, 5, 4, 3, 2, 1, 123456), + 1700000000, + ] + fast = self._encode(type_name, vals) + assert fast == self._encode(type_name, tuple(vals), len(vals)) + assert fast == self._encode(type_name, _NdarrayLikeColumn(vals), len(vals)) + + def test_datetime64_string_fallback_matches_generic(self): + vals = ["2024-01-15T12:34:56.789000+00:00", 5] + fast = self._encode("DateTime64(3)", vals) + assert fast == self._encode("DateTime64(3)", _NdarrayLikeColumn(vals), 2) + + def test_date_inputs_match_generic_and_helper(self): + vals = [dt.date(2024, 1, 2), 19737, dt.date(1970, 1, 1)] + fast = self._encode("Date", vals) + assert fast == self._encode("Date", _NdarrayLikeColumn(vals), 3) + expected_days = [dt.date(2024, 1, 2).toordinal() - 719163, 19737, 0] + assert fast == build_native_block([("t", "Date", expected_days)]) + + def test_nullable_datetime_with_none(self): + vals = [dt.datetime(2024, 1, 15, 12, 0, 0, tzinfo=dt.timezone.utc), None, 5] + fast = self._encode("Nullable(DateTime)", vals) + assert fast == self._encode("Nullable(DateTime)", _NdarrayLikeColumn(vals), 3) + + def test_out_of_range_errors_unchanged(self): + with pytest.raises(ValueError, match="outside UInt32 range"): + self._encode("DateTime", [2**32]) + with pytest.raises(ValueError, match="outside UInt16 range"): + self._encode("Date", [65536]) + with pytest.raises(ValueError, match="row 0 is None but DateTime is not Nullable"): + self._encode("DateTime", [None]) + + def test_negative_ints(self): + # Negative is out of range for DateTime (falls back to the specific + # range error) but a valid pre-epoch value for Date32. + with pytest.raises(ValueError, match="outside UInt32 range"): + self._encode("DateTime", [-1]) + fast = self._encode("Date32", [-100, 0, 100]) + assert fast == self._encode("Date32", _NdarrayLikeColumn([-100, 0, 100]), 3) + decoded = list(_ch_core.ColBatch.decode_native(fast).column_data(0)) + assert decoded == [dt.date(1969, 9, 23), dt.date(1970, 1, 1), dt.date(1970, 4, 11)] + + +class TestTimeInsert: + def _encode(self, type_name, vals, n=None): + n = len(vals) if n is None else n + return _ch_core.encode_native_block(["t"], [type_name], [vals], n) + + def _decode(self, encoded): + return list(_ch_core.ColBatch.decode_native(encoded).column_data(0)) + + def test_time_accepted_values_and_fast_raw_ints(self): + values = [ + 13, + dt.timedelta(seconds=-1, microseconds=-500_000), + dt.time(1, 2, 3, 999_999), + "002:03:04.999", + 79.9, + ] + ticks = [13, -1, 3_723, 7_384, 79] + encoded = self._encode("Time", values) + assert encoded == build_native_block([("t", "Time", ticks)]) + assert encoded == self._encode("Time", tuple(values), len(values)) + assert encoded == self._encode( + "Time", _NdarrayLikeColumn(values), len(values) + ) + assert self._decode(encoded) == [dt.timedelta(seconds=v) for v in ticks] + + def test_established_inputs_skip_numpy_scalar_probe(self): + probes = {"count": 0} + + def dtype_probe(): + probes["count"] += 1 + return "not-a-numpy-dtype" + + class StringValue(str): + @property + def dtype(self): + return dtype_probe() + + class FloatValue(float): + @property + def dtype(self): + return dtype_probe() + + class IntValue(int): + @property + def dtype(self): + return dtype_probe() + + values = [ + dt.timedelta(seconds=1, microseconds=234_567), + dt.time(0, 0, 1, 234_567), + StringValue("000:00:01.234567"), + FloatValue(79), + IntValue(-13), + ] + expected_ticks = [1_234_567, 1_234_567, 1_234_567, 79, -13] + direct = self._encode("Nullable(Time64(6))", [*values, None]) + generic = self._encode( + "Nullable(Time64(6))", _NdarrayLikeColumn([*values, None]), 6 + ) + assert direct == generic == build_native_block( + [ + ( + "t", + "Nullable(Time64(6))", + [*expected_ticks, None], + ) + ] + ) + self._encode("Array(Time64(6))", [values]) + self._encode( + "LowCardinality(Time)", + [values[0], values[1], values[2], values[3], values[4]], + ) + assert probes["count"] == 0 + + def test_delta_and_time_subclasses_probe_then_fall_back(self): + # Subclasses go through the numpy scalar probe (pd.Timedelta needs it); + # a non-dtype attribute falls back to the struct-field conversion. + class DeltaValue(dt.timedelta): + @property + def dtype(self): + return "not-a-numpy-dtype" + + class TimeValue(dt.time): + @property + def dtype(self): + return "not-a-numpy-dtype" + + values = [ + DeltaValue(seconds=1, microseconds=234_567), + TimeValue(0, 0, 1, 234_567), + ] + encoded = self._encode("Nullable(Time64(6))", values) + assert encoded == build_native_block( + [("t", "Nullable(Time64(6))", [1_234_567, 1_234_567])] + ) + + @pytest.mark.parametrize( + "precision,values,ticks", + [ + ( + 0, + [ + dt.timedelta(seconds=13, microseconds=999_999), + dt.time(1, 2, 3, 999_999), + "-002:03:04.9", + ], + [13, 3_723, -7_384], + ), + ( + 3, + [ + dt.timedelta(seconds=1, microseconds=234_567), + dt.time(1, 2, 3, 456_789), + "-002:03:04.56789", + ], + [1_234, 3_723_456, -7_384_567], + ), + ( + 6, + [ + dt.timedelta(microseconds=79), + dt.time(0, 0, 1, 234_567), + "000:00:01.2", + ], + [79, 1_234_567, 1_200_000], + ), + ( + 9, + [ + dt.timedelta(microseconds=79), + dt.time(0, 0, 1, 234_567), + "000:00:01.000000079", + ], + [79_000, 1_234_567_000, 1_000_000_079], + ), + ], + ) + def test_time64_precisions(self, precision, values, ticks): + type_name = f"Time64({precision})" + encoded = self._encode(type_name, values) + assert encoded == build_native_block([("t", type_name, ticks)]) + assert self._decode(encoded) == [ + dt.timedelta( + microseconds=(abs(v) * 1_000_000 // (10**precision)) + * (-1 if v < 0 else 1) + ) + for v in ticks + ] + + @pytest.mark.parametrize( + "type_name,value", + [ + ("Time64(0)", dt.timedelta(milliseconds=-999)), + ("Time64(3)", dt.timedelta(microseconds=-999)), + ], + ) + def test_negative_timedelta_sub_tick_truncates_toward_zero( + self, type_name, value + ): + encoded = self._encode(type_name, [value]) + assert encoded == build_native_block([("t", type_name, [0])]) + + @pytest.mark.parametrize( + "type_name,value", + [ + ("Time64(0)", ("timedelta64[ms]", -999)), + ("Time64(3)", ("timedelta64[us]", -999)), + ("Time64(6)", ("timedelta64[ns]", -999)), + ("Time64(9)", ("timedelta64[ps]", -999)), + ], + ) + def test_numpy_scalar_negative_sub_tick_truncates_toward_zero( + self, type_name, value + ): + np = pytest.importorskip("numpy") + dtype, raw = value + scalar = np.array(raw, dtype=dtype)[()] + encoded = self._encode(type_name, [scalar]) + assert encoded == build_native_block([("t", type_name, [0])]) + + @pytest.mark.parametrize( + "type_name,expected_ticks", + [ + ("Time", [1, -1, 0]), + ("Time64(3)", [1_234, -1_234, 0]), + ("Time64(6)", [1_234_567, -1_234_567, 0]), + ("Time64(9)", [1_234_567_890, -1_234_567_890, 0]), + ], + ) + def test_numpy_timedelta64_ndarray_bulk_path(self, type_name, expected_ticks): + np = pytest.importorskip("numpy") + values = np.array( + [1_234_567_890, -1_234_567_890, 0], dtype="timedelta64[ns]" + ) + encoded = self._encode(type_name, values) + assert encoded == build_native_block([("t", type_name, expected_ticks)]) + + def test_numpy_timedelta64_two_dimensional_array_rejected(self): + np = pytest.importorskip("numpy") + values = np.array([13, 79], dtype="timedelta64[ns]").reshape(2, 1) + with pytest.raises(ValueError, match="must be one-dimensional"): + self._encode("Time64(9)", values) + + def test_numpy_timedelta64_byte_order_multiplier_and_general_ratios(self): + np = pytest.importorskip("numpy") + + big_endian = np.array([1_234, -1_234], dtype=">m8[ms]") + assert self._encode("Time64(3)", big_endian) == build_native_block( + [("t", "Time64(3)", [1_234, -1_234])] + ) + + multiplied = np.array([13, -13], dtype="timedelta64[10us]") + assert self._encode("Time64(9)", multiplied) == build_native_block( + [("t", "Time64(9)", [130_000, -130_000])] + ) + + general = np.array([334, -334], dtype="timedelta64[3ps]") + assert self._encode("Time64(9)", general) == build_native_block( + [("t", "Time64(9)", [1, -1])] + ) + + def test_numpy_timedelta64_nullable_nat_array_and_scalars(self): + np = pytest.importorskip("numpy") + values = np.array([1_234_567, "NaT", -1_234_567], dtype="timedelta64[us]") + encoded = self._encode("Nullable(Time64(6))", values) + assert encoded == build_native_block( + [("t", "Nullable(Time64(6))", [1_234_567, None, -1_234_567])] + ) + assert self._decode(encoded) == [ + dt.timedelta(microseconds=1_234_567), + None, + dt.timedelta(microseconds=-1_234_567), + ] + + scalar_values = [np.timedelta64(13, "ms"), np.timedelta64("NaT")] + scalar_encoded = self._encode("Nullable(Time64(3))", scalar_values) + assert scalar_encoded == build_native_block( + [("t", "Nullable(Time64(3))", [13, None])] + ) + + def test_pandas_timedelta_series_bulk_path(self): + pd = pytest.importorskip("pandas") + values = pd.Series( + pd.to_timedelta(["1.234567890s", None, "-1.234567890s"]) + ) + encoded = self._encode("Nullable(Time64(9))", values) + assert encoded == build_native_block( + [ + ( + "t", + "Nullable(Time64(9))", + [1_234_567_890, None, -1_234_567_890], + ) + ] + ) + + def test_pandas_timedelta_scalar_ns_precision_in_list(self): + pd = pytest.importorskip("pandas") + values = [pd.Timedelta("1s 123456789ns"), pd.Timedelta("-1s")] + encoded = self._encode("Time64(9)", values) + assert encoded == build_native_block( + [("t", "Time64(9)", [1_123_456_789, -1_000_000_000])] + ) + sub_tick = self._encode("Time64(0)", [pd.Timedelta("-999ms")]) + assert sub_tick == build_native_block([("t", "Time64(0)", [0])]) + + def test_pandas_nat_nullable_and_non_nullable(self): + pd = pytest.importorskip("pandas") + encoded = self._encode("Nullable(Time64(9))", [pd.Timedelta("1us"), pd.NaT]) + assert encoded == build_native_block( + [("t", "Nullable(Time64(9))", [1_000, None])] + ) + with pytest.raises(ValueError, match="row 0 is NaT but Time is not Nullable"): + self._encode("Time", [pd.NaT]) + with pytest.raises(ValueError, match="row 1 is NaT"): + self._encode("Time64(3)", _NdarrayLikeColumn([13, pd.NaT]), 2) + + def test_numpy_nat_non_nullable_and_range_errors(self): + np = pytest.importorskip("numpy") + with pytest.raises(ValueError, match="row 0 is NaT.*not Nullable"): + self._encode("Time", np.array(["NaT"], dtype="timedelta64[ns]")) + with pytest.raises(ValueError, match="outside logical range"): + self._encode("Time", np.array([1_000], dtype="timedelta64[h]")) + + def test_late_day_time64_nested_value_avoids_i64_overflow(self): + value = dt.time(23, 59, 59, 999_999) + type_name = "Array(Tuple(Time64(9)))" + encoded = self._encode(type_name, [[(value,)]]) + expected_ticks = 86_399_999_999_000 + assert encoded == build_native_block( + [("t", type_name, [[(expected_ticks,)]])] + ) + assert self._decode(encoded) == [ + [(dt.timedelta(seconds=86_399, microseconds=999_999),)] + ] + + def test_nullable_and_recursive_shapes(self): + nullable = [ + dt.timedelta(seconds=1, microseconds=250_000), + None, + "-000:00:00.001", + ] + encoded = self._encode("Nullable(Time64(3))", nullable) + assert self._decode(encoded) == [ + dt.timedelta(milliseconds=1_250), + None, + dt.timedelta(milliseconds=-1), + ] + + array_rows = [[dt.timedelta(milliseconds=13), "000:00:00.079"], [], [-1]] + encoded = self._encode("Array(Time64(3))", array_rows) + assert self._decode(encoded) == [ + [dt.timedelta(milliseconds=13), dt.timedelta(milliseconds=79)], + [], + [dt.timedelta(milliseconds=-1)], + ] + + tuple_rows = [ + (dt.time(0, 0, 13), dt.timedelta(microseconds=79)), + (-13, None), + ] + encoded = self._encode("Tuple(Time, Nullable(Time64(6)))", tuple_rows) + assert self._decode(encoded) == [ + (dt.timedelta(seconds=13), dt.timedelta(microseconds=79)), + (dt.timedelta(seconds=-13), None), + ] + + array_tuple_rows = [ + [("000:00:13", 79_000)], + [], + [(dt.time(0, 1, 19), -1_000)], + ] + encoded = self._encode( + "Array(Tuple(Time, Time64(6)))", array_tuple_rows + ) + assert self._decode(encoded) == [ + [(dt.timedelta(seconds=13), dt.timedelta(microseconds=79_000))], + [], + [(dt.timedelta(seconds=79), dt.timedelta(microseconds=-1_000))], + ] + + def test_low_cardinality_time_and_time64_rejection(self): + encoded = self._encode("LowCardinality(Time)", [13, 79, 13, -1]) + assert self._decode(encoded) == [ + dt.timedelta(seconds=13), + dt.timedelta(seconds=79), + dt.timedelta(seconds=13), + dt.timedelta(seconds=-1), + ] + with pytest.raises( + NotImplementedError, match="unsupported LowCardinality inner type" + ): + self._encode("LowCardinality(Time64(3))", [13]) + + @pytest.mark.parametrize( + "type_name,value", + [ + ("Time", 3_600_000), + ("Time", -3_600_000), + ("Time64(3)", 3_600_000_000), + ("Time64(3)", -3_600_000_000), + ("Time", "1000:00:00"), + ("Time64(6)", "001:60:00"), + ("Time64(6)", float("inf")), + ], + ) + def test_invalid_values(self, type_name, value): + with pytest.raises(ValueError, match="column.*row 0.*Time"): + self._encode(type_name, [value]) + + def test_none_requires_nullable(self): + with pytest.raises(ValueError, match="row 0 is None but Time is not Nullable"): + self._encode("Time", [None]) + + +# --------------------------------------------------------------------------- +# Temporal types +# --------------------------------------------------------------------------- + + +class TestDecodeDate: + def test_date(self): + days = [0, 19737, 100, 65535] + data = build_native_block([("d", "Date", days)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["Date"] + expected = [_EPOCH_DATE + dt.timedelta(days=x) for x in days] + assert list(batch.column_data(0)) == expected + + def test_date32_pre_epoch(self): + days = [-25567, -1, 0, 19737] # -25567 ~= 1900-01-01, signed days + data = build_native_block([("d", "Date32", days)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["Date32"] + expected = [_EPOCH_DATE + dt.timedelta(days=x) for x in days] + assert list(batch.column_data(0)) == expected + + +class TestDecodeDateTime: + def test_datetime_naive(self): + secs = [0, 1705322096, 961056000] + data = build_native_block([("dt", "DateTime", secs)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["DateTime"] + expected = [_EPOCH_NAIVE + dt.timedelta(seconds=s) for s in secs] + result = list(batch.column_data(0)) + assert result == expected + assert all(v.tzinfo is None for v in result) + + def test_datetime_utc_is_naive(self): + # A UTC-equivalent timezone renders naive, matching clickhouse-connect. + secs = [1705322096] + data = build_native_block([("dt", "DateTime('UTC')", secs)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["DateTime('UTC')"] + v = list(batch.column_data(0))[0] + assert v == _EPOCH_NAIVE + dt.timedelta(seconds=secs[0]) + assert v.tzinfo is None + + def test_datetime_named_zone_is_aware(self): + secs = [1705322096] # 2024-01-15 12:34:56 UTC + data = build_native_block([("dt", "DateTime('America/New_York')", secs)]) + batch = _ch_core.ColBatch.decode_native(data) + v = list(batch.column_data(0))[0] + assert v.tzinfo == ZoneInfo("America/New_York") + # Same instant as the source UTC seconds, expressed in New York time. + assert v == dt.datetime(2024, 1, 15, 12, 34, 56, tzinfo=dt.timezone.utc) + assert (v.hour, v.minute, v.second) == (7, 34, 56) + + +class TestDateTimeTzSubSecond: + _NY = ZoneInfo("America/New_York") + + def _decode_one(self, type_name, tick): + data = build_native_block([("ts", type_name, [tick])]) + return list(_ch_core.ColBatch.decode_native(data).column_data(0))[0] + + def test_dst_fall_back_fold(self): + # America/New_York 2020-11-01: 1:00:00.5 wall time exists twice. The + # epoch values pin each side of the fold; sub-second micros must not + # disturb the UTC offset. + for epoch_secs, offset_hours in ((1604206800, -4), (1604210400, -5)): + v = self._decode_one( + "DateTime64(6, 'America/New_York')", epoch_secs * 1_000_000 + 500_000 + ) + assert v.microsecond == 500_000 + assert v.utcoffset() == dt.timedelta(hours=offset_hours) + assert (v.hour, v.minute, v.second) == (1, 0, 0) + + def test_microsecond_exactness(self): + # Sweep awkward microsecond values at a recent epoch and assert exact + # round-trips through the tz-aware path. + for secs in (1705322096, -1, -86_400): + for micros in (1, 3, 333_333, 499_999, 500_000, 500_001, 999_999): + v = self._decode_one( + "DateTime64(6, 'America/New_York')", secs * 1_000_000 + micros + ) + expected = dt.datetime.fromtimestamp(secs, self._NY).replace( + microsecond=micros + ) + assert v == expected + assert v.microsecond == micros + + def test_far_future_exactness(self): + # Year 2200 is beyond f64 sub-microsecond precision for epoch seconds; + # the exact path must still produce the precise microsecond. + secs = 7_258_204_800 + micros = 123_457 + v = self._decode_one("DateTime64(6, 'America/New_York')", secs * 1_000_000 + micros) + expected = dt.datetime.fromtimestamp(secs, self._NY).replace(microsecond=micros) + assert v == expected + assert v.microsecond == micros + + +class TestDecodeDateTime64: + def test_dt64_millis(self): + ticks = [0, 1705322096789] # milliseconds + data = build_native_block([("ts", "DateTime64(3)", ticks)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["DateTime64(3)"] + expected = [_EPOCH_NAIVE + dt.timedelta(milliseconds=t) for t in ticks] + assert list(batch.column_data(0)) == expected + + def test_dt64_nanos_truncate_to_micros(self): + # Precision 9 (ns); Python datetime resolves to microseconds, so the + # sub-microsecond digits are truncated. + ticks = [1705322096789012345] + data = build_native_block([("ts", "DateTime64(9)", ticks)]) + batch = _ch_core.ColBatch.decode_native(data) + v = list(batch.column_data(0))[0] + assert v == dt.datetime(2024, 1, 15, 12, 34, 56, 789012) + + def test_dt64_nullable(self): + ticks = [1705322096789, None, 0] + data = build_native_block([("ts", "Nullable(DateTime64(3))", ticks)]) + batch = _ch_core.ColBatch.decode_native(data) + assert batch.column_type_names == ["Nullable(DateTime64(3))"] + result = list(batch.column_data(0)) + assert result[1] is None + assert result[0] == _EPOCH_NAIVE + dt.timedelta(milliseconds=ticks[0]) + + def test_temporal_rows_and_columns(self): + days = [0, 19737] + ticks = [0, 1705322096789] + data = build_native_block([("d", "Date", days), ("ts", "DateTime64(3)", ticks)]) + batch = _ch_core.ColBatch.decode_native(data) + rows = list(batch.to_python_rows()) + assert rows[1] == ( + _EPOCH_DATE + dt.timedelta(days=days[1]), + _EPOCH_NAIVE + dt.timedelta(milliseconds=ticks[1]), + ) + cols = list(batch.to_python_columns()) + assert list(cols[0])[0] == _EPOCH_DATE + + +class TestDecodeTime: + @pytest.mark.parametrize( + "type_name,ticks,expected", + [ + ( + "Time", + [-3_599_999, -13, 0, 3_599_999], + [ + dt.timedelta(seconds=v) + for v in [-3_599_999, -13, 0, 3_599_999] + ], + ), + ( + "Time64(0)", + [-13, 0, 79], + [dt.timedelta(seconds=v) for v in [-13, 0, 79]], + ), + ( + "Time64(3)", + [-1_500, -1, 0, 79_999], + [ + dt.timedelta(milliseconds=-1_500), + dt.timedelta(milliseconds=-1), + dt.timedelta(0), + dt.timedelta(milliseconds=79_999), + ], + ), + ( + "Time64(6)", + [-1_500_001, -1, 0, 79_999_999], + [ + dt.timedelta(microseconds=-1_500_001), + dt.timedelta(microseconds=-1), + dt.timedelta(0), + dt.timedelta(microseconds=79_999_999), + ], + ), + ( + "Time64(9)", + [-1_999, -1_001, -999, 1_999], + [ + dt.timedelta(microseconds=-1), + dt.timedelta(microseconds=-1), + dt.timedelta(0), + dt.timedelta(microseconds=1), + ], + ), + ], + ) + def test_plain_precisions_and_truncation(self, type_name, ticks, expected): + batch = _ch_core.ColBatch.decode_native( + build_native_block([("t", type_name, ticks)]) + ) + assert batch.column_type_names == [type_name] + assert list(batch.column_data(0)) == expected + + def test_nullable_and_all_object_exits(self): + ticks = [-1_001, None, 1_999] + expected = [ + dt.timedelta(microseconds=-1), + None, + dt.timedelta(microseconds=1), + ] + batch = _ch_core.ColBatch.decode_native( + build_native_block([("t", "Nullable(Time64(9))", ticks)]) + ) + assert batch.column_type_names == ["Nullable(Time64(9))"] + assert list(batch.column_data(0)) == expected + assert list(batch.to_python_columns()[0]) == expected + assert [row[0] for row in batch.to_python_rows()] == expected + + def test_raw_time_ticks_scalar_nullable_and_low_cardinality(self): + data = build_native_block( + [ + ("t", "Time", [-13, 0, 79]), + ("t64", "Nullable(Time64(9))", [-1_001, None, 1_999]), + ("lc", "LowCardinality(Time)", [13, 79, 13]), + ] + ) + batch = _ch_core.ColBatch.decode_native(data) + assert list(batch.column_data(0, raw_time_ticks=True)) == [-13, 0, 79] + assert list(batch.column_data(1, raw_time_ticks=True)) == [ + -1_001, + None, + 1_999, + ] + assert list(batch.column_data(2, raw_time_ticks=True)) == [13, 79, 13] + assert list(batch.column_data(0)) == [ + dt.timedelta(seconds=-13), + dt.timedelta(0), + dt.timedelta(seconds=79), + ] + + def test_raw_time_ticks_recursive_array_tuple_and_low_cardinality(self): + type_name = ( + "Array(Tuple(Time, Nullable(Time64(9)), " + "Array(LowCardinality(Time))))" + ) + rows = [ + [(13, -1_001, [13, 13, 79]), (-79, None, [])], + [], + [(0, 1_999, [-1])], + ] + batch = _ch_core.ColBatch.decode_native( + build_native_block([("v", type_name, rows)]) + ) + assert list(batch.column_data(0, True)) == rows + assert list(batch.column_data(0)) == [ + [ + ( + dt.timedelta(seconds=13), + dt.timedelta(microseconds=-1), + [dt.timedelta(seconds=13)] * 2 + [dt.timedelta(seconds=79)], + ), + (dt.timedelta(seconds=-79), None, []), + ], + [], + [ + ( + dt.timedelta(0), + dt.timedelta(microseconds=1), + [dt.timedelta(seconds=-1)], + ) + ], + ] + + def test_raw_time_ticks_map_duplicate_key_last_value_wins(self): + type_name = "Map(Time, String)" + body = bytearray(struct.pack(">= 7 + if value: + buf.append(byte | 0x80) + else: + buf.append(byte) + return + + +def _empty_batch_from_meta(client, query, settings): + """Schema-bearing empty batch for a columns-only query. + + Mirrors the v1 LIMIT 0 branch: the server sends zero Native bytes for + these, so v1 fetches FORMAT JSON metadata instead (httpclient.py:263). + The server's own name and type strings are framed as a zero-row Native + header and decoded by the core, so the type parsing and the unsupported + type boundary stay in Rust. + """ + meta = json.loads(client.raw_query(query, settings=settings, fmt="JSON"))["meta"] + has_info = bool(client.protocol_version) + buf = bytearray() + if has_info: + buf += bytes((1, 0, 2, 0xFF, 0xFF, 0xFF, 0xFF, 0)) + _write_varint(buf, len(meta)) + _write_varint(buf, 0) + for col in meta: + for text in (col["name"], col["type"]): + encoded = text.encode() + _write_varint(buf, len(encoded)) + buf += encoded + return _ch_core.ColBatch.decode_native(bytes(buf), has_block_info=has_info) + + +def _put(q, item, stop): + """Bounded put that gives up once the consumer has signaled stop.""" + while not stop.is_set(): + try: + q.put(item, timeout=PUT_TIMEOUT) + return True + except queue.Full: + pass + return False + + +def query_rust(client, query, settings=None): + """Run a SELECT down the Rust decode path. Returns RustQueryResult. + + Streams the response on a producer thread into a bounded queue while the + caller thread decodes. On any error the response connection is closed, + not returned to the pool, and the original exception propagates. + """ + if columns_only_re.search(remove_sql_comments(query)): + return RustQueryResult(_empty_batch_from_meta(client, query, settings)) + response = _mirrored_request(client, query, settings) + q = queue.Queue(maxsize=QUEUE_MAX) + stop = threading.Event() + + def produce(): + try: + for chunk in _decompressed_chunks(response): + if not _put(q, ("data", chunk), stop): + return + except BaseException as exc: # noqa: BLE001 - propagated to consumer + _put(q, ("error", exc), stop) + finally: + _put(q, ("eof", None), stop) + + producer = threading.Thread(target=produce, daemon=True) + producer.start() + + decoder = _ch_core.StreamDecoder(has_block_info=bool(client.protocol_version)) + batches = [] + try: + while True: + tag, payload = q.get() + if tag == "data": + batches.extend(decoder.feed(payload)) + elif tag == "error": + raise payload + else: + batches.extend(decoder.finish()) + break + except BaseException: + stop.set() + response.close() + producer.join(timeout=JOIN_TIMEOUT) + raise + producer.join(timeout=JOIN_TIMEOUT) + response.release_conn() + + if not batches: + # The server sent zero bytes (empty result without a LIMIT 0 suffix). + # v1 query() yields a schema-less empty result for these too. + return RustQueryResult(None) + batch = _ch_core.ColBatch.from_batches(batches) + del batches + return RustQueryResult(batch) + + +class RustQueryResult: + """Materialized query result over a merged ColBatch. + + batch is None for a schema-less empty result, matching v1 behavior when + the server sends no Native block at all. + """ + + def __init__(self, batch): + self.batch = batch + + @property + def column_names(self): + if self.batch is None: + return () + return tuple(self.batch.column_names) + + @property + def result_rows(self): + if self.batch is None: + return [] + return self.batch.to_python_rows() + + @property + def result_columns(self): + if self.batch is None: + return [] + return self.batch.to_python_columns() + + def arrow_table(self): + import pyarrow as pa + + if self.batch is None: + return pa.table({}) + return pa.RecordBatchReader.from_stream(self.batch).read_all() + + def to_pandas(self): + return self.arrow_table().to_pandas() + + def to_polars(self): + import polars as pl + + return pl.from_arrow(self.arrow_table()) diff --git a/rust/streaming_demo.py b/rust/streaming_demo.py new file mode 100644 index 00000000..b4da4559 --- /dev/null +++ b/rust/streaming_demo.py @@ -0,0 +1,249 @@ +"""Streaming transport prototype for the rust core — sync AND async. + +Demonstrates pull/parse overlap, following the same architecture as v1's +aiohttp async client: + + * SYNC : a producer thread reads HTTP chunks into a bounded queue.Queue; + the main thread pulls chunks and feeds _ch_core.StreamDecoder. + * ASYNC: an aiohttp task reads chunks on the event loop into v1's + AsyncSyncQueue (the bidirectional bridge); a thread executor + pulls the sync side and feeds the same StreamDecoder. + +The decode engine is identical in both cases — _ch_core.StreamDecoder.feed(), +which releases the GIL during decode so pull and parse genuinely overlap. + +Run against the local ClickHouse on localhost:8123. +""" +import asyncio +import os +import queue +import sys +import threading +import time +from urllib.parse import quote + +sys.path.insert(0, os.environ.get("CHC_BASELINE_PATH", os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))) + +import _ch_core # noqa: E402 +import aiohttp # noqa: E402 +import urllib3 # noqa: E402 + +# v1's proven async<->sync bridge queue — the same one the aiohttp client uses. +from clickhouse_connect.driver.asyncqueue import EOF_SENTINEL, AsyncSyncQueue # noqa: E402 + +HOST, PORT = "localhost", 8123 +CHUNK = 1 << 20 # 1 MiB network reads +QUEUE_MAX = 16 # bounded -> backpressure + + +def _url(query: str) -> str: + return f"http://{HOST}:{PORT}/?query={quote(query + ' FORMAT Native')}" + + +_http = urllib3.PoolManager() + + +# --------------------------------------------------------------------------- +# Sequential baseline (fetch fully, then decode) — what bench.py measured. +# --------------------------------------------------------------------------- +def fetch_full(query: str) -> bytes: + resp = _http.request("GET", _url(query), preload_content=True) + return resp.data + + +# --------------------------------------------------------------------------- +# SYNC streaming: producer thread -> queue.Queue -> main thread feeds decoder +# --------------------------------------------------------------------------- +def stream_sync(query: str): + resp = _http.request("GET", _url(query), preload_content=False) + q: queue.Queue = queue.Queue(maxsize=QUEUE_MAX) + + def producer(): + try: + for chunk in resp.stream(CHUNK): + q.put(chunk) + finally: + q.put(None) # EOF sentinel + resp.release_conn() + + t = threading.Thread(target=producer, daemon=True) + t.start() + + dec = _ch_core.StreamDecoder() + while True: + chunk = q.get() + if chunk is None: + break + for block in dec.feed(chunk): + yield block + for block in dec.finish(): + yield block + t.join() + + +# --------------------------------------------------------------------------- +# ASYNC streaming: aiohttp loop producer -> AsyncSyncQueue -> executor decoder +# --------------------------------------------------------------------------- +async def stream_async(query: str): + bridge: AsyncSyncQueue = AsyncSyncQueue(maxsize=QUEUE_MAX) + loop = asyncio.get_running_loop() + + # Request identity encoding: aiohttp defaults to Accept-Encoding: gzip, + # and ClickHouse will honor it — then aiohttp burns CPU decompressing in + # Python. v1's async client disables this for the same reason. + async with aiohttp.ClientSession(headers={"Accept-Encoding": "identity"}) as session: + async with session.get(_url(query)) as resp: + + async def producer(): + try: + async for chunk in resp.content.iter_chunked(CHUNK): + await bridge.async_q.put(chunk) + finally: + bridge.shutdown() # wakes the sync consumer with EOF + + def consume(): + # Runs in a thread executor: blocking pulls + GIL-releasing decode, + # overlapping with the aiohttp producer on the event loop. + dec = _ch_core.StreamDecoder() + blocks = [] + while True: + chunk = bridge.sync_q.get() + if chunk is EOF_SENTINEL: + break + blocks.extend(dec.feed(chunk)) + blocks.extend(dec.finish()) + return blocks + + prod = asyncio.create_task(producer()) + blocks = await loop.run_in_executor(None, consume) + await prod + return blocks + + +# --------------------------------------------------------------------------- +# ASYNC streaming, executor-driven (#2): no aiohttp at all. +# +# Drive the proven SYNC pipeline (urllib3 producer thread + GIL-releasing +# decode) inside a thread executor and await it. Overlap happens between the +# two worker threads (both release the GIL); the event loop stays free the +# whole time, so it remains genuinely non-blocking AND overlaps fetch+decode. +# --------------------------------------------------------------------------- +async def stream_async_executor(query: str): + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, lambda: list(stream_sync(query))) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def rows_of(blocks) -> list: + out = [] + for b in blocks: + out.extend(tuple(r) for r in b.to_python_rows()) + return out + + +def ms(s): + return f"{s * 1000:7.1f} ms" + + +# --------------------------------------------------------------------------- +# Correctness: streamed (sync + async) must equal a single full decode. +# --------------------------------------------------------------------------- +def correctness(): + print("== correctness (deterministic ORDER BY, exact equality) ==") + q = "SELECT id, val, name, flag, small_int, big_uint FROM bench_types ORDER BY id LIMIT 100000" + raw = fetch_full(q) + baseline = rows_of([_ch_core.ColBatch.decode_native(raw, has_block_info=False)]) + + sync_rows = rows_of(list(stream_sync(q))) + async_rows = rows_of(asyncio.run(stream_async(q))) + async_exec_rows = rows_of(asyncio.run(stream_async_executor(q))) + + print(f" baseline rows : {len(baseline)}") + print(f" sync == baseline: {sync_rows == baseline}") + print(f" async(aiohttp) == baseline: {async_rows == baseline}") + print(f" async(executor)== baseline: {async_exec_rows == baseline}") + + +# --------------------------------------------------------------------------- +# Overlap: streaming total should approach max(fetch, decode), not the sum. +# --------------------------------------------------------------------------- +def overlap(): + print("\n== overlap (streaming total vs sequential fetch+decode) ==") + workloads = { + "mixed_6col_10M": "SELECT id, val, name, flag, small_int, big_uint FROM bench_types", + "string_1col_10M": "SELECT name FROM bench_types", + "int_3col_10M": "SELECT id, small_int, big_uint FROM bench_types", + } + for name, q in workloads.items(): + # sequential baseline + t0 = time.perf_counter() + raw = fetch_full(q) + t_fetch = time.perf_counter() - t0 + t0 = time.perf_counter() + base = _ch_core.ColBatch.decode_native(raw, has_block_info=False) + t_dec = time.perf_counter() - t0 + n = base.num_rows + seq = t_fetch + t_dec + + # sync streaming (count rows to force full consumption) + t0 = time.perf_counter() + rows_sync = sum(b.num_rows for b in stream_sync(q)) + t_sync = time.perf_counter() - t0 + + # async streaming, aiohttp + bridge (GIL-bound loop — doesn't overlap) + t0 = time.perf_counter() + rows_async = sum(b.num_rows for b in asyncio.run(stream_async(q))) + t_async = time.perf_counter() - t0 + + # async streaming, executor-driven sync pipeline (#2 — should match sync) + t0 = time.perf_counter() + rows_aexec = sum(b.num_rows for b in asyncio.run(stream_async_executor(q))) + t_aexec = time.perf_counter() - t0 + + assert rows_sync == n and rows_async == n and rows_aexec == n, "row count mismatch" + print(f"\n {name} ({n:,} rows)") + print(f" sequential : fetch {ms(t_fetch)} + decode {ms(t_dec)} = {ms(seq)}") + print(f" sync stream : {ms(t_sync)} ({(seq / t_sync):.2f}x vs sequential)") + print(f" async (aiohttp) : {ms(t_async)} ({(seq / t_async):.2f}x vs sequential)") + print(f" async (executor) : {ms(t_aexec)} ({(seq / t_aexec):.2f}x vs sequential)") + + +# --------------------------------------------------------------------------- +# Non-blocking proof: the event loop must keep running other tasks while the +# executor-driven query overlaps fetch+decode on worker threads. +# --------------------------------------------------------------------------- +def non_blocking(): + print("\n== non-blocking event loop (executor-driven async) ==") + q = "SELECT id, val, name, flag, small_int, big_uint FROM bench_types" + + async def run(): + ticks = 0 + + async def heartbeat(): + nonlocal ticks + while True: + await asyncio.sleep(0.005) # 5ms heartbeat + ticks += 1 + + hb = asyncio.create_task(heartbeat()) + t0 = time.perf_counter() + blocks = await stream_async_executor(q) + dt = time.perf_counter() - t0 + hb.cancel() + rows = sum(b.num_rows for b in blocks) + return dt, rows, ticks + + dt, rows, ticks = asyncio.run(run()) + expected = int(dt / 0.005) + print(f" query: {rows:,} rows in {ms(dt)}") + print(f" heartbeat ticked {ticks} times during the query " + f"(~{expected} expected if loop never blocked)") + print(f" loop stayed responsive: {ticks > 0.5 * expected}") + + +if __name__ == "__main__": + correctness() + overlap() + non_blocking() diff --git a/rust/temporal_e2e_check.py b/rust/temporal_e2e_check.py new file mode 100644 index 00000000..1a13835d --- /dev/null +++ b/rust/temporal_e2e_check.py @@ -0,0 +1,106 @@ +"""End-to-end temporal check: live ClickHouse -> Native -> _ch_core -> verify. + +Validates the binding's temporal value policy (Date/Date32/DateTime/DateTime64, +naive vs tz-aware, DateTime64 precision) against: + 1. clickhouse-connect v1's own decode (value-level ground truth) + 2. pyarrow (via the zero-copy __arrow_c_stream__ export) + +The v1 baseline is pinned to the compiled source tree and we hard-assert the C +extension loaded, so we never silently compare against pure-Python v1. +""" +import os +import sys + +sys.path.insert(0, os.environ.get("CHC_BASELINE_PATH", os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))) + +import _ch_core # noqa: E402 +import pyarrow as pa # noqa: E402 + +import clickhouse_connect # noqa: E402 +from clickhouse_connect.driverc import dataconv # noqa: E402,F401 (hard-load C ext) + +client = clickhouse_connect.get_client(host="localhost", port=8123) + +# Raw HTTP `FORMAT Native` with no protocol version (revision 0) makes the server +# drop the timezone from a bare DateTime type string, so it would decode naive. +# Requesting this protocol version (clickhouse-connect's PROTOCOL_VERSION_WITH_LOW_CARD) +# preserves the timezone and adds the block-info preamble, while staying below the +# custom-serialization revision (54454) that the bool `has_block_info` flag cannot +# express. Decode the resulting stream with has_block_info=True. +PROTOCOL_VERSION = 54405 + +# A spread of every temporal shape the core decodes: naive Date/Date32 (incl. +# pre-epoch), bare/UTC/named-zone DateTime, DateTime64 at precision 3/6/9 with +# and without a zone, and a nullable DateTime64. +QUERY = """ +SELECT + toDate('1970-01-01') + (number % 60000) AS d, + toDate32('1925-01-01') + (number % 130000) AS d32, + toDateTime('1970-01-01 00:00:00') + number AS dt, + toDateTime('2000-01-01 00:00:00', 'UTC') + number AS dt_utc, + toDateTime('2000-01-01 00:00:00', 'America/New_York') + number AS dt_ny, + addMilliseconds(toDateTime64('2000-01-01 00:00:00.000', 3), number) AS dt64_3, + addMicroseconds(toDateTime64('2000-01-01 00:00:00.000000', 6, 'UTC'), number) AS dt64_6_utc, + toDateTime64('2000-01-01 00:00:00.000000000', 9, 'America/New_York') + number AS dt64_9_ny, + if(number % 7 = 0, NULL, addMilliseconds(toDateTime64('2010-06-15 08:00:00.000', 3), number)) AS dt64_nil +FROM numbers({n}) +""" + +N = 50_000 + + +def main(): + q = QUERY.format(n=N) + + raw = client.raw_query(q, fmt="Native", settings={"client_protocol_version": PROTOCOL_VERSION}) + print(f"fetched {len(raw):,} bytes of Native data") + + batch = _ch_core.ColBatch.decode_native(raw, has_block_info=True) + print(f"_ch_core decoded: {batch.num_rows} rows x {batch.num_columns} cols") + print(f" types: {batch.column_type_names}") + assert batch.num_rows == N, f"row count {batch.num_rows} != {N}" + + v1 = client.query(q, column_oriented=True) + v1_cols = v1.result_columns + v1_names = v1.column_names + assert list(batch.column_names) == list(v1_names), ( + f"name mismatch: {batch.column_names} vs {v1_names}" + ) + + # Value-level comparison: rust to_python_columns vs v1, cell by cell. + chc_cols = batch.to_python_columns() + mismatches = 0 + for ci, name in enumerate(v1_names): + rust_col = list(chc_cols[ci]) + v1_col = list(v1_cols[ci]) + if rust_col == v1_col: + print(f"OK {name:12s} {len(rust_col):,} values match v1 (sample {rust_col[1]!r})") + continue + mismatches += 1 + for ri, (a, b) in enumerate(zip(rust_col, v1_col)): + if a != b: + print(f"FAIL {name:12s} row {ri}: rust={a!r} v1={b!r}") + break + + # to_python_rows must agree with the columnar view it is built from. + rows = batch.to_python_rows() + assert len(rows) == N + assert tuple(chc_cols[ci][13] for ci in range(batch.num_columns)) == tuple(rows[13]) + + # Arrow zero-copy export sanity: Date32 maps to a real Arrow date32 and + # round-trips to the same Python dates v1 produced. + table = pa.table(batch) + assert table.num_rows == N + assert table.schema.field("d32").type == pa.date32() + assert table.column("d32").to_pylist() == list(v1_cols[v1_names.index("d32")]) + print("OK arrow date32 export matches v1") + + if mismatches == 0: + print(f"\nTEMPORAL END-TO-END: PASS ({batch.num_columns} columns match v1 exactly)") + else: + print(f"\nTEMPORAL END-TO-END: FAIL ({mismatches} columns differ)") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 0a83088c..d1de647d 100644 --- a/setup.py +++ b/setup.py @@ -82,6 +82,7 @@ def run_setup(try_c: bool = True): "tzdata": ["tzdata"], "async": ["aiohttp>=3.9.0"], "chdb": ["chdb>=4.1.7"], + "rust": ["clickhouse-connect-core>=0.1.0,<0.2"], }, tests_require=["pytest"], entry_points={ diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index 44c4c62f..f66896fd 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -8,11 +8,12 @@ from subprocess import PIPE, Popen from typing import NamedTuple +import pytest import pytest_asyncio from pytest import fixture from clickhouse_connect import common, get_async_client -from clickhouse_connect.driver import Client, create_client +from clickhouse_connect.driver import AsyncClient, Client, create_client from clickhouse_connect.driver.common import coerce_bool from clickhouse_connect.driver.exceptions import OperationalError from clickhouse_connect.driver.httpclient import HttpClient @@ -31,6 +32,7 @@ class TestConfig(NamedTuple): compress: str insert_quorum: int proxy_address: str + native_codec: str __test__ = False @@ -38,6 +40,15 @@ class TestException(BaseException): pass +def type_available(client: Client | AsyncClient, data_type: str) -> None: + if client.get_client_setting(f"allow_experimental_{data_type}_type") is None: + return + setting_def = client.server_settings.get(f"allow_experimental_{data_type}_type", None) + if setting_def is not None and setting_def.value == "1": + return + pytest.skip(f"New {data_type.upper()} type not available in this version: {client.server_version}") + + def make_client_config(test_config: TestConfig, **kwargs): """Helper to build client config dict from test_config with optional overrides.""" settings = kwargs.pop("settings", {}).copy() @@ -75,7 +86,11 @@ def test_config_fixture() -> Iterator[TestConfig]: compress = os.environ.get("CLICKHOUSE_CONNECT_TEST_COMPRESS", "True") insert_quorum = int(os.environ.get("CLICKHOUSE_CONNECT_TEST_INSERT_QUORUM", "0")) proxy_address = os.environ.get("CLICKHOUSE_CONNECT_TEST_PROXY_ADDR", "") - yield TestConfig(host, port, username, password, docker, test_database, cloud, compress, insert_quorum, proxy_address) + # The native_codec common-setting default is env-seeded at import time, so fixture-created + # clients already pick up this env value without any per-fixture wiring. This field is for + # tests that want to read the configured codec. + native_codec = os.environ.get("CLICKHOUSE_CONNECT_NATIVE_CODEC", "python") + yield TestConfig(host, port, username, password, docker, test_database, cloud, compress, insert_quorum, proxy_address, native_codec) @fixture(scope="session", name="test_db") diff --git a/tests/integration_tests/test_geometric.py b/tests/integration_tests/test_geometric.py index a4f33c3e..0cb63536 100644 --- a/tests/integration_tests/test_geometric.py +++ b/tests/integration_tests/test_geometric.py @@ -1,6 +1,21 @@ from collections.abc import Callable +import pytest + +from clickhouse_connect.datatypes.dynamic import typed_variant from clickhouse_connect.driver import Client +from clickhouse_connect.driver.exceptions import DatabaseError, DataError + + +def _require_geometry(client: Client, call) -> None: + try: + resolved_type = call(client.command, "SELECT toTypeName(defaultValueOfTypeName('Geometry'))") + except DatabaseError as ex: + if ex.name != "UNKNOWN_TYPE": + raise + pytest.skip(f"Geometry is not supported by server {client.server_version}") + if resolved_type != "Geometry": + pytest.skip(f"Geometry is not supported by server {client.server_version}") def test_point_column(param_client: Client, call, table_context: Callable): @@ -32,3 +47,75 @@ def test_polygon_column(param_client: Client, call, table_context: Callable): call(param_client.insert, "polygon_column_test", [(1, pg), (4, pg)]) query_result = call(param_client.query, "SELECT key, polygon FROM polygon_column_test WHERE key = 4") assert query_result.first_row[1] == pg + + +def test_geometry_python_codec_round_trip(client_factory, call, client_mode): + client = client_factory(native_codec="python") + _require_geometry(client, call) + table = f"geometry_python_codec_{client_mode}" + values = [ + ("LineString", [(13.0, 23.0), (14.0, 24.0)]), + ("MultiLineString", [[(31.0, 41.0), (32.0, 42.0)]]), + ("MultiPolygon", [[[(51.0, 61.0)]]]), + ("Point", (71.0, 81.0)), + ("Polygon", [[(91.0, 101.0), (92.0, 102.0)]]), + ("Ring", [(111.0, 121.0)]), + ] + rows = [] + expected = [] + for index, (type_name, value) in enumerate(values): + tagged = typed_variant(value, type_name) + rows.append( + [ + index, + tagged, + [tagged, None], + (tagged, index), + [(tagged, index), (None, 79)], + {"value": tagged, "null": None}, + ] + ) + expected.append((index, value, [value, None], (value, index), [(value, index), (None, 79)], {"value": value, "null": None})) + rows.append([len(values), None, [None], (None, len(values)), [(None, len(values))], {"null": None}]) + expected.append((len(values), None, [None], (None, len(values)), [(None, len(values))], {"null": None})) + schema = "id UInt8, g Geometry, a Array(Geometry), t Tuple(Geometry, UInt8), at Array(Tuple(Geometry, UInt8)), m Map(String, Geometry)" + + try: + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE MergeTree ORDER BY id") + call(client.insert, table, rows, column_names=["id", "g", "a", "t", "at", "m"]) + result = call(client.query, f"SELECT * FROM {table} ORDER BY id").result_rows + assert result == expected + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +def test_geometry_python_codec_rejects_ambiguous_values(client_factory, call, client_mode): + client = client_factory(native_codec="python") + _require_geometry(client, call) + table = f"geometry_python_codec_errors_{client_mode}" + + try: + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} (id UInt8, g Geometry) ENGINE MergeTree ORDER BY id") + with pytest.raises(DataError, match="Cannot map Python type list"): + call(client.insert, table, [[0, [(13.0, 23.0)]]], column_names=["id", "g"]) + with pytest.raises(DataError, match="Type 'String' is not a member"): + call(client.insert, table, [[1, typed_variant("bad", "String")]], column_names=["id", "g"]) + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +@pytest.mark.parametrize("point", [(13.0,), (13.0, 23.0, 79.0)]) +def test_geometry_python_codec_rejects_malformed_point(client_factory, call, client_mode, point): + client = client_factory(native_codec="python") + _require_geometry(client, call) + table = f"geometry_python_codec_point_error_{client_mode}" + + try: + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} (id UInt8, g Geometry) ENGINE MergeTree ORDER BY id") + with pytest.raises(DataError, match=r"Tuple\(Float64, Float64\).*row 0"): + call(client.insert, table, [[0, typed_variant(point, "Point")]], column_names=["id", "g"]) + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") diff --git a/tests/integration_tests/test_inserts.py b/tests/integration_tests/test_inserts.py index 6a279198..a30bd7ca 100644 --- a/tests/integration_tests/test_inserts.py +++ b/tests/integration_tests/test_inserts.py @@ -1,9 +1,11 @@ import os -import time +import time as time_module import zoneinfo from collections.abc import Callable -from datetime import datetime +from datetime import date, datetime, time, timedelta from decimal import Decimal +from ipaddress import IPv4Address, IPv6Address +from uuid import UUID import pytest @@ -11,7 +13,7 @@ from clickhouse_connect.driver.client import Client from clickhouse_connect.driver.exceptions import DataError -HAS_TZSET = hasattr(time, "tzset") +HAS_TZSET = hasattr(time_module, "tzset") @pytest.fixture @@ -21,7 +23,7 @@ def naive_datetime_server_environment(monkeypatch): original_tz = os.environ.get("TZ") try: monkeypatch.setenv("TZ", "America/New_York") - time.tzset() + time_module.tzset() common.set_setting("naive_datetime_insert", "server") common.set_setting("naive_datetime_binding", "wall") yield @@ -32,7 +34,7 @@ def naive_datetime_server_environment(monkeypatch): monkeypatch.delenv("TZ", raising=False) else: monkeypatch.setenv("TZ", original_tz) - time.tzset() + time_module.tzset() def test_insert(param_client: Client, call, test_table_engine: str): @@ -125,6 +127,302 @@ def test_float_decimal_conv(param_client: Client, call, table_context: Callable) assert result == [(Decimal("0.492917"), Decimal("0.492917"), Decimal("0.492917"), Decimal("0.492917"))] +def test_rust_codec_insert(client_factory, call, table_context: Callable): + pytest.importorskip("_ch_core") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + + columns = [ + "id UInt32", + "b Bool", + "i Int32", + "u UInt64", + "f Float64", + "s String", + "fs FixedString(4)", + "n Nullable(Int32)", + "d Date", + "dt DateTime", + "ts DateTime64(3)", + "e Enum8('red' = 1, 'green' = 2)", + "dec Decimal(18, 4)", + "uuid_col UUID", + "ip4 IPv4", + "ip6 IPv6", + "lc LowCardinality(String)", + "lcn LowCardinality(Nullable(String))", + ] + data = [ + [ + 1, + True, + -13, + 79, + 1.25, + "user_1", + "abcd", + 13, + date(2024, 1, 15), + 1705322096, + 1705322096789, + "red", + Decimal("123.4567"), + UUID("00112233-4455-6677-8899-aabbccddeeff"), + "192.0.2.1", + IPv6Address("2001:db8::1"), + "x", + "nx", + ], + [ + 2, + False, + -79, + 500, + -2.5, + "user_2", + "xy", + None, + 19738, + 1705322097, + 1705322097790, + 2, + "-1.5", + "11111111-2222-3333-4444-555555555555", + IPv4Address("198.51.100.7"), + "2001:db8::2", + "x", + None, + ], + ] + expected = [ + ( + 1, + True, + -13, + 79, + 1.25, + "user_1", + b"abcd", + 13, + 19737, + 1705322096, + 1705322096789, + "red", + Decimal("123.4567"), + UUID("00112233-4455-6677-8899-aabbccddeeff"), + IPv4Address("192.0.2.1"), + IPv6Address("2001:db8::1"), + "x", + "nx", + ), + ( + 2, + False, + -79, + 500, + -2.5, + "user_2", + b"xy\x00\x00", + None, + 19738, + 1705322097, + 1705322097790, + "green", + Decimal("-1.5000"), + UUID("11111111-2222-3333-4444-555555555555"), + IPv4Address("198.51.100.7"), + IPv6Address("2001:db8::2"), + "x", + None, + ), + ] + + with table_context("test_rust_native_insert", columns): + call( + rust_client.insert, + "test_rust_native_insert", + data, + ) + result = call( + python_client.query, + "SELECT * FROM test_rust_native_insert ORDER BY id", + query_formats={"Date": "int", "DateTime": "int", "DateTime64": "int"}, + ).result_rows + assert result == expected + + +def test_rust_codec_time_insert(client_factory, call, test_config, test_table_engine: str): + pytest.importorskip("_ch_core") + if test_config.cloud: + pytest.skip("Time/Time64 settings are locked in ClickHouse Cloud") + + version_client = client_factory(native_codec="python") + if not version_client.min_version("25.6"): + pytest.skip("Time and Time64 require ClickHouse 25.6+") + + settings = {"enable_time_time64_type": 1} + rust_client = client_factory(native_codec="rust_strict", settings=settings) + python_client = client_factory(native_codec="python", settings=settings) + + table = "test_rust_native_time_insert" + columns = ( + "id UInt32, t Time, nt Nullable(Time), t64 Time64(6), " + "nt64 Nullable(Time64(6)), a Array(Time64(6)), " + "pair Tuple(Time, Nullable(Time64(6)))" + ) + rows = [ + [ + 13, + timedelta(seconds=-5), + None, + timedelta(seconds=-5, microseconds=-500_000), + None, + [timedelta(microseconds=1), timedelta(hours=1, minutes=2, seconds=3, microseconds=123_456)], + (timedelta(seconds=-13), None), + ], + [ + 79, + time(1, 2, 3), + "030:00:00", + time(1, 2, 3, 123_456), + "002:00:00.000079", + ["-000:00:05.500000", 79], + (79, 79), + ], + ] + expected = [ + ( + 13, + timedelta(seconds=-5), + None, + timedelta(seconds=-5, microseconds=-500_000), + None, + [timedelta(microseconds=1), timedelta(hours=1, minutes=2, seconds=3, microseconds=123_456)], + (timedelta(seconds=-13), None), + ), + ( + 79, + timedelta(hours=1, minutes=2, seconds=3), + timedelta(hours=30), + timedelta(hours=1, minutes=2, seconds=3, microseconds=123_456), + timedelta(hours=2, microseconds=79), + [timedelta(seconds=-5, microseconds=-500_000), timedelta(microseconds=79)], + (timedelta(seconds=79), timedelta(microseconds=79)), + ), + ] + + call(rust_client.command, f"DROP TABLE IF EXISTS {table}") + try: + call(rust_client.command, f"CREATE TABLE {table} ({columns}) ENGINE {test_table_engine} ORDER BY id") + call(rust_client.insert, table, rows) + result = call(python_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows + assert result == expected + + parity_rows = [ + [13, timedelta(microseconds=-1), timedelta(seconds=-5, microseconds=-500_000)], + [79, timedelta(seconds=79), timedelta(microseconds=-1)], + ] + parity_expected = [ + (13, timedelta(0), timedelta(seconds=-5, microseconds=-500_000)), + (79, timedelta(seconds=79), timedelta(microseconds=-1)), + ] + results = [] + for insert_client in (rust_client, python_client): + call(rust_client.command, f"TRUNCATE TABLE {table}") + call(insert_client.insert, table, parity_rows, column_names=["id", "t", "t64"]) + results.append(call(python_client.query, f"SELECT id, t, t64 FROM {table} ORDER BY id").result_rows) + assert results == [parity_expected, parity_expected] + + np = pytest.importorskip("numpy") + numpy_columns = [ + np.array([13, 79], dtype="uint32"), + np.array([13, "NaT"], dtype="timedelta64[s]"), + np.array([1, "NaT"], dtype="timedelta64[us]"), + ] + numpy_expected = [(13, timedelta(seconds=13), timedelta(microseconds=1)), (79, None, None)] + results = [] + for insert_client in (rust_client, python_client): + call(rust_client.command, f"TRUNCATE TABLE {table}") + call( + insert_client.insert, + table, + numpy_columns, + column_names=["id", "nt", "nt64"], + column_oriented=True, + ) + results.append(call(python_client.query, f"SELECT id, nt, nt64 FROM {table} ORDER BY id").result_rows) + assert results == [numpy_expected, numpy_expected] + + pd = pytest.importorskip("pandas") + frame = pd.DataFrame( + { + "id": [13, 79], + "t": [timedelta(seconds=5), timedelta(hours=1)], + "t64": [ + timedelta(seconds=1, microseconds=79), + timedelta(seconds=-5, microseconds=-500_000), + ], + "nt": pd.to_timedelta(["13s", None]), + "nt64": pd.to_timedelta([None, "0.000079s"]), + } + ) + frame_expected = [ + (13, timedelta(seconds=5), timedelta(seconds=1, microseconds=79), timedelta(seconds=13), None), + (79, timedelta(hours=1), timedelta(seconds=-5, microseconds=-500_000), None, timedelta(microseconds=79)), + ] + results = [] + for insert_client in (rust_client, python_client): + call(rust_client.command, f"TRUNCATE TABLE {table}") + call(insert_client.insert_df, table, frame) + results.append(call(python_client.query, f"SELECT id, t, t64, nt, nt64 FROM {table} ORDER BY id").result_rows) + assert results == [frame_expected, frame_expected] + finally: + call(rust_client.command, f"DROP TABLE IF EXISTS {table}") + + +def test_rust_codec_insert_dataframe(client_factory, call, table_context: Callable): + pytest.importorskip("_ch_core") + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + + data = pd.DataFrame( + { + "id": [13, 79], + "name": ["user_1", "user_2"], + "score": [Decimal("12.30"), Decimal("45.60")], + } + ) + + with table_context("test_rust_native_insert_df", ["id Int32", "name String", "score Decimal(9, 2)"]): + call( + rust_client.insert_df, + "test_rust_native_insert_df", + data, + ) + result = call(python_client.query, "SELECT * FROM test_rust_native_insert_df ORDER BY id").result_rows + assert result == [(13, "user_1", Decimal("12.30")), (79, "user_2", Decimal("45.60"))] + + +def test_rust_codec_insert_numpy(client_factory, call, table_context: Callable): + pytest.importorskip("_ch_core") + np = pytest.importorskip("numpy") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + + data = np.array([(13, 1.25), (79, 2.5)], dtype=[("id", " toUInt128(x) + toUInt128('18446744073709551616'), range(number % 4))", + "tuple_wide_int": "tuple(toInt128(number) - 2, toUInt256(number) + toUInt256('340282366920938463463374607431768211456'))", + "array_tuple_wide_int": ("arrayMap(x -> tuple(toInt256(x) - 2, toUInt128(x) + toUInt128('18446744073709551616')), range(number % 4))"), + "nullable_int": "CAST(if(number % 3 = 0, NULL, toInt32(number)) AS Nullable(Int32))", + "low_card_string": "CAST(toString(number % 3) AS LowCardinality(String))", + "low_card_nullable_string": "CAST(if(number % 2 = 0, NULL, toString(number)) AS LowCardinality(Nullable(String)))", + "enum8": "CAST(if(number % 2 = 0, 'a', 'b') AS Enum8('a' = 1, 'b' = 2))", + "enum16": "CAST(if(number % 2 = 0, 'x', 'y') AS Enum16('x' = 100, 'y' = 200))", + "datetime64": "toDateTime64(number, 3)", + "datetime64_utc": "toDateTime64(number, 3, 'UTC')", + "fixed_string": "CAST(leftPad(toString(number), 6, '0') AS FixedString(6))", + "datetime_utc": "toDateTime(number, 'UTC')", + "datetime_named_tz": "toDateTime(number, 'America/New_York')", + "float64": "toFloat64(number) / 2", + "uuid": "toUUID(concat(leftPad(lower(hex(number)), 8, '0'), '-1122-3344-5566-778899aabbcc'))", + "nullable_uuid": ( + "CAST(if(number % 3 = 0, NULL, toUUID(concat(leftPad(lower(hex(number)), 8, '0'), " + "'-1122-3344-5566-778899aabbcc'))) AS Nullable(UUID))" + ), + "decimal32": "toDecimal32(number / 3 - 1, 2)", + "decimal64": "toDecimal64(number / 3 - 2, 4)", + "decimal128": "toDecimal128(number / 3 - 2, 10)", + "decimal256": "toDecimal256(number / 3 - 2, 20)", + "nullable_decimal": "CAST(if(number % 3 = 0, NULL, toDecimal64(number / 3, 4)) AS Nullable(Decimal(18, 4)))", + "ipv4": "toIPv4(toUInt32(number * 16909060))", + "nullable_ipv4": "CAST(if(number % 3 = 0, NULL, toIPv4(toUInt32(number * 16909060))) AS Nullable(IPv4))", + "ipv6": "toIPv6(concat('2001:db8::', lower(hex(toUInt16(number + 1)))))", + "ipv6_v4_mapped": "toIPv6(toIPv4(toUInt32(number + 1)))", + "array_int": "range(number % 4)", + "array_string": "arrayMap(x -> toString(x), range(number % 4))", + "array_nullable_int": "arrayMap(x -> if(x % 2 = 0, NULL, toInt64(x)), range(number % 4))", + "array_low_card_string": "CAST(arrayMap(x -> toString(x % 3), range(number % 4)) AS Array(LowCardinality(String)))", + "array_nested": "arrayMap(x -> range(x % 3), range(number % 4))", + "array_uuid": "arrayMap(x -> toUUID(concat(leftPad(lower(hex(x)), 8, '0'), '-1122-3344-5566-778899aabbcc')), range(number % 4))", + "array_datetime": "arrayMap(x -> toDateTime(x), range(number % 4))", + "array_decimal": "arrayMap(x -> toDecimal64(x, 4), range(number % 4))", + "tuple_unnamed": "tuple(number, toString(number))", + "tuple_named": "CAST((toInt64(number), toString(number)), 'Tuple(a Int64, b String)')", + "tuple_low_card": "CAST((toString(number % 3), number), 'Tuple(LowCardinality(String), UInt64)')", + "tuple_nullable_element": "tuple(if(number % 2 = 0, NULL, toString(number)))", + "map_string_int": "mapFromArrays(arrayMap(x -> concat('k', toString(x)), range(number % 4)), range(number % 4))", + "map_array_value": "CAST(map('a', range(number % 4)), 'Map(String, Array(UInt64))')", + "map_low_card_key": "CAST(map(toString(number % 3), number), 'Map(LowCardinality(String), UInt64)')", + "array_of_tuple": "arrayMap(x -> (x, toString(x)), range(number % 4))", + "map_of_tuple_value": "CAST(map('a', (number, toString(number))), 'Map(String, Tuple(UInt64, String))')", + "point": "(toFloat64(number), toFloat64(number) / 2)::Point", + "ring": "[(toFloat64(number), toFloat64(number) / 2), (toFloat64(number) + 1, toFloat64(number))]::Ring", + "linestring": "[(toFloat64(number), toFloat64(number)), (toFloat64(number) + 1, 0.)]::LineString", + "polygon": "[[(toFloat64(number), 0.), (0., toFloat64(number)), (1., 1.)]]::Polygon", + "multilinestring": "[[(toFloat64(number), 0.), (1., toFloat64(number))]]::MultiLineString", + "multipolygon": "[[[(toFloat64(number), 0.), (0., 1.), (1., 0.)]]]::MultiPolygon", + "simple_agg_uint64": "CAST(number AS SimpleAggregateFunction(sum, UInt64))", + "simple_agg_string": "CAST(toString(number) AS SimpleAggregateFunction(anyLast, String))", + "simple_agg_low_card": "CAST(toString(number % 3) AS SimpleAggregateFunction(anyLast, LowCardinality(String)))", + "simple_agg_array": "CAST(range(number % 4) AS SimpleAggregateFunction(groupArrayArray, Array(UInt64)))", +} + + +def test_rust_codec_ab_parity(client_factory, call): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + + rust_result = call(rust_client.query, SCALAR_QUERY) + python_result = call(python_client.query, SCALAR_QUERY) + + assert rust_result.result_rows == python_result.result_rows + assert rust_result.column_names == python_result.column_names + assert [t.name for t in rust_result.column_types] == [t.name for t in python_result.column_types] + + +def test_rust_codec_qbit_cross_codec_round_trip(client_factory, call, client_mode): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + if not python_client.min_version("25.10"): + pytest.skip("QBit requires ClickHouse 25.10+") + + setting = python_client.server_settings.get("allow_experimental_qbit_type") + if setting is not None and setting.value != "1": + try: + rust_client.set_client_setting("allow_experimental_qbit_type", "1") + python_client.set_client_setting("allow_experimental_qbit_type", "1") + call(python_client.command, "SELECT 1") + except DatabaseError: + pytest.skip("QBit is experimental and cannot be enabled on this server") + + rust_table = f"rc_qbit_rust_{client_mode}" + python_table = f"rc_qbit_python_{client_mode}" + schema = ( + "id UInt8, f32 QBit(Float32, 9), f64 QBit(Float64, 5), bf Nullable(QBit(BFloat16, 3)), nested Array(Tuple(QBit(Float32, 9), UInt8))" + ) + names = ["id", "f32", "f64", "bf", "nested"] + rows = [ + [ + 1, + [-1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, -8.0, -9.0], + [-0.99105519, 1.28887844, -0.43526649, -0.98520696, 0.66154391], + [0.1, -2.5, 3.25], + [([-1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, -8.0, -9.0], 13)], + ], + [ + 2, + [0.25, 1.5, -2.75, 3.0, 4.125, -5.5, 6.75, 7.0, 8.5], + [float("inf"), float("-inf"), -0.0, 13.0, 79.0], + None, + [], + ], + ] + + def assert_rows(actual, expected): + assert [row[0] for row in actual] == [row[0] for row in expected] + for actual_row, expected_row in zip(actual, expected): + assert actual_row[1] == pytest.approx(expected_row[1], rel=1e-6) + assert actual_row[2] == pytest.approx(expected_row[2]) + if expected_row[3] is None: + assert actual_row[3] is None + else: + assert actual_row[3] == pytest.approx(expected_row[3], rel=1e-2, abs=1e-2) + assert len(actual_row[4]) == len(expected_row[4]) + for (actual_vector, actual_tag), (expected_vector, expected_tag) in zip(actual_row[4], expected_row[4]): + assert actual_tag == expected_tag + assert actual_vector == pytest.approx(expected_vector, rel=1e-6) + + def cross_read(insert_client, table): + call(insert_client.command, f"CREATE TABLE {table} ({schema}) ENGINE Memory") + call(insert_client.insert, table, rows, column_names=names) + query = f"SELECT * FROM {table} ORDER BY id" + rust_rows = call(rust_client.query, query).result_rows + python_rows = call(python_client.query, query).result_rows + assert rust_rows == python_rows + assert_rows(rust_rows, rows) + distance = call( + python_client.query, + f"SELECT L2DistanceTransposed(f32, %(reference)s, 32) FROM {table} WHERE id = 1", + parameters={"reference": rows[0][1]}, + ).result_rows[0][0] + assert distance == pytest.approx(0.0, abs=1e-7) + + try: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {python_table}") + cross_read(rust_client, rust_table) + cross_read(python_client, python_table) + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {python_table}") + + +def test_rust_codec_variant_round_trip_parity(client_factory, call, client_mode): + pytest.importorskip("pandas") + probe = client_factory(native_codec="python") + type_available(probe, "variant") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + rust_table = f"rc_variant_rust_{client_mode}" + python_table = f"rc_variant_python_{client_mode}" + schema = ( + "id UInt8, v Variant(String, UInt64), " + "a Array(Variant(Bool, Int32, String)), " + "m Map(String, Variant(String, UInt64)), " + "va Variant(Array(String), Array(UInt32))" + ) + rows = [ + [ + 0, + None, + [True, 13, "a", None], + {"a": "x", "b": 13, "c": "y"}, + typed_variant([13, 79], "Array(UInt32)"), + ], + [ + 1, + "user_1", + [], + {}, + typed_variant(["a", "b"], "Array(String)"), + ], + [ + 2, + 79, + [False, -13, "b"], + {"d": 79, "e": "z"}, + typed_variant([], "Array(UInt32)"), + ], + ] + names = ["id", "v", "a", "m", "va"] + expected = [ + (0, None, [True, 13, "a", None], {"a": "x", "b": 13, "c": "y"}, [13, 79]), + (1, "user_1", [], {}, ["a", "b"]), + (2, 79, [False, -13, "b"], {"d": 79, "e": "z"}, []), + ] + + def round_trip(insert_client, table): + call(insert_client.command, f"DROP TABLE IF EXISTS {table}") + call(insert_client.command, f"CREATE TABLE {table} ({schema}) ENGINE Memory") + call(insert_client.insert, table, rows, column_names=names) + query = f"SELECT * FROM {table} ORDER BY id" + rust_result = call(rust_client.query, query) + python_result = call(python_client.query, query) + assert rust_result.result_rows == python_result.result_rows == expected + assert [ch_type.name for ch_type in rust_result.column_types] == [ch_type.name for ch_type in python_result.column_types] + rust_np = call(rust_client.query_np, query) + python_np = call(python_client.query_np, query) + assert rust_np.dtype == python_np.dtype + assert rust_np.tolist() == python_np.tolist() + rust_df = call(rust_client.query_df, query) + python_df = call(python_client.query_df, query) + # Cells are objects; cell-type parity is a known divergence, so compare values only. + assert list(rust_df.dtypes) == list(python_df.dtypes) + assert rust_df.equals(python_df) + + try: + round_trip(rust_client, rust_table) + round_trip(python_client, python_table) + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {python_table}") + + +DYNAMIC_QUERY = """ + SELECT id, d, s, [d, CAST('tail', 'Dynamic')] AS a, + tuple(d, id) AS t, map('v', d) AS m + FROM + ( + SELECT toUInt8(0) AS id, CAST(NULL, 'Dynamic') AS d, + CAST(NULL, 'Dynamic(max_types=0)') AS s + UNION ALL + SELECT 1, CAST('user_1', 'Dynamic'), + CAST(toDate('2024-01-02'), 'Dynamic(max_types=0)') + UNION ALL + SELECT 2, CAST(toUInt64(79), 'Dynamic'), + CAST([toUInt8(1), 2, 3], 'Dynamic(max_types=0)') + UNION ALL + SELECT 3, CAST(toInt32(-13), 'Dynamic'), + CAST('hello', 'Dynamic(max_types=0)') + UNION ALL + SELECT 4, CAST(toFloat64(2.5), 'Dynamic'), + CAST(toUInt64(79), 'Dynamic(max_types=0)') + UNION ALL + SELECT 5, CAST(true, 'Dynamic'), + CAST(toFloat64(-0.5), 'Dynamic(max_types=0)') + UNION ALL + SELECT 6, CAST(toInt64(7), 'Dynamic'), + CAST(true, 'Dynamic(max_types=0)') + ) + ORDER BY id +""" + +# Typed SharedVariant decode under rust: date/list/str/int/float/bool. +DYNAMIC_SHARED_TYPED = [None, date(2024, 1, 2), [1, 2, 3], "hello", 79, -0.5, True] +# Known divergence: the python codec's shared-cell heuristic only decodes +# int/float/str/bool and returns raw wire bytes for Date and Array cells +# (FINDINGS.md finding 4). +DYNAMIC_SHARED_PYTHON = [None, b"\x0f\x0cM", b"\x1e\x01\x03\x01\x02\x03", "hello", 79, -0.5, True] +DYNAMIC_DIRECT = [None, "user_1", 79, -13, 2.5, True, 7] + + +def test_rust_codec_dynamic_decode_parity(client_factory, call, consume_stream): + np = pytest.importorskip("numpy") + pytest.importorskip("pandas") + probe = client_factory(native_codec="python") + type_available(probe, "dynamic") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + + rust_result = call(rust_client.query, DYNAMIC_QUERY, settings={"max_block_size": 1}) + python_result = call(python_client.query, DYNAMIC_QUERY) + rust_rows = rust_result.result_rows + python_rows = python_result.result_rows + assert [t.name for t in rust_result.column_types] == [t.name for t in python_result.column_types] + + # The direct Dynamic column and every container built from it are full parity. + for ix in (0, 1, 3, 4, 5): # id, d, a, t, m + assert [row[ix] for row in rust_rows] == [row[ix] for row in python_rows] + assert [row[1] for row in rust_rows] == DYNAMIC_DIRECT + assert [row[3] for row in rust_rows] == [[value, "tail"] for value in DYNAMIC_DIRECT] + assert [row[4] for row in rust_rows] == list(zip(DYNAMIC_DIRECT, range(7))) + assert [row[5] for row in rust_rows] == [{"v": value} for value in DYNAMIC_DIRECT] + + assert [row[2] for row in rust_rows] == DYNAMIC_SHARED_TYPED + assert [row[2] for row in python_rows] == DYNAMIC_SHARED_PYTHON + + # max_block_size=1 forces one row per block, exercising block-local child unification. + for settings in (None, {"max_block_size": 1}): + rust_np = call(rust_client.query_np, DYNAMIC_QUERY, settings=settings) + python_np = call(python_client.query_np, DYNAMIC_QUERY, settings=settings) + assert rust_np.dtype == python_np.dtype + for name in ("id", "d", "a", "t", "m"): + assert rust_np[name].tolist() == python_np[name].tolist() + assert rust_np["s"].tolist() == DYNAMIC_SHARED_TYPED + assert python_np["s"].tolist() == DYNAMIC_SHARED_PYTHON + + # np-scalar residue (FINDINGS.md finding 4): rust yields python-native + # scalars in object cells where python yields value-equal numpy scalars. + assert type(rust_np["d"].tolist()[2]) is int + assert isinstance(python_np["d"].tolist()[2], np.unsignedinteger) + + rust_df = call(rust_client.query_df, DYNAMIC_QUERY, settings=settings) + python_df = call(python_client.query_df, DYNAMIC_QUERY, settings=settings) + assert list(rust_df.dtypes) == list(python_df.dtypes) + for name in ("id", "d", "a", "t", "m"): + assert rust_df[name].tolist() == python_df[name].tolist() + assert rust_df["s"].tolist() == DYNAMIC_SHARED_TYPED + # Known divergence: the python codec's pandas exit stringifies every + # non-null shared cell, including the heuristically decoded ones. + assert python_df["s"].tolist() == [None, "\x0f\x0cM", "\x1e\x01\x03\x01\x02\x03", "hello", "79", "-0.5", "True"] + + # np-scalar residue inside a container cell (FINDINGS.md finding 4). + assert type(rust_df["a"].tolist()[2][0]) is int + assert isinstance(python_df["a"].tolist()[2][0], np.unsignedinteger) + + streamed = [] + consume_stream(call(rust_client.query_rows_stream, DYNAMIC_QUERY, settings={"max_block_size": 1}), streamed.append) + assert [row[1] for row in streamed] == DYNAMIC_DIRECT + assert [row[2] for row in streamed] == DYNAMIC_SHARED_TYPED + + +def test_rust_codec_json_round_trip_parity(client_factory, call, consume_stream, client_mode): + pytest.importorskip("numpy") + pytest.importorskip("pandas") + probe = client_factory(native_codec="python") + type_available(probe, "json") + if not probe.min_version("25.3"): + pytest.skip("Nullable(JSON) requires ClickHouse 25.3+") + + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + rust_table = f"rc_json_rust_{client_mode}" + python_table = f"rc_json_python_{client_mode}" + schema = ( + "id UInt8, " + "payload json(max_dynamic_paths=2, `typed.value` Int64), " + "shared_payload JSON(max_dynamic_paths=0), " + "nullable_payload Nullable(JSON), " + "items Array(JSON), " + "wrapped Tuple(payload JSON)" + ) + names = ["id", "payload", "shared_payload", "nullable_payload", "items", "wrapped"] + rows = [ + [ + 0, + {"typed": {"value": 13}, "name": "user_1", "active": True, "score": 17}, + {"items": [{"kind": "shared"}]}, + None, + [{"kind": "first"}, {"count": 13}], + ({"nested": {"value": "a"}},), + ], + [ + 1, + {"typed": {"value": 79}, "ratio": 2.5, "deep": {"value": "x"}, "extra": -13}, + {"mix": [1, "a", None]}, + {"nullable": "present"}, + [], + ({"nested": {"value": "b"}},), + ], + [ + 2, + {"typed": {"value": 5}, "other": 79, "flag": False, "tail": "shared"}, + {"empty": []}, + {"array": [1, None, 3]}, + [{"kind": "last"}], + ({"nested": {"value": "c"}},), + ], + ] + expected = [ + ( + 0, + {"typed": {"value": 13}, "name": "user_1", "active": True, "score": 17}, + {"items": [{"kind": "shared"}]}, + None, + [{"kind": "first"}, {"count": 13}], + {"payload": {"nested": {"value": "a"}}}, + ), + ( + 1, + {"typed": {"value": 79}, "ratio": 2.5, "deep": {"value": "x"}, "extra": -13}, + {"mix": [1, "a", None]}, + {"nullable": "present"}, + [], + {"payload": {"nested": {"value": "b"}}}, + ), + ( + 2, + {"typed": {"value": 5}, "other": 79, "flag": False, "tail": "shared"}, + {"empty": []}, + {"array": [1, None, 3]}, + [{"kind": "last"}], + {"payload": {"nested": {"value": "c"}}}, + ), + ] + + def create_and_insert(client, table): + call(python_client.command, f"DROP TABLE IF EXISTS {table}") + call(python_client.command, f"CREATE TABLE {table} ({schema}) ENGINE Memory") + call(client.insert, table, rows, column_names=names) + + try: + create_and_insert(rust_client, rust_table) + create_and_insert(python_client, python_table) + rust_query = f"SELECT * FROM {rust_table} ORDER BY id" + python_query = f"SELECT * FROM {python_table} ORDER BY id" + + # max_block_size=1 varies the block-local dynamic path set and forces + # shared-data overflow under max_dynamic_paths=2. + settings = {"max_block_size": 1} + rust_result = call(rust_client.query, rust_query, settings=settings) + python_result = call(python_client.query, python_query, settings=settings) + assert rust_result.result_rows == python_result.result_rows == expected + assert [ch_type.name for ch_type in rust_result.column_types] == [ch_type.name for ch_type in python_result.column_types] + + rust_np = call(rust_client.query_np, rust_query, settings=settings) + python_np = call(python_client.query_np, python_query, settings=settings) + assert rust_np.dtype == python_np.dtype + assert rust_np.tolist() == rust_result.result_rows + + rust_df = call(rust_client.query_df, rust_query, settings=settings) + assert list(rust_df.itertuples(index=False, name=None)) == rust_result.result_rows + + streamed = [] + consume_stream( + call(rust_client.query_rows_stream, rust_query, settings=settings), + streamed.append, + ) + assert streamed == rust_result.result_rows + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {python_table}") + + +def test_rust_codec_json_skip_regexp_round_trip(client_factory, call, client_mode): + probe = client_factory(native_codec="python") + type_available(probe, "json") + + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + table = f"rc_json_skip_regexp_{client_mode}" + # SQL literal renders the regex skip'me\..* with both a single quote and a backslash. + schema = "id UInt8, payload JSON(SKIP REGEXP 'skip\\'me\\\\..*')" + rows = [ + [0, {"kept": 13, "skip'me": {"inner": 79}}], + [1, {"kept": "user_1", "skip'me": {"inner": "x"}, "other": 2.5}], + ] + expected = [ + (0, {"kept": 13}), + (1, {"kept": "user_1", "other": 2.5}), + ] + try: + call(python_client.command, f"DROP TABLE IF EXISTS {table}") + call(python_client.command, f"CREATE TABLE {table} ({schema}) ENGINE Memory") + call(rust_client.insert, table, rows, column_names=["id", "payload"]) + query = f"SELECT * FROM {table} ORDER BY id" + rust_result = call(rust_client.query, query) + python_result = call(python_client.query, query) + assert rust_result.result_rows == python_result.result_rows == expected + assert [ch_type.name for ch_type in rust_result.column_types] == [ch_type.name for ch_type in python_result.column_types] + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {table}") + + +@pytest.mark.parametrize("native_codec", ["rust", "rust_strict"]) +def test_rust_codec_dynamic_insert_parity(client_factory, call, client_mode, native_codec): + probe = client_factory(native_codec="python") + type_available(probe, "dynamic") + python_client = client_factory(native_codec="python") + rows = [[0, None], [1, True], [2, 13], [3, 2.5], [4, "user_1"], [5, [1, 2]]] + + def roundtrip(codec): + client = client_factory(native_codec=codec) + table = f"rc_ins_dynamic_{codec}_{client_mode}" + call(python_client.command, f"DROP TABLE IF EXISTS {table}") + try: + call(python_client.command, f"CREATE TABLE {table} (id UInt8, d Dynamic) ENGINE Memory") + call(client.insert, table, rows, column_names=["id", "d"]) + return call(python_client.query, f"SELECT id, d, dynamicType(d) FROM {table} ORDER BY id").result_rows + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {table}") + + # Dynamic inserts send str(value) with None as the literal "NULL"; the rust + # encoder now builds that String column natively with exact wire parity, so + # both rust and rust_strict insert without the python fallback. + expected = roundtrip("python") + assert roundtrip(native_codec) == expected + assert [row[1] for row in expected] == ["NULL", "True", "13", "2.5", "user_1", "[1, 2]"] + + +def test_rust_codec_time_parity(client_factory, call, test_config): + if test_config.cloud: + pytest.skip("Time/Time64 settings are locked in ClickHouse Cloud") + + version_client = client_factory(native_codec="python") + if not version_client.min_version("25.6"): + pytest.skip("Time and Time64 require ClickHouse 25.6+") + + settings = {"allow_suspicious_low_cardinality_types": 1, "enable_time_time64_type": 1} + rust_client = client_factory(native_codec="rust_strict", settings=settings) + python_client = client_factory(native_codec="python", settings=settings) + + cases = [ + ( + "CAST(v AS Time)", + ["-000:00:05", "000:00:00", "030:00:00"], + [timedelta(seconds=-5), timedelta(0), timedelta(hours=30)], + ), + ( + "CAST(v AS Time64(9))", + ["-000:00:05.500000000", "000:00:00.000000001", "001:02:03.123456789"], + [timedelta(seconds=-5, microseconds=-500_000), timedelta(0), timedelta(seconds=3_723, microseconds=123_456)], + ), + ] + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + for expression, values, expected in cases: + rows = ", ".join(f"('{value}')" for value in values) + query = f"SELECT {expression} AS c FROM values('v String', {rows})" + + rust_result = call(rust_client.query, query) + python_result = call(python_client.query, query) + assert rust_result.result_rows == python_result.result_rows == [(value,) for value in expected] + + rust_np = call(rust_client.query_np, query) + python_np = call(python_client.query_np, query) + assert rust_np.dtype == python_np.dtype + np.testing.assert_array_equal(rust_np, python_np) + + for use_extended_dtypes in (False, True): + rust_df = call(rust_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + python_df = call(python_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + assert rust_df["c"].dtype == python_df["c"].dtype + pd.testing.assert_frame_equal(rust_df, python_df) + nullable_cases = [ + ("Time", ["000:00:13", None, "-000:00:05"]), + ("Time64(6)", ["000:00:13.000079", None, "-000:00:05.500000"]), + ("Time64(9)", ["-000:00:00.000000001", None, "000:00:00.000000001"]), + ] + for type_name, values in nullable_cases: + rows = ", ".join("(NULL)" if value is None else f"('{value}')" for value in values) + query = f"SELECT CAST(v AS Nullable({type_name})) AS c FROM values('v Nullable(String)', {rows})" + + rust_np = call(rust_client.query_np, query) + python_np = call(python_client.query_np, query) + assert rust_np.dtype == python_np.dtype + np.testing.assert_array_equal(rust_np, python_np) + + for use_extended_dtypes in (False, True): + rust_df = call(rust_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + python_df = call(python_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + assert rust_df["c"].dtype == python_df["c"].dtype + pd.testing.assert_frame_equal(rust_df, python_df) + + low_card_queries = [ + ("SELECT CAST(v AS LowCardinality(Time)) AS c FROM values('v String', ('000:00:13'), ('-000:00:05'), ('000:00:13'))"), + ("SELECT CAST(v AS LowCardinality(Nullable(Time))) AS c FROM values('v Nullable(String)', ('000:00:13'), (NULL), ('-000:00:05'))"), + ] + for query in low_card_queries: + rust_np = call(rust_client.query_np, query) + python_np = call(python_client.query_np, query) + assert rust_np.dtype == python_np.dtype + np.testing.assert_array_equal(rust_np, python_np) + rust_df = call(rust_client.query_df, query) + python_df = call(python_client.query_df, query) + assert rust_df["c"].dtype == python_df["c"].dtype + pd.testing.assert_frame_equal(rust_df, python_df) + + nested_queries = [ + ("SELECT [CAST('-000:00:00.000000001' AS Time64(9)), CAST('000:00:00.000000001' AS Time64(9))] AS c"), + ("SELECT tuple(CAST('-000:00:05' AS Time), CAST('000:00:00.000000001' AS Time64(9))) AS c"), + ( + "SELECT [tuple(CAST('-000:00:05' AS Time), " + "CAST('-000:00:00.000000001' AS Time64(9))), " + "tuple(CAST('000:00:13' AS Time), " + "CAST('000:00:00.000000001' AS Time64(9)))] AS c" + ), + ("SELECT map('key', CAST('000:00:00.000000001' AS Time64(9)), 'key', CAST('000:00:00.000000002' AS Time64(9))) AS c"), + ("SELECT [CAST(NULL AS Nullable(Time64(9))), CAST('000:00:00.000000001' AS Nullable(Time64(9)))] AS c"), + ] + for query in nested_queries: + rust_np = call(rust_client.query_np, query) + python_np = call(python_client.query_np, query) + assert rust_np.dtype == python_np.dtype + np.testing.assert_array_equal(rust_np, python_np) + + for use_extended_dtypes in (False, True): + rust_df = call(rust_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + python_df = call(python_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + pd.testing.assert_frame_equal(rust_df, python_df) + + +def test_rust_codec_interval_parity(client_factory, call, client_mode): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = ( + "SELECT toIntervalYear(-13), toIntervalQuarter(79), toIntervalMonth(-13), " + "toIntervalWeek(79), toIntervalDay(-13), toIntervalHour(79), " + "toIntervalMinute(-13), toIntervalSecond(79), toIntervalMillisecond(-13), " + "toIntervalMicrosecond(79), toIntervalNanosecond(-13), 'sentinel'" + ) + + rust_result = call(rust_client.query, query) + python_result = call(python_client.query, query) + assert rust_result.result_rows == python_result.result_rows + assert [ch_type.name for ch_type in rust_result.column_types] == [ch_type.name for ch_type in python_result.column_types] + + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + parity_queries = [ + "SELECT toIntervalDay(v) AS c FROM values('v Int64', (-13), (0), (79))", + "SELECT toIntervalSecond(v) AS c FROM values('v Int64', (-13), (0), (79))", + "SELECT toIntervalDay(v) AS c FROM values('v Nullable(Int64)', (-13), (NULL), (79))", + "SELECT toIntervalNanosecond(v) AS c FROM values('v Nullable(Int64)', (-13), (NULL), (79))", + "SELECT [toIntervalMinute(-13), toIntervalMinute(79)] AS c", + "SELECT tuple(toIntervalSecond(-13), 'x') AS c", + "SELECT [CAST(NULL AS Nullable(IntervalHour)), toIntervalHour(79)] AS c", + "SELECT map(toIntervalDay(-13), 'x') AS c", + ] + for parity_query in parity_queries: + rust_np = call(rust_client.query_np, parity_query) + python_np = call(python_client.query_np, parity_query) + assert rust_np.dtype == python_np.dtype + np.testing.assert_array_equal(rust_np, python_np) + + for use_extended_dtypes in (False, True): + rust_df = call(rust_client.query_df, parity_query, use_extended_dtypes=use_extended_dtypes) + python_df = call(python_client.query_df, parity_query, use_extended_dtypes=use_extended_dtypes) + assert rust_df["c"].dtype == python_df["c"].dtype + pd.testing.assert_frame_equal(rust_df, python_df) + + lc_query = "SELECT toLowCardinality(toIntervalHour(v)) AS c FROM values('v Int64', (13), (79), (13))" + rust_df = call(rust_client.query_df, lc_query) + python_df = call(python_client.query_df, lc_query) + pd.testing.assert_frame_equal(rust_df, python_df) + + table = f"rc_ins_interval_{client_mode}" + schema = ( + "id UInt32, d IntervalDay, n Nullable(IntervalHour), a Array(IntervalMinute), " + "t Tuple(IntervalSecond, String), " + "at Array(Tuple(IntervalMillisecond, IntervalMonth)), m Map(IntervalDay, String)" + ) + rows = [ + [0, -13, None, [-13, 79], (-13, "x"), [(-13, 1), (79, -2)], {-13: "x"}], + [1, 79, 13, [], (79, "y"), [], {}], + ] + try: + call(python_client.command, f"DROP TABLE IF EXISTS {table}") + call(python_client.command, f"CREATE TABLE {table} ({schema}) ENGINE Memory") + call(rust_client.insert, table, rows, column_names=["id", "d", "n", "a", "t", "at", "m"]) + + expected = [tuple(row) for row in rows] + assert call(rust_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows == expected + assert call(python_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows == expected + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {table}") + + +@pytest.mark.parametrize( + "expr", + [ + "CAST((CAST('000:00:13' AS Time), if(number = 0, NULL, toInt64(13))), 'Tuple(t Time, v Nullable(Int64))')", + "CAST((CAST('000:00:13' AS Time), toDate(13)), 'Tuple(t Time, d Date)')", + ], + ids=["nullable_int", "date"], +) +def test_rust_codec_nested_time_sibling_df_parity(client_factory, call, test_config, expr): + if test_config.cloud: + pytest.skip("Time/Time64 settings are locked in ClickHouse Cloud") + + version_client = client_factory(native_codec="python") + if not version_client.min_version("25.6"): + pytest.skip("Time and Time64 require ClickHouse 25.6+") + + pd = pytest.importorskip("pandas") + settings = {"enable_time_time64_type": 1} + rust_client = client_factory(native_codec="rust_strict", settings=settings) + python_client = client_factory(native_codec="python", settings=settings) + query = f"SELECT {expr} AS c FROM numbers(2)" + + rust_df = call(rust_client.query_df, query, use_extended_dtypes=True) + python_df = call(python_client.query_df, query, use_extended_dtypes=True) + pd.testing.assert_frame_equal(rust_df, python_df) + if "Nullable(Int64)" in expr: + assert rust_df.iloc[0, 0]["v"] is pd.NA + + +@pytest.mark.parametrize("expr", DECODE_MATRIX.values(), ids=list(DECODE_MATRIX)) +def test_rust_codec_decode_matrix_parity(client_factory, call, expr): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = f"SELECT {expr} AS c FROM numbers(13)" + + rust_result = call(rust_client.query, query) + python_result = call(python_client.query, query) + + assert rust_result.result_rows == python_result.result_rows + assert [t.name for t in rust_result.column_types] == [t.name for t in python_result.column_types] + + +def test_rust_codec_nothing_insert_parity(client_factory, call, client_mode): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + schema = "id UInt8, tn Tuple(Nullable(Nothing), UInt8), at Array(Tuple(Nullable(Nothing), UInt8)), an Tuple(Array(Nothing), UInt8)" + rows = [ + [0, (None, 13), [(None, 5), (None, 7)], ([], 79)], + [1, (None, 79), [], ([], 13)], + ] + names = ["id", "tn", "at", "an"] + rust_table = f"rc_ins_nothing_rust_{client_mode}" + py_table = f"rc_ins_nothing_py_{client_mode}" + + def roundtrip(client, table): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE Memory") + call(client.insert, table, rows, column_names=names) + select = f"SELECT * FROM {table} ORDER BY id" + assert call(rust_client.query, select).result_rows == call(python_client.query, select).result_rows + return call(rust_client.query, select).result_rows + + try: + expected = [tuple(row) for row in rows] + assert roundtrip(rust_client, rust_table) == expected + assert roundtrip(python_client, py_table) == expected + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {py_table}") + + +def test_rust_codec_nothing_np_df_parity(client_factory, call): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + + # Sibling scalars inside tuples/maps are the accepted value-equal residue + # (rust python int vs numpy scalar); assert_array_equal and + # assert_frame_equal compare by value. + parity_queries = [ + "SELECT NULL AS c FROM numbers(3)", + "SELECT [NULL] AS c FROM numbers(3)", + "SELECT tuple(NULL, toUInt8(number)) AS c FROM numbers(3)", + "SELECT map() AS c FROM numbers(3)", + "SELECT mapFromArrays([toUInt8(number)], [NULL]) AS c FROM numbers(3)", + ] + for query in parity_queries: + rust_np = call(rust_client.query_np, query) + python_np = call(python_client.query_np, query) + assert rust_np.dtype == python_np.dtype + np.testing.assert_array_equal(rust_np, python_np) + + for use_extended_dtypes in (False, True): + rust_df = call(rust_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + python_df = call(python_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + assert rust_df["c"].dtype == python_df["c"].dtype + pd.testing.assert_frame_equal(rust_df, python_df) + + # Empty flat Array runs break query_np identically in both codecs, so these + # shapes are df-only. + df_only_queries = [ + "SELECT array() AS c FROM numbers(3)", + "SELECT tuple([], toUInt8(number)) AS c FROM numbers(3)", + ] + for query in df_only_queries: + for use_extended_dtypes in (False, True): + rust_df = call(rust_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + python_df = call(python_client.query_df, query, use_extended_dtypes=use_extended_dtypes) + assert rust_df["c"].dtype == python_df["c"].dtype + pd.testing.assert_frame_equal(rust_df, python_df) + + +def test_rust_codec_streaming_parity(client_factory, call, consume_stream): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = "SELECT number AS n, toString(number) AS s FROM numbers(5000)" + + def totals(client): + block_rows = 0 + + def on_block(block): + nonlocal block_rows + block_rows += len(block[0]) + + consume_stream(call(client.query_column_block_stream, query, settings={"max_block_size": 1000}), on_block) + + row_count = 0 + + def on_row(_row): + nonlocal row_count + row_count += 1 + + consume_stream(call(client.query_rows_stream, query), on_row) + return block_rows, row_count + + rust_totals = totals(rust_client) + assert rust_totals == totals(python_client) + assert rust_totals == (5000, 5000) + + +def test_rust_codec_numeric_column_blocks_keep_typed_arrays(client_factory, call, consume_stream): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + expressions = [ + "toInt8(number) AS i8", + "toInt16(number) AS i16", + "toInt32(number) AS i32", + "toInt64(number) AS i64", + "toUInt8(number) AS u8", + "toUInt16(number) AS u16", + "toUInt32(number) AS u32", + "toUInt64(number) AS u64", + "toFloat32(number / 2) AS f32", + "toFloat64(number) / 2 AS f64", + "CAST(if(number = 1, NULL, number), 'Nullable(Int32)') AS nullable_i32", + ] + query = f"SELECT {', '.join(expressions)} FROM numbers(3)" + + def blocks(client): + output = [] + consume_stream(call(client.query_column_block_stream, query), output.append) + return output + + rust_blocks = blocks(rust_client) + python_blocks = blocks(python_client) + assert len(rust_blocks) == len(python_blocks) == 1 + expected_typecodes = ["b", "h", "i", "q", "B", "H", "I", "Q", "f", "d"] + for rust_column, python_column, typecode in zip(rust_blocks[0][:-1], python_blocks[0][:-1], expected_typecodes): + assert isinstance(rust_column, array.array) + assert rust_column.typecode == python_column.typecode == typecode + assert list(rust_column) == list(python_column) + assert isinstance(rust_blocks[0][-1], list) + assert rust_blocks[0][-1] == python_blocks[0][-1] + + +def test_rust_codec_column_block_container_parity(client_factory, call, consume_stream): + from clickhouse_connect.driver.ctypes import data_conv + + if "driverc" not in data_conv.__name__: + pytest.skip("container parity targets the C-accelerated python readers") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + expressions = [ + "toString(number) AS s", + "CAST(if(number % 2 = 0, NULL, toString(number)) AS Nullable(String)) AS ns", + "CAST(leftPad(toString(number), 6, '0') AS FixedString(6)) AS fs", + "toDate(number) AS d", + "toDateTime(number) AS dt", + "toDateTime64(number, 3, 'UTC') AS dt64", + "CAST(if(number % 2 = 0, NULL, number) AS Nullable(UInt64)) AS n", + "CAST(if(number % 2 = 0, NULL, toDateTime(number)) AS Nullable(DateTime)) AS ndt", + "CAST(toString(number % 3) AS LowCardinality(String)) AS lc", + "CAST(if(number % 2 = 0, NULL, toString(number)) AS LowCardinality(Nullable(String))) AS lcn", + "toUUID(concat(leftPad(lower(hex(number)), 8, '0'), '-1122-3344-5566-778899aabbcc')) AS u", + "toIPv4(toUInt32(number)) AS ip", + "tuple(number, toString(number)) AS tup", + "CAST((toInt64(number), toString(number)), 'Tuple(a Int64, b String)') AS ntup", + "range(number % 4) AS arr", + "number AS num", + ] + query = f"SELECT {', '.join(expressions)} FROM numbers(10)" + + def blocks(client): + output = [] + consume_stream(call(client.query_column_block_stream, query), output.append) + return output + + rust_blocks = blocks(rust_client) + python_blocks = blocks(python_client) + assert len(rust_blocks) == len(python_blocks) == 1 + for rust_column, python_column in zip(rust_blocks[0], python_blocks[0]): + assert type(rust_column) is type(python_column) + assert list(rust_column) == list(python_column) + + +@pytest.mark.parametrize("native_codec", ["rust", "rust_strict"]) +def test_rust_codec_aggregate_function_decode(client_factory, call, native_codec): + client = client_factory(native_codec=native_codec) + result = call( + client.query, + "SELECT countState() AS count_state, " + "sumState(toUInt64(number)) AS sum_state, " + "sumState(CAST(number, 'Nullable(UInt64)')) AS nullable_sum_state, " + "sumState(CAST(NULL, 'Nullable(UInt64)')) AS empty_nullable_sum_state " + "FROM numbers(13)", + ) + + assert result.result_rows == [ + ( + b"\x0d", + (78).to_bytes(8, "little"), + b"\x01" + (78).to_bytes(8, "little"), + b"\x00", + ) + ] + assert [column_type.name for column_type in result.column_types] == [ + "AggregateFunction(count)", + "AggregateFunction(sum, UInt64)", + "AggregateFunction(sum, Nullable(UInt64))", + "AggregateFunction(sum, Nullable(UInt64))", + ] + + +def test_rust_codec_aggregate_function_streaming(client_factory, call, consume_stream): + client = client_factory(native_codec="rust_strict") + query = ( + "SELECT number, sumState(toUInt64(number)) AS state, " + "sumState(CAST(number, 'Nullable(UInt64)')) AS nullable_state " + "FROM numbers(3) GROUP BY number ORDER BY number" + ) + expected = [ + ( + number, + number.to_bytes(8, "little"), + b"\x01" + number.to_bytes(8, "little"), + ) + for number in range(3) + ] + rows = [] + consume_stream(call(client.query_rows_stream, query), rows.append) + + assert rows == expected + + +def test_rust_codec_aggregate_function_np_df(client_factory, call): + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + client = client_factory(native_codec="rust_strict") + query = ( + "SELECT number, sumState(toUInt64(number)) AS state, " + "sumState(CAST(number, 'Nullable(UInt64)')) AS nullable_state " + "FROM numbers(3) GROUP BY number ORDER BY number" + ) + expected_states = [number.to_bytes(8, "little") for number in range(3)] + expected_nullable_states = [b"\x01" + state for state in expected_states] + + array = call(client.query_np, query) + assert array.dtype["state"] == np.dtype("O") + assert array.dtype["nullable_state"] == np.dtype("O") + assert array["state"].tolist() == expected_states + assert array["nullable_state"].tolist() == expected_nullable_states + + frame = call(client.query_df, query) + assert frame["state"].dtype == np.dtype("O") + assert frame["nullable_state"].dtype == np.dtype("O") + pd.testing.assert_series_equal( + frame["state"], + pd.Series(expected_states, name="state", dtype=object), + ) + pd.testing.assert_series_equal( + frame["nullable_state"], + pd.Series(expected_nullable_states, name="nullable_state", dtype=object), + ) + + +@pytest.mark.parametrize("native_codec", ["rust", "rust_strict"]) +def test_rust_codec_unsupported_aggregate_function_decode(client_factory, call, native_codec): + client = client_factory(native_codec=native_codec) + with pytest.raises(NotSupportedError): + call(client.query, "SELECT avgState(number) AS agg FROM numbers(3)") + + +def test_rust_codec_nullable_tuple_decode(client_factory, call): + # The python codec cannot parse Nullable(Tuple), so the rust path is the + # reference here rather than a parity target. + client = client_factory(native_codec="rust_strict") + query = "SELECT if(number % 2 = 0, CAST((number, 'x'), 'Nullable(Tuple(UInt64, String))'), NULL) AS t FROM numbers(4)" + result = call(client.query, query, settings={"enable_nullable_tuple_type": 1}) + assert result.result_rows == [((0, "x"),), (None,), ((2, "x"),), (None,)] + + +def test_rust_codec_nullable_tuple_aggregate_function_decode(client_factory, call): + # Null rows carry server-written placeholder states under the null mask, + # so boundary recovery must stay correct across the masked rows. + client = client_factory(native_codec="rust_strict") + query = "SELECT if(number % 2 = 0, tuple(initializeAggregation('countState', number)), NULL) AS t FROM numbers(4)" + result = call(client.query, query, settings={"enable_nullable_tuple_type": 1}) + assert result.result_rows == [((b"\x01",),), (None,), ((b"\x01",),), (None,)] + + +def test_rust_codec_eligibility_routing(client_factory, call): + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust") + strict_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + + python_df = call(python_client.query_df, "SELECT number FROM numbers(10)") + # numpy/pandas output is now served by the rust codec, so both rust and rust_strict route through it. + pd.testing.assert_frame_equal(call(rust_client.query_df, "SELECT number FROM numbers(10)"), python_df) + pd.testing.assert_frame_equal(call(strict_client.query_df, "SELECT number FROM numbers(10)"), python_df) + + # Timezone contexts remain a Python fallback until type/tz coverage lands. + with pytest.raises(NotSupportedError): + call(strict_client.query, "SELECT toDateTime(number) FROM numbers(3)", query_tz="America/New_York") + + +def test_rust_codec_insert_refusal_surfaces_directly(client_factory, call, caplog): + client = client_factory(native_codec="rust_strict") + context = InsertContext( + "unused_table", + ["v"], + [get_from_name("UInt8")], + [(13,)], + query_formats={"UInt8": "string"}, + ) + + with caplog.at_level(logging.DEBUG, logger="clickhouse_connect.driver.streaming"): + with pytest.raises(NotSupportedError, match="per-column or per-type write formats"): + call(client.data_insert, context) + + # An expected refusal is quiet: no ERROR-level records and no rebuilt second attempt. + streaming_records = [r for r in caplog.records if r.name == "clickhouse_connect.driver.streaming"] + assert not [r for r in streaming_records if r.levelno >= logging.ERROR] + assert sum("Insert producer error" in r.getMessage() for r in streaming_records) == 1 + + +def test_rust_codec_midstream_error_parity(client_factory, call, consume_stream): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = "SELECT number, throwIf(number = 100000) FROM numbers(200000)" + + def run(client): + consume_stream(call(client.query_row_block_stream, query), lambda _block: None) + + with pytest.raises(StreamFailureError): + run(rust_client) + with pytest.raises(StreamFailureError): + run(python_client) + + +def test_rust_codec_columns_only_limit_zero(client_factory, call): + strict_client = client_factory(native_codec="rust_strict") + + result = call(strict_client.query, "SELECT number AS n, toString(number) AS s FROM numbers(10) LIMIT 0") + + assert result.column_names == ("n", "s") + assert [t.name for t in result.column_types] == ["UInt64", "String"] + assert result.result_rows == [] + + +# numpy/pandas dtype parity matrix. Values chosen to exercise nulls, tz materialization, and dtype width. +NP_DF_MATRIX = { + "int128": "toInt128(number) - toInt128('170141183460469231731687303715884105000')", + "uint256": "toUInt256('115792089237316195423570985008687907853269984665640564039457584007913129639000') + number", + "int32": "toInt32(number) - 5", + "uint64": "number", + "float64": "toFloat64(number) / 2", + "bool": "CAST(number % 2 AS Bool)", + "string": "toString(number)", + "fixed_string": "CAST(leftPad(toString(number), 6, '0') AS FixedString(6))", + "date": "toDate(number)", + "date32": "toDate32(number)", + "datetime": "toDateTime(number)", + "datetime_utc": "toDateTime(number, 'UTC')", + "datetime_named_tz": "toDateTime(number, 'America/New_York')", + "datetime64_3": "toDateTime64(number, 3)", + "datetime64_6": "toDateTime64(number, 6)", + "datetime64_9": "toDateTime64(number, 9)", + "datetime64_named_tz": "toDateTime64(number, 3, 'America/New_York')", + "nullable_int": "CAST(if(number % 3 = 0, NULL, toInt32(number)) AS Nullable(Int32))", + "nullable_uint64": "CAST(if(number % 3 = 0, NULL, number) AS Nullable(UInt64))", + "nullable_float": "CAST(if(number % 3 = 0, NULL, toFloat64(number) / 2) AS Nullable(Float64))", + "nullable_string": "CAST(if(number % 3 = 0, NULL, toString(number)) AS Nullable(String))", + "nullable_datetime": "CAST(if(number % 3 = 0, NULL, toDateTime(number)) AS Nullable(DateTime))", + "nullable_datetime64": "CAST(if(number % 3 = 0, NULL, toDateTime64(number, 3)) AS Nullable(DateTime64(3)))", + "low_card_string": "CAST(toString(number % 3) AS LowCardinality(String))", + "low_card_nullable_string": "CAST(if(number % 2 = 0, NULL, toString(number)) AS LowCardinality(Nullable(String)))", + "enum8": "CAST(if(number % 2 = 0, 'a', 'b') AS Enum8('a' = 1, 'b' = 2))", + "enum16": "CAST(if(number % 2 = 0, 'x', 'y') AS Enum16('x' = 100, 'y' = 200))", + "uuid": "toUUID(concat(leftPad(lower(hex(number)), 8, '0'), '-1122-3344-5566-778899aabbcc'))", + "nullable_uuid": ( + "CAST(if(number % 3 = 0, NULL, toUUID(concat(leftPad(lower(hex(number)), 8, '0'), " + "'-1122-3344-5566-778899aabbcc'))) AS Nullable(UUID))" + ), + "decimal64": "toDecimal64(number / 3 - 2, 4)", + "decimal128": "toDecimal128(number / 3 - 2, 10)", + "ipv4": "toIPv4(toUInt32(number * 16909060))", + "ipv6": "toIPv6(concat('2001:db8::', lower(hex(toUInt16(number + 1)))))", + # Non-null nested scalars stay python-native under rust (value-equal to python's numpy scalars). + "array_int": "range(number % 4)", + "array_string": "arrayMap(x -> toString(x), range(number % 4))", + # Nested nulls are refinalized to match the python codec per leaf type (pd.NA, NaN, NaT, numpy scalars). + "array_nullable_int": "arrayMap(x -> if(x % 2 = 0, NULL, toInt64(x)), range(number % 4))", + "array_nullable_string": "arrayMap(x -> if(x % 2 = 0, NULL, toString(x)), range(number % 4))", + "array_nullable_float": "arrayMap(x -> if(x % 2 = 0, NULL, toFloat64(x) / 2), range(number % 4))", + "array_nullable_date": "arrayMap(x -> if(x % 2 = 0, NULL, toDate(x)), range(number % 4))", + "array_nullable_datetime": "arrayMap(x -> if(x % 2 = 0, NULL, toDateTime(x, 'UTC')), range(number % 4))", + "tuple_unnamed": "tuple(number, toString(number))", + "tuple_named": "CAST((toInt64(number), toString(number)), 'Tuple(a Int64, b String)')", + "tuple_nullable_int": "tuple(if(number % 2 = 0, NULL, toInt64(number)), toString(number))", + "map_string_int": "mapFromArrays(arrayMap(x -> concat('k', toString(x)), range(number % 4)), range(number % 4))", + "map_nullable_int": "CAST(map('k', if(number % 2 = 0, NULL, toInt64(number))), 'Map(String, Nullable(Int64))')", + "array_tuple_nullable": "arrayMap(x -> tuple(if(x % 2 = 0, NULL, toInt64(x)), toString(x)), range(number % 4))", + # Non-nullable and null-free temporal leaves are still rewrapped to numpy datetime64 to match python. + "array_date_nonnull": "arrayMap(x -> toDate(x), range(number % 4))", + "tuple_nullable_int_date": "CAST((if(number % 2 = 0, NULL, toInt64(number)), toDate(number)), 'Tuple(a Nullable(Int64), b Date)')", + # SimpleAggregateFunction converts as its element type. Geo aliases take the object exit on both codecs. + "simple_agg_uint64": "CAST(number AS SimpleAggregateFunction(sum, UInt64))", + "simple_agg_string": "CAST(toString(number) AS SimpleAggregateFunction(anyLast, String))", + "simple_agg_date": "CAST(toDate(number + 13) AS SimpleAggregateFunction(anyLast, Date))", + "simple_agg_datetime": "CAST(toDateTime(number, 'UTC') AS SimpleAggregateFunction(anyLast, DateTime('UTC')))", + "point": "(toFloat64(number), toFloat64(number) / 2)::Point", + "ring": "[(toFloat64(number), toFloat64(number) / 2), (toFloat64(number) + 1, toFloat64(number))]::Ring", +} + + +@pytest.mark.parametrize("expr", NP_DF_MATRIX.values(), ids=list(NP_DF_MATRIX)) +def test_rust_codec_np_df_parity(client_factory, call, expr): + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = f"SELECT {expr} AS c FROM numbers(13)" + + rust_np = call(rust_client.query_np, query) + python_np = call(python_client.query_np, query) + assert rust_np.dtype == python_np.dtype + np.testing.assert_array_equal(rust_np, python_np) + + rust_df = call(rust_client.query_df, query) + python_df = call(python_client.query_df, query) + assert rust_df["c"].dtype == python_df["c"].dtype + pd.testing.assert_frame_equal(rust_df, python_df) + + +def test_rust_codec_bfloat16_parity(client_factory, call, client_mode): + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + settings = {"allow_suspicious_low_cardinality_types": 1} + rust_client = client_factory(native_codec="rust_strict", settings=settings) + python_client = client_factory(native_codec="python", settings=settings) + if not python_client.min_version("24.11"): + pytest.skip("BFloat16 requires ClickHouse 24.11+") + + query = ( + "SELECT " + "CAST(toFloat32(number) + 0.1 AS BFloat16) AS bf, " + "CAST(if(number = 1, NULL, toFloat32(number) + 0.1) AS Nullable(BFloat16)) AS nbf, " + "[CAST(toFloat32(number) + 0.1 AS BFloat16)] AS abf, " + "tuple(CAST(toFloat32(number) + 0.1 AS BFloat16), toString(number)) AS tbf, " + "[tuple(CAST(toFloat32(number) + 0.1 AS BFloat16), toUInt8(number))] AS atbf, " + "CAST(toFloat32(number % 2) + 0.1 AS LowCardinality(BFloat16)) AS lcbf " + "FROM numbers(3)" + ) + rust_result = call(rust_client.query, query) + python_result = call(python_client.query, query) + assert rust_result.result_rows == python_result.result_rows + assert [ch_type.name for ch_type in rust_result.column_types] == [ch_type.name for ch_type in python_result.column_types] + + numpy_query = ( + "SELECT " + "CAST(toFloat32(number) + 0.1 AS BFloat16) AS bf, " + "CAST(if(number = 1, NULL, toFloat32(number) + 0.1) AS Nullable(BFloat16)) AS nbf " + "FROM numbers(3)" + ) + rust_np = call(rust_client.query_np, numpy_query) + python_np = call(python_client.query_np, numpy_query) + assert rust_np.dtype == python_np.dtype + np.testing.assert_equal(rust_np, python_np) + + for use_extended_dtypes in (False, True): + rust_df = call(rust_client.query_df, numpy_query, use_extended_dtypes=use_extended_dtypes) + python_df = call(python_client.query_df, numpy_query, use_extended_dtypes=use_extended_dtypes) + pd.testing.assert_frame_equal(rust_df, python_df) + + # Nested nullable BFloat16 leaves densify to np.float32 with NaN in numpy output and pd.NA in + # extended output. The Nullable(Float32) tuple sibling keeps python float/None in numpy output. + nested_query = ( + "SELECT " + "CAST(multiIf(number = 0, [toFloat32(1.5), NULL, toFloat32(-0.0)], number = 1, [NULL], number = 2, [], " + "[toFloat32(nan), toFloat32(79)]), 'Array(Nullable(BFloat16))') AS anbf, " + "CAST((if(number = 0, NULL, toFloat32(1.5)), if(number = 1, NULL, toFloat32(2.5))), " + "'Tuple(Nullable(BFloat16), Nullable(Float32))') AS tnbf " + "FROM numbers(4)" + ) + + def assert_nested_equal(rust_value, python_value): + assert type(rust_value) is type(python_value) + if isinstance(rust_value, (list, tuple)): + assert len(rust_value) == len(python_value) + for rust_item, python_item in zip(rust_value, python_value): + assert_nested_equal(rust_item, python_item) + elif rust_value is python_value: + pass + elif isinstance(rust_value, (float, np.floating)) and np.isnan(rust_value): + assert np.isnan(python_value) + else: + assert rust_value == python_value + + rust_nested_np = call(rust_client.query_np, nested_query) + python_nested_np = call(python_client.query_np, nested_query) + assert_nested_equal(rust_nested_np.tolist(), python_nested_np.tolist()) + assert [type(leaf) for leaf in rust_nested_np[0, 0]] == [np.float32, np.float32, np.float32] + assert np.isnan(rust_nested_np[0, 0][1]) + assert isinstance(rust_nested_np[0, 1][0], np.float32) and np.isnan(rust_nested_np[0, 1][0]) + assert rust_nested_np[1, 1][1] is None + + for use_extended_dtypes in (False, True): + rust_df = call(rust_client.query_df, nested_query, use_extended_dtypes=use_extended_dtypes) + python_df = call(python_client.query_df, nested_query, use_extended_dtypes=use_extended_dtypes) + pd.testing.assert_frame_equal(rust_df, python_df) + cell = rust_df["anbf"].iloc[0] + if use_extended_dtypes: + assert cell[1] is pd.NA + else: + assert np.isnan(cell[1]) + + schema = ( + "id UInt8, bf BFloat16, nbf Nullable(BFloat16), abf Array(BFloat16), " + "tbf Tuple(BFloat16, String), atbf Array(Tuple(BFloat16, UInt8)), " + "lcbf LowCardinality(BFloat16)" + ) + rows = [ + [0, 1.1, None, [1.1, -1.1], (1.1, "user_1"), [(1.1, 13)], 1.1], + [1, -1.1, -1.1, [], (-1.1, "user_2"), [], -1.1], + [2, 13.0, 79.0, [13.0], (13.0, "user_3"), [(13.0, 79)], 1.1], + ] + names = ["id", "bf", "nbf", "abf", "tbf", "atbf", "lcbf"] + rust_table = f"rc_ins_bf16_rust_{client_mode}" + python_table = f"rc_ins_bf16_python_{client_mode}" + + def roundtrip(client, table): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE MergeTree ORDER BY id") + call(client.insert, table, rows, column_names=names) + return call(python_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows + + try: + assert roundtrip(rust_client, rust_table) == roundtrip(python_client, python_table) + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {python_table}") + + +@pytest.mark.parametrize( + "expr", + [ + "arrayMap(x -> if(x % 2 = 0, NULL, toDateTime(x, 'America/Denver')), range(number % 4))", + "tuple(if(number % 2 = 0, NULL, toDateTime64(number, 3, 'America/Denver')), toString(number))", + ], + ids=["array_nullable_datetime_named_tz", "tuple_nullable_datetime64_named_tz"], +) +def test_rust_codec_nested_named_tz_df_parity(client_factory, call, expr): + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = f"SELECT {expr} AS c FROM numbers(13)" + + rust_df = call(rust_client.query_df, query, use_extended_dtypes=True) + python_df = call(python_client.query_df, query, use_extended_dtypes=True) + pd.testing.assert_frame_equal(rust_df, python_df) + + +@pytest.mark.parametrize("type_name", ["DateTime('America/Denver')", "DateTime64(3, 'America/Denver')"]) +def test_rust_codec_nullable_named_tz_df_parity(client_factory, call, type_name): + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = f"SELECT CAST(if(number = 0, NULL, number) AS Nullable({type_name})) AS c FROM numbers(3)" + + rust_df = call(rust_client.query_df, query) + python_df = call(python_client.query_df, query) + pd.testing.assert_frame_equal(rust_df, python_df) + + +@pytest.mark.parametrize("type_name", ["DateTime('America/Denver')", "DateTime64(3, 'America/Denver')"]) +def test_rust_codec_nullable_named_tz_np_object_cells(client_factory, call, type_name): + # query_np object columns keep stdlib datetime cells; the pd.Timestamp wrap is pandas-only. + # The python codec renders these cells as naive-UTC np.datetime64, a pre-existing scalar-type + # divergence, so parity here is on the UTC instants. + np = pytest.importorskip("numpy") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = f"SELECT CAST(if(number = 0, NULL, number) AS Nullable({type_name})) AS c FROM numbers(3)" + + rust_values = call(rust_client.query_np, query).flatten().tolist() + python_values = call(python_client.query_np, query).flatten().tolist() + + assert rust_values[0] is None and python_values[0] is None + for rust_value, python_value in zip(rust_values[1:], python_values[1:]): + assert type(rust_value) is datetime + assert isinstance(python_value, np.datetime64) + assert np.datetime64(rust_value.astimezone(timezone.utc).replace(tzinfo=None)) == python_value + + +def test_rust_codec_np_df_stream_parity(client_factory, call, consume_stream): + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = "SELECT number AS n, toString(number) AS s, toDateTime64(number, 3) AS dt FROM numbers(5000)" + + def np_blocks(client): + blocks = [] + consume_stream(call(client.query_np_stream, query, settings={"max_block_size": 1000}), blocks.append) + return blocks + + def df_blocks(client): + parts = [] + consume_stream(call(client.query_df_stream, query, settings={"max_block_size": 1000}), parts.append) + return parts + + rust_np, python_np = np_blocks(rust_client), np_blocks(python_client) + assert len(rust_np) > 1 # max_block_size forced multiple blocks + for name in ("n", "s", "dt"): + np.testing.assert_array_equal(np.concatenate([b[name] for b in rust_np]), np.concatenate([b[name] for b in python_np])) + + rust_df = pd.concat(df_blocks(rust_client), ignore_index=True) + python_df = pd.concat(df_blocks(python_client), ignore_index=True) + pd.testing.assert_frame_equal(rust_df, python_df) + + +def test_rust_codec_empty_df_parity(client_factory, call): + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = "SELECT number AS n, toString(number) AS s FROM numbers(0)" + pd.testing.assert_frame_equal(call(rust_client.query_df, query), call(python_client.query_df, query)) + + +def test_rust_codec_dt64_unsupported_precision_parity(client_factory, call): + pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = "SELECT toDateTime64(number, 1) AS c FROM numbers(3)" + with pytest.raises(ProgrammingError): + call(rust_client.query_df, query) + with pytest.raises(ProgrammingError): + call(python_client.query_df, query) + + +def test_rust_codec_uuid_df_parity(client_factory, call): + pd = pytest.importorskip("pandas") + python_client = client_factory(native_codec="python") + query = "SELECT toUUID(concat(leftPad(lower(hex(number)), 8, '0'), '-1122-3344-5566-778899aabbcc')) AS u FROM numbers(3)" + python_df = call(python_client.query_df, query) + for codec in ("rust", "rust_strict"): + client = client_factory(native_codec=codec) + pd.testing.assert_frame_equal(call(client.query_df, query), python_df) + + +def test_rust_codec_abandoned_stream_no_read_ahead_thread(client_factory, call, client_mode): + rust_client = client_factory(native_codec="rust_strict") + # The result must exceed the read-ahead queue capacity so the producer thread is still blocked at + # abandonment. A small result the producer fully buffers would exit on its own and hide a close() leak. + stream = call(rust_client.query_column_block_stream, "SELECT number FROM numbers(20000000)") + + if client_mode == "sync": + with stream as blocks: + for _ in blocks: + break + else: + + async def abandon(): + async with stream as blocks: + async for _ in blocks: + break + + call(abandon) + + def read_ahead_threads(): + return [t for t in threading.enumerate() if t.name == "clickhouse-read-ahead" and t.is_alive()] + + deadline = time.time() + 2.0 + while time.time() < deadline and read_ahead_threads(): + time.sleep(0.05) + assert not read_ahead_threads() + + +def _insert_df_roundtrip(client, python_client, call, table, schema, df): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE MergeTree ORDER BY id") + call(client.insert_df, table, df) + return call(python_client.query_df, f"SELECT * FROM {table} ORDER BY id") + + +def test_rust_codec_insert_df_parity(client_factory, call, client_mode): + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + schema = "id UInt32, i Int64, f Float64, b Bool, s String, dt DateTime" + # Exercises _convert_pandas shapes: int list, float to_numpy array, np.bool_ array, string object array, + # and the datetime "int" tick list injected post-init. + df = pd.DataFrame( + { + "id": pd.Series([0, 1, 2], dtype="uint32"), + "i": pd.Series([13, 79, -5], dtype="int64"), + "f": pd.Series([1.5, 2.5, 3.5], dtype="float64"), + "b": pd.Series([True, False, True], dtype="bool"), + "s": pd.Series(["user_1", "user_2", "user_3"]), + "dt": pd.to_datetime(["2020-01-01 00:00:13", "2021-06-15 12:00:00", "2022-12-31 23:59:59"]), + } + ) + rust_table = f"rc_ins_basic_rust_{client_mode}" + py_table = f"rc_ins_basic_py_{client_mode}" + try: + rust_back = _insert_df_roundtrip(rust_client, python_client, call, rust_table, schema, df) + python_back = _insert_df_roundtrip(python_client, python_client, call, py_table, schema, df) + pd.testing.assert_frame_equal(rust_back, python_back) + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {py_table}") + + +def test_rust_codec_insert_df_nullable_parity(client_factory, call, client_mode): + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + schema = "id UInt32, ni Nullable(Int64), nf Nullable(Float64), ns Nullable(String), ndt Nullable(DateTime)" + # Exercises nullable object arrays, NaN->None float lists, and NaT->None datetime tick lists. + df = pd.DataFrame( + { + "id": pd.Series([0, 1, 2], dtype="uint32"), + "ni": pd.array([13, None, 79], dtype="Int64"), + "nf": pd.Series([1.5, np.nan, 3.5]), + "ns": pd.Series(["user_1", None, "user_2"]), + "ndt": pd.to_datetime(["2020-01-01 00:00:13", None, "2022-12-31 23:59:59"]), + } + ) + rust_table = f"rc_ins_nul_rust_{client_mode}" + py_table = f"rc_ins_nul_py_{client_mode}" + try: + rust_back = _insert_df_roundtrip(rust_client, python_client, call, rust_table, schema, df) + python_back = _insert_df_roundtrip(python_client, python_client, call, py_table, schema, df) + pd.testing.assert_frame_equal(rust_back, python_back) + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {py_table}") + + +@pytest.mark.parametrize("native_codec", ["python", "rust_strict"]) +def test_native_codec_insert_df_enum_nan(client_factory, call, client_mode, native_codec): + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + client = client_factory(native_codec=native_codec) + table = f"rc_ins_enum_nan_{native_codec}_{client_mode}" + df = pd.DataFrame({"id": [0, 1, 2], "e": [1, np.nan, 2]}) + + call(client.command, f"DROP TABLE IF EXISTS {table}") + try: + call( + client.command, + f"CREATE TABLE {table} (id UInt8, e Enum8('missing' = 0, 'one' = 1, 'two' = 2)) ENGINE Memory", + ) + call(client.insert_df, table, df) + assert call(client.query, f"SELECT * FROM {table} ORDER BY id").result_rows == [ + (0, "one"), + (1, "missing"), + (2, "two"), + ] + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +@pytest.mark.parametrize("native_codec", ["python", "rust_strict"]) +def test_native_codec_insert_df_enum_int_dtype(client_factory, call, client_mode, native_codec): + pd = pytest.importorskip("pandas") + client = client_factory(native_codec=native_codec) + table = f"fixa_ins_enum_int_{native_codec}_{client_mode}" + df = pd.DataFrame({"id": [0, 1, 2], "e": pd.Series([1, 2, 1], dtype="int64")}) + + call(client.command, f"DROP TABLE IF EXISTS {table}") + try: + call( + client.command, + f"CREATE TABLE {table} (id UInt8, e Enum8('missing' = 0, 'one' = 1, 'two' = 2)) ENGINE Memory", + ) + call(client.insert_df, table, df) + assert call(client.query, f"SELECT * FROM {table} ORDER BY id").result_rows == [ + (0, "one"), + (1, "two"), + (2, "one"), + ] + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +def test_rust_codec_integer_coercion_matrix(client_factory, call, client_mode): + client = client_factory(native_codec="rust_strict") + table = f"rc_ins_integer_coercion_{client_mode}" + schema = ( + "id UInt8, scalar Int32, nullable Nullable(Int32), array Array(Int32), " + "tuple Tuple(Int32, Int32), array_tuple Array(Tuple(Int32, Int32))" + ) + rows = [ + [0, 13.0, None, [13.0, Decimal("79.000")], (13.0, Decimal("79.000")), [(13.0, Decimal("79.000"))]], + [1, Decimal("-13.000"), 79.0, [], (Decimal("-13.000"), 79.0), []], + ] + expected = [ + (0, 13, None, [13, 79], (13, 79), [(13, 79)]), + (1, -13, 79, [], (-13, 79), []), + ] + + call(client.command, f"DROP TABLE IF EXISTS {table}") + try: + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE Memory") + call( + client.insert, + table, + rows, + column_names=["id", "scalar", "nullable", "array", "tuple", "array_tuple"], + ) + assert call(client.query, f"SELECT * FROM {table} ORDER BY id").result_rows == expected + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +@pytest.mark.parametrize( + "value,message", + [ + (13.5, "would lose fractional data; pass an integer value"), + (Decimal("79.5"), "would lose fractional data; pass an integer value"), + ("13", "strings are not accepted; pass an int instead"), + ], +) +def test_rust_codec_integer_coercion_errors(client_factory, call, client_mode, value, message): + client = client_factory(native_codec="rust_strict") + table = f"rc_ins_integer_coercion_error_{client_mode}" + + call(client.command, f"DROP TABLE IF EXISTS {table}") + try: + call(client.command, f"CREATE TABLE {table} (v Int32) ENGINE Memory") + with pytest.raises(DataError, match=message): + call(client.insert, table, [[value]], column_names=["v"]) + assert call(client.query, f"SELECT count() FROM {table}").result_rows == [(0,)] + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +def test_rust_codec_wide_integer_insert_parity(client_factory, call, client_mode): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + schema = ( + "id UInt8, i128 Int128, u128 UInt128, i256 Int256, u256 UInt256, " + "ni Nullable(Int256), a Array(UInt128), t Tuple(Int128, UInt256), " + "at Array(Tuple(Int256, UInt128)), m Map(UInt8, Int256)" + ) + i128_min, i128_max = -(2**127), 2**127 - 1 + u128_max = 2**128 - 1 + i256_min, i256_max = -(2**255), 2**255 - 1 + u256_max = 2**256 - 1 + rows = [ + [0, i128_min, 0, i256_min, 0, None, [], (i128_min, u256_max), [], {}], + [ + 1, + -1, + 2**127, + -1, + 2**255, + i256_max, + [0, u128_max], + (-1, 2**255), + [(i256_min, 0), (-1, u128_max)], + {1: i256_min, 2: -1}, + ], + [ + 2, + i128_max, + u128_max, + i256_max, + u256_max, + i256_min, + [79], + (i128_max, 79), + [(i256_max, 79)], + {13: i256_max}, + ], + ] + names = ["id", "i128", "u128", "i256", "u256", "ni", "a", "t", "at", "m"] + rust_table = f"rc_ins_wide_rust_{client_mode}" + py_table = f"rc_ins_wide_py_{client_mode}" + + def roundtrip(client, table): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE MergeTree ORDER BY id") + call(client.insert, table, rows, column_names=names) + return call(python_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows + + try: + assert roundtrip(rust_client, rust_table) == roundtrip(python_client, py_table) + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {py_table}") + + +def test_rust_codec_wide_integer_string_insert_rejected(client_factory, call, client_mode): + client = client_factory(native_codec="rust_strict") + table = f"rc_ins_wide_str_{client_mode}" + + call(client.command, f"DROP TABLE IF EXISTS {table}") + try: + call(client.command, f"CREATE TABLE {table} (v Int256) ENGINE Memory") + with pytest.raises(DataError, match="strings are not accepted; pass an int instead"): + call(client.insert, table, [["79"]], column_names=["v"]) + assert call(client.query, f"SELECT count() FROM {table}").result_rows == [(0,)] + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +def test_rust_codec_tuple_map_insert_parity(client_factory, call, client_mode): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + schema = ( + "id UInt32, tu Tuple(Int64, String), tn Tuple(a Int64, b Nullable(String)), " + "ms Map(String, Int64), ma Map(String, Array(Int64)), at Array(Tuple(Int64, String))" + ) + rows = [ + [0, (1, "x"), {"a": 5, "b": "named"}, {"k1": 1, "k2": 2}, {"m": [1, 2]}, [(1, "a"), (2, "b")]], + [1, (2, "y"), {"a": 6}, {}, {"n": []}, []], + [2, (3, "z"), {"a": 7, "b": None}, {"k": -1}, {"p": [3]}, [(3, "c")]], + ] + names = ["id", "tu", "tn", "ms", "ma", "at"] + rust_table = f"rc_ins_nested_rust_{client_mode}" + py_table = f"rc_ins_nested_py_{client_mode}" + + def roundtrip(client, table): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE MergeTree ORDER BY id") + call(client.insert, table, rows, column_names=names) + return call(python_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows + + try: + assert roundtrip(rust_client, rust_table) == roundtrip(python_client, py_table) + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {py_table}") + + +def test_rust_codec_nullable_tuple_insert(client_factory, call, client_mode): + # The python codec cannot insert Nullable(Tuple) at all, so the rust path is + # the reference. Requires the true type name to reach the encoder because + # Tuple.insert_name drops the Nullable wrapper. + client = client_factory(native_codec="rust_strict") + table = f"rc_ins_ntup_{client_mode}" + call(client.command, f"DROP TABLE IF EXISTS {table}") + try: + call( + client.command, + f"CREATE TABLE {table} (id UInt32, t Nullable(Tuple(a Int64, b String))) ENGINE Memory", + settings={"enable_nullable_tuple_type": 1}, + ) + call(client.insert, table, [[0, (1, "x")], [1, None], [2, (3, "z")]], column_names=["id", "t"]) + result = call(client.query, f"SELECT * FROM {table} ORDER BY id") + assert result.result_rows == [(0, {"a": 1, "b": "x"}), (1, None), (2, {"a": 3, "b": "z"})] + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +def test_rust_codec_geo_insert_parity(client_factory, call, client_mode): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + schema = "id UInt32, pt Point, rg Ring, ls LineString, pg Polygon, mls MultiLineString, mpg MultiPolygon" + rows = [ + [ + 0, + (3.55, 3.55), + [(5.522, 58.472), (3.55, 3.55)], + [(1.0, 2.0), (3.0, 4.0)], + [[(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)]], + [[(0.0, 0.0), (1.0, 1.0)], [(2.0, 2.0), (3.0, 3.0)]], + [[[(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)]]], + ], + [1, (4.55, 4.55), [(4.55, 4.55)], [(7.0, 8.0)], [[(2.0, 2.0)]], [[(9.0, 9.0)]], [[[(5.0, 5.0)]]]], + ] + names = ["id", "pt", "rg", "ls", "pg", "mls", "mpg"] + rust_table = f"rc_ins_geo_rust_{client_mode}" + py_table = f"rc_ins_geo_py_{client_mode}" + + def roundtrip(client, table): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE MergeTree ORDER BY id") + call(client.insert, table, rows, column_names=names) + return call(python_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows + + try: + assert roundtrip(rust_client, rust_table) == roundtrip(python_client, py_table) == [tuple(r) for r in rows] + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {py_table}") + + +def test_rust_codec_geometry_round_trip_parity(client_factory, call, client_mode): + probe = client_factory(native_codec="python") + try: + resolved_type = call(probe.command, "SELECT toTypeName(defaultValueOfTypeName('Geometry'))") + except DatabaseError as ex: + if ex.name != "UNKNOWN_TYPE": + raise + pytest.skip(f"Geometry is not supported by server {probe.server_version}") + if resolved_type != "Geometry": + pytest.skip(f"Geometry is not supported by server {probe.server_version}") + + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + rust_table = f"rc_geometry_rust_{client_mode}" + python_table = f"rc_geometry_python_{client_mode}" + tagged = [ + typed_variant([(13.0, 23.0), (14.0, 24.0)], "LineString"), + typed_variant([[(31.0, 41.0), (32.0, 42.0)]], "MultiLineString"), + typed_variant([[[(51.0, 61.0)]]], "MultiPolygon"), + typed_variant((71.0, 81.0), "Point"), + typed_variant([[(91.0, 101.0), (92.0, 102.0)]], "Polygon"), + typed_variant([(111.0, 121.0)], "Ring"), + None, + ] + rows = [[index, value] for index, value in enumerate(tagged)] + expected = [(index, value.value if value is not None else None) for index, value in enumerate(tagged)] + expanded_geometry = "Variant(LineString, MultiLineString, MultiPolygon, Point, Polygon, Ring)" + expected_variant_types = [ + ("LineString",), + ("MultiLineString",), + ("MultiPolygon",), + ("Point",), + ("Polygon",), + ("Ring",), + ("None",), + ] + + def create_and_insert(client, table): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} (id UInt8, g Geometry) ENGINE MergeTree ORDER BY id") + call(client.insert, table, rows, column_names=["id", "g"]) + + try: + create_and_insert(rust_client, rust_table) + assert call(rust_client.query, f"SELECT * FROM {rust_table} ORDER BY id").result_rows == expected + assert call(python_client.query, f"SELECT * FROM {rust_table} ORDER BY id").result_rows == expected + assert call(rust_client.query, f"SELECT variantType(g) FROM {rust_table} ORDER BY id").result_rows == expected_variant_types + expanded_result = call( + rust_client.query, + f"SELECT CAST(g, '{expanded_geometry}') FROM {rust_table} ORDER BY id", + settings={"allow_suspicious_variant_types": 1}, + ) + assert expanded_result.column_types[0].name == expanded_geometry + assert expanded_result.result_rows == [(value,) for _, value in expected] + + create_and_insert(python_client, python_table) + assert call(python_client.query, f"SELECT * FROM {python_table} ORDER BY id").result_rows == expected + assert call(rust_client.query, f"SELECT * FROM {python_table} ORDER BY id").result_rows == expected + assert call(python_client.query, f"SELECT variantType(g) FROM {python_table} ORDER BY id").result_rows == expected_variant_types + finally: + call(rust_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {python_table}") + + +def test_rust_codec_nested_insert_parity(client_factory, call, client_mode): + # A single Nested(...) typed column only appears with flatten_nested=0; the + # default splits it into sibling Array columns. Nested reads as a list of + # dicts keyed by the field names in both codecs. + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + schema = "id UInt32, n Nested(sku String, qty UInt32)" + rows = [ + [0, []], + [1, [{"sku": "sku_1", "qty": 5}, {"sku": "sku_2", "qty": 77}]], + [2, [{"sku": "sku_3", "qty": 13}]], + ] + expected = [ + (0, []), + (1, [{"sku": "sku_1", "qty": 5}, {"sku": "sku_2", "qty": 77}]), + (2, [{"sku": "sku_3", "qty": 13}]), + ] + rust_table = f"rc_ins_nested_col_rust_{client_mode}" + py_table = f"rc_ins_nested_col_py_{client_mode}" + + def roundtrip(client, read_client, table): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE MergeTree ORDER BY id", settings={"flatten_nested": 0}) + call(client.insert, table, rows, column_names=["id", "n"]) + return call(read_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows + + try: + # rust->rust exercises encode and decode end to end; python->python is the reference. + assert roundtrip(rust_client, rust_client, rust_table) == expected + assert roundtrip(python_client, python_client, py_table) == expected + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {py_table}") + + +def test_rust_codec_simple_agg_insert_parity(client_factory, call, client_mode): + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + schema = ( + "id UInt32, s SimpleAggregateFunction(anyLast, String), " + "lc SimpleAggregateFunction(anyLast, LowCardinality(String)), " + "n SimpleAggregateFunction(sum, UInt64), f SimpleAggregateFunction(max, Float64), " + "arr SimpleAggregateFunction(groupArrayArray, Array(UInt64))" + ) + rows = [ + [0, "first", "lc_1", 100, 3.5, [1, 2, 3]], + [1, "second", "lc_2", 79, -1.5, []], + ] + names = ["id", "s", "lc", "n", "f", "arr"] + rust_table = f"rc_ins_saf_rust_{client_mode}" + py_table = f"rc_ins_saf_py_{client_mode}" + + def roundtrip(client, table): + call(client.command, f"DROP TABLE IF EXISTS {table}") + call(client.command, f"CREATE TABLE {table} ({schema}) ENGINE MergeTree ORDER BY id") + call(client.insert, table, rows, column_names=names) + return call(python_client.query, f"SELECT * FROM {table} ORDER BY id").result_rows + + try: + assert roundtrip(rust_client, rust_table) == roundtrip(python_client, py_table) == [tuple(r) for r in rows] + finally: + call(python_client.command, f"DROP TABLE IF EXISTS {rust_table}") + call(python_client.command, f"DROP TABLE IF EXISTS {py_table}") + + +def test_rust_codec_aggregate_function_insert(client_factory, call, client_mode): + client = client_factory(native_codec="rust_strict") + table = f"rc_ins_aggregate_function_{client_mode}" + rows = [ + [0, b"\x00", (13).to_bytes(8, "little"), b"\x00"], + [ + 1, + b"\x0d", + (79).to_bytes(8, "little"), + b"\x01" + (-13).to_bytes(8, "little", signed=True), + ], + [ + 2, + b"\x80\x01", + (258).to_bytes(8, "little"), + b"\x80" + (79).to_bytes(8, "little"), + ], + ] + call(client.command, f"DROP TABLE IF EXISTS {table}") + try: + call( + client.command, + f"CREATE TABLE {table} (" + "id UInt8, count_state AggregateFunction(count), " + "sum_state AggregateFunction(sum, UInt64), " + "nullable_sum_state AggregateFunction(sum, Nullable(Int32))) ENGINE Memory", + ) + call( + client.insert, + table, + rows, + column_names=["id", "count_state", "sum_state", "nullable_sum_state"], + ) + result = call( + client.query, + f"SELECT id, finalizeAggregation(count_state), finalizeAggregation(sum_state), " + f"finalizeAggregation(nullable_sum_state) FROM {table} ORDER BY id", + ) + assert result.result_rows == [ + (0, 0, 13, None), + (1, 13, 79, -13), + (2, 128, 258, 79), + ] + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +def test_rust_codec_aggregate_function_insert_failure(client_factory, call, client_mode): + client = client_factory(native_codec="rust_strict") + table = f"rc_ins_aggregate_function_bad_{client_mode}" + call(client.command, f"DROP TABLE IF EXISTS {table}") + try: + call( + client.command, + f"CREATE TABLE {table} (s AggregateFunction(sum, UInt64)) ENGINE Memory", + ) + with pytest.raises(DataError, match="not exactly one valid serialized"): + call(client.insert, table, [[b"\x00" * 7]], column_names=["s"]) + assert call(client.query, f"SELECT count() FROM {table}").result_rows == [(0,)] + finally: + call(client.command, f"DROP TABLE IF EXISTS {table}") + + +def test_rust_codec_midstream_error_df_parity(client_factory, call, consume_stream): + pytest.importorskip("pandas") + rust_client = client_factory(native_codec="rust_strict") + python_client = client_factory(native_codec="python") + query = "SELECT number, throwIf(number = 100000) FROM numbers(200000)" + + def run(client): + consume_stream(call(client.query_df_stream, query), lambda _df: None) + + with pytest.raises(StreamFailureError): + run(rust_client) + with pytest.raises(StreamFailureError): + run(python_client) diff --git a/tests/integration_tests/test_sqlalchemy/test_reflect.py b/tests/integration_tests/test_sqlalchemy/test_reflect.py index 472b00ea..854d0017 100644 --- a/tests/integration_tests/test_sqlalchemy/test_reflect.py +++ b/tests/integration_tests/test_sqlalchemy/test_reflect.py @@ -5,8 +5,9 @@ from sqlalchemy.exc import NoResultFound from clickhouse_connect import common -from clickhouse_connect.cc_sqlalchemy.datatypes.sqltypes import Point, SimpleAggregateFunction, UInt32 +from clickhouse_connect.cc_sqlalchemy.datatypes.sqltypes import Geometry, Point, SimpleAggregateFunction, UInt32 from clickhouse_connect.datatypes.format import clear_all_formats, set_default_formats +from clickhouse_connect.driver.exceptions import DatabaseError def test_basic_reflection(test_engine: Engine): @@ -50,6 +51,31 @@ def test_types_reflection(test_engine: Engine, test_db: str): assert "MergeTree" in table.engine.name +def test_geometry_reflection(test_engine: Engine, test_db: str, test_client): + try: + resolved_type = test_client.command("SELECT toTypeName(defaultValueOfTypeName('Geometry'))") + except DatabaseError as ex: + if ex.name != "UNKNOWN_TYPE": + raise + pytest.skip(f"Geometry is not supported by server {test_client.server_version}") + if resolved_type != "Geometry": + pytest.skip(f"Geometry is not supported by server {test_client.server_version}") + + common.set_setting("invalid_setting_action", "drop") + with test_engine.begin() as conn: + try: + conn.execute(text(f"DROP TABLE IF EXISTS {test_db}.sqlalchemy_geometry_test")) + conn.execute( + text(f"CREATE TABLE {test_db}.sqlalchemy_geometry_test (key UInt32, geometry Geometry) ENGINE MergeTree ORDER BY key") + ) + metadata = db.MetaData(schema=test_db) + table = db.Table("sqlalchemy_geometry_test", metadata, autoload_with=test_engine) + assert table.columns.geometry.type.__class__ == Geometry + assert table.columns.geometry.type.name == "Geometry" + finally: + conn.execute(text(f"DROP TABLE IF EXISTS {test_db}.sqlalchemy_geometry_test")) + + def test_table_exists(test_engine: Engine): common.set_setting("invalid_setting_action", "drop") inspector = inspect(test_engine) diff --git a/tests/unit_tests/test_backend_orchestration.py b/tests/unit_tests/test_backend_orchestration.py index 4fd4b8d2..1a67bcc2 100644 --- a/tests/unit_tests/test_backend_orchestration.py +++ b/tests/unit_tests/test_backend_orchestration.py @@ -264,7 +264,10 @@ def test_client_execute_operation_dispatch(): assert result is client.command.return_value client.command.assert_called_once_with("SELECT version()", settings=None, use_database=False) Client._execute_operation(client, QueryOp("SELECT 1", settings={"max_threads": "4"})) - client.query.assert_called_once_with("SELECT 1", settings={"max_threads": "4"}, query_formats={"String": "string"}) + client.create_query_context.assert_called_once_with(query="SELECT 1", settings={"max_threads": "4"}, query_formats={"String": "string"}) + context = client.create_query_context.return_value + assert context.internal is True + client._query_with_context.assert_called_once_with(context) Client._execute_operation(client, RawQueryOp("SELECT 1", fmt="Native")) client.raw_query.assert_called_once_with("SELECT 1", settings=None, fmt="Native") with pytest.raises(TypeError, match="Unsupported operation type"): @@ -276,8 +279,12 @@ def test_async_client_execute_operation_dispatch(): result = asyncio.run(AsyncClient._execute_operation(client, CommandOp("SELECT version()", use_database=False))) assert result is client.command.return_value client.command.assert_awaited_once_with("SELECT version()", settings=None, use_database=False) + client.create_query_context = Mock() asyncio.run(AsyncClient._execute_operation(client, QueryOp("SELECT 1", settings={"max_threads": "4"}))) - client.query.assert_awaited_once_with("SELECT 1", settings={"max_threads": "4"}, query_formats={"String": "string"}) + client.create_query_context.assert_called_once_with(query="SELECT 1", settings={"max_threads": "4"}, query_formats={"String": "string"}) + context = client.create_query_context.return_value + assert context.internal is True + client._query_with_context.assert_awaited_once_with(context) asyncio.run(AsyncClient._execute_operation(client, RawQueryOp("SELECT 1", fmt="Native"))) client.raw_query.assert_awaited_once_with("SELECT 1", settings=None, fmt="Native") with pytest.raises(TypeError, match="Unsupported operation type"): diff --git a/tests/unit_tests/test_chtypes.py b/tests/unit_tests/test_chtypes.py index ab29c788..3b4443ea 100644 --- a/tests/unit_tests/test_chtypes.py +++ b/tests/unit_tests/test_chtypes.py @@ -1,7 +1,24 @@ from datetime import timedelta, timezone +import pytest + from clickhouse_connect.datatypes.container import Nested from clickhouse_connect.datatypes.registry import get_from_name as gfn +from clickhouse_connect.driver.query import QueryContext + +INTERVAL_TYPES = ( + "IntervalYear", + "IntervalQuarter", + "IntervalMonth", + "IntervalWeek", + "IntervalDay", + "IntervalHour", + "IntervalMinute", + "IntervalSecond", + "IntervalMillisecond", + "IntervalMicrosecond", + "IntervalNanosecond", +) def test_enum_parse(): @@ -58,3 +75,19 @@ def test_datetime_fixed_offset_negative_timezone(): def test_datetime64_fixed_offset_timezone(): dt64_type = gfn("DateTime64(3, 'Fixed/UTC+05:30:00')") assert dt64_type.tzinfo == timezone(timedelta(hours=5, minutes=30)) + + +def test_intervals_use_signed_i64_storage(): + for type_name in INTERVAL_TYPES: + interval_type = gfn(type_name) + assert interval_type.byte_size == 8 + assert interval_type.np_type == " 0 + + +def test_convert_pandas_enum_column_paths(): + np = pytest.importorskip("numpy") + pd = pytest.importorskip("pandas") + enum_type = get_from_name("Enum8('missing' = 0, 'one' = 1, 'two' = 2)") + + # NA-free int dtypes keep the vectorized object-array conversion. + ctx = InsertContext("fake_table", ["e"], [enum_type], pd.DataFrame({"e": pd.Series([1, 2, 1], dtype="int64")})) + column = ctx._block_columns[0] + assert isinstance(column, np.ndarray) + assert list(column) == [1, 2, 1] + + # Float dtypes take the per-row path that maps NaN to the zero code. + ctx = InsertContext("fake_table", ["e"], [enum_type], pd.DataFrame({"e": [1.0, float("nan"), 2.0]})) + column = ctx._block_columns[0] + assert isinstance(column, list) + assert column == [1, 0, 2] diff --git a/tests/unit_tests/test_driver/test_native_read.py b/tests/unit_tests/test_driver/test_native_read.py index 6ee10b87..f4cf87e7 100644 --- a/tests/unit_tests/test_driver/test_native_read.py +++ b/tests/unit_tests/test_driver/test_native_read.py @@ -2,6 +2,7 @@ from uuid import UUID from clickhouse_connect.datatypes import registry +from clickhouse_connect.datatypes.dynamic import typed_variant from clickhouse_connect.driver.insert import InsertContext from clickhouse_connect.driver.query import QueryContext, QueryResult from clickhouse_connect.driver.transform import NativeTransform @@ -106,6 +107,29 @@ def test_point(): assert tuple(python) == tuple(point for point in points) +def test_geometry(): + tagged = [ + typed_variant([(13.0, 23.0), (14.0, 24.0)], "LineString"), + typed_variant([[(31.0, 41.0), (32.0, 42.0)]], "MultiLineString"), + typed_variant([[[(51.0, 61.0)]]], "MultiPolygon"), + typed_variant((71.0, 81.0), "Point"), + typed_variant([[(91.0, 101.0), (92.0, 102.0)]], "Polygon"), + typed_variant([(111.0, 121.0)], "Ring"), + None, + ] + expected = [value.value if value is not None else None for value in tagged] + geometry_type = registry.get_from_name("Geometry") + dest = bytearray() + geometry_type.write_column(tagged, dest, InsertContext("", [], [])) + + source = bytes_source(bytes(dest)) + ctx = QueryContext() + read_state = geometry_type.read_column_prefix(source, ctx) + assert geometry_type.name == "Geometry" + assert registry.get_from_name("GEOMETRY").name == "Geometry" + assert geometry_type.read_column_data(source, len(tagged), ctx, read_state) == expected + + def test_nested(): result = parse_response(bytes_source(NESTED_BINARY)) check_result(result, [{"str1": "one", "int32": 5}, {"str1": "two", "int32": 55}], 2, 0) diff --git a/tests/unit_tests/test_driver/test_native_write.py b/tests/unit_tests/test_driver/test_native_write.py index 7f641776..6bad5361 100644 --- a/tests/unit_tests/test_driver/test_native_write.py +++ b/tests/unit_tests/test_driver/test_native_write.py @@ -1,5 +1,8 @@ +import pytest + from clickhouse_connect.datatypes.registry import get_from_name -from clickhouse_connect.driver.exceptions import ProgrammingError +from clickhouse_connect.driver.exceptions import DataError, ProgrammingError +from clickhouse_connect.driver.insert import InsertContext from tests.helpers import native_insert_block, to_bytes from tests.unit_tests.test_driver.binary import NESTED_BINARY @@ -89,6 +92,22 @@ def test_point(): assert bytes(output) == bytes.fromhex(POINT_OUTPUT) +@pytest.mark.parametrize( + ("points", "message"), + [ + ([(13.0,)], "got 1 at row 0"), + ([(13.0, 23.0, 79.0)], "got 3 at row 0"), + ([(13.0, 23.0), (79.0,)], "got 1 at row 1"), + ([(13.0, 23.0), (13.0, 23.0, 79.0)], "got 3 at row 1"), + ], +) +def test_point_rejects_wrong_coordinate_count(points, message): + point_type = get_from_name("Point") + context = InsertContext("table", ["value"], [point_type], [[(13.0, 23.0)]]) + with pytest.raises(DataError, match=message): + point_type.write_column(points, bytearray(), context) + + def test_nested(): data = [ ([],), diff --git a/tests/unit_tests/test_driver/test_numeric.py b/tests/unit_tests/test_driver/test_numeric.py index 8dec0caa..da3b78df 100644 --- a/tests/unit_tests/test_driver/test_numeric.py +++ b/tests/unit_tests/test_driver/test_numeric.py @@ -1,4 +1,5 @@ import unittest +from math import copysign, isnan from unittest.mock import Mock import numpy as np @@ -93,6 +94,24 @@ def test_numpy_fastpath(self): expected = np.array(expected, dtype=np.float32) self.assertTrue(np.all(out == expected)) + def test_nullable_special_values_are_preserved(self): + data_in = [None, -0.0, float("inf"), float("-inf"), float("nan")] + null_map = bytes([1, 0, 0, 0, 0]) + dest = bytearray(null_map) + self.null_bf16._write_column_binary(data_in, dest, self.ins_ctx) + + # Little-endian bf16 words: null 0x0000, -0.0 0x8000, inf 0x7F80, -inf 0xFF80, nan 0x7FC0. + expected_words = b"\x00\x00\x00\x80\x80\x7f\x80\xff\xc0\x7f" + self.assertEqual(bytes(dest), null_map + expected_words) + + source = create_test_source(bytes(dest)) + out = self.null_bf16._read_nullable_column(source, len(data_in), self.qry_ctx, None) + + self.assertIsNone(out[0]) + self.assertEqual(copysign(1.0, out[1]), -1.0) + self.assertEqual(out[2:4], [float("inf"), float("-inf")]) + self.assertTrue(isnan(out[4])) + def test_encoded_size(self): "Test we only encode 2 bytes per float" sample = [10.0, 20.0, 30.0] diff --git a/tests/unit_tests/test_driver/test_rustcodec.py b/tests/unit_tests/test_driver/test_rustcodec.py new file mode 100644 index 00000000..4f8c0a81 --- /dev/null +++ b/tests/unit_tests/test_driver/test_rustcodec.py @@ -0,0 +1,848 @@ +import logging +import struct +import sys +from datetime import datetime, timezone +from zoneinfo import ZoneInfo + +import pytest + +from clickhouse_connect import common +from clickhouse_connect.common import _native_codec_env_default +from clickhouse_connect.datatypes.registry import get_from_name +from clickhouse_connect.driver import rustcodec +from clickhouse_connect.driver.exceptions import ( + DataError, + NotSupportedError, + ProgrammingError, + StreamClosedError, + StreamFailureError, +) +from clickhouse_connect.driver.insert import InsertContext +from clickhouse_connect.driver.query import QueryContext +from clickhouse_connect.driver.rustcodec import ( + _make_native_transform, + _rust_query_ineligible_reason, + _RustNativeTransform, + resolve_native_codec, +) +from clickhouse_connect.driver.transform import NativeTransform +from tests.helpers import TAGGED_EXCEPTION_BODY, TAGGED_EXCEPTION_TAG + + +class FakeSource: + def __init__(self, chunks, exception_tag=None): + self.gen = iter(chunks) + self.exception_tag = exception_tag + self.closed = False + + def close(self): + self.closed = True + + +class _PresentCore: + """Stand-in for the compiled module with a compatible binding API.""" + + __version__ = "0.1.0" + BINDING_API_VERSION = rustcodec.REQUIRED_BINDING_API_VERSION + + +def eligible_ctx(**kwargs) -> QueryContext: + """Baseline rust-eligible context (mirrors a real UTC-server client).""" + return QueryContext(apply_server_tz=True, server_tz=timezone.utc, **kwargs) + + +def _uint64_ctx(data, block_size=None, query_formats=None, column_formats=None) -> InsertContext: + return InsertContext( + "fake_table", + ["key"], + [get_from_name("UInt64")], + data, + block_size=block_size, + query_formats=query_formats, + column_formats=column_formats, + ) + + +def _json_ctx(data, type_name="JSON") -> InsertContext: + return InsertContext("fake_table", ["payload"], [get_from_name(type_name)], data) + + +@pytest.fixture(name="restore_native_codec") +def restore_native_codec_fixture(): + original = common.get_setting("native_codec") + yield + common.set_setting("native_codec", original) + + +@pytest.fixture(name="restore_naive_datetime_insert") +def restore_naive_datetime_insert_fixture(): + original = common.get_setting("naive_datetime_insert") + yield + common.set_setting("naive_datetime_insert", original) + + +@pytest.fixture(name="clean_formats") +def clean_formats_fixture(): + from clickhouse_connect.datatypes.base import ch_read_formats, ch_write_formats + + read_snapshot = dict(ch_read_formats) + write_snapshot = dict(ch_write_formats) + yield + ch_read_formats.clear() + ch_read_formats.update(read_snapshot) + ch_write_formats.clear() + ch_write_formats.update(write_snapshot) + + +# --- Resolution -------------------------------------------------------------- + + +def test_resolve_native_codec_kwarg_beats_common_setting(monkeypatch, restore_native_codec): + monkeypatch.setitem(sys.modules, "_ch_core", _PresentCore) + common.set_setting("native_codec", "rust") + assert resolve_native_codec("python") == "python" + assert resolve_native_codec(None) == "rust" + + +@pytest.mark.parametrize( + ("env_value", "expected"), + [("rust", "rust"), (" RUST ", "rust")], +) +def test_native_codec_env_default_valid(monkeypatch, env_value, expected): + monkeypatch.setenv("CLICKHOUSE_CONNECT_NATIVE_CODEC", env_value) + assert _native_codec_env_default() == expected + + +def test_native_codec_env_default_unset(monkeypatch): + monkeypatch.delenv("CLICKHOUSE_CONNECT_NATIVE_CODEC", raising=False) + assert _native_codec_env_default() == "python" + + +def test_native_codec_env_default_invalid_warns(monkeypatch, caplog): + monkeypatch.setenv("CLICKHOUSE_CONNECT_NATIVE_CODEC", "bogus") + with caplog.at_level(logging.WARNING): + assert _native_codec_env_default() == "python" + assert any("CLICKHOUSE_CONNECT_NATIVE_CODEC" in record.getMessage() for record in caplog.records) + + +def test_resolve_native_codec_invalid_kwarg(): + with pytest.raises(ProgrammingError): + resolve_native_codec("bogus") + + +# --- Availability ------------------------------------------------------------ + + +@pytest.fixture(name="reset_version_log") +def reset_version_log_fixture(monkeypatch): + monkeypatch.setattr(rustcodec, "_versions_logged", False) + + +@pytest.mark.parametrize("codec", ["rust", "rust_strict"]) +def test_resolve_rust_raises_when_module_missing(monkeypatch, codec): + monkeypatch.setitem(sys.modules, "_ch_core", None) + with pytest.raises(NotSupportedError) as excinfo: + resolve_native_codec(codec) + message = str(excinfo.value) + assert "compiled _ch_core extension module" in message + assert 'pip install "clickhouse-connect[rust]"' in message + + +class _StaleCore: + __version__ = "0.0.9" + BINDING_API_VERSION = rustcodec.REQUIRED_BINDING_API_VERSION - 1 + + +class _NoApiVersionCore: + __version__ = "0.0.5" + + +@pytest.mark.parametrize("core", [_StaleCore, _NoApiVersionCore], ids=["api_stale", "api_missing"]) +@pytest.mark.parametrize("codec", ["rust", "rust_strict"]) +def test_resolve_rust_raises_when_binding_api_too_old(monkeypatch, codec, core): + monkeypatch.setitem(sys.modules, "_ch_core", core) + with pytest.raises(NotSupportedError) as excinfo: + resolve_native_codec(codec) + message = str(excinfo.value) + assert f"clickhouse-connect-core version {core.__version__}" in message + assert "pip install --upgrade clickhouse-connect-core" in message + + +@pytest.mark.parametrize( + "api_version", + [rustcodec.REQUIRED_BINDING_API_VERSION, rustcodec.REQUIRED_BINDING_API_VERSION + 1], + ids=["required", "newer"], +) +def test_resolve_rust_accepts_compatible_binding_api(monkeypatch, reset_version_log, api_version): + class _Core: + __version__ = "0.1.0" + BINDING_API_VERSION = api_version + + monkeypatch.setitem(sys.modules, "_ch_core", _Core) + assert resolve_native_codec("rust") == "rust" + assert resolve_native_codec("rust_strict") == "rust_strict" + + +def test_built_binding_matches_required_api_version(): + """The compiled extension and the driver ship from one tree, so their + binding API versions must be bumped together. Enforced wherever the + extension is built; skipped on pure-Python environments.""" + core = pytest.importorskip("_ch_core") + assert core.BINDING_API_VERSION == rustcodec.REQUIRED_BINDING_API_VERSION + + +def test_resolve_rust_logs_versions_once(monkeypatch, caplog, reset_version_log): + monkeypatch.setitem(sys.modules, "_ch_core", _PresentCore) + with caplog.at_level(logging.INFO, logger="clickhouse_connect"): + resolve_native_codec("rust") + resolve_native_codec("rust_strict") + records = [r for r in caplog.records if "clickhouse-connect-core" in r.getMessage()] + assert len(records) == 1 + message = records[0].getMessage() + assert "native_codec=rust" in message + assert _PresentCore.__version__ in message + assert f"binding API {_PresentCore.BINDING_API_VERSION}" in message + assert common.version() in message + + +@pytest.mark.parametrize(("codec", "strict"), [("rust", False), ("rust_strict", True)]) +def test__make_native_transform_rust_variants(monkeypatch, codec, strict): + monkeypatch.setitem(sys.modules, "_ch_core", _PresentCore) + transform = _make_native_transform(codec) + assert isinstance(transform, _RustNativeTransform) + assert transform.strict is strict + assert transform.threaded_insert is True + + +def test__make_native_transform_python(): + transform = _make_native_transform("python") + assert isinstance(transform, NativeTransform) + assert transform.threaded_insert is False + + +# --- Eligibility ------------------------------------------------------------- + + +def _ctx_response_tz(): + ctx = eligible_ctx() + ctx.set_response_tz(ZoneInfo("America/New_York")) + return ctx + + +@pytest.mark.parametrize( + ("builder", "reason"), + [ + pytest.param(lambda: eligible_ctx(query_formats={"Int*": "string"}), "query_formats", id="query_formats"), + pytest.param(lambda: eligible_ctx(column_formats={"x": "string"}), "column_formats", id="column_formats"), + pytest.param(lambda: eligible_ctx(use_none=False), "use_none=False", id="use_none"), + pytest.param(lambda: eligible_ctx(encoding="latin-1"), "custom encoding", id="encoding"), + pytest.param(lambda: eligible_ctx(query_tz="America/New_York"), "query_tz", id="query_tz"), + pytest.param(lambda: eligible_ctx(column_tzs={"x": "America/New_York"}), "column_tzs", id="column_tzs"), + pytest.param(lambda: eligible_ctx(tz_mode="aware"), "tz_mode", id="tz_mode_aware"), + pytest.param(lambda: eligible_ctx(tz_mode="schema"), "tz_mode", id="tz_mode_schema"), + pytest.param(_ctx_response_tz, "server timezone header", id="response_tz"), + pytest.param( + lambda: QueryContext(apply_server_tz=True, server_tz=ZoneInfo("America/New_York")), + "ambient timezone", + id="non_utc_server", + ), + ], +) +def test_rust_query_ineligible(builder, reason): + assert _rust_query_ineligible_reason(builder()) == reason + + +@pytest.mark.parametrize( + "builder", + [ + pytest.param(eligible_ctx, id="baseline"), + pytest.param(lambda: eligible_ctx(column_oriented=True), id="column_oriented"), + pytest.param(lambda: eligible_ctx(streaming=True), id="streaming"), + pytest.param(lambda: eligible_ctx(rename_response_column="remove_prefix"), id="renamer"), + pytest.param(lambda: eligible_ctx(use_numpy=True), id="use_numpy"), + pytest.param(lambda: eligible_ctx(use_numpy=True, as_pandas=True), id="as_pandas"), + ], +) +def test_rust_query_eligible(builder): + assert _rust_query_ineligible_reason(builder()) is None + + +def test_rust_query_ineligible_pyarrow_missing(monkeypatch): + monkeypatch.setattr(rustcodec.options, "arrow", None, raising=False) + assert _rust_query_ineligible_reason(eligible_ctx(use_numpy=True)) == "pyarrow not installed" + # Non-numpy queries never touch the Arrow exit, so a missing pyarrow is irrelevant. + assert _rust_query_ineligible_reason(eligible_ctx()) is None + + +def test_rust_query_ineligible_global_read_format(clean_formats): + from clickhouse_connect.datatypes.format import set_read_format + + set_read_format("IPv4", "string") + assert _rust_query_ineligible_reason(eligible_ctx()) == "global read format override" + + +# --- Routing ----------------------------------------------------------------- + + +def test_strict_ineligible_raises_and_closes_source(): + src = FakeSource([]) + with pytest.raises(NotSupportedError): + _RustNativeTransform(strict=True).parse_response(src, eligible_ctx(use_none=False)) + assert src.closed is True + + +def test_strict_global_read_format_raises_and_closes_source(clean_formats): + from clickhouse_connect.datatypes.format import set_read_format + + set_read_format("IPv4", "string") + src = FakeSource([]) + with pytest.raises(NotSupportedError): + _RustNativeTransform(strict=True).parse_response(src, eligible_ctx()) + assert src.closed is True + + +def test_non_strict_ineligible_delegates_to_python_and_logs_reason(monkeypatch, caplog): + sentinel = object() + monkeypatch.setattr(NativeTransform, "parse_response", staticmethod(lambda source, context: sentinel)) + src = FakeSource([]) + with caplog.at_level(logging.INFO, logger="clickhouse_connect"): + result = _RustNativeTransform(strict=False).parse_response(src, eligible_ctx(use_none=False)) + assert result is sentinel + assert src.closed is False + assert "fallback to Python for query: use_none=False" in caplog.text + + +# --- Insert build (FakeCore) ------------------------------------------------- + + +def test_build_insert_probe_not_implemented_non_strict(monkeypatch): + calls = [] + + class FakeCore: + @staticmethod + def encode_native_block(*args): + calls.append(args) + raise NotImplementedError("unsupported") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + rust_ctx = _uint64_ctx([(13,), (79,)], block_size=1) + python_ctx = _uint64_ctx([(13,), (79,)], block_size=1) + + rust_out = b"".join(_RustNativeTransform(strict=False).build_insert(rust_ctx)) + python_out = b"".join(NativeTransform.build_insert(python_ctx)) + + assert rust_out == python_out + assert len(calls) == 1 # only the probe ran + assert rust_ctx.insert_exception is None + + +def test_build_insert_probe_not_implemented_strict(monkeypatch): + class FakeCore: + @staticmethod + def encode_native_block(*args): + raise NotImplementedError("unsupported") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + with pytest.raises(NotSupportedError): + _RustNativeTransform(strict=True).build_insert(_uint64_ctx([(13,)])) + + +def test_build_insert_mid_block_failure(monkeypatch): + calls = {"n": 0} + + class FakeCore: + @staticmethod + def encode_native_block(names, type_names, columns, row_count, prefix): + calls["n"] += 1 + if row_count == 0: + return b"" + if calls["n"] == 2: + return b"rust_block" + raise ValueError("late") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + ctx = _uint64_ctx([(13,), (79,)], block_size=1) + + chunks = list(_RustNativeTransform(strict=False).build_insert(ctx)) + + assert chunks == [b"rust_block", b"INTERNAL EXCEPTION WHILE SERIALIZING"] + # Binding ValueErrors surface as DataError with the binding's message. + assert isinstance(ctx.insert_exception, DataError) + assert str(ctx.insert_exception) == "late" + assert isinstance(ctx.insert_exception.__cause__, ValueError) + + +def test_build_insert_global_write_format_strict_raises(monkeypatch, clean_formats): + from clickhouse_connect.datatypes.format import set_write_format + + class FakeCore: + @staticmethod + def encode_native_block(*args): + raise AssertionError("encoder must not run when a write format override is set") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + set_write_format("IPv4", "string") + with pytest.raises(NotSupportedError): + _RustNativeTransform(strict=True).build_insert(_uint64_ctx([(13,)])) + + +def test_build_insert_global_write_format_non_strict_falls_back(monkeypatch, clean_formats): + from clickhouse_connect.datatypes.format import set_write_format + + class FakeCore: + @staticmethod + def encode_native_block(*args): + raise AssertionError("encoder must not run when a write format override is set") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + set_write_format("IPv4", "string") + rust_ctx = _uint64_ctx([(13,), (79,)], block_size=1) + python_ctx = _uint64_ctx([(13,), (79,)], block_size=1) + + rust_out = b"".join(_RustNativeTransform(strict=False).build_insert(rust_ctx)) + python_out = b"".join(NativeTransform.build_insert(python_ctx)) + assert rust_out == python_out + + +@pytest.mark.parametrize( + "type_name", + [ + "DateTime", + "Nullable(DateTime64(6))", + "Array(DateTime)", + "Tuple(DateTime64(6), String)", + "Map(String, DateTime)", + "Variant(DateTime, String)", + "Nested(ts DateTime64(6))", + ], +) +def test_build_insert_naive_datetime_server_strict_raises(monkeypatch, restore_naive_datetime_insert, type_name): + class FakeCore: + @staticmethod + def encode_native_block(*args): + raise AssertionError("encoder must not run for server-timezone datetime inserts") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + common.set_setting("naive_datetime_insert", "server") + context = InsertContext("fake_table", ["value"], [get_from_name(type_name)]) + + with pytest.raises(NotSupportedError, match='naive_datetime_insert="server"'): + _RustNativeTransform(strict=True).build_insert(context) + + +def test_build_insert_naive_datetime_server_non_strict_falls_back(monkeypatch, restore_naive_datetime_insert): + class FakeCore: + @staticmethod + def encode_native_block(*args): + raise AssertionError("encoder must not run for server-timezone datetime inserts") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + common.set_setting("naive_datetime_insert", "server") + value = datetime(2025, 7, 15, 12, 34, 56, 250306) + column_names = ["dt", "nested"] + column_types = [get_from_name("DateTime"), get_from_name("Array(Tuple(DateTime64(6), String))")] + rows = [(value, [(value, "user_1")])] + rust_ctx = InsertContext("fake_table", column_names, column_types, rows, server_tz=ZoneInfo("America/Denver")) + python_ctx = InsertContext("fake_table", column_names, column_types, rows, server_tz=ZoneInfo("America/Denver")) + + rust_out = b"".join(_RustNativeTransform(strict=False).build_insert(rust_ctx)) + python_out = b"".join(NativeTransform.build_insert(python_ctx)) + + assert rust_out == python_out + + +@pytest.mark.parametrize("strict", [False, True]) +def test_build_insert_json_uses_rust(monkeypatch, strict): + calls = [] + + class FakeCore: + @staticmethod + def encode_native_block(names, type_names, columns, row_count, prefix): + calls.append(row_count) + return b"rust_block" + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + rows = [({"id": 13},), ({"name": "user_1"},)] + + chunks = list(_RustNativeTransform(strict=strict).build_insert(_json_ctx(rows))) + + assert chunks == [b"rust_block"] + assert calls == [0, 2] # probe then one block + + +def _qbit_ctx(data, type_name) -> InsertContext: + return InsertContext("fake_table", ["vec"], [get_from_name(type_name)], data) + + +@pytest.mark.parametrize("type_name", ["QBit(Float32, 8)", "Array(QBit(Float32, 8))"]) +def test_build_insert_probe_fallback_qbit(monkeypatch, type_name): + class FakeCore: + @staticmethod + def encode_native_block(*args): + raise NotImplementedError("unsupported") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + vector = [0.5, 1.5, -2.0, 13.0, 79.0, 0.0, -1.25, 3.75] + value = vector if type_name.startswith("QBit") else [vector] + rows = [(value,), (value,)] + + rust_out = b"".join(_RustNativeTransform(strict=False).build_insert(_qbit_ctx(rows, type_name))) + python_out = b"".join(NativeTransform.build_insert(_qbit_ctx(rows, type_name))) + assert rust_out == python_out + + with pytest.raises(NotSupportedError): + _RustNativeTransform(strict=True).build_insert(_qbit_ctx(rows, type_name)) + + +_USER_FORMATS = [ + pytest.param({"column_formats": {"absent_col": "string"}}, id="column_format"), + pytest.param({"query_formats": {"IPv4": "string"}}, id="query_format"), +] + + +@pytest.mark.parametrize("fmt", _USER_FORMATS) +def test_build_insert_user_format_strict_raises(monkeypatch, fmt): + class FakeCore: + @staticmethod + def encode_native_block(*args): + raise AssertionError("encoder must not run when a user format is set") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + with pytest.raises(NotSupportedError): + _RustNativeTransform(strict=True).build_insert(_uint64_ctx([(13,)], **fmt)) + + +@pytest.mark.parametrize("fmt", _USER_FORMATS) +def test_build_insert_user_format_non_strict_falls_back(monkeypatch, fmt): + class FakeCore: + @staticmethod + def encode_native_block(*args): + raise AssertionError("encoder must not run when a user format is set") + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + rust_out = b"".join(_RustNativeTransform(strict=False).build_insert(_uint64_ctx([(13,), (79,)], block_size=1, **fmt))) + python_out = b"".join(NativeTransform.build_insert(_uint64_ctx([(13,), (79,)], block_size=1, **fmt))) + assert rust_out == python_out + + +def test_build_insert_datetime_dataframe_routes_rust(monkeypatch): + pd = pytest.importorskip("pandas") + calls = [] + + class FakeCore: + @staticmethod + def encode_native_block(names, type_names, columns, row_count, prefix): + calls.append((names, type_names, columns, row_count, prefix)) + return b"rust_block" + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + df = pd.DataFrame({"ts": pd.to_datetime(["2020-01-01 00:00:13", "2020-01-01 00:01:19"])}) + ctx = InsertContext("fake_table", ["ts"], [get_from_name("DateTime")], df) + + chunks = list(_RustNativeTransform(strict=True).build_insert(ctx)) + + # _convert_pandas injects the "int" hint into column_formats, but the compiled dicts stay empty, + # so rust_strict still serves the insert instead of raising. + assert ctx.column_formats == {"ts": "int"} + assert ctx.col_simple_formats == {} + assert chunks == [b"rust_block"] + assert len(calls) == 2 # probe + one block + + +def test_build_insert_success(monkeypatch): + calls = [] + + class FakeCore: + @staticmethod + def encode_native_block(names, type_names, columns, row_count, prefix): + calls.append((names, type_names, columns, row_count, prefix)) + return b"rust_block" + + monkeypatch.setitem(sys.modules, "_ch_core", FakeCore) + ctx = _uint64_ctx([(13,), (79,)]) + + chunks = list(_RustNativeTransform(strict=False).build_insert(ctx)) + + assert calls[0][3] == 0 # probe carries row_count 0 + assert chunks == [b"rust_block"] + assert ctx.insert_exception is None + + +# --- Decode error taxonomy (fake decoder) ------------------------------------ + + +class _FakeBatch: + def __init__(self, names, type_names, columns=None, error=None): + self.column_names = names + self.column_type_names = type_names + self._columns = columns + self._error = error + + def to_python_columns(self, typed_numeric=False, tuple_columns=None): + del typed_numeric + if self._error is not None: + raise self._error + if tuple_columns is None: + return self._columns + return [tuple(col) if flag else col for flag, col in zip(tuple_columns, self._columns)] + + def to_python_rows(self): + return list(zip(*self.to_python_columns())) + + +def _fake_decoder_core(feed_batches=(), finish_batches=(), feed_error=None): + class _FakeStreamDecoder: + def __init__(self, has_block_info=False): + self._feed = list(feed_batches) + self._finish = list(finish_batches) + + def feed(self, chunk): + if feed_error is not None: + raise feed_error + out, self._feed = self._feed, [] + return out + + def finish(self): + return self._finish + + class _FakeCore: + StreamDecoder = _FakeStreamDecoder + + return _FakeCore + + +@pytest.mark.parametrize( + ("error", "exception_tag", "chunk", "expected"), + [ + (NotImplementedError("Unsupported ClickHouse type 'X'"), None, b"random block bytes", NotSupportedError), + (ValueError("Malformed payload: bad bytes for column 'x'"), None, b"random block bytes", DataError), + ( + NotImplementedError("Unsupported ClickHouse type 'X'"), + None, + b"prefix Code: 62 DB::Exception boom", + StreamFailureError, + ), + ( + ValueError("Malformed payload: bad bytes for column 'x'"), + None, + b"prefix Code: 62 DB::Exception boom", + StreamFailureError, + ), + ( + NotImplementedError("Unsupported ClickHouse type 'X'"), + "T", + b"prefix Code: 62 DB::Exception boom", + NotSupportedError, + ), + ( + ValueError("Malformed payload: bad bytes for column 'x'"), + "T", + b"prefix Code: 62 DB::Exception boom", + DataError, + ), + ], + ids=[ + "unsupported_plain", + "malformed_plain", + "unsupported_untagged_server_error", + "malformed_untagged_server_error", + "unsupported_tagged_no_false_positive", + "malformed_tagged_no_false_positive", + ], +) +def test_decode_feed_error_disambiguation(monkeypatch, error, exception_tag, chunk, expected): + core = _fake_decoder_core(feed_error=error) + monkeypatch.setitem(sys.modules, "_ch_core", core) + src = FakeSource([chunk], exception_tag=exception_tag) + with pytest.raises(expected) as excinfo: + _RustNativeTransform(strict=True).parse_response(src, eligible_ctx()) + if expected is DataError: + assert str(excinfo.value) == str(error) + assert src.closed is True + + +def test_decode_feed_not_implemented_maps_to_not_supported(monkeypatch): + core = _fake_decoder_core(feed_error=NotImplementedError("python object exit")) + monkeypatch.setitem(sys.modules, "_ch_core", core) + src = FakeSource([b"random block bytes"]) + with pytest.raises(NotSupportedError): + _RustNativeTransform(strict=True).parse_response(src, eligible_ctx()) + assert src.closed is True + + +def test_decode_later_block_unsupported_closes_source(monkeypatch): + batch0 = _FakeBatch(["a"], ["Int32"], columns=[[13]]) + batch1 = _FakeBatch(["a"], ["Int32"], error=NotImplementedError("python object exit")) + monkeypatch.setitem(sys.modules, "_ch_core", _fake_decoder_core(feed_batches=[batch0, batch1])) + src = FakeSource([b"chunk"]) + result = _RustNativeTransform(strict=True).parse_response(src, eligible_ctx(streaming=True)) + with pytest.raises(NotSupportedError), result.column_block_stream as stream: + list(stream) + assert src.closed is True + + +def test_decode_early_abandonment_closes_source(monkeypatch): + batch0 = _FakeBatch(["a"], ["Int32"], columns=[[13]]) + batch1 = _FakeBatch(["a"], ["Int32"], columns=[[79]]) + monkeypatch.setitem(sys.modules, "_ch_core", _fake_decoder_core(feed_batches=[batch0, batch1])) + src = FakeSource([b"chunk"]) + result = _RustNativeTransform(strict=True).parse_response(src, eligible_ctx(streaming=True)) + with result.column_block_stream as stream: + for _ in stream: + break + assert src.closed is True + + +def test_decode_buffered_unsupported_raises(monkeypatch): + batch0 = _FakeBatch(["a"], ["Int32"], columns=[[13]]) + batch1 = _FakeBatch(["a"], ["Int32"], error=NotImplementedError("python object exit")) + monkeypatch.setitem(sys.modules, "_ch_core", _fake_decoder_core(feed_batches=[batch0, batch1])) + src = FakeSource([b"chunk"]) + with pytest.raises(NotSupportedError): + _RustNativeTransform(strict=True).parse_response(src, eligible_ctx()) + assert src.closed is True + + +def test_decode_buffered_malformed_fill_raises_data_error(monkeypatch): + batch0 = _FakeBatch(["a"], ["Int32"], columns=[[13]]) + batch1 = _FakeBatch(["a"], ["Int32"], error=ValueError("Malformed payload: bad cell")) + monkeypatch.setitem(sys.modules, "_ch_core", _fake_decoder_core(feed_batches=[batch0, batch1])) + src = FakeSource([b"chunk"]) + with pytest.raises(DataError, match="Malformed payload: bad cell"): + _RustNativeTransform(strict=True).parse_response(src, eligible_ctx()) + assert src.closed is True + + +def test_decode_streaming_malformed_fill_raises_data_error(monkeypatch): + batch0 = _FakeBatch(["a"], ["Int32"], columns=[[13]]) + batch1 = _FakeBatch(["a"], ["Int32"], error=ValueError("Malformed payload: bad cell")) + monkeypatch.setitem(sys.modules, "_ch_core", _fake_decoder_core(feed_batches=[batch0, batch1])) + src = FakeSource([b"chunk"]) + result = _RustNativeTransform(strict=True).parse_response(src, eligible_ctx(streaming=True)) + with pytest.raises(DataError, match="Malformed payload: bad cell"), result.column_block_stream as stream: + list(stream) + assert src.closed is True + + +def test_decode_buffered_result_is_materialized(monkeypatch): + batch0 = _FakeBatch(["a"], ["Int32"], columns=[[13]]) + batch1 = _FakeBatch(["a"], ["Int32"], columns=[[79]]) + monkeypatch.setitem(sys.modules, "_ch_core", _fake_decoder_core(feed_batches=[batch0, batch1])) + src = FakeSource([b"chunk"]) + result = _RustNativeTransform(strict=True).parse_response(src, eligible_ctx()) + assert src.closed is True + assert result.result_rows == [(13,), (79,)] + assert result.result_columns == [[13, 79]] + with pytest.raises(StreamClosedError): + _ = result.column_block_stream + + +def test_decode_buffered_column_oriented(monkeypatch): + batch0 = _FakeBatch(["a", "b"], ["Int32", "Int32"], columns=[[13], [1]]) + batch1 = _FakeBatch(["a", "b"], ["Int32", "Int32"], columns=[[79], [2]]) + monkeypatch.setitem(sys.modules, "_ch_core", _fake_decoder_core(feed_batches=[batch0, batch1])) + src = FakeSource([b"chunk"]) + result = _RustNativeTransform(strict=True).parse_response(src, eligible_ctx(column_oriented=True)) + assert result.result_columns == [[13, 79], [1, 2]] + assert result.result_rows == [(13, 1), (79, 2)] + # Repeated access in either order stays stable across the multi-batch buffers. + assert result.result_columns == [[13, 79], [1, 2]] + assert result.result_rows == [(13, 1), (79, 2)] + + +def test_decode_buffered_result_columns_uses_direct_column_exit(monkeypatch): + class ColumnsOnlyBatch(_FakeBatch): + def to_python_rows(self): + raise AssertionError("row exit must not run for first result_columns access") + + batch = ColumnsOnlyBatch(["a", "b"], ["Int32", "Int32"], columns=[[13, 79], [1, 2]]) + monkeypatch.setitem(sys.modules, "_ch_core", _fake_decoder_core(feed_batches=[batch])) + + result = _RustNativeTransform(strict=True).parse_response(FakeSource([b"chunk"]), eligible_ctx()) + + assert result.result_columns == [[13, 79], [1, 2]] + + +# --- Decode against real _ch_core -------------------------------------------- + + +@pytest.fixture(name="ch_core") +def ch_core_fixture(): + return pytest.importorskip("_ch_core") + + +def test_decode_basic(ch_core): + data = ch_core.encode_native_block(["a", "b"], ["Int32", "String"], [[13, 79], ["user_1", "user_2"]], 2, None) + chunks = [data[i : i + 8] for i in range(0, len(data), 8)] + src = FakeSource(chunks) + + result = _RustNativeTransform(strict=True).parse_response(src, eligible_ctx()) + + assert result.result_rows == [(13, "user_1"), (79, "user_2")] + assert result.column_names == ("a", "b") + assert [t.name for t in result.column_types] == ["Int32", "String"] + + +def test_decode_empty_stream(ch_core): + result = _RustNativeTransform(strict=True).parse_response(FakeSource([]), eligible_ctx()) + assert result.result_rows == [] + + +def test_decode_truncated(ch_core): + data = ch_core.encode_native_block(["a", "b"], ["Int32", "String"], [[13, 79], ["user_1", "user_2"]], 2, None) + with pytest.raises(StreamFailureError): + _RustNativeTransform(strict=True).parse_response(FakeSource([data[:-4]]), eligible_ctx()) + + +def test_decode_tagged_exception(ch_core): + src = FakeSource([TAGGED_EXCEPTION_BODY], exception_tag=TAGGED_EXCEPTION_TAG) + with pytest.raises(StreamFailureError) as excinfo: + _RustNativeTransform(strict=True).parse_response(src, eligible_ctx()) + assert str(excinfo.value) == "Big bam occurred right while reading the data" + + +def _varint_str(value: str) -> bytes: + encoded = value.encode() + return bytes([len(encoded)]) + encoded + + +def _dynamic_shared_block(cells: list) -> bytes: + """V2 Dynamic block whose only child is SharedVariant.""" + body = bytearray(struct.pack("= 3 keeps at the scalar unit. + default_dtype = "timedelta64[ns]" if int(pd.__version__.split(".", 1)[0]) < 3 else dtype + assert default_result.dtype == np.dtype(default_dtype) + np.testing.assert_array_equal(default_result, np.array([-5, "NaT", 79], dtype=dtype).astype(default_dtype)) + + +def test_nullable_time64_query_np_preserves_nanosecond_scalars(monkeypatch): + np = pytest.importorskip("numpy") + pa = pytest.importorskip("pyarrow") + ch_type = get_from_name("Nullable(Time64(9))") + wire_values = pa.array([-1, None, 1], type=pa.int64()) + monkeypatch.setattr(rustnumpy, "_arrow_column", lambda _table, _index: wire_values) + + result = rustnumpy._make_time_convert(ch_type)(None, None, 0) + + assert isinstance(result, list) + assert isinstance(result[0], np.timedelta64) + assert result == [np.timedelta64(-1, "ns"), None, np.timedelta64(1, "ns")] + + +@pytest.mark.parametrize( + ("type_name", "rows", "unit"), + [ + ("Array(Time)", [[], [-5], [13, 79]], "s"), + ("Array(Time64(3))", [[], [-5], [13, 79]], "ms"), + ("Array(Time64(6))", [[], [-5], [13, 79]], "us"), + ("Array(Time64(9))", [[], [-5], [13, 79]], "ns"), + ("Array(Array(Time))", [[[13], []], [[-5, 79]]], "s"), + ], +) +def test_array_time_converter_slices_flat_values(monkeypatch, type_name, rows, unit): + np = pytest.importorskip("numpy") + pa = pytest.importorskip("pyarrow") + ch_type = get_from_name(type_name) + depth, leaf = rustnumpy._array_time_leaf(ch_type) + arrow_type = pa.int32() if unit == "s" else pa.int64() + for _ in range(depth): + arrow_type = pa.large_list(arrow_type) + wire = pa.array(rows, type=arrow_type) + monkeypatch.setattr(rustnumpy, "_arrow_column", lambda _table, _index: wire) + context = type("Context", (), {"as_pandas": False, "use_extended_dtypes": False})() + + result = rustnumpy._make_array_time_convert(leaf, depth, context)(None, None, 0) + + def expect(node): + if isinstance(node, int): + return np.timedelta64(node, unit) + return [expect(value) for value in node] + + assert result == expect(rows) + assert all(isinstance(row, list) for row in result) + + +def test_array_nullable_time_converter_null_policy(monkeypatch): + np = pytest.importorskip("numpy") + pa = pytest.importorskip("pyarrow") + ch_type = get_from_name("Array(Nullable(Time64(9)))") + depth, leaf = rustnumpy._array_time_leaf(ch_type) + wire = pa.array([[None, 1], [-1]], type=pa.large_list(pa.int64())) + monkeypatch.setattr(rustnumpy, "_arrow_column", lambda _table, _index: wire) + + np_context = type("Context", (), {"as_pandas": False, "use_extended_dtypes": False})() + result = rustnumpy._make_array_time_convert(leaf, depth, np_context)(None, None, 0) + assert result == [[None, np.timedelta64(1, "ns")], [np.timedelta64(-1, "ns")]] + + ext_context = type("Context", (), {"as_pandas": True, "use_extended_dtypes": True})() + result = rustnumpy._make_array_time_convert(leaf, depth, ext_context)(None, None, 0) + assert np.isnat(result[0][0]) + assert result[0][0].dtype == np.dtype("timedelta64[ns]") + assert result[0][1] == np.timedelta64(1, "ns") + + +def test_low_card_time_converter_decodes_dictionary(monkeypatch): + np = pytest.importorskip("numpy") + pa = pytest.importorskip("pyarrow") + ch_type = get_from_name("LowCardinality(Nullable(Time))") + wire = pa.DictionaryArray.from_arrays(pa.array([1, None, 0], type=pa.int32()), pa.array([13, -5], type=pa.int32())) + monkeypatch.setattr(rustnumpy, "_arrow_column", lambda _table, _index: wire) + + result = rustnumpy._make_low_card_time_convert(ch_type, as_pandas=False)(None, None, 0) + assert result == [np.timedelta64(-5, "s"), None, np.timedelta64(13, "s")] + + values = rustnumpy._make_low_card_time_convert(ch_type, as_pandas=True)(None, None, 0) + assert values.dtype == np.dtype("timedelta64[s]") + np.testing.assert_array_equal(values, np.array([-5, "NaT", 13], dtype="timedelta64[s]")) + + +def test_nested_time64_converter_preserves_nanoseconds(): + np = pytest.importorskip("numpy") + ch_type = get_from_name("Array(Tuple(Time, Nullable(Time64(9))))") + + class _Batch: + @staticmethod + def column_data(_index, *, raw_time_ticks): + assert raw_time_ticks is True + return [[(-5, -1), (13, None), (79, 1)]] + + context = type("Context", (), {"as_pandas": False, "use_extended_dtypes": False, "use_numpy": True})() + result = rustnumpy._make_nested_time_convert(ch_type, context)(None, _Batch(), 0) + + assert result == [ + [ + (np.timedelta64(-5, "s"), np.timedelta64(-1, "ns")), + (np.timedelta64(13, "s"), None), + (np.timedelta64(79, "s"), np.timedelta64(1, "ns")), + ] + ] + + +@pytest.mark.parametrize( + ("type_name", "expected"), + [ + ("Array(Date)", True), + ("Array(DateTime('UTC'))", False), + ("Array(DateTime64(3, 'UTC'))", False), + ("Array(Nullable(DateTime64(3, 'UTC')))", True), + ("Tuple(Time, Nullable(Int64))", True), + ("Tuple(String, Nullable(Time))", False), + ], +) +def test_needs_refinalize_only_gates_transforming_leaves(type_name, expected): + assert rustnumpy._needs_refinalize(get_from_name(type_name)) is expected + + +@pytest.mark.parametrize( + ("type_name", "column"), + [ + ("Tuple(Nullable(Int64), String)", [(None, "untouched"), (13, "also untouched")]), + ("Map(String, Nullable(Int64))", [{"untouched": None}, {"also untouched": 13}]), + ], +) +def test_refinalize_skips_unaffected_sibling_leaves(monkeypatch, type_name, column): + pd = pytest.importorskip("pandas") + + def fail_finalize(*_args): + raise AssertionError("unaffected String leaf was finalized") + + monkeypatch.setattr(String, "_finalize_column", fail_finalize) + context = QueryContext(use_numpy=True, as_pandas=True, use_extended_dtypes=True) + + result = rustnumpy._refinalize_leaves(get_from_name(type_name), column, context) + + leaf = result[0][0] if isinstance(result[0], tuple) else result[0]["untouched"] + assert leaf is pd.NA + + +class _ObjectExitBatch: + def __init__(self, column): + self._column = column + + def column_data(self, _index): + return list(self._column) + + +@pytest.mark.parametrize("type_name", ["DateTime('America/Denver')", "DateTime64(3, 'America/Denver')"]) +def test_object_convert_numpy_keeps_stdlib_aware_datetimes(type_name): + pytest.importorskip("numpy") + tz = ZoneInfo("America/Denver") + values = [None, datetime.fromtimestamp(13, tz), datetime.fromtimestamp(79, tz)] + + converter = rustnumpy._build_converter(get_from_name(f"Nullable({type_name})"), QueryContext(use_numpy=True)) + result = converter(None, _ObjectExitBatch(values), 0) + + assert list(result) == values + assert type(result[1]) is datetime + + +@pytest.mark.parametrize("type_name", ["DateTime('America/Denver')", "DateTime64(3, 'America/Denver')"]) +def test_object_convert_numpy_aware_datetimes_without_pandas(monkeypatch, type_name): + pytest.importorskip("numpy") + monkeypatch.setattr(rustnumpy.options, "pd", None) + tz = ZoneInfo("America/Denver") + values = [None, datetime.fromtimestamp(13, tz)] + + converter = rustnumpy._build_converter(get_from_name(f"Nullable({type_name})"), QueryContext(use_numpy=True)) + result = converter(None, _ObjectExitBatch(values), 0) + + assert list(result) == values + + +@pytest.mark.parametrize("type_name", ["DateTime('America/Denver')", "DateTime64(3, 'America/Denver')"]) +def test_refinalize_nullable_named_timezone_datetimes(type_name): + pd = pytest.importorskip("pandas") + timezone = ZoneInfo("America/Denver") + value = datetime.fromtimestamp(1, timezone) + ch_type = get_from_name(f"Array(Nullable({type_name}))") + context = QueryContext(use_numpy=True, as_pandas=True, use_extended_dtypes=True) + + result = rustnumpy._refinalize_leaves(ch_type, [[None, value]], context) + + assert result[0][0] is pd.NaT + assert result[0][1] == pd.Timestamp(value) diff --git a/tests/unit_tests/test_driver/test_temporal.py b/tests/unit_tests/test_driver/test_temporal.py index c8942ce8..f3edd4cc 100644 --- a/tests/unit_tests/test_driver/test_temporal.py +++ b/tests/unit_tests/test_driver/test_temporal.py @@ -1,6 +1,6 @@ import array import unittest -from datetime import time, timedelta +from datetime import datetime, time, timedelta, timezone from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -8,7 +8,7 @@ import pandas as pd from clickhouse_connect.datatypes.base import TypeDef -from clickhouse_connect.datatypes.temporal import Time, Time64 +from clickhouse_connect.datatypes.temporal import DateTime64, Time, Time64 from clickhouse_connect.driver.exceptions import ProgrammingError @@ -205,6 +205,40 @@ def test_write_time_format(self): self.assertEqual(dest, array.array("i", expected).tobytes()) +class TestDateTime64DataType(unittest.TestCase): + def setUp(self): + self.dt64_type = DateTime64(TypeDef(values=(3,))) + self.insert_ctx = SimpleNamespace(column_name="test_datetime64_col") + + def test_write_pre_epoch_fractional_datetime_ticks(self): + column = [ + datetime(1969, 12, 31, 23, 59, 59, 500000, tzinfo=timezone.utc), + datetime(1969, 12, 31, 23, 59, 59, 999999, tzinfo=timezone.utc), + datetime(1970, 1, 1, 0, 0, 0, 999000, tzinfo=timezone.utc), + ] + expected_ticks = [-500, -1, 999] + dest = bytearray() + + with patch.object(self.dt64_type, "write_format", return_value="native"): + self.dt64_type._write_column_binary(column, dest, self.insert_ctx) + + self.assertEqual(dest, array.array("q", expected_ticks).tobytes()) + + def test_write_pre_epoch_fractional_string_ticks(self): + column = [ + "1969-12-31T23:59:59.500000+00:00", + "1969-12-31T23:59:59.999999+00:00", + "1970-01-01T00:00:00.999000+00:00", + ] + expected_ticks = [-500, -1, 999] + dest = bytearray() + + with patch.object(self.dt64_type, "write_format", return_value="native"): + self.dt64_type._write_column_binary(column, dest, self.insert_ctx) + + self.assertEqual(dest, array.array("q", expected_ticks).tobytes()) + + class TestTime64DataType(unittest.TestCase): """Tests for the Time64 type.""" @@ -270,10 +304,25 @@ def test_timedelta_negative_and_out_of_range(self): t6 = self.make(6) td_neg = -timedelta(seconds=2) self.assertEqual(t6._timedelta_to_ticks(td_neg), -2_000_000) + self.assertEqual(t6._timedelta_to_ticks(timedelta(seconds=-5, microseconds=-500_000)), -5_500_000) + self.assertEqual(self.make(3)._timedelta_to_ticks(timedelta(microseconds=-1)), 0) + self.assertEqual(t6._timedelta_to_ticks(timedelta(microseconds=-1)), -1) + self.assertEqual(self.make(9)._timedelta_to_ticks(timedelta(microseconds=-1)), -1_000) td_big = timedelta(hours=1000) with self.assertRaises(ValueError): t6._timedelta_to_ticks(td_big) + def test_numpy_timedelta_rescales_to_ticks(self): + self.assertEqual(self.make(3)._timedelta_to_ticks(np.timedelta64(-1, "ns")), 0) + self.assertEqual(self.make(6)._timedelta_to_ticks(np.timedelta64(-1, "ns")), 0) + self.assertEqual(self.make(9)._timedelta_to_ticks(np.timedelta64(-1, "ns")), -1) + self.assertEqual(self.make(9)._timedelta_to_ticks(np.timedelta64(-1, "ps")), 0) + self.assertEqual(self.make(6)._timedelta_to_ticks(np.timedelta64(1_000_079, "us")), 1_000_079) + with self.assertRaises(ValueError): + self.make(9)._timedelta_to_ticks(np.timedelta64(18_446_744_073, "s")) + with self.assertRaisesRegex(ValueError, "Unsupported NumPy timedelta64 unit"): + self.make(9)._timedelta_to_ticks(np.timedelta64(1, "M")) + # ------------------------------------------------------------------ # Binary write/read and numpy # ------------------------------------------------------------------ @@ -292,6 +341,28 @@ def test_round_trip_write_read_nullable_native(self): res = t6._read_column_binary(mock_source, 2, self.base_read_ctx, None) self.assertEqual(res, ["000:00:01.000001", "000:00:00.000000"]) + def test_nullable_numpy_nat_writes_null_map(self): + t6 = self.make(6, nullable=True) + for column in ( + np.array([1, "NaT"], dtype="timedelta64[us]"), + [pd.Timedelta(microseconds=1), pd.NaT], + ): + dest = bytearray() + with patch(f"{Time64.__module__}.write_array", _dummy_write_array): + t6.write_column_data(column, dest, SimpleNamespace(column_name="c")) + self.assertEqual(dest, bytes([0, 1]) + array.array("q", [1, 0]).tobytes()) + + def test_nullable_timedelta_then_nat_writes_null_map(self): + t6 = self.make(6, nullable=True) + for column in ( + [timedelta(microseconds=1), pd.NaT], + [timedelta(microseconds=1), np.timedelta64("NaT")], + ): + dest = bytearray() + with patch(f"{Time64.__module__}.write_array", _dummy_write_array): + t6.write_column_data(column, dest, SimpleNamespace(column_name="c")) + self.assertEqual(dest, bytes([0, 1]) + array.array("q", [1, 0]).tobytes()) + def test_read_numpy_format(self): t6 = self.make(6) ticks = [0, 10**6] diff --git a/tests/unit_tests/test_qbit_type.py b/tests/unit_tests/test_qbit_type.py index d4063486..c0205c6a 100644 --- a/tests/unit_tests/test_qbit_type.py +++ b/tests/unit_tests/test_qbit_type.py @@ -238,6 +238,53 @@ def test_transpose_large_dimension(): assert result == pytest.approx(vector, rel=1e-6) +def test_transpose_reverses_native_byte_groups(monkeypatch): + """Native stores eight-element groups from the end of each plane.""" + qbit = get_from_name("QBit(Float32, 9)") + vector = [-1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, -8.0, -9.0] + + # Plane zero is the sign bit. Element 8 occupies bit zero of the first + # byte; elements 0 and 7 occupy bits zero and seven of the second byte. + assert qbit._transpose_row(vector)[0] == b"\x01\x81" + + monkeypatch.setattr(options, "np", None) + transposed = qbit._transpose_row(vector) + assert transposed[0] == b"\x01\x81" + assert qbit._untranspose_row(transposed) == pytest.approx(vector) + + +def test_transpose_reverses_three_native_byte_groups(monkeypatch): + """A three-byte plane pins complete group reversal, not a two-byte swap.""" + qbit = get_from_name("QBit(Float32, 17)") + vector = [0.0] * 17 + for index in (0, 7, 8, 9, 16): + vector[index] = -0.0 + + assert qbit._transpose_row(vector)[0] == b"\x01\x03\x81" + monkeypatch.setattr(options, "np", None) + transposed = qbit._transpose_row(vector) + assert transposed[0] == b"\x01\x03\x81" + assert [math.copysign(1.0, value) for value in qbit._untranspose_row(transposed)] == [ + -1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + -1.0, + -1.0, + -1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + -1.0, + ] + + def test_untranspose_manual_bit_pattern(): """Test untranspose with a manually constructed bit pattern""" qbit = get_from_name("QBit(Float32, 2)") diff --git a/tests/unit_tests/test_sqlalchemy/test_types.py b/tests/unit_tests/test_sqlalchemy/test_types.py index 9a81307d..8415fd9b 100644 --- a/tests/unit_tests/test_sqlalchemy/test_types.py +++ b/tests/unit_tests/test_sqlalchemy/test_types.py @@ -12,6 +12,7 @@ Decimal, Enum, Float64, + Geometry, Int64, LowCardinality, Nullable, @@ -39,6 +40,15 @@ def test_sqla(): assert "Enum8('value2' = 5, 'value1' = 7)" == enum._compiler_dispatch(None) +def test_geometry(): + geometry = sqla_type_from_name("Geometry") + assert geometry.__class__ == Geometry + assert geometry.python_type is object + assert geometry.name == "Geometry" + assert geometry._compiler_dispatch(None) == "Geometry" + assert sqla_type_from_name("GEOMETRY").name == "Geometry" + + def test_nullable(): nullable = Nullable(Int64) assert nullable.__class__ == Int64 diff --git a/tests/unit_tests/test_streaming_source.py b/tests/unit_tests/test_streaming_source.py index bbbc1195..d9fd61fa 100644 --- a/tests/unit_tests/test_streaming_source.py +++ b/tests/unit_tests/test_streaming_source.py @@ -1,5 +1,6 @@ import asyncio import gzip +import logging import threading import time import zlib @@ -9,10 +10,12 @@ import pytest from clickhouse_connect.driver.compression import _zstd_compress -from clickhouse_connect.driver.exceptions import OperationalError +from clickhouse_connect.driver.exceptions import NotSupportedError, OperationalError from clickhouse_connect.driver.streaming import ( + ReadAheadSource, StreamingInsertSource, StreamingResponseSource, + _SyncStreamingInsertSource, ) @@ -459,6 +462,74 @@ async def test_streaming_insert_error_propagation(): # Should have received first chunk before error assert chunks == [b"chunk1"] + assert isinstance(context.insert_exception, ValueError) + assert str(context.insert_exception) == "Serialization error" + + +def test_sync_streaming_insert_error_propagation(): + """Test that insert producer errors are propagated to sync consumer.""" + transform = FailingTransform() + context = MockContext() + + source = _SyncStreamingInsertSource(transform, context) + source.start_producer() + + chunks = [] + with pytest.raises(ValueError, match="Serialization error"): + for chunk in source.gen: + chunks.append(chunk) + + assert isinstance(context.insert_exception, ValueError) + + assert chunks == [b"chunk1"] + + +class RefusingTransform: + """Mock transform whose build_insert raises a deterministic driver refusal at call time.""" + + @staticmethod + def build_insert(context): + raise NotSupportedError("strict refusal") + + +def _streaming_error_records(caplog): + records = [r for r in caplog.records if r.name == "clickhouse_connect.driver.streaming"] + return [r for r in records if r.levelno >= logging.ERROR] + + +@pytest.mark.asyncio +async def test_streaming_insert_driver_error_logs_debug(caplog): + """Deterministic driver refusals propagate without ERROR-level noise.""" + context = MockContext() + loop = asyncio.get_running_loop() + + source = StreamingInsertSource(RefusingTransform(), context, loop) + with caplog.at_level(logging.DEBUG, logger="clickhouse_connect.driver.streaming"): + source.start_producer() + with pytest.raises(NotSupportedError, match="strict refusal"): + async for _chunk in source.async_generator(): + pass + await source.close() + + assert isinstance(context.insert_exception, NotSupportedError) + assert not _streaming_error_records(caplog) + assert any("Insert producer error" in r.getMessage() for r in caplog.records) + + +def test_sync_streaming_insert_driver_error_logs_debug(caplog): + """Deterministic driver refusals propagate without ERROR-level noise.""" + context = MockContext() + + source = _SyncStreamingInsertSource(RefusingTransform(), context) + with caplog.at_level(logging.DEBUG, logger="clickhouse_connect.driver.streaming"): + source.start_producer() + with pytest.raises(NotSupportedError, match="strict refusal"): + for _chunk in source.gen: + pass + + assert isinstance(context.insert_exception, NotSupportedError) + assert not _streaming_error_records(caplog) + assert any("Insert producer error" in r.getMessage() for r in caplog.records) @pytest.mark.asyncio @@ -485,5 +556,78 @@ async def test_streaming_insert_backpressure(): assert received == chunks +class MockByteSource: + """Mock ByteSource for ReadAheadSource tests.""" + + def __init__(self, chunks, exception_tag=None, error=None): + self._chunks = list(chunks) + self._error = error + self.exception_tag = exception_tag + self.closed = False + + @property + def gen(self): + yield from self._chunks + if self._error is not None: + raise self._error + + def close(self): + self.closed = True + + +def test_read_ahead_chunk_order(): + src = MockByteSource([b"a", b"b", b"c"]) + read_source = ReadAheadSource(src) + assert list(read_source.gen) == [b"a", b"b", b"c"] + read_source.close() + assert src.closed is True + + +def test_read_ahead_gen_cached(): + read_source = ReadAheadSource(MockByteSource([b"a"])) + assert read_source.gen is read_source.gen + read_source.close() + + +def test_read_ahead_error_forwarded_verbatim(): + err = ValueError("boom") + src = MockByteSource([b"a", b"b"], error=err) + read_source = ReadAheadSource(src) + collected = [] + with pytest.raises(ValueError, match="boom") as excinfo: + for chunk in read_source.gen: + collected.append(chunk) + assert collected == [b"a", b"b"] # forwarded in stream order, error last + assert excinfo.value is err # verbatim, not re-wrapped + read_source.close() + + +def test_read_ahead_tagged_exception_chunk_unchanged(): + # exception_tag is delegated and the chunk passes through verbatim so the codec scanner can find it. + src = MockByteSource([b"prefix __exception__T\r\nboom\r\n"], exception_tag="T") + read_source = ReadAheadSource(src) + assert read_source.exception_tag == "T" + assert list(read_source.gen) == [b"prefix __exception__T\r\nboom\r\n"] + read_source.close() + + +def test_read_ahead_close_during_block_terminates_thread(): + # A large producer fills the bounded queue while the consumer reads nothing, so the producer blocks in _put. + src = MockByteSource([bytes([i % 256]) for i in range(500)]) + read_source = ReadAheadSource(src, maxsize=2) + assert next(read_source.gen) == b"\x00" + read_source.close() + assert src.closed is True + assert read_source._thread.is_alive() is False + + +def test_read_ahead_join_on_close(): + src = MockByteSource([b"a", b"b"]) + read_source = ReadAheadSource(src) + read_source.close() + assert read_source._thread.is_alive() is False + assert src.closed is True + + if __name__ == "__main__": pytest.main([__file__, "-v"])