diff --git a/packages/google-cloud-spanner-arrow/.coveragerc b/packages/google-cloud-spanner-arrow/.coveragerc new file mode 100644 index 000000000000..fa1733575f90 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/.coveragerc @@ -0,0 +1,13 @@ +[run] +branch = True +source = + google_cloud_spanner_arrow + +[report] +show_missing = True +exclude_lines = + pragma: no cover + def __repr__ + raise AssertionError + raise NotImplementedError + if __name__ == .__main__.: diff --git a/packages/google-cloud-spanner-arrow/.flake8 b/packages/google-cloud-spanner-arrow/.flake8 new file mode 100644 index 000000000000..bba7a13cbf4f --- /dev/null +++ b/packages/google-cloud-spanner-arrow/.flake8 @@ -0,0 +1,10 @@ +[flake8] +exclude = + .git, + __pycache__, + build, + dist, + .nox, + venv* +max-line-length = 88 +extend-ignore = E203, E501 diff --git a/packages/google-cloud-spanner-arrow/.repo-metadata.json b/packages/google-cloud-spanner-arrow/.repo-metadata.json new file mode 100644 index 000000000000..60365aa99503 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/.repo-metadata.json @@ -0,0 +1,9 @@ +{ + "client_documentation": "https://googleapis.dev/python/google-cloud-spanner-arrow/latest", + "distribution_name": "google-cloud-spanner-arrow", + "language": "python", + "library_type": "OTHER", + "name": "google-cloud-spanner-arrow", + "release_level": "preview", + "repo": "googleapis/google-cloud-python" +} diff --git a/packages/google-cloud-spanner-arrow/BUILDING.md b/packages/google-cloud-spanner-arrow/BUILDING.md new file mode 100644 index 000000000000..e39b2235fb65 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/BUILDING.md @@ -0,0 +1,33 @@ +# Building `google-cloud-spanner-arrow` + +## Requirements + +- Python >= 3.10 +- C99 compatible compiler (`gcc`, `clang`, or MSVC on Windows) +- `setuptools >= 64.0.0`, `wheel` +- `pyarrow >= 14.0.0` + +## Local Development & Installation + +### Editable install with C extension: +```bash +pip install -e . +``` + +### Pure Python fallback build: +```bash +SPANNER_ARROW_PURE_PYTHON=1 pip install -e . +``` + +## Running Tests + +```bash +pytest tests +``` + +## Multi-Platform Wheels + +Multi-platform wheels are built using the scripts in `scripts/`: +- Linux: `scripts/manylinux/build.sh` +- macOS: `scripts/osx/build.sh` +- Windows: `scripts\windows\build.bat` diff --git a/packages/google-cloud-spanner-arrow/CHANGELOG.md b/packages/google-cloud-spanner-arrow/CHANGELOG.md new file mode 100644 index 000000000000..2d70a218dd8b --- /dev/null +++ b/packages/google-cloud-spanner-arrow/CHANGELOG.md @@ -0,0 +1,8 @@ +# Changelog + +## 0.1.0 (2026-08-16) + +- Initial release of `google-cloud-spanner-arrow`. +- Native C extension leveraging `nanoarrow` and the Arrow C Data Interface (`RecordBatch._import_from_c`). +- GIL-free parsing of Spanner rows and partial result set streams into Arrow RecordBatches. +- Dynamic integration with `google-cloud-spanner`'s `StreamedResultSet`. diff --git a/packages/google-cloud-spanner-arrow/CONTRIBUTING.md b/packages/google-cloud-spanner-arrow/CONTRIBUTING.md new file mode 100644 index 000000000000..f16a38266a67 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/CONTRIBUTING.md @@ -0,0 +1,6 @@ +# Contributing to google-cloud-spanner-arrow + +1. Fork the repository and create your branch from `main`. +2. Ensure you have standard build tools (C compiler, Python >= 3.10). +3. Run tests with `pytest tests`. +4. Submit a pull request. diff --git a/packages/google-cloud-spanner-arrow/DESIGN.md b/packages/google-cloud-spanner-arrow/DESIGN.md new file mode 100644 index 000000000000..6ba1dbfe127e --- /dev/null +++ b/packages/google-cloud-spanner-arrow/DESIGN.md @@ -0,0 +1,137 @@ +# Python Spanner Apache Arrow Accelerator + +## Overview +This document describes the design, architecture, and benchmark results for `google-cloud-spanner-arrow`, an optional C-extension companion package for the Google Cloud Spanner Python client. The package provides direct, accelerated conversion from incoming protobuf `PartialResultSet` stream messages into Apache Arrow record batches and tables without intermediate Python object instantiation. + +Project implementation and long-term maintenance are straightforward and low-risk because the library directly follows established patterns in the Google Cloud Python client ecosystem: +- It reuses the proven companion package architecture, CI/CD build scripts, and multi-OS binary wheel distribution infrastructure of `google-crc32c`. +- It implements the standard Arrow and DataFrame query result API conventions established by `google-cloud-bigquery`. +- It maintains a complete pure-Python fallback ensuring zero breakage on unsupported platforms or environments without a C compiler. + +--- + +## Problem Description + +Data engineering, analytics, and machine learning workloads in Python (such as those using Pandas, Polars, DuckDB, and Ray) increasingly require reading large result sets from Cloud Spanner into Apache Arrow columnar formats. + +Currently, applications attempting to consume Spanner data into Arrow or DataFrames face significant performance and scalability constraints: + +1. **Intermediate Python Object Instantiation Overhead**: + The standard Spanner Python client decodes gRPC messages into Python `PartialResultSet` protobuf structures, instantiating individual `google.protobuf.Value` Python objects on the heap for every column of every row. For a query returning 1,000,000 rows across 12 columns, this generates 12,000,000 intermediate Python objects, resulting in substantial CPU overhead, memory consumption (peaking at hundreds of megabytes), and garbage collection pauses. + +2. **Inefficient Client-Side Conversion Paths**: + Without native Arrow support, customers must manually iterate through row tuples, convert them to Python dictionaries or lists, and pass them to `pyarrow.Table.from_pylist()` or `pandas.DataFrame()`. This approach incurs double-conversion overhead: first decoding protobuf to Python objects, and then converting Python objects into Arrow column buffers. + +3. **Global Interpreter Lock (GIL) Contention in Concurrent Workloads**: + High-throughput workloads typically use multi-threading to parallelize reads. In pure Python, all worker threads compete for the Global Interpreter Lock (GIL) while deserializing protobuf objects and allocating heap memory. As a result, client-side ingestion throughput plateaus at approximately 28,000 to 35,000 rows/second total, failing to scale with additional CPU cores or worker threads. + +To address these limitations, a solution is needed that decodes incoming gRPC stream data directly into Arrow columnar memory in native code, bypassing Python object creation and releasing the GIL during parsing. + +--- + +## Key Requirements and Architecture + +### 1. Direct Conversion to Apache Arrow +- **Direct Wire Ingestion**: The extension decodes Spanner `PartialResultSet` protobuf wire bytes directly into native Apache Arrow columnar buffers in C using the header-only `nanoarrow` library. +- **gRPC Raw Byte Interception**: In standard client execution, the gRPC transport stub configures `response_deserializer=PartialResultSet.deserialize`, which deserializes messages into Python protobuf objects. When streaming query results into Arrow, the client overrides the response deserializer (`response_deserializer=lambda raw_bytes: raw_bytes`) on the underlying `ExecuteStreamingSql` callable. This delivers raw network byte chunks directly to the C extension without invoking Python protobuf deserialization. +- **Zero Python Object Overhead**: Bypasses intermediate `google.protobuf.Value` Python object instantiation on the Python heap. +- **GIL Release**: Byte parsing and buffer construction run within `Py_BEGIN_ALLOW_THREADS` / `Py_END_ALLOW_THREADS` blocks, allowing concurrent queries across multiple threads to execute on separate CPU cores without contention on the Global Interpreter Lock (GIL). + +### 2. Alignment with Google Cloud Python Architecture +This project intentionally follows established patterns in the Google Cloud Python client ecosystem rather than inventing new mechanisms: + +#### A. The Companion Accelerator Pattern (`google-crc32c`) +- **Precedent**: The `google-crc32c` companion package in this monorepo provides a mature, production-tested blueprint for C-accelerated optional extensions (used by `google-cloud-storage` and `google-resumable-media`). +- **Shared Infrastructure**: Reuses the exact same build scripts, `cibuildwheel` configurations, and multi-OS wheel generation matrix (`manylinux`, macOS universal2/arm64/x86_64, Windows AMD64). +- **Release Integration**: `google-crc32c` is actively maintained as part of the monorepo's standard automated release cycle (managed by `release-please`), with regular updates tracking new Python releases (including Python 3.12, 3.13, and 3.14). The Arrow accelerator plugs directly into this existing release machinery. +- **Transparent Consumer Integration**: `google-cloud-spanner` dynamically checks for the presence of `google_cloud_spanner_arrow`. If installed, it delegates to the C extension; otherwise, it operates using the pure-Python implementation without error. + +#### B. Consistency with BigQuery's Arrow API (`google-cloud-bigquery`) +- **API Symmetry**: `google-cloud-bigquery` established the standard for returning analytical query results as Arrow and DataFrames in Google Cloud Python SDKs. The Spanner implementation matches this interface: + - `StreamedResultSet.to_arrow()`: Returns a complete `pyarrow.Table`. + - `StreamedResultSet.to_arrow_batches()`: Yields an iterator of `pyarrow.RecordBatch` chunks as they arrive from the network (matching BigQuery's `to_arrow_iterable()`). + - `StreamedResultSet.to_dataframe()`: Returns a `pandas.DataFrame` created directly from Arrow record batches. +- **Downstream Tool Interoperability**: + - **Polars & DuckDB**: Directly consume `pyarrow.Table` and batch iterators with zero memory copy via the standard Arrow C Data Interface (`polars.from_arrow()`, `duckdb.arrow()`). + - **Ray Data**: Uses `pyarrow.Table` as its native distributed in-memory block format (`ray.data.from_arrow()`). + - **Pandas 2.0+**: Supports backing DataFrame columns directly with Arrow storage (`dtype_backend="pyarrow"`), avoiding legacy NumPy object-array conversions. + +### 3. Opt-in and Graceful Fallback +- **Opt-in Dependency**: The core `google-cloud-spanner` package does not depend on `google-cloud-spanner-arrow`. Applications opt in by installing the accelerator package explicitly or via an extra (`google-cloud-spanner[arrow]`). +- **Pure-Python Fallback**: If `google-cloud-spanner-arrow` is not installed, or if C extension compilation fails on a given platform, the library automatically falls back to the pure-Python implementation in `google_cloud_spanner_arrow.python` / `google.cloud.spanner_v1._arrow`. +- **API Parity**: The pure-Python fallback and the C-accelerated implementation implement identical function signatures and schema type mappings. + +### 4. Build and Distribution Strategy +- **Minimal Dependencies**: The extension relies on `nanoarrow` and the standard C library, avoiding a compile-time dependency on the full C++ `libarrow`. +- **Wheel Matrix**: Reuses the build automation and platform scripts from `google-crc32c` (`noxfile.py` and GitHub Actions `cibuildwheel` configurations) to produce pre-compiled binary wheels for: + - Linux (x86_64, aarch64 via manylinux) + - macOS (x86_64, arm64) + - Windows (AMD64) +- **Source Distribution (sdist)**: Includes a fallback in `setup.py` that allows installation to succeed as pure-Python if a C compiler is unavailable. + +### 5. Interaction with Future Shared Native Core +- If a shared native core is introduced for Python and Node.js to manage gRPC transport and stream processing, wire-to-Arrow decoding will be integrated directly into that core. +- The Arrow C Data Interface (`ArrowArray` / `ArrowSchema`) will remain the handoff boundary to Python (`pyarrow.RecordBatch._import_from_c`), ensuring the public API in `google-cloud-spanner` remains backward-compatible without code changes. +- The standalone `google-cloud-spanner-arrow` package would then be superseded by the shared native core. + +--- + +## Memory and Resilience Considerations + +### Memory Utilization +- In pure-Python parsing, each row value is instantiated as a Python `Value` object, generating approximately 120–160 MB of temporary heap allocations per 100,000 rows. +- The C wire parser writes values directly to contiguous Arrow column memory buffers, eliminating intermediate Python object creation and associated garbage collection overhead. + +| Workload (50,000 rows $\times$ 12 columns) | Pure-Python Conversion | Direct Wire C Extension | Reduction | +| :--- | :---: | :---: | :---: | +| **Python Heap Peak Allocation (`tracemalloc`)** | 36.28 MB | **0.00 MB** | **100% reduction** | +| **Intermediate Python Objects** | ~600,000 objects | **0 objects** | **Zero GC tracking overhead** | + +### Stream Resumption and Chunk Merging State Machine +- **Chunk Merging**: When Spanner splits a large field across `PartialResultSet` boundaries (`chunked_value = true`), the C parser maintains the active column buffer offset across messages and appends subsequent chunk bytes before finalizing the Arrow element. +- **Stateful Stream Decoding**: In production, the converter operates as a stateful decoder (`SpannerArrowStreamDecoder`). It buffers incoming wire chunks and tracks `last_seen_resume_token` across message boundaries. +- **Batch Handoff and Retry**: When `max_chunk_size` is reached (e.g., 65,536 rows), the decoder finalizes the `RecordBatch` and returns it alongside the `resume_token` corresponding to that completed row boundary. If a gRPC connection breaks with an `UNAVAILABLE` error between batches, the Python client restarts the stream using that recorded token without duplicate rows or data loss. + +--- + +## Benchmark Results + +Benchmarks were executed against a provisioned Google Cloud Spanner instance (`4 node, regional-europe-north1`) using a schema of 12 diverse column types (`BOOL`, `BYTES`, `DATE`, `FLOAT32`, `FLOAT64`, `INTERVAL`, `JSON`, `INT64`, `NUMERIC`, `STRING`, `TIMESTAMP`, `UUID`). + +### Comparison Implementations +The benchmarks compare three implementations: +1. **Traditional Rows**: The standard Spanner Python client row iteration (`for row in results:`), returning rows as lists/tuples of Python objects. +2. **Pure-Python Arrow**: The baseline pure-Python conversion path in `google-cloud-spanner` (`to_arrow_batches()`), converting Spanner row objects into Arrow tables in Python without C acceleration. +3. **Direct Wire C Extension (`google-cloud-spanner-arrow`)**: The proposed companion accelerator, parsing raw gRPC protobuf wire bytes directly into Apache Arrow memory buffers with the GIL released. + +--- + +### 1. Co-located End-to-End Concurrency Benchmark +- **Environment**: GCE `n2-standard-8` VM (8 vCPUs, Debian 12) located in `europe-north1-a` (same zone as the Spanner instance, < 1 ms RTT). +- **Workload**: 30,000 rows per query, 3 iterations per concurrency tier. +- **Metric**: Full wall-clock time from `ExecuteStreamingSql` RPC dispatch to complete result set ingestion. + +| Concurrency | Total Rows | Total Wire Size | Traditional Rows | Pure-Python Arrow | Direct Wire C Extension | Speedup vs Traditional | Speedup vs Pure-Python | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| **1 Thread** | 30,000 | 9.1 MB | 940.8 ms *(31.9k rows/s)* | 755.5 ms *(39.7k rows/s)* | **238.5 ms (125.8k rows/s)** | 3.94x | 3.17x | +| **4 Threads** | 120,000 | 36.5 MB | 3,961.9 ms *(30.3k rows/s)* | 3,141.3 ms *(38.2k rows/s)* | **275.8 ms (435.1k rows/s)** | 14.37x | 11.39x | +| **8 Threads** | 240,000 | 73.0 MB | 8,301.2 ms *(28.9k rows/s)* | 6,536.7 ms *(36.7k rows/s)* | **507.5 ms (472.9k rows/s)** | 16.36x | 12.88x | +| **16 Threads** | 480,000 | 146.0 MB | 16,871.3 ms *(28.5k rows/s)* | 12,220.9 ms *(39.3k rows/s)* | **987.9 ms (485.9k rows/s)** | 17.08x | 12.37x | +| **32 Threads** | 960,000 | 292.0 MB | 33,914.3 ms *(28.3k rows/s)* | 27,405.7 ms *(35.0k rows/s)* | **1,957.4 ms (490.4k rows/s)** | 17.33x | 14.00x | + +*Notes on table metrics:* +- **Speedup vs Traditional**: Ratio of execution time of standard Python row iteration over the Direct Wire C Extension. +- **Speedup vs Pure-Python**: Ratio of execution time of pure-Python Arrow conversion over the Direct Wire C Extension. + +--- + +### 2. Pure CPU Parsing Throughput +- **Environment**: GCE `n2-standard-8` VM in `europe-north1-a`. +- **Metric**: Time to decode pre-fetched protobuf wire payloads into Arrow record batches (network latency excluded). + +| Workload | Pure-Python Arrow | Direct Wire C Extension | Speedup vs Pure-Python | +| :--- | :---: | :---: | :---: | +| **10,000 rows (Single Thread)** | 87.54 ms *(114k rows/s)* | **7.43 ms (1.35M rows/s)** | 11.8x | +| **50,000 rows (Single Thread)** | 447.41 ms *(112k rows/s)* | **35.69 ms (1.40M rows/s)** | 12.5x | +| **100,000 rows (Single Thread)** | 889.09 ms *(112k rows/s)* | **77.34 ms (1.29M rows/s)** | 11.5x | +| **800,000 rows (8 Threads Parallel)** | 7,236.89 ms *(111k rows/s)* | **165.70 ms (4.83M rows/s)** | 43.7x | diff --git a/packages/google-cloud-spanner-arrow/LICENSE b/packages/google-cloud-spanner-arrow/LICENSE new file mode 100644 index 000000000000..f433b1a53f5b --- /dev/null +++ b/packages/google-cloud-spanner-arrow/LICENSE @@ -0,0 +1,177 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS diff --git a/packages/google-cloud-spanner-arrow/MANIFEST.in b/packages/google-cloud-spanner-arrow/MANIFEST.in new file mode 100644 index 000000000000..2da97051688e --- /dev/null +++ b/packages/google-cloud-spanner-arrow/MANIFEST.in @@ -0,0 +1,4 @@ +include README.md LICENSE +recursive-include src/google_cloud_spanner_arrow/nanoarrow *.h *.c +include src/google_cloud_spanner_arrow/*.c +include src/google_cloud_spanner_arrow/py.typed diff --git a/packages/google-cloud-spanner-arrow/README.md b/packages/google-cloud-spanner-arrow/README.md new file mode 100644 index 000000000000..2be3171c9b0a --- /dev/null +++ b/packages/google-cloud-spanner-arrow/README.md @@ -0,0 +1,52 @@ +# Google Cloud Spanner Apache Arrow Accelerator + +[![PyPI version](https://badge.fury.io/py/google-cloud-spanner-arrow.svg)](https://badge.fury.io/py/google-cloud-spanner-arrow) + +High-performance native C extension and Apache Arrow accelerator for the Google Cloud Spanner Python client library (`google-cloud-spanner`). + +## Overview + +`google-cloud-spanner-arrow` accelerates ingestion of Cloud Spanner `PartialResultSet` protobuf streams into Apache Arrow `RecordBatch` / `Table` / `DataFrame` structures. + +### Highlights +- **Native C Data Ingestion**: Zero-copy parsing of Spanner primitive types (INT64, FLOAT64, BOOL, STRING, BYTES, DATE, TIMESTAMP, NUMERIC) into Arrow memory buffers using `nanoarrow` and the Arrow C Data Interface (`RecordBatch._import_from_c`). +- **GIL Release for Multi-Threaded Partitioned Queries**: Releases CPython's Global Interpreter Lock (GIL) during parsing for linear scaling across multi-core CPU architectures. +- **Seamless Drop-in Integration**: Automatically discovered by `google-cloud-spanner`'s `StreamedResultSet.to_arrow_batches()`, `to_arrow()`, and `to_dataframe()`, with pure-Python fallback if not installed. + +## Quick Start + +### Installation + +```bash +pip install google-cloud-spanner-arrow +``` + +### Usage with `google-cloud-spanner` + +When `google-cloud-spanner-arrow` is installed in your Python environment, `google-cloud-spanner` automatically uses the native C accelerator: + +```python +from google.cloud import spanner + +client = spanner.Client() +instance = client.instance("my-instance") +database = instance.database("my-database") + +with database.snapshot() as snapshot: + results = snapshot.execute_sql("SELECT * FROM large_table") + # Native C-accelerated Arrow batch iterator + for batch in results.to_arrow_batches(max_chunk_size=65536): + print(f"Batch rows: {batch.num_rows}") + + # Or directly as a PyArrow Table or Pandas DataFrame + table = results.to_arrow() + df = results.to_dataframe() +``` + +## Pure-Python Fallback Mode + +To explicitly disable the C extension and force pure-Python fallback, set the environment variable: + +```bash +export SPANNER_ARROW_PURE_PYTHON=1 +``` diff --git a/packages/google-cloud-spanner-arrow/mypy.ini b/packages/google-cloud-spanner-arrow/mypy.ini new file mode 100644 index 000000000000..abe31502b47d --- /dev/null +++ b/packages/google-cloud-spanner-arrow/mypy.ini @@ -0,0 +1,3 @@ +[mypy] +python_version = 3.12 +ignore_missing_imports = True diff --git a/packages/google-cloud-spanner-arrow/noxfile.py b/packages/google-cloud-spanner-arrow/noxfile.py new file mode 100644 index 000000000000..3019c2b70d28 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/noxfile.py @@ -0,0 +1,119 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import nox + +HERE = os.path.dirname(__file__) + +DEFAULT_PYTHON_VERSION = "3.12" +UNIT_TEST_PYTHON_VERSIONS = ["3.10", "3.11", "3.12", "3.13", "3.14"] +ALL_PYTHON = list(UNIT_TEST_PYTHON_VERSIONS) + +FLAKE8_VERSION = "flake8==6.1.0" +BLACK_VERSION = "black[jupyter]==23.7.0" +RUFF_VERSION = "ruff==0.14.14" +ISORT_VERSION = "isort==5.11.0" +LINT_PATHS = ["src", "tests", "noxfile.py", "setup.py"] + +nox.options.sessions = [ + "check", + "lint", + "blacken", + "format", + "lint_setup_py", + "mypy", + "unit", +] + + +@nox.session(python=UNIT_TEST_PYTHON_VERSIONS) +def check(session): + """Run tests against built wheels.""" + session.install("pytest", "pyarrow>=14.0.0") + session.install("--no-index", f"--find-links={HERE}/wheels", "google-cloud-spanner-arrow") + session.run("pytest", "tests") + session.run("python", f"{HERE}/scripts/check_spanner_arrow_extension.py", *session.posargs) + + +@nox.session(python=UNIT_TEST_PYTHON_VERSIONS) +def unit(session): + """Run all unit tests.""" + session.install("pytest", "pyarrow>=14.0.0", "protobuf") + session.install("-e", ".") + session.run("pytest", "tests") + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def mypy(session): + """Verify type hints are mypy compatible.""" + session.install( + "mypy", + "types-mock", + "types-setuptools", + "pyarrow", + ) + session.env["MYPYPATH"] = "src" + session.run("mypy", "src/google_cloud_spanner_arrow/", "tests/") + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def lint(session): + """Run linters.""" + session.install(FLAKE8_VERSION, BLACK_VERSION) + session.run( + "black", + "--check", + *LINT_PATHS, + ) + session.run("flake8", *LINT_PATHS) + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def blacken(session): + """Run black formatting.""" + session.install(BLACK_VERSION) + session.run( + "black", + *LINT_PATHS, + ) + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def format(session): + """Run ruff to sort imports and format code.""" + session.install(RUFF_VERSION) + session.run( + "ruff", + "check", + "--select", + "I", + "--fix", + "--line-length=88", + *LINT_PATHS, + ) + session.run( + "ruff", + "format", + "--line-length=88", + *LINT_PATHS, + ) + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def lint_setup_py(session): + """Verify setup.py validity.""" + session.install("docutils", "pygments", "setuptools") + session.run("python", "setup.py", "check", "--strict") diff --git a/packages/google-cloud-spanner-arrow/pyproject.toml b/packages/google-cloud-spanner-arrow/pyproject.toml new file mode 100644 index 000000000000..9683b4b23df5 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/pyproject.toml @@ -0,0 +1,28 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[build-system] +requires = ["setuptools>=64.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "google-cloud-spanner-arrow" +version = "0.1.0" +description = "High-performance Apache Arrow accelerator for Google Cloud Spanner" +readme = "README.md" +requires-python = ">=3.10" +license = {text = "Apache-2.0"} +dependencies = [ + "pyarrow >= 14.0.0", +] diff --git a/packages/google-cloud-spanner-arrow/scripts/check_spanner_arrow_extension.py b/packages/google-cloud-spanner-arrow/scripts/check_spanner_arrow_extension.py new file mode 100644 index 000000000000..d838f7ae723a --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/check_spanner_arrow_extension.py @@ -0,0 +1,27 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from google_cloud_spanner_arrow import _spanner_arrow +import google_cloud_spanner_arrow as sa + + +def main(): + print("_spanner_arrow: {}".format(_spanner_arrow)) + print("dir(_spanner_arrow): {}".format(dir(_spanner_arrow))) + print("implementation: {}".format(sa.implementation)) + assert sa.implementation == "c", "Expected C extension implementation" + + +if __name__ == "__main__": + main() diff --git a/packages/google-cloud-spanner-arrow/scripts/dev-requirements.txt b/packages/google-cloud-spanner-arrow/scripts/dev-requirements.txt new file mode 100644 index 000000000000..4a5ea1268a11 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/dev-requirements.txt @@ -0,0 +1,6 @@ +auditwheel +delocate +pytest +setuptools >= 64.0.0 +wheel +pyarrow >= 14.0.0 diff --git a/packages/google-cloud-spanner-arrow/scripts/local-linux/build_spanner_arrow.sh b/packages/google-cloud-spanner-arrow/scripts/local-linux/build_spanner_arrow.sh new file mode 100644 index 000000000000..9278beca55e6 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/local-linux/build_spanner_arrow.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e -x + +PY_BIN=${PY_BIN:-python3.12} +REPO_ROOT=${REPO_ROOT:-$(pwd)} + +VENV=${REPO_ROOT}/venv +${PY_BIN} -m venv ${VENV} +${VENV}/bin/python -m pip install --upgrade setuptools pip wheel +${VENV}/bin/python -m pip install --requirement ${REPO_ROOT}/scripts/dev-requirements.txt + +cd ${REPO_ROOT} +${VENV}/bin/python -m pip wheel . --wheel-dir=wheels + +rm -fr ${VENV} diff --git a/packages/google-cloud-spanner-arrow/scripts/manylinux/build.sh b/packages/google-cloud-spanner-arrow/scripts/manylinux/build.sh new file mode 100644 index 000000000000..f6c3ed800f0d --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/manylinux/build.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e -x +echo "BUILDING ON LINUX" +export BUILD_PYTHON=${BUILD_PYTHON} + +MANYLINUX_DIR=$(echo $(cd $(dirname ${0}); pwd)) +SCRIPTS_DIR=$(dirname ${MANYLINUX_DIR}) +REPO_ROOT=$(dirname ${SCRIPTS_DIR}) + +cd $REPO_ROOT +git config --global --add safe.directory '*' + +docker pull quay.io/pypa/manylinux2014_x86_64 +docker run \ + --rm \ + --interactive \ + --volume ${REPO_ROOT}:/var/code/spanner-arrow/ \ + --env BUILD_PYTHON=${BUILD_PYTHON} \ + quay.io/pypa/manylinux2014_x86_64 \ + /var/code/spanner-arrow/scripts/manylinux/build_on_centos.sh + +docker run --rm --privileged hypriot/qemu-register || true +docker pull quay.io/pypa/manylinux2014_aarch64 || true +docker run \ + --rm \ + --interactive \ + --volume ${REPO_ROOT}:/var/code/spanner-arrow/ \ + --env BUILD_PYTHON=${BUILD_PYTHON} \ + quay.io/pypa/manylinux2014_aarch64 \ + /var/code/spanner-arrow/scripts/manylinux/build_on_centos.sh || true + +if [[ "${PUBLISH_WHEELS}" == "true" ]]; then + . /${MANYLINUX_DIR}/publish_python_wheel.sh +fi diff --git a/packages/google-cloud-spanner-arrow/scripts/manylinux/build_on_centos.sh b/packages/google-cloud-spanner-arrow/scripts/manylinux/build_on_centos.sh new file mode 100644 index 000000000000..c624146d6759 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/manylinux/build_on_centos.sh @@ -0,0 +1,76 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e -x +MAIN_PYTHON_BIN="/opt/python/cp310-cp310/bin/" +echo "BUILD_PYTHON: ${BUILD_PYTHON}" +REPO_ROOT=/var/code/spanner-arrow/ + +${MAIN_PYTHON_BIN}/python -m pip install --upgrade pip +${MAIN_PYTHON_BIN}/python -m pip install \ + --requirement ${REPO_ROOT}/scripts/dev-requirements.txt + +PYTHON_VERSIONS="" +if [[ -z ${BUILD_PYTHON} ]]; then + for PYTHON_BIN in /opt/python/*/bin; do + if [[ "${PYTHON_BIN}" == *"310"* ]]; then + PYTHON_VERSIONS="${PYTHON_VERSIONS} ${PYTHON_BIN}" + elif [[ "${PYTHON_BIN}" == *"311"* ]]; then + PYTHON_VERSIONS="${PYTHON_VERSIONS} ${PYTHON_BIN}" + elif [[ "${PYTHON_BIN}" == *"312"* ]]; then + PYTHON_VERSIONS="${PYTHON_VERSIONS} ${PYTHON_BIN}" + elif [[ "${PYTHON_BIN}" == *"313"* && "${PYTHON_BIN}" != *"313t"* ]]; then + PYTHON_VERSIONS="${PYTHON_VERSIONS} ${PYTHON_BIN}" + elif [[ "${PYTHON_BIN}" == *"314"* && "${PYTHON_BIN}" != *"314t"* ]]; then + PYTHON_VERSIONS="${PYTHON_VERSIONS} ${PYTHON_BIN}" + fi + done +else + STRIPPED_PYTHON=$(echo ${BUILD_PYTHON} | sed -e "s/\.//g" | sed -e "s/-dev$//") + for PYTHON_BIN in /opt/python/*/bin; do + if [[ "${PYTHON_BIN}" == *"${STRIPPED_PYTHON}"* ]]; then + PYTHON_VERSIONS="${PYTHON_VERSIONS} ${PYTHON_BIN}" + fi + done +fi + +# Build wheels +cd ${REPO_ROOT} +mkdir -p dist_wheels +for PYTHON_BIN in ${PYTHON_VERSIONS}; do + ${PYTHON_BIN}/python -m pip install --upgrade pip + ${PYTHON_BIN}/python -m pip install \ + --requirement ${REPO_ROOT}/scripts/dev-requirements.txt + ${PYTHON_BIN}/python -m pip wheel . --wheel-dir dist_wheels/ +done + +# Audit wheels +mkdir -p wheels +for whl in dist_wheels/google_cloud_spanner_arrow*.whl; do + "${MAIN_PYTHON_BIN}/auditwheel" repair "${whl}" --wheel-dir wheels/ || cp "${whl}" wheels/ +done + +# Install and test wheels +for PYTHON_BIN in ${PYTHON_VERSIONS}; do + ABI_TAG=$(basename $(dirname ${PYTHON_BIN})) + ARCH=$(uname -m) + ${PYTHON_BIN}/python -m venv /tmp/venv + WHEEL_FILE=$(ls ${REPO_ROOT}/wheels/google_cloud_spanner_arrow-*-${ABI_TAG}-*${ARCH}*.whl | head -n 1) + /tmp/venv/bin/pip install "${WHEEL_FILE}" + /tmp/venv/bin/python ${REPO_ROOT}/scripts/check_spanner_arrow_extension.py + rm -rf /tmp/venv +done + +rm -rf ${REPO_ROOT}/dist_wheels/ diff --git a/packages/google-cloud-spanner-arrow/scripts/manylinux/publish_python_wheel.sh b/packages/google-cloud-spanner-arrow/scripts/manylinux/publish_python_wheel.sh new file mode 100644 index 000000000000..694f645a1a76 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/manylinux/publish_python_wheel.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e -x + +python3 -m pip install twine +twine upload --skip-existing ${REPO_ROOT}/wheels/* diff --git a/packages/google-cloud-spanner-arrow/scripts/osx/build.sh b/packages/google-cloud-spanner-arrow/scripts/osx/build.sh new file mode 100644 index 000000000000..49f62fee1d00 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/osx/build.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e -x +echo "BUILDING FOR OSX" + +export MACOSX_DEPLOYMENT_TARGET=12 + +SCRIPT_FI=$(python3 -c "import os; print(os.path.realpath('${0}'))") +OSX_DIR=$(dirname ${SCRIPT_FI}) +SCRIPTS_DIR=$(dirname ${OSX_DIR}) +export REPO_ROOT=$(dirname ${SCRIPTS_DIR}) + +cd ${REPO_ROOT} +git config --global --add safe.directory '*' + +SUPPORTED_PYTHON_VERSIONS=("3.10" "3.11" "3.12" "3.13" "3.14") + +for PYTHON_VERSION in ${SUPPORTED_PYTHON_VERSIONS[@]}; do + echo "Build wheel for Python ${PYTHON_VERSION}" + export PY_BIN=$PYTHON_VERSION + export PY_TAG="cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}" + . ${OSX_DIR}/build_python_wheel.sh +done diff --git a/packages/google-cloud-spanner-arrow/scripts/osx/build_python_wheel.sh b/packages/google-cloud-spanner-arrow/scripts/osx/build_python_wheel.sh new file mode 100644 index 000000000000..575b359af86f --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/osx/build_python_wheel.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e -x + +if [[ -z "${REPO_ROOT}" ]]; then + echo "REPO_ROOT environment variable should be set by the caller." + exit 1 +fi +if [[ -z "${PY_BIN}" ]]; then + echo "PY_BIN environment variable should be set by the caller." + exit 1 +fi + +VENV=${REPO_ROOT}/venv${PY_BIN} +"python${PY_BIN}" -m venv ${VENV} +${VENV}/bin/python -m pip install --upgrade pip setuptools wheel +${VENV}/bin/python -m pip install --requirement ${REPO_ROOT}/scripts/dev-requirements.txt + +DIST_WHEELS="${REPO_ROOT}/dist_wheels" +mkdir -p ${DIST_WHEELS} +cd ${REPO_ROOT} +${VENV}/bin/python -m pip wheel ${REPO_ROOT} --wheel-dir ${DIST_WHEELS} + +FIXED_WHEELS="${REPO_ROOT}/wheels" +mkdir -p ${FIXED_WHEELS} +cp ${DIST_WHEELS}/google_cloud_spanner_arrow*${PY_TAG}*.whl ${FIXED_WHEELS}/ 2>/dev/null || cp ${DIST_WHEELS}/*.whl ${FIXED_WHEELS}/ + +# Test wheel +${VENV}/bin/pip install --no-index --find-links=${FIXED_WHEELS} google-cloud-spanner-arrow --force-reinstall +${VENV}/bin/pip install pytest +${VENV}/bin/pytest ${REPO_ROOT}/tests +${VENV}/bin/python ${REPO_ROOT}/scripts/check_spanner_arrow_extension.py + +rm -rf ${DIST_WHEELS} +rm -rf ${VENV} diff --git a/packages/google-cloud-spanner-arrow/scripts/osx/publish_python_wheel.sh b/packages/google-cloud-spanner-arrow/scripts/osx/publish_python_wheel.sh new file mode 100644 index 000000000000..694f645a1a76 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/osx/publish_python_wheel.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e -x + +python3 -m pip install twine +twine upload --skip-existing ${REPO_ROOT}/wheels/* diff --git a/packages/google-cloud-spanner-arrow/scripts/requirements.in b/packages/google-cloud-spanner-arrow/scripts/requirements.in new file mode 100644 index 000000000000..361a45740dcc --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/requirements.in @@ -0,0 +1 @@ +pyarrow >= 14.0.0 diff --git a/packages/google-cloud-spanner-arrow/scripts/requirements.txt b/packages/google-cloud-spanner-arrow/scripts/requirements.txt new file mode 100644 index 000000000000..a79af95afb8e --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/requirements.txt @@ -0,0 +1 @@ +pyarrow>=14.0.0 diff --git a/packages/google-cloud-spanner-arrow/scripts/windows/build.bat b/packages/google-cloud-spanner-arrow/scripts/windows/build.bat new file mode 100644 index 000000000000..9461675f3af7 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/windows/build.bat @@ -0,0 +1,35 @@ +@rem Copyright 2026 Google LLC +@rem +@rem Licensed under the Apache License, Version 2.0 (the "License"); +@rem you may not use this file except in compliance with the License. +@rem You may obtain a copy of the License at +@rem +@rem https://www.apache.org/licenses/LICENSE-2.0 +@rem +@rem Unless required by applicable law or agreed to in writing, software +@rem distributed under the License is distributed on an "AS IS" BASIS, +@rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +@rem See the License for the specific language governing permissions and +@rem limitations under the License. + +setlocal ENABLEDELAYEDEXPANSION + +FOR %%P IN (3.10, 3.11, 3.12, 3.13, 3.14) DO ( + echo "Building for Python version %%P" + set python_version=%%P + set python_version_trimmed=!python_version:~0,4! + + py -!python_version_trimmed!-64 -m pip install --upgrade pip setuptools wheel + py -!python_version_trimmed!-64 -m pip install -r scripts\dev-requirements.txt + + mkdir wheels 2>nul + py -!python_version_trimmed!-64 -m pip wheel . --wheel-dir wheels\ + + call %~dp0\test.bat !python_version_trimmed! || goto :error +) + +goto :EOF + +:error +echo Failed with error #%errorlevel%. +exit /b %errorlevel% diff --git a/packages/google-cloud-spanner-arrow/scripts/windows/test.bat b/packages/google-cloud-spanner-arrow/scripts/windows/test.bat new file mode 100644 index 000000000000..5aaa0e5161af --- /dev/null +++ b/packages/google-cloud-spanner-arrow/scripts/windows/test.bat @@ -0,0 +1,28 @@ +@rem Copyright 2026 Google LLC +@rem +@rem Licensed under the Apache License, Version 2.0 (the "License"); +@rem you may not use this file except in compliance with the License. +@rem You may obtain a copy of the License at +@rem +@rem https://www.apache.org/licenses/LICENSE-2.0 +@rem +@rem Unless required by applicable law or agreed to in writing, software +@rem distributed under the License is distributed on an "AS IS" BASIS, +@rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +@rem See the License for the specific language governing permissions and +@rem limitations under the License. + +set python_version=%1 + +py -%python_version%-64 -m venv test_venv +call test_venv\Scripts\activate.bat + +python -m pip install --upgrade pip +python -m pip install --no-index --find-links=wheels google-cloud-spanner-arrow +python -m pip install pytest + +pytest tests +python scripts\check_spanner_arrow_extension.py + +call deactivate +rmdir /s /q test_venv diff --git a/packages/google-cloud-spanner-arrow/setup.cfg b/packages/google-cloud-spanner-arrow/setup.cfg new file mode 100644 index 000000000000..1fa2c16b9c49 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/setup.cfg @@ -0,0 +1,51 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[metadata] +name = google-cloud-spanner-arrow +version = 0.1.0 +description = High-performance Apache Arrow accelerator for Google Cloud Spanner +url = https://github.com/googleapis/google-cloud-python/tree/main/packages/google-cloud-spanner-arrow +long_description = file: README.md +long_description_content_type = text/markdown +author = Google LLC +author_email = googleapis-packages@google.com + +license = Apache 2.0 +license_files = LICENSE +platforms = Posix, MacOS X, Windows +classifiers = + Development Status :: 4 - Beta + Intended Audience :: Developers + Operating System :: OS Independent + Programming Language :: Python :: 3 + 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 + +[options] +zip_safe = False +python_requires = >=3.10 +install_requires = + pyarrow >= 14.0.0 + +[options.extras_require] +testing = + pytest + +[options.package_data] +google_cloud_spanner_arrow = + py.typed diff --git a/packages/google-cloud-spanner-arrow/setup.py b/packages/google-cloud-spanner-arrow/setup.py new file mode 100644 index 000000000000..a73a7e2ab8c6 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/setup.py @@ -0,0 +1,74 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import os +import setuptools + +# Explicit environment variable disables pure-Python fallback +SPANNER_ARROW_PURE_PYTHON_EXPLICIT = "SPANNER_ARROW_PURE_PYTHON" in os.environ +_FALSE_OPTIONS = ("0", "false", "no", "False", "No", None) +SPANNER_ARROW_PURE_PYTHON = os.getenv("SPANNER_ARROW_PURE_PYTHON") not in _FALSE_OPTIONS + + +def build_pure_python(): + setuptools.setup( + packages=["google_cloud_spanner_arrow"], + package_dir={"": "src"}, + ext_modules=[], + ) + + +def build_c_extension(): + module_sources = [ + os.path.normcase(os.path.join("src", "google_cloud_spanner_arrow", "_spanner_arrow.c")), + os.path.normcase(os.path.join("src", "google_cloud_spanner_arrow", "nanoarrow", "nanoarrow.c")), + ] + include_dirs = [ + os.path.normcase(os.path.join("src", "google_cloud_spanner_arrow")), + os.path.normcase(os.path.join("src", "google_cloud_spanner_arrow", "nanoarrow")), + ] + + module = setuptools.Extension( + "google_cloud_spanner_arrow._spanner_arrow", + sources=module_sources, + include_dirs=include_dirs, + ) + + setuptools.setup( + packages=["google_cloud_spanner_arrow"], + package_dir={"": "src"}, + ext_modules=[module], + ) + + +if SPANNER_ARROW_PURE_PYTHON: + build_pure_python() +else: + try: + build_c_extension() + except SystemExit: + if SPANNER_ARROW_PURE_PYTHON_EXPLICIT: + logging.error( + "Compiling the C Extension for google-cloud-spanner-arrow failed. " + "To enable building / installing a pure-Python-only version, " + "set 'SPANNER_ARROW_PURE_PYTHON=1' in the environment." + ) + raise + + logging.info( + "Compiling the C Extension for google-cloud-spanner-arrow failed. " + "Falling back to pure Python build." + ) + build_pure_python() diff --git a/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/__config__.py b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/__config__.py new file mode 100644 index 000000000000..edf31873e55a --- /dev/null +++ b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/__config__.py @@ -0,0 +1,23 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +MODIFY_CORE = False + +SPANNER_ARROW_PURE_PYTHON = os.getenv("SPANNER_ARROW_PURE_PYTHON", "0").lower() in ( + "1", + "true", + "yes", +) diff --git a/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/__init__.py b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/__init__.py new file mode 100644 index 000000000000..191af6ce9d66 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/__init__.py @@ -0,0 +1,49 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""High-performance Apache Arrow accelerator for Google Cloud Spanner.""" + +import os +import warnings + +_SLOW_SPANNER_ARROW_WARNING = ( + "As the C extension couldn't be imported, `google-cloud-spanner-arrow` is using a " + "pure Python implementation that is significantly slower. If possible, " + "please compile the C extension for maximum throughput." +) + +_DISABLE_CEXT = os.getenv("SPANNER_ARROW_PURE_PYTHON", "0").lower() in ("1", "true", "yes") + +if not _DISABLE_CEXT: + try: + from google_cloud_spanner_arrow import cext as impl + implementation = "c" + except ImportError: + from google_cloud_spanner_arrow import python as impl # type: ignore + warnings.warn(_SLOW_SPANNER_ARROW_WARNING, RuntimeWarning) + implementation = "python" +else: + from google_cloud_spanner_arrow import python as impl + implementation = "python" + +rows_to_arrow_batch = impl.rows_to_arrow_batch +fields_to_arrow_schema = impl.fields_to_arrow_schema +spanner_type_to_arrow_type = impl.spanner_type_to_arrow_type + +__all__ = [ + "rows_to_arrow_batch", + "fields_to_arrow_schema", + "spanner_type_to_arrow_type", + "implementation", +] diff --git a/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/_spanner_arrow.c b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/_spanner_arrow.c new file mode 100644 index 000000000000..5792bc732bb5 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/_spanner_arrow.c @@ -0,0 +1,969 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#define PY_SSIZE_T_CLEAN +#include +#include "nanoarrow/nanoarrow.h" +#include +#include +#include + +// Spanner TypeCode enum constants (matching google.cloud.spanner_v1.types.TypeCode) +#define SPANNER_TYPE_UNSPECIFIED 0 +#define SPANNER_TYPE_BOOL 1 +#define SPANNER_TYPE_INT64 2 +#define SPANNER_TYPE_FLOAT64 3 +#define SPANNER_TYPE_TIMESTAMP 4 +#define SPANNER_TYPE_DATE 5 +#define SPANNER_TYPE_STRING 6 +#define SPANNER_TYPE_BYTES 7 +#define SPANNER_TYPE_ARRAY 8 +#define SPANNER_TYPE_STRUCT 9 +#define SPANNER_TYPE_NUMERIC 10 +#define SPANNER_TYPE_JSON 11 +#define SPANNER_TYPE_PROTO 13 +#define SPANNER_TYPE_ENUM 14 +#define SPANNER_TYPE_FLOAT32 15 +#define SPANNER_TYPE_INTERVAL 16 +#define SPANNER_TYPE_UUID 17 + +static const int8_t base64_decode_table[256] = { + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 62, -1, -1, -1, 63, + 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, -1, -1, -1, 0, -1, -1, + -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, + 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, -1, -1, -1, -1, -1, + -1, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, + 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 +}; + +static size_t base64_decode(const char* src, size_t src_len, uint8_t* dst) { + if (src_len == 0 || src == NULL || dst == NULL) return 0; + size_t out_len = 0; + uint32_t buf = 0; + int bits = 0; + + for (size_t i = 0; i < src_len; i++) { + unsigned char c = (unsigned char)src[i]; + if (c == '=') break; + int8_t val = base64_decode_table[c]; + if (val < 0) continue; + buf = (buf << 6) | (uint32_t)val; + bits += 6; + if (bits >= 8) { + bits -= 8; + dst[out_len++] = (uint8_t)((buf >> bits) & 0xFF); + } + } + return out_len; +} + +static int32_t parse_date32_fast(const char* str, size_t len) { + if (len < 10 || str == NULL) return 0; + int year = (str[0]-'0')*1000 + (str[1]-'0')*100 + (str[2]-'0')*10 + (str[3]-'0'); + int month = (str[5]-'0')*10 + (str[6]-'0'); + int day = (str[8]-'0')*10 + (str[9]-'0'); + + year -= (month <= 2); + const int era = (year >= 0 ? year : year - 399) / 400; + const unsigned yoe = (unsigned)(year - era * 400); + const unsigned doy = (153 * (month + (month > 2 ? -3 : 9)) + 2) / 5 + day - 1; + const unsigned doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; + return (int32_t)(era * 146097 + (int)doe - 719468); +} + +static int64_t parse_timestamp_us_fast(const char* str, size_t len) { + if (len < 19 || str == NULL) return 0; + int32_t days = parse_date32_fast(str, len); + int hour = (str[11]-'0')*10 + (str[12]-'0'); + int min = (str[14]-'0')*10 + (str[15]-'0'); + int sec = (str[17]-'0')*10 + (str[18]-'0'); + + int64_t total_us = ((int64_t)days * 86400LL + (int64_t)hour * 3600LL + (int64_t)min * 60LL + (int64_t)sec) * 1000000LL; + + if (len > 19 && str[19] == '.') { + size_t idx = 20; + int64_t frac_us = 0; + int digits = 0; + while (idx < len && isdigit((unsigned char)str[idx]) && digits < 6) { + frac_us = frac_us * 10 + (str[idx] - '0'); + digits++; + idx++; + } + while (digits < 6) { + frac_us *= 10; + digits++; + } + total_us += frac_us; + } + return total_us; +} + +static void parse_decimal128_fast(const char* str, size_t len, struct ArrowDecimal128* out) { + memset(out->bytes, 0, 16); + if (len == 0 || str == NULL) return; + + int sign = 1; + size_t i = 0; + if (str[0] == '-') { + sign = -1; + i = 1; + } else if (str[0] == '+') { + i = 1; + } + +#if defined(__SIZEOF_INT128__) + unsigned __int128 whole = 0; + unsigned __int128 frac = 0; + int frac_digits = 0; + bool in_frac = false; + + for (; i < len; i++) { + char c = str[i]; + if (c == '.') { + in_frac = true; + continue; + } + if (isdigit((unsigned char)c)) { + if (!in_frac) { + whole = whole * 10 + (c - '0'); + } else if (frac_digits < 9) { + frac = frac * 10 + (c - '0'); + frac_digits++; + } + } + } + while (frac_digits < 9) { + frac *= 10; + frac_digits++; + } + unsigned __int128 scale_multiplier = 1000000000ULL; + unsigned __int128 total = whole * scale_multiplier + frac; + __int128 final_val = (sign < 0) ? -((__int128)total) : ((__int128)total); + memcpy(out->bytes, &final_val, 16); +#else + int64_t whole = 0; + int64_t frac = 0; + int frac_digits = 0; + bool in_frac = false; + for (; i < len; i++) { + char c = str[i]; + if (c == '.') { + in_frac = true; + continue; + } + if (isdigit((unsigned char)c)) { + if (!in_frac) { + whole = whole * 10 + (c - '0'); + } else if (frac_digits < 9) { + frac = frac * 10 + (c - '0'); + frac_digits++; + } + } + } + while (frac_digits < 9) { + frac *= 10; + frac_digits++; + } + int64_t total = (whole * 1000000000LL + frac) * sign; + memcpy(out->bytes, &total, 8); + if (sign < 0) { + memset(out->bytes + 8, 0xFF, 8); + } +#endif +} + +static void configure_field_schema(struct ArrowSchema* schema, PyObject* f_obj) { + const char* col_name = "col"; + int type_code = SPANNER_TYPE_STRING; + PyObject* children_obj = NULL; + + if (PyTuple_Check(f_obj) && PyTuple_Size(f_obj) >= 2) { + PyObject* name_obj = PyTuple_GET_ITEM(f_obj, 0); + PyObject* type_obj = PyTuple_GET_ITEM(f_obj, 1); + if (PyUnicode_Check(name_obj)) { + col_name = PyUnicode_AsUTF8(name_obj); + } + if (PyLong_Check(type_obj)) { + type_code = (int)PyLong_AsLong(type_obj); + } + if (PyTuple_Size(f_obj) >= 3) { + children_obj = PyTuple_GET_ITEM(f_obj, 2); + } + } + + ArrowSchemaSetName(schema, col_name); + schema->flags = ARROW_FLAG_NULLABLE; + schema->release = &ArrowSchemaRelease; + + switch (type_code) { + case SPANNER_TYPE_BOOL: + ArrowSchemaSetFormat(schema, "b"); + break; + case SPANNER_TYPE_INT64: + case SPANNER_TYPE_ENUM: + ArrowSchemaSetFormat(schema, "l"); + break; + case SPANNER_TYPE_FLOAT32: + ArrowSchemaSetFormat(schema, "f"); + break; + case SPANNER_TYPE_FLOAT64: + ArrowSchemaSetFormat(schema, "g"); + break; + case SPANNER_TYPE_STRING: + case SPANNER_TYPE_JSON: + case SPANNER_TYPE_INTERVAL: + case SPANNER_TYPE_UUID: + ArrowSchemaSetFormat(schema, "u"); + break; + case SPANNER_TYPE_BYTES: + case SPANNER_TYPE_PROTO: + ArrowSchemaSetFormat(schema, "z"); + break; + case SPANNER_TYPE_DATE: + ArrowSchemaSetFormat(schema, "tdD"); + break; + case SPANNER_TYPE_TIMESTAMP: + ArrowSchemaSetFormat(schema, "tsu:UTC"); + break; + case SPANNER_TYPE_NUMERIC: + ArrowSchemaSetFormat(schema, "d:38,9"); + break; + case SPANNER_TYPE_ARRAY: + ArrowSchemaSetFormat(schema, "+l"); + ArrowSchemaAllocateChildren(schema, 1); + if (children_obj != NULL) { + configure_field_schema(schema->children[0], children_obj); + } else { + ArrowSchemaSetName(schema->children[0], "item"); + ArrowSchemaSetFormat(schema->children[0], "u"); + } + break; + case SPANNER_TYPE_STRUCT: + ArrowSchemaSetFormat(schema, "+s"); + if (children_obj != NULL && PySequence_Check(children_obj)) { + Py_ssize_t n_sub = PySequence_Size(children_obj); + ArrowSchemaAllocateChildren(schema, (int64_t)n_sub); + for (Py_ssize_t s = 0; s < n_sub; s++) { + PyObject* sub_item = PySequence_GetItem(children_obj, s); + configure_field_schema(schema->children[s], sub_item); + Py_XDECREF(sub_item); + } + } + break; + default: + ArrowSchemaSetFormat(schema, "u"); + break; + } +} + +// -------------------------------------------------------------------------- +// Python Object Cell Ingestion +// -------------------------------------------------------------------------- + +static int append_python_cell(struct ArrowArray* col_array, PyObject* cell, PyObject* f_obj) { + int type_code = SPANNER_TYPE_STRING; + PyObject* children_obj = NULL; + + if (PyTuple_Check(f_obj) && PyTuple_Size(f_obj) >= 2) { + PyObject* type_obj = PyTuple_GET_ITEM(f_obj, 1); + if (PyLong_Check(type_obj)) { + type_code = (int)PyLong_AsLong(type_obj); + } + if (PyTuple_Size(f_obj) >= 3) { + children_obj = PyTuple_GET_ITEM(f_obj, 2); + } + } + + if (cell == NULL || cell == Py_None) { + return ArrowArrayAppendNull(col_array, 1); + } + + if (PyObject_HasAttrString(cell, "WhichOneof")) { + PyObject* kind_obj = PyObject_CallMethod(cell, "WhichOneof", "s", "kind"); + if (kind_obj == NULL || kind_obj == Py_None) { + Py_XDECREF(kind_obj); + return ArrowArrayAppendNull(col_array, 1); + } + const char* kind = PyUnicode_AsUTF8(kind_obj); + if (kind == NULL || strcmp(kind, "null_value") == 0) { + Py_DECREF(kind_obj); + return ArrowArrayAppendNull(col_array, 1); + } + if (strcmp(kind, "bool_value") == 0) { + PyObject* val_obj = PyObject_GetAttrString(cell, "bool_value"); + int b = PyObject_IsTrue(val_obj); + Py_XDECREF(val_obj); + Py_DECREF(kind_obj); + return ArrowArrayAppendBool(col_array, (uint8_t)b); + } + if (strcmp(kind, "number_value") == 0) { + PyObject* val_obj = PyObject_GetAttrString(cell, "number_value"); + double d = PyFloat_AsDouble(val_obj); + Py_XDECREF(val_obj); + Py_DECREF(kind_obj); + if (type_code == SPANNER_TYPE_FLOAT32) { + return ArrowArrayAppendFloat(col_array, (float)d); + } + return ArrowArrayAppendDouble(col_array, d); + } + if (strcmp(kind, "string_value") == 0) { + PyObject* val_obj = PyObject_GetAttrString(cell, "string_value"); + Py_ssize_t str_len = 0; + const char* str_val = PyUnicode_AsUTF8AndSize(val_obj, &str_len); + int ret = 0; + + if (str_val == NULL) { + Py_XDECREF(val_obj); + Py_DECREF(kind_obj); + return ArrowArrayAppendNull(col_array, 1); + } + + switch (type_code) { + case SPANNER_TYPE_INT64: + case SPANNER_TYPE_ENUM: { + int64_t val = (int64_t)strtoll(str_val, NULL, 10); + ret = ArrowArrayAppendInt(col_array, val); + break; + } + case SPANNER_TYPE_FLOAT64: { + double val = 0.0; + if (strcmp(str_val, "NaN") == 0) { + val = NAN; + } else if (strcmp(str_val, "Infinity") == 0) { + val = INFINITY; + } else if (strcmp(str_val, "-Infinity") == 0) { + val = -INFINITY; + } else { + val = strtod(str_val, NULL); + } + ret = ArrowArrayAppendDouble(col_array, val); + break; + } + case SPANNER_TYPE_FLOAT32: { + float val = 0.0f; + if (strcmp(str_val, "NaN") == 0) { + val = (float)NAN; + } else if (strcmp(str_val, "Infinity") == 0) { + val = (float)INFINITY; + } else if (strcmp(str_val, "-Infinity") == 0) { + val = (float)-INFINITY; + } else { + val = strtof(str_val, NULL); + } + ret = ArrowArrayAppendFloat(col_array, val); + break; + } + case SPANNER_TYPE_BYTES: + case SPANNER_TYPE_PROTO: { + size_t max_decoded = (size_t)(str_len * 3 / 4 + 4); + uint8_t* decode_buf = (uint8_t*)malloc(max_decoded); + if (decode_buf != NULL) { + size_t decoded_len = base64_decode(str_val, (size_t)str_len, decode_buf); + struct ArrowBufferView view = {decode_buf, (int64_t)decoded_len}; + ret = ArrowArrayAppendBytes(col_array, view); + free(decode_buf); + } else { + ret = ArrowArrayAppendNull(col_array, 1); + } + break; + } + case SPANNER_TYPE_DATE: { + int32_t days = parse_date32_fast(str_val, (size_t)str_len); + ret = ArrowArrayAppendInt(col_array, days); + break; + } + case SPANNER_TYPE_TIMESTAMP: { + int64_t ts_us = parse_timestamp_us_fast(str_val, (size_t)str_len); + ret = ArrowArrayAppendInt(col_array, ts_us); + break; + } + case SPANNER_TYPE_NUMERIC: { + struct ArrowDecimal128 dec; + parse_decimal128_fast(str_val, (size_t)str_len, &dec); + ret = ArrowArrayAppendDecimal128(col_array, dec); + break; + } + default: { + struct ArrowStringView view = {str_val, (int64_t)str_len}; + ret = ArrowArrayAppendString(col_array, view); + break; + } + } + Py_XDECREF(val_obj); + Py_DECREF(kind_obj); + return ret; + } + if (strcmp(kind, "list_value") == 0) { + PyObject* val_obj = PyObject_GetAttrString(cell, "list_value"); + PyObject* values_list = val_obj ? PyObject_GetAttrString(val_obj, "values") : NULL; + if (values_list && PySequence_Check(values_list)) { + Py_ssize_t list_len = PySequence_Size(values_list); + for (Py_ssize_t li = 0; li < list_len; li++) { + PyObject* elem = PySequence_GetItem(values_list, li); + if (col_array->n_children > 0) { + append_python_cell(col_array->children[0], elem, children_obj); + } + Py_XDECREF(elem); + } + } + Py_XDECREF(values_list); + Py_XDECREF(val_obj); + Py_DECREF(kind_obj); + return ArrowArrayAppendList(col_array); + } + if (strcmp(kind, "struct_value") == 0) { + PyObject* val_obj = PyObject_GetAttrString(cell, "struct_value"); + PyObject* fields_dict = val_obj ? PyObject_GetAttrString(val_obj, "fields") : NULL; + if (fields_dict && children_obj && PySequence_Check(children_obj)) { + Py_ssize_t n_sub = PySequence_Size(children_obj); + for (Py_ssize_t s = 0; s < n_sub; s++) { + PyObject* sub_info = PySequence_GetItem(children_obj, s); + PyObject* sub_name = PyTuple_GET_ITEM(sub_info, 0); + PyObject* sub_val = PyObject_GetItem(fields_dict, sub_name); + if (sub_val == NULL) { + PyErr_Clear(); + sub_val = Py_None; + Py_INCREF(sub_val); + } + if (s < col_array->n_children) { + append_python_cell(col_array->children[s], sub_val, sub_info); + } + Py_XDECREF(sub_val); + Py_XDECREF(sub_info); + } + } + Py_XDECREF(fields_dict); + Py_XDECREF(val_obj); + Py_DECREF(kind_obj); + return ArrowArrayAppendStruct(col_array); + } + Py_DECREF(kind_obj); + return ArrowArrayAppendNull(col_array, 1); + } + + if (PyBool_Check(cell)) { + return ArrowArrayAppendBool(col_array, cell == Py_True ? 1 : 0); + } + if (PyLong_Check(cell)) { + int64_t v = (int64_t)PyLong_AsLongLong(cell); + if (type_code == SPANNER_TYPE_DATE) { + return ArrowArrayAppendInt(col_array, (int32_t)v); + } + return ArrowArrayAppendInt(col_array, v); + } + if (PyFloat_Check(cell)) { + double d = PyFloat_AS_DOUBLE(cell); + if (type_code == SPANNER_TYPE_FLOAT32) { + return ArrowArrayAppendFloat(col_array, (float)d); + } + return ArrowArrayAppendDouble(col_array, d); + } + if (PyBytes_Check(cell)) { + Py_ssize_t b_len = PyBytes_GET_SIZE(cell); + const char* b_data = PyBytes_AS_STRING(cell); + struct ArrowBufferView view = {b_data, (int64_t)b_len}; + return ArrowArrayAppendBytes(col_array, view); + } + if (PyUnicode_Check(cell)) { + Py_ssize_t str_len = 0; + const char* str_val = PyUnicode_AsUTF8AndSize(cell, &str_len); + if (str_val == NULL) { + return ArrowArrayAppendNull(col_array, 1); + } + switch (type_code) { + case SPANNER_TYPE_INT64: + case SPANNER_TYPE_ENUM: { + int64_t val = (int64_t)strtoll(str_val, NULL, 10); + return ArrowArrayAppendInt(col_array, val); + } + case SPANNER_TYPE_FLOAT64: { + double val = strtod(str_val, NULL); + return ArrowArrayAppendDouble(col_array, val); + } + case SPANNER_TYPE_FLOAT32: { + float val = strtof(str_val, NULL); + return ArrowArrayAppendFloat(col_array, val); + } + case SPANNER_TYPE_BYTES: + case SPANNER_TYPE_PROTO: { + size_t max_decoded = (size_t)(str_len * 3 / 4 + 4); + uint8_t* decode_buf = (uint8_t*)malloc(max_decoded); + if (decode_buf != NULL) { + size_t decoded_len = base64_decode(str_val, (size_t)str_len, decode_buf); + struct ArrowBufferView view = {decode_buf, (int64_t)decoded_len}; + int ret = ArrowArrayAppendBytes(col_array, view); + free(decode_buf); + return ret; + } + return ArrowArrayAppendNull(col_array, 1); + } + case SPANNER_TYPE_DATE: { + int32_t days = parse_date32_fast(str_val, (size_t)str_len); + return ArrowArrayAppendInt(col_array, days); + } + case SPANNER_TYPE_TIMESTAMP: { + int64_t ts_us = parse_timestamp_us_fast(str_val, (size_t)str_len); + return ArrowArrayAppendInt(col_array, ts_us); + } + case SPANNER_TYPE_NUMERIC: { + struct ArrowDecimal128 dec; + parse_decimal128_fast(str_val, (size_t)str_len, &dec); + return ArrowArrayAppendDecimal128(col_array, dec); + } + default: { + struct ArrowStringView view = {str_val, (int64_t)str_len}; + return ArrowArrayAppendString(col_array, view); + } + } + } + + return ArrowArrayAppendNull(col_array, 1); +} + +static PyObject* py_rows_to_c_batch(PyObject* self, PyObject* args) { + PyObject* py_fields; + PyObject* py_rows; + + if (!PyArg_ParseTuple(args, "OO", &py_fields, &py_rows)) { + return NULL; + } + + if (!PySequence_Check(py_fields) || !PySequence_Check(py_rows)) { + PyErr_SetString(PyExc_TypeError, "fields and rows must be sequences"); + return NULL; + } + + Py_ssize_t num_cols = PySequence_Size(py_fields); + Py_ssize_t num_rows = PySequence_Size(py_rows); + + struct ArrowSchema* out_schema = (struct ArrowSchema*)calloc(1, sizeof(struct ArrowSchema)); + struct ArrowArray* out_array = (struct ArrowArray*)calloc(1, sizeof(struct ArrowArray)); + struct ArrowError error; + + if (out_schema == NULL || out_array == NULL) { + if (out_schema) free(out_schema); + if (out_array) free(out_array); + PyErr_NoMemory(); + return NULL; + } + + ArrowSchemaInit(out_schema, NANOARROW_TYPE_STRUCT); + ArrowSchemaAllocateChildren(out_schema, (int64_t)num_cols); + + for (Py_ssize_t i = 0; i < num_cols; i++) { + PyObject* f = PySequence_GetItem(py_fields, i); + configure_field_schema(out_schema->children[i], f); + Py_XDECREF(f); + } + + if (ArrowArrayInitFromSchema(out_array, out_schema, &error) != 0) { + ArrowSchemaRelease(out_schema); + free(out_schema); + free(out_array); + PyErr_Format(PyExc_RuntimeError, "Failed to init ArrowArray: %s", error.message); + return NULL; + } + + ArrowArrayStartAppending(out_array); + + for (Py_ssize_t r = 0; r < num_rows; r++) { + PyObject* row = PySequence_GetItem(py_rows, r); + if (row != NULL && PySequence_Check(row)) { + Py_ssize_t row_len = PySequence_Size(row); + for (Py_ssize_t c = 0; c < num_cols; c++) { + PyObject* f_info = PySequence_GetItem(py_fields, c); + PyObject* cell = (c < row_len) ? PySequence_GetItem(row, c) : NULL; + append_python_cell(out_array->children[c], cell, f_info); + Py_XDECREF(cell); + Py_XDECREF(f_info); + } + out_array->length++; + } + Py_XDECREF(row); + } + + if (ArrowArrayFinishBuildingDefault(out_array, &error) != 0) { + ArrowArrayRelease(out_array); + ArrowSchemaRelease(out_schema); + free(out_array); + free(out_schema); + PyErr_Format(PyExc_RuntimeError, "Failed to finish ArrowArray: %s", error.message); + return NULL; + } + + uintptr_t array_ptr = (uintptr_t)out_array; + uintptr_t schema_ptr = (uintptr_t)out_schema; + + return Py_BuildValue("(KK)", (unsigned long long)array_ptr, (unsigned long long)schema_ptr); +} + +// -------------------------------------------------------------------------- +// Direct Protobuf Wire Parser Helpers (No Python Object Allocations) +// -------------------------------------------------------------------------- + +static inline uint64_t decode_varint(const uint8_t** ptr, const uint8_t* end) { + uint64_t result = 0; + int shift = 0; + const uint8_t* p = *ptr; + while (p < end && shift < 64) { + uint8_t byte = *p++; + result |= ((uint64_t)(byte & 0x7F)) << shift; + if ((byte & 0x80) == 0) { + *ptr = p; + return result; + } + shift += 7; + } + *ptr = p; + return result; +} + +static inline const uint8_t* read_length_delimited(const uint8_t** ptr, const uint8_t* end, uint64_t* out_len) { + uint64_t len = decode_varint(ptr, end); + *out_len = len; + const uint8_t* slice = *ptr; + *ptr += len; + if (*ptr > end) { + *ptr = end; + } + return slice; +} + +static inline void skip_wire_field(const uint8_t** ptr, const uint8_t* end, int wire_type) { + switch (wire_type) { + case 0: + decode_varint(ptr, end); + break; + case 1: + *ptr += 8; + if (*ptr > end) *ptr = end; + break; + case 2: { + uint64_t len = decode_varint(ptr, end); + *ptr += len; + if (*ptr > end) *ptr = end; + break; + } + case 5: + *ptr += 4; + if (*ptr > end) *ptr = end; + break; + default: + *ptr = end; + break; + } +} + +static int append_wire_value(struct ArrowArray* col_array, int type_code, const uint8_t* p, const uint8_t* val_end) { + if (p >= val_end) { + return ArrowArrayAppendNull(col_array, 1); + } + + while (p < val_end) { + uint64_t tag = decode_varint(&p, val_end); + int field_num = (int)(tag >> 3); + int wire_type = (int)(tag & 0x07); + + if (field_num == 1 && wire_type == 0) { + decode_varint(&p, val_end); + return ArrowArrayAppendNull(col_array, 1); + } else if (field_num == 2 && wire_type == 1) { + if (p + 8 <= val_end) { + double d; + memcpy(&d, p, 8); + p += 8; + if (type_code == SPANNER_TYPE_FLOAT32) { + return ArrowArrayAppendFloat(col_array, (float)d); + } + return ArrowArrayAppendDouble(col_array, d); + } + return ArrowArrayAppendNull(col_array, 1); + } else if (field_num == 3 && wire_type == 2) { + uint64_t str_len = 0; + const uint8_t* str_data = read_length_delimited(&p, val_end, &str_len); + const char* str_val = (const char*)str_data; + + switch (type_code) { + case SPANNER_TYPE_INT64: + case SPANNER_TYPE_ENUM: { + int64_t val = (int64_t)strtoll(str_val, NULL, 10); + return ArrowArrayAppendInt(col_array, val); + } + case SPANNER_TYPE_FLOAT64: { + double val = 0.0; + if (str_len == 3 && strncmp(str_val, "NaN", 3) == 0) { + val = NAN; + } else if (str_len == 8 && strncmp(str_val, "Infinity", 8) == 0) { + val = INFINITY; + } else if (str_len == 9 && strncmp(str_val, "-Infinity", 9) == 0) { + val = -INFINITY; + } else { + val = strtod(str_val, NULL); + } + return ArrowArrayAppendDouble(col_array, val); + } + case SPANNER_TYPE_FLOAT32: { + float val = 0.0f; + if (str_len == 3 && strncmp(str_val, "NaN", 3) == 0) { + val = (float)NAN; + } else if (str_len == 8 && strncmp(str_val, "Infinity", 8) == 0) { + val = (float)INFINITY; + } else if (str_len == 9 && strncmp(str_val, "-Infinity", 9) == 0) { + val = (float)-INFINITY; + } else { + val = strtof(str_val, NULL); + } + return ArrowArrayAppendFloat(col_array, val); + } + case SPANNER_TYPE_BYTES: + case SPANNER_TYPE_PROTO: { + size_t max_decoded = (size_t)(str_len * 3 / 4 + 4); + uint8_t* decode_buf = (uint8_t*)malloc(max_decoded); + if (decode_buf != NULL) { + size_t decoded_len = base64_decode(str_val, (size_t)str_len, decode_buf); + struct ArrowBufferView view = {decode_buf, (int64_t)decoded_len}; + int ret = ArrowArrayAppendBytes(col_array, view); + free(decode_buf); + return ret; + } + return ArrowArrayAppendNull(col_array, 1); + } + case SPANNER_TYPE_DATE: { + int32_t days = parse_date32_fast(str_val, (size_t)str_len); + return ArrowArrayAppendInt(col_array, days); + } + case SPANNER_TYPE_TIMESTAMP: { + int64_t ts_us = parse_timestamp_us_fast(str_val, (size_t)str_len); + return ArrowArrayAppendInt(col_array, ts_us); + } + case SPANNER_TYPE_NUMERIC: { + struct ArrowDecimal128 dec; + parse_decimal128_fast(str_val, (size_t)str_len, &dec); + return ArrowArrayAppendDecimal128(col_array, dec); + } + default: { + struct ArrowStringView view = {str_val, (int64_t)str_len}; + return ArrowArrayAppendString(col_array, view); + } + } + } else if (field_num == 4 && wire_type == 0) { + uint64_t b = decode_varint(&p, val_end); + return ArrowArrayAppendBool(col_array, b ? 1 : 0); + } else if (field_num == 6 && wire_type == 2) { + uint64_t list_len = 0; + const uint8_t* list_data = read_length_delimited(&p, val_end, &list_len); + const uint8_t* lp = list_data; + const uint8_t* lend = list_data + list_len; + while (lp < lend) { + uint64_t ltag = decode_varint(&lp, lend); + if ((ltag >> 3) == 1 && (ltag & 0x07) == 2) { + uint64_t elem_len = 0; + const uint8_t* elem_data = read_length_delimited(&lp, lend, &elem_len); + if (col_array->n_children > 0) { + append_wire_value(col_array->children[0], SPANNER_TYPE_STRING, elem_data, elem_data + elem_len); + } + } else { + skip_wire_field(&lp, lend, (int)(ltag & 0x07)); + } + } + return ArrowArrayAppendList(col_array); + } else { + skip_wire_field(&p, val_end, wire_type); + } + } + return ArrowArrayAppendNull(col_array, 1); +} + +static void parse_single_wire_prs( + const uint8_t* p, + const uint8_t* end, + int num_cols, + const int* col_type_codes, + struct ArrowArray* out_array, + int* current_col_idx +) { + while (p < end) { + uint64_t tag = decode_varint(&p, end); + int field_num = (int)(tag >> 3); + int wire_type = (int)(tag & 0x07); + + if (field_num == 2 && wire_type == 2) { + uint64_t val_len = 0; + const uint8_t* val_data = read_length_delimited(&p, end, &val_len); + int col_idx = *current_col_idx; + append_wire_value(out_array->children[col_idx], col_type_codes[col_idx], val_data, val_data + val_len); + col_idx++; + if (col_idx == num_cols) { + col_idx = 0; + out_array->length++; + } + *current_col_idx = col_idx; + } else { + skip_wire_field(&p, end, wire_type); + } + } +} + +static PyObject* py_wire_prs_to_c_batch(PyObject* self, PyObject* args) { + PyObject* py_fields; + PyObject* py_wire_chunks; + + if (!PyArg_ParseTuple(args, "OO", &py_fields, &py_wire_chunks)) { + return NULL; + } + + if (!PySequence_Check(py_fields) || !PySequence_Check(py_wire_chunks)) { + PyErr_SetString(PyExc_TypeError, "fields and wire_chunks must be sequences"); + return NULL; + } + + Py_ssize_t num_cols = PySequence_Size(py_fields); + Py_ssize_t num_chunks = PySequence_Size(py_wire_chunks); + + int* type_codes = (int*)malloc(num_cols * sizeof(int)); + if (type_codes == NULL) { + PyErr_NoMemory(); + return NULL; + } + + struct ArrowSchema* out_schema = (struct ArrowSchema*)calloc(1, sizeof(struct ArrowSchema)); + struct ArrowArray* out_array = (struct ArrowArray*)calloc(1, sizeof(struct ArrowArray)); + struct ArrowError error; + + if (out_schema == NULL || out_array == NULL) { + if (out_schema) free(out_schema); + if (out_array) free(out_array); + free(type_codes); + PyErr_NoMemory(); + return NULL; + } + + ArrowSchemaInit(out_schema, NANOARROW_TYPE_STRUCT); + ArrowSchemaAllocateChildren(out_schema, (int64_t)num_cols); + + for (Py_ssize_t i = 0; i < num_cols; i++) { + PyObject* f = PySequence_GetItem(py_fields, i); + configure_field_schema(out_schema->children[i], f); + type_codes[i] = SPANNER_TYPE_STRING; + if (PyTuple_Check(f) && PyTuple_Size(f) >= 2) { + PyObject* t_obj = PyTuple_GET_ITEM(f, 1); + if (PyLong_Check(t_obj)) { + type_codes[i] = (int)PyLong_AsLong(t_obj); + } + } + Py_XDECREF(f); + } + + if (ArrowArrayInitFromSchema(out_array, out_schema, &error) != 0) { + ArrowSchemaRelease(out_schema); + free(out_schema); + free(out_array); + free(type_codes); + PyErr_Format(PyExc_RuntimeError, "Failed to init ArrowArray: %s", error.message); + return NULL; + } + + ArrowArrayStartAppending(out_array); + + typedef struct { + const uint8_t* ptr; + size_t len; + } RawBuf; + + RawBuf* raw_buffers = (RawBuf*)malloc(num_chunks * sizeof(RawBuf)); + if (raw_buffers == NULL) { + ArrowArrayRelease(out_array); + ArrowSchemaRelease(out_schema); + free(out_array); + free(out_schema); + free(type_codes); + PyErr_NoMemory(); + return NULL; + } + + for (Py_ssize_t i = 0; i < num_chunks; i++) { + PyObject* chunk_obj = PySequence_GetItem(py_wire_chunks, i); + if (chunk_obj != NULL && PyBytes_Check(chunk_obj)) { + raw_buffers[i].ptr = (const uint8_t*)PyBytes_AS_STRING(chunk_obj); + raw_buffers[i].len = (size_t)PyBytes_GET_SIZE(chunk_obj); + } else { + raw_buffers[i].ptr = NULL; + raw_buffers[i].len = 0; + } + Py_XDECREF(chunk_obj); + } + + Py_BEGIN_ALLOW_THREADS + int current_col_idx = 0; + for (Py_ssize_t i = 0; i < num_chunks; i++) { + if (raw_buffers[i].ptr != NULL && raw_buffers[i].len > 0) { + parse_single_wire_prs( + raw_buffers[i].ptr, + raw_buffers[i].ptr + raw_buffers[i].len, + (int)num_cols, + type_codes, + out_array, + ¤t_col_idx + ); + } + } + Py_END_ALLOW_THREADS + + free(raw_buffers); + free(type_codes); + + if (ArrowArrayFinishBuildingDefault(out_array, &error) != 0) { + ArrowArrayRelease(out_array); + ArrowSchemaRelease(out_schema); + free(out_array); + free(out_schema); + PyErr_Format(PyExc_RuntimeError, "Failed to finish ArrowArray: %s", error.message); + return NULL; + } + + uintptr_t array_ptr = (uintptr_t)out_array; + uintptr_t schema_ptr = (uintptr_t)out_schema; + + return Py_BuildValue("(KK)", (unsigned long long)array_ptr, (unsigned long long)schema_ptr); +} + +static PyMethodDef SpannerArrowMethods[] = { + {"rows_to_c_batch", py_rows_to_c_batch, METH_VARARGS, + "Convert sequence of Spanner rows into Arrow C Data Interface pointers."}, + {"wire_prs_to_c_batch", py_wire_prs_to_c_batch, METH_VARARGS, + "Convert raw protobuf wire PartialResultSet bytes directly into Arrow C Data Interface pointers without allocating Python objects."}, + {NULL, NULL, 0, NULL} +}; + +static struct PyModuleDef spannerarrowmodule = { + PyModuleDef_HEAD_INIT, + "_spanner_arrow", + "High-performance native Apache Arrow accelerator for Google Cloud Spanner", + -1, + SpannerArrowMethods +}; + +PyMODINIT_FUNC PyInit__spanner_arrow(void) { + return PyModule_Create(&spannerarrowmodule); +} diff --git a/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/cext.py b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/cext.py new file mode 100644 index 000000000000..ae3d28f121fc --- /dev/null +++ b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/cext.py @@ -0,0 +1,96 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""C-extension accelerated implementation of Spanner Arrow converter.""" + +from typing import Any, Optional, Sequence +import pyarrow as pa + +from google_cloud_spanner_arrow import _spanner_arrow +from google_cloud_spanner_arrow.python import ( + fields_to_arrow_schema, + spanner_type_to_arrow_type, + _get_type_code, + SPANNER_TYPE_ARRAY, + SPANNER_TYPE_STRUCT, +) + + +def _normalize_single_field(f: Any) -> tuple: + if isinstance(f, tuple): + name = f[0] + type_obj = f[1] + type_code = _get_type_code(type_obj) + if len(f) >= 3: + sub = f[2] + if type_code == SPANNER_TYPE_ARRAY: + return (name, type_code, _normalize_single_field(sub)) + elif type_code == SPANNER_TYPE_STRUCT and isinstance(sub, (list, tuple)): + return (name, type_code, tuple(_normalize_single_field(sf) for sf in sub)) + return (name, type_code, sub) + else: + name = getattr(f, "name", "col") + type_obj = getattr(f, "type_", 6) + type_code = _get_type_code(type_obj) + + if type_code == SPANNER_TYPE_ARRAY: + elem_type = getattr(type_obj, "array_element_type", 6) + child_tuple = _normalize_single_field(("item", elem_type)) + return (name, type_code, child_tuple) + elif type_code == SPANNER_TYPE_STRUCT: + struct_type = getattr(type_obj, "struct_type", None) + sub_fields = getattr(struct_type, "fields", ()) + children = tuple(_normalize_single_field(sub_f) for sub_f in sub_fields) + return (name, type_code, children) + return (name, type_code) + + +def _normalize_fields(fields: Sequence[Any]): + return [_normalize_single_field(f) for f in fields] + + +def rows_to_arrow_batch( + fields: Sequence[Any], + rows: Sequence[Sequence[Any]], + schema: Optional[pa.Schema] = None, +) -> pa.RecordBatch: + """Convert sequence of Spanner rows to pyarrow.RecordBatch using native C acceleration.""" + if not rows: + if schema is None: + schema = fields_to_arrow_schema(fields) + return pa.RecordBatch.from_arrays( + [pa.array([], type=f.type) for f in schema], schema=schema + ) + + field_tuples = _normalize_fields(fields) + array_ptr, schema_ptr = _spanner_arrow.rows_to_c_batch(field_tuples, rows) + return pa.RecordBatch._import_from_c(array_ptr, schema_ptr) + + +def wire_prs_to_arrow_batch( + fields: Sequence[Any], + wire_chunks: Sequence[bytes], + schema: Optional[pa.Schema] = None, +) -> pa.RecordBatch: + """Convert sequence of raw PartialResultSet protobuf wire bytes to pyarrow.RecordBatch in C.""" + if not wire_chunks: + if schema is None: + schema = fields_to_arrow_schema(fields) + return pa.RecordBatch.from_arrays( + [pa.array([], type=f.type) for f in schema], schema=schema + ) + + field_tuples = _normalize_fields(fields) + array_ptr, schema_ptr = _spanner_arrow.wire_prs_to_c_batch(field_tuples, wire_chunks) + return pa.RecordBatch._import_from_c(array_ptr, schema_ptr) diff --git a/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/nanoarrow/nanoarrow.c b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/nanoarrow/nanoarrow.c new file mode 100644 index 000000000000..cd4d3605d68e --- /dev/null +++ b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/nanoarrow/nanoarrow.c @@ -0,0 +1,698 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "nanoarrow.h" +#include +#include + +void ArrowErrorSet(struct ArrowError* error, const char* fmt, ...) { + if (error == NULL) return; + va_list args; + va_start(args, fmt); + vsnprintf(error->message, sizeof(error->message), fmt, args); + va_end(args); +} + +void ArrowBufferInit(struct ArrowBuffer* buffer) { + buffer->data = NULL; + buffer->size_bytes = 0; + buffer->capacity_bytes = 0; +} + +int ArrowBufferReserve(struct ArrowBuffer* buffer, int64_t additional_bytes) { + int64_t target_capacity = buffer->size_bytes + additional_bytes; + if (target_capacity <= buffer->capacity_bytes) { + return 0; + } + int64_t new_capacity = buffer->capacity_bytes == 0 ? 64 : buffer->capacity_bytes * 2; + while (new_capacity < target_capacity) { + new_capacity *= 2; + } + uint8_t* new_data = (uint8_t*)realloc(buffer->data, (size_t)new_capacity); + if (new_data == NULL) { + return -1; + } + buffer->data = new_data; + buffer->capacity_bytes = new_capacity; + return 0; +} + +int ArrowBufferAppend(struct ArrowBuffer* buffer, const void* data, int64_t size_bytes) { + if (ArrowBufferReserve(buffer, size_bytes) != 0) { + return -1; + } + if (data != NULL && size_bytes > 0) { + memcpy(buffer->data + buffer->size_bytes, data, (size_t)size_bytes); + } + buffer->size_bytes += size_bytes; + return 0; +} + +int ArrowBufferAppendFill(struct ArrowBuffer* buffer, uint8_t value, int64_t size_bytes) { + if (ArrowBufferReserve(buffer, size_bytes) != 0) { + return -1; + } + if (size_bytes > 0) { + memset(buffer->data + buffer->size_bytes, value, (size_t)size_bytes); + } + buffer->size_bytes += size_bytes; + return 0; +} + +void ArrowBufferReset(struct ArrowBuffer* buffer) { + if (buffer->data != NULL) { + free(buffer->data); + buffer->data = NULL; + } + buffer->size_bytes = 0; + buffer->capacity_bytes = 0; +} + +void ArrowBitmapInit(struct ArrowBitmap* bitmap) { + ArrowBufferInit(&bitmap->buffer); + bitmap->null_count = 0; +} + +int ArrowBitmapReserve(struct ArrowBitmap* bitmap, int64_t additional_elements) { + int64_t current_elements = bitmap->buffer.size_bytes * 8; + int64_t target_elements = current_elements + additional_elements; + int64_t target_bytes = (target_elements + 7) / 8; + return ArrowBufferReserve(&bitmap->buffer, target_bytes - bitmap->buffer.size_bytes); +} + +int ArrowBitmapAppend(struct ArrowBitmap* bitmap, uint8_t is_valid, int64_t count) { + for (int64_t i = 0; i < count; i++) { + int64_t bit_index = bitmap->buffer.size_bytes * 8; + if (bit_index % 8 == 0) { + uint8_t zero = 0; + if (ArrowBufferAppend(&bitmap->buffer, &zero, 1) != 0) { + return -1; + } + } + int64_t byte_pos = bitmap->buffer.size_bytes - 1; + int bit_pos = (int)((bit_index) % 8); + if (is_valid) { + bitmap->buffer.data[byte_pos] |= (uint8_t)(1 << bit_pos); + } else { + bitmap->null_count++; + } + } + return 0; +} + +void ArrowBitmapReset(struct ArrowBitmap* bitmap) { + ArrowBufferReset(&bitmap->buffer); + bitmap->null_count = 0; +} + +static char* nanoarrow_strdup(const char* s) { + if (s == NULL) return NULL; + size_t len = strlen(s); + char* copy = (char*)malloc(len + 1); + if (copy) { + memcpy(copy, s, len + 1); + } + return copy; +} + +void ArrowSchemaRelease(struct ArrowSchema* schema) { + if (schema == NULL || schema->release == NULL) { + return; + } + if (schema->format != NULL) { + free((void*)schema->format); + schema->format = NULL; + } + if (schema->name != NULL) { + free((void*)schema->name); + schema->name = NULL; + } + if (schema->metadata != NULL) { + free((void*)schema->metadata); + schema->metadata = NULL; + } + if (schema->children != NULL) { + for (int64_t i = 0; i < schema->n_children; i++) { + if (schema->children[i] != NULL) { + if (schema->children[i]->release != NULL) { + schema->children[i]->release(schema->children[i]); + } + free(schema->children[i]); + } + } + free(schema->children); + schema->children = NULL; + } + if (schema->dictionary != NULL) { + if (schema->dictionary->release != NULL) { + schema->dictionary->release(schema->dictionary); + } + free(schema->dictionary); + schema->dictionary = NULL; + } + schema->release = NULL; +} + +void ArrowSchemaInit(struct ArrowSchema* schema, enum ArrowType type) { + schema->format = NULL; + schema->name = NULL; + schema->metadata = NULL; + schema->flags = ARROW_FLAG_NULLABLE; + schema->n_children = 0; + schema->children = NULL; + schema->dictionary = NULL; + schema->release = &ArrowSchemaRelease; + schema->private_data = NULL; + + switch (type) { + case NANOARROW_TYPE_BOOL: + ArrowSchemaSetFormat(schema, "b"); + break; + case NANOARROW_TYPE_INT8: + ArrowSchemaSetFormat(schema, "c"); + break; + case NANOARROW_TYPE_UINT8: + ArrowSchemaSetFormat(schema, "C"); + break; + case NANOARROW_TYPE_INT16: + ArrowSchemaSetFormat(schema, "s"); + break; + case NANOARROW_TYPE_UINT16: + ArrowSchemaSetFormat(schema, "S"); + break; + case NANOARROW_TYPE_INT32: + ArrowSchemaSetFormat(schema, "i"); + break; + case NANOARROW_TYPE_UINT32: + ArrowSchemaSetFormat(schema, "I"); + break; + case NANOARROW_TYPE_INT64: + ArrowSchemaSetFormat(schema, "l"); + break; + case NANOARROW_TYPE_UINT64: + ArrowSchemaSetFormat(schema, "L"); + break; + case NANOARROW_TYPE_FLOAT: + ArrowSchemaSetFormat(schema, "f"); + break; + case NANOARROW_TYPE_DOUBLE: + ArrowSchemaSetFormat(schema, "g"); + break; + case NANOARROW_TYPE_STRING: + ArrowSchemaSetFormat(schema, "u"); + break; + case NANOARROW_TYPE_BINARY: + ArrowSchemaSetFormat(schema, "z"); + break; + case NANOARROW_TYPE_DATE32: + ArrowSchemaSetFormat(schema, "tdD"); + break; + case NANOARROW_TYPE_TIMESTAMP: + ArrowSchemaSetFormat(schema, "tsu:UTC"); + break; + case NANOARROW_TYPE_DECIMAL128: + ArrowSchemaSetFormat(schema, "d:38,9"); + break; + case NANOARROW_TYPE_LIST: + ArrowSchemaSetFormat(schema, "+l"); + break; + case NANOARROW_TYPE_STRUCT: + ArrowSchemaSetFormat(schema, "+s"); + break; + case NANOARROW_TYPE_NA: + ArrowSchemaSetFormat(schema, "n"); + break; + default: + ArrowSchemaSetFormat(schema, "n"); + break; + } +} + +int ArrowSchemaSetFormat(struct ArrowSchema* schema, const char* format) { + if (schema->format != NULL) { + free((void*)schema->format); + } + schema->format = nanoarrow_strdup(format); + return schema->format == NULL ? -1 : 0; +} + +int ArrowSchemaSetName(struct ArrowSchema* schema, const char* name) { + if (schema->name != NULL) { + free((void*)schema->name); + } + schema->name = nanoarrow_strdup(name); + return schema->name == NULL ? -1 : 0; +} + +int ArrowSchemaAllocateChildren(struct ArrowSchema* schema, int64_t n_children) { + schema->n_children = n_children; + schema->children = (struct ArrowSchema**)calloc((size_t)n_children, sizeof(struct ArrowSchema*)); + if (schema->children == NULL) { + return -1; + } + for (int64_t i = 0; i < n_children; i++) { + schema->children[i] = (struct ArrowSchema*)calloc(1, sizeof(struct ArrowSchema)); + if (schema->children[i] == NULL) { + return -1; + } + ArrowSchemaInit(schema->children[i], NANOARROW_TYPE_UNINITIALIZED); + } + return 0; +} + +void ArrowArrayRelease(struct ArrowArray* array) { + if (array == NULL || array->release == NULL) { + return; + } + if (array->private_data != NULL) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + ArrowBitmapReset(&private_data->bitmap); + ArrowBufferReset(&private_data->buffer1); + ArrowBufferReset(&private_data->buffer2); + free(private_data); + array->private_data = NULL; + } + if (array->buffers != NULL) { + free((void*)array->buffers); + array->buffers = NULL; + } + if (array->children != NULL) { + for (int64_t i = 0; i < array->n_children; i++) { + if (array->children[i] != NULL) { + if (array->children[i]->release != NULL) { + array->children[i]->release(array->children[i]); + } + free(array->children[i]); + } + } + free(array->children); + array->children = NULL; + } + if (array->dictionary != NULL) { + if (array->dictionary->release != NULL) { + array->dictionary->release(array->dictionary); + } + free(array->dictionary); + array->dictionary = NULL; + } + array->release = NULL; +} + +static enum ArrowType type_from_format(const char* format) { + if (format == NULL) return NANOARROW_TYPE_UNINITIALIZED; + if (strcmp(format, "b") == 0) return NANOARROW_TYPE_BOOL; + if (strcmp(format, "c") == 0) return NANOARROW_TYPE_INT8; + if (strcmp(format, "C") == 0) return NANOARROW_TYPE_UINT8; + if (strcmp(format, "s") == 0) return NANOARROW_TYPE_INT16; + if (strcmp(format, "S") == 0) return NANOARROW_TYPE_UINT16; + if (strcmp(format, "i") == 0) return NANOARROW_TYPE_INT32; + if (strcmp(format, "I") == 0) return NANOARROW_TYPE_UINT32; + if (strcmp(format, "l") == 0) return NANOARROW_TYPE_INT64; + if (strcmp(format, "L") == 0) return NANOARROW_TYPE_UINT64; + if (strcmp(format, "f") == 0) return NANOARROW_TYPE_FLOAT; + if (strcmp(format, "g") == 0) return NANOARROW_TYPE_DOUBLE; + if (strcmp(format, "u") == 0) return NANOARROW_TYPE_STRING; + if (strcmp(format, "z") == 0) return NANOARROW_TYPE_BINARY; + if (strcmp(format, "tdD") == 0) return NANOARROW_TYPE_DATE32; + if (strncmp(format, "ts", 2) == 0) return NANOARROW_TYPE_TIMESTAMP; + if (strncmp(format, "d:", 2) == 0) return NANOARROW_TYPE_DECIMAL128; + if (strcmp(format, "+l") == 0) return NANOARROW_TYPE_LIST; + if (strcmp(format, "+s") == 0) return NANOARROW_TYPE_STRUCT; + if (strcmp(format, "n") == 0) return NANOARROW_TYPE_NA; + return NANOARROW_TYPE_UNINITIALIZED; +} + +int ArrowArrayInitFromSchema(struct ArrowArray* array, struct ArrowSchema* schema, struct ArrowError* error) { + array->length = 0; + array->null_count = 0; + array->offset = 0; + array->n_buffers = 0; + array->n_children = 0; + array->buffers = NULL; + array->children = NULL; + array->dictionary = NULL; + array->release = &ArrowArrayRelease; + array->private_data = NULL; + + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)calloc(1, sizeof(struct ArrowArrayPrivateData)); + if (private_data == NULL) { + ArrowErrorSet(error, "Failed to allocate private data for ArrowArray"); + return -1; + } + ArrowBitmapInit(&private_data->bitmap); + ArrowBufferInit(&private_data->buffer1); + ArrowBufferInit(&private_data->buffer2); + private_data->type = type_from_format(schema->format); + array->private_data = private_data; + + switch (private_data->type) { + case NANOARROW_TYPE_BOOL: + case NANOARROW_TYPE_INT8: + case NANOARROW_TYPE_UINT8: + case NANOARROW_TYPE_INT16: + case NANOARROW_TYPE_UINT16: + case NANOARROW_TYPE_INT32: + case NANOARROW_TYPE_UINT32: + case NANOARROW_TYPE_INT64: + case NANOARROW_TYPE_UINT64: + case NANOARROW_TYPE_FLOAT: + case NANOARROW_TYPE_DOUBLE: + case NANOARROW_TYPE_DATE32: + case NANOARROW_TYPE_TIMESTAMP: + case NANOARROW_TYPE_DECIMAL128: + array->n_buffers = 2; + array->buffers = (const void**)calloc(2, sizeof(void*)); + break; + case NANOARROW_TYPE_STRING: + case NANOARROW_TYPE_BINARY: + array->n_buffers = 3; + array->buffers = (const void**)calloc(3, sizeof(void*)); + break; + case NANOARROW_TYPE_STRUCT: + array->n_buffers = 1; + array->buffers = (const void**)calloc(1, sizeof(void*)); + if (schema->n_children > 0) { + array->n_children = schema->n_children; + array->children = (struct ArrowArray**)calloc((size_t)schema->n_children, sizeof(struct ArrowArray*)); + for (int64_t i = 0; i < schema->n_children; i++) { + array->children[i] = (struct ArrowArray*)calloc(1, sizeof(struct ArrowArray)); + if (ArrowArrayInitFromSchema(array->children[i], schema->children[i], error) != 0) { + return -1; + } + } + } + break; + case NANOARROW_TYPE_LIST: + array->n_buffers = 2; + array->buffers = (const void**)calloc(2, sizeof(void*)); + if (schema->n_children == 1) { + array->n_children = 1; + array->children = (struct ArrowArray**)calloc(1, sizeof(struct ArrowArray*)); + array->children[0] = (struct ArrowArray*)calloc(1, sizeof(struct ArrowArray)); + if (ArrowArrayInitFromSchema(array->children[0], schema->children[0], error) != 0) { + return -1; + } + } + break; + case NANOARROW_TYPE_NA: + array->n_buffers = 0; + break; + default: + array->n_buffers = 2; + array->buffers = (const void**)calloc(2, sizeof(void*)); + break; + } + return 0; +} + +int ArrowArrayStartAppending(struct ArrowArray* array) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + + if (private_data->type == NANOARROW_TYPE_STRING || + private_data->type == NANOARROW_TYPE_BINARY || + private_data->type == NANOARROW_TYPE_LIST) { + int32_t zero_offset = 0; + if (ArrowBufferAppend(&private_data->buffer1, &zero_offset, sizeof(int32_t)) != 0) { + return -1; + } + } + if (array->n_children > 0 && array->children != NULL) { + for (int64_t i = 0; i < array->n_children; i++) { + if (ArrowArrayStartAppending(array->children[i]) != 0) { + return -1; + } + } + } + return 0; +} + +int ArrowArrayAppendNull(struct ArrowArray* array, int64_t n) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + + for (int64_t i = 0; i < n; i++) { + if (ArrowBitmapAppend(&private_data->bitmap, 0, 1) != 0) return -1; + array->length++; + + switch (private_data->type) { + case NANOARROW_TYPE_BOOL: { + int64_t bit_idx = private_data->buffer1.size_bytes * 8; + if (bit_idx % 8 == 0) { + uint8_t zero = 0; + ArrowBufferAppend(&private_data->buffer1, &zero, 1); + } + break; + } + case NANOARROW_TYPE_INT8: + case NANOARROW_TYPE_UINT8: { + uint8_t zero = 0; + ArrowBufferAppend(&private_data->buffer1, &zero, 1); + break; + } + case NANOARROW_TYPE_INT16: + case NANOARROW_TYPE_UINT16: { + int16_t zero = 0; + ArrowBufferAppend(&private_data->buffer1, &zero, 2); + break; + } + case NANOARROW_TYPE_INT32: + case NANOARROW_TYPE_UINT32: + case NANOARROW_TYPE_DATE32: { + int32_t zero = 0; + ArrowBufferAppend(&private_data->buffer1, &zero, 4); + break; + } + case NANOARROW_TYPE_INT64: + case NANOARROW_TYPE_UINT64: + case NANOARROW_TYPE_TIMESTAMP: { + int64_t zero = 0; + ArrowBufferAppend(&private_data->buffer1, &zero, 8); + break; + } + case NANOARROW_TYPE_FLOAT: { + float zero = 0.0f; + ArrowBufferAppend(&private_data->buffer1, &zero, 4); + break; + } + case NANOARROW_TYPE_DOUBLE: { + double zero = 0.0; + ArrowBufferAppend(&private_data->buffer1, &zero, 8); + break; + } + case NANOARROW_TYPE_DECIMAL128: { + uint8_t zero[16] = {0}; + ArrowBufferAppend(&private_data->buffer1, zero, 16); + break; + } + case NANOARROW_TYPE_STRING: + case NANOARROW_TYPE_BINARY: { + int32_t current_offset = (int32_t)private_data->buffer2.size_bytes; + ArrowBufferAppend(&private_data->buffer1, ¤t_offset, sizeof(int32_t)); + break; + } + case NANOARROW_TYPE_LIST: { + int32_t child_len = array->n_children > 0 ? (int32_t)array->children[0]->length : 0; + ArrowBufferAppend(&private_data->buffer1, &child_len, sizeof(int32_t)); + break; + } + case NANOARROW_TYPE_STRUCT: { + // For struct null, append null to children + if (array->n_children > 0 && array->children != NULL) { + for (int64_t c = 0; c < array->n_children; c++) { + ArrowArrayAppendNull(array->children[c], 1); + } + } + break; + } + default: + break; + } + } + return 0; +} + +int ArrowArrayAppendInt(struct ArrowArray* array, int64_t value) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + array->length++; + + switch (private_data->type) { + case NANOARROW_TYPE_INT8: { + int8_t v = (int8_t)value; + return ArrowBufferAppend(&private_data->buffer1, &v, sizeof(int8_t)); + } + case NANOARROW_TYPE_INT16: { + int16_t v = (int16_t)value; + return ArrowBufferAppend(&private_data->buffer1, &v, sizeof(int16_t)); + } + case NANOARROW_TYPE_INT32: + case NANOARROW_TYPE_DATE32: { + int32_t v = (int32_t)value; + return ArrowBufferAppend(&private_data->buffer1, &v, sizeof(int32_t)); + } + case NANOARROW_TYPE_INT64: + case NANOARROW_TYPE_TIMESTAMP: { + return ArrowBufferAppend(&private_data->buffer1, &value, sizeof(int64_t)); + } + default: + return ArrowBufferAppend(&private_data->buffer1, &value, sizeof(int64_t)); + } +} + +int ArrowArrayAppendDouble(struct ArrowArray* array, double value) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + array->length++; + return ArrowBufferAppend(&private_data->buffer1, &value, sizeof(double)); +} + +int ArrowArrayAppendFloat(struct ArrowArray* array, float value) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + array->length++; + return ArrowBufferAppend(&private_data->buffer1, &value, sizeof(float)); +} + +int ArrowArrayAppendBool(struct ArrowArray* array, uint8_t value) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + + int64_t bit_index = array->length; + array->length++; + if (bit_index % 8 == 0) { + uint8_t zero = 0; + if (ArrowBufferAppend(&private_data->buffer1, &zero, 1) != 0) { + return -1; + } + } + int64_t byte_pos = private_data->buffer1.size_bytes - 1; + int bit_pos = (int)(bit_index % 8); + if (value) { + private_data->buffer1.data[byte_pos] |= (uint8_t)(1 << bit_pos); + } + return 0; +} + +int ArrowArrayAppendString(struct ArrowArray* array, struct ArrowStringView value) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + array->length++; + + if (value.data != NULL && value.size_bytes > 0) { + if (ArrowBufferAppend(&private_data->buffer2, value.data, value.size_bytes) != 0) { + return -1; + } + } + int32_t new_offset = (int32_t)private_data->buffer2.size_bytes; + return ArrowBufferAppend(&private_data->buffer1, &new_offset, sizeof(int32_t)); +} + +int ArrowArrayAppendBytes(struct ArrowArray* array, struct ArrowBufferView value) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + array->length++; + + if (value.data != NULL && value.size_bytes > 0) { + if (ArrowBufferAppend(&private_data->buffer2, value.data, value.size_bytes) != 0) { + return -1; + } + } + int32_t new_offset = (int32_t)private_data->buffer2.size_bytes; + return ArrowBufferAppend(&private_data->buffer1, &new_offset, sizeof(int32_t)); +} + +int ArrowArrayAppendDecimal128(struct ArrowArray* array, struct ArrowDecimal128 value) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + array->length++; + return ArrowBufferAppend(&private_data->buffer1, value.bytes, sizeof(value.bytes)); +} + +int ArrowArrayAppendList(struct ArrowArray* array) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + array->length++; + int32_t child_len = array->n_children > 0 ? (int32_t)array->children[0]->length : 0; + return ArrowBufferAppend(&private_data->buffer1, &child_len, sizeof(int32_t)); +} + +int ArrowArrayAppendStruct(struct ArrowArray* array) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return -1; + if (ArrowBitmapAppend(&private_data->bitmap, 1, 1) != 0) return -1; + array->length++; + return 0; +} + +int ArrowArrayFinishBuildingDefault(struct ArrowArray* array, struct ArrowError* error) { + struct ArrowArrayPrivateData* private_data = (struct ArrowArrayPrivateData*)array->private_data; + if (private_data == NULL) return 0; + + array->null_count = private_data->bitmap.null_count; + + switch (private_data->type) { + case NANOARROW_TYPE_BOOL: + case NANOARROW_TYPE_INT8: + case NANOARROW_TYPE_UINT8: + case NANOARROW_TYPE_INT16: + case NANOARROW_TYPE_UINT16: + case NANOARROW_TYPE_INT32: + case NANOARROW_TYPE_UINT32: + case NANOARROW_TYPE_INT64: + case NANOARROW_TYPE_UINT64: + case NANOARROW_TYPE_FLOAT: + case NANOARROW_TYPE_DOUBLE: + case NANOARROW_TYPE_DATE32: + case NANOARROW_TYPE_TIMESTAMP: + case NANOARROW_TYPE_DECIMAL128: + array->buffers[0] = (array->null_count > 0) ? private_data->bitmap.buffer.data : NULL; + array->buffers[1] = private_data->buffer1.data; + break; + case NANOARROW_TYPE_STRING: + case NANOARROW_TYPE_BINARY: + array->buffers[0] = (array->null_count > 0) ? private_data->bitmap.buffer.data : NULL; + array->buffers[1] = private_data->buffer1.data; + array->buffers[2] = private_data->buffer2.data; + break; + case NANOARROW_TYPE_STRUCT: + array->buffers[0] = (array->null_count > 0) ? private_data->bitmap.buffer.data : NULL; + for (int64_t i = 0; i < array->n_children; i++) { + if (ArrowArrayFinishBuildingDefault(array->children[i], error) != 0) { + return -1; + } + } + break; + case NANOARROW_TYPE_LIST: + array->buffers[0] = (array->null_count > 0) ? private_data->bitmap.buffer.data : NULL; + array->buffers[1] = private_data->buffer1.data; + if (array->n_children == 1) { + if (ArrowArrayFinishBuildingDefault(array->children[0], error) != 0) { + return -1; + } + } + break; + default: + break; + } + return 0; +} diff --git a/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/nanoarrow/nanoarrow.h b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/nanoarrow/nanoarrow.h new file mode 100644 index 000000000000..1700006b24e2 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/nanoarrow/nanoarrow.h @@ -0,0 +1,169 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef NANOARROW_H_INCLUDED +#define NANOARROW_H_INCLUDED + +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Arrow C Data Interface specification definitions +#ifndef ARROW_C_DATA_INTERFACE +#define ARROW_C_DATA_INTERFACE + +#define ARROW_FLAG_DICTIONARY_ORDERED 1 +#define ARROW_FLAG_NULLABLE 2 +#define ARROW_FLAG_MAP_KEYS_SORTED 4 + +struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; +}; + +struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; +}; + +#endif // ARROW_C_DATA_INTERFACE + +enum ArrowType { + NANOARROW_TYPE_UNINITIALIZED = 0, + NANOARROW_TYPE_NA = 1, + NANOARROW_TYPE_BOOL = 2, + NANOARROW_TYPE_INT8 = 3, + NANOARROW_TYPE_UINT8 = 4, + NANOARROW_TYPE_INT16 = 5, + NANOARROW_TYPE_UINT16 = 6, + NANOARROW_TYPE_INT32 = 7, + NANOARROW_TYPE_UINT32 = 8, + NANOARROW_TYPE_INT64 = 9, + NANOARROW_TYPE_UINT64 = 10, + NANOARROW_TYPE_FLOAT = 11, + NANOARROW_TYPE_DOUBLE = 12, + NANOARROW_TYPE_STRING = 13, + NANOARROW_TYPE_BINARY = 14, + NANOARROW_TYPE_DATE32 = 15, + NANOARROW_TYPE_TIMESTAMP = 16, + NANOARROW_TYPE_DECIMAL128 = 17, + NANOARROW_TYPE_LIST = 18, + NANOARROW_TYPE_STRUCT = 19 +}; + +struct ArrowError { + char message[1024]; +}; + +struct ArrowStringView { + const char* data; + int64_t size_bytes; +}; + +struct ArrowBufferView { + const void* data; + int64_t size_bytes; +}; + +struct ArrowDecimal128 { + uint8_t bytes[16]; +}; + +struct ArrowBuffer { + uint8_t* data; + int64_t size_bytes; + int64_t capacity_bytes; +}; + +struct ArrowBitmap { + struct ArrowBuffer buffer; + int64_t null_count; +}; + +struct ArrowArrayPrivateData { + enum ArrowType type; + struct ArrowBitmap bitmap; + struct ArrowBuffer buffer1; // Offsets (for string/binary/list) or data (for primitive) + struct ArrowBuffer buffer2; // Data (for string/binary) +}; + +// Buffer functions +void ArrowBufferInit(struct ArrowBuffer* buffer); +int ArrowBufferReserve(struct ArrowBuffer* buffer, int64_t additional_bytes); +int ArrowBufferAppend(struct ArrowBuffer* buffer, const void* data, int64_t size_bytes); +int ArrowBufferAppendFill(struct ArrowBuffer* buffer, uint8_t value, int64_t size_bytes); +void ArrowBufferReset(struct ArrowBuffer* buffer); + +// Bitmap functions +void ArrowBitmapInit(struct ArrowBitmap* bitmap); +int ArrowBitmapReserve(struct ArrowBitmap* bitmap, int64_t additional_elements); +int ArrowBitmapAppend(struct ArrowBitmap* bitmap, uint8_t is_valid, int64_t count); +static inline void ArrowBitmapSet(struct ArrowBitmap* bitmap, int64_t index, uint8_t is_valid) { + if (is_valid) { + bitmap->buffer.data[index / 8] |= (uint8_t)(1 << (index % 8)); + } else { + bitmap->buffer.data[index / 8] &= (uint8_t)~(1 << (index % 8)); + bitmap->null_count++; + } +} +void ArrowBitmapReset(struct ArrowBitmap* bitmap); + +// Schema functions +void ArrowSchemaInit(struct ArrowSchema* schema, enum ArrowType type); +int ArrowSchemaSetFormat(struct ArrowSchema* schema, const char* format); +int ArrowSchemaSetName(struct ArrowSchema* schema, const char* name); +int ArrowSchemaAllocateChildren(struct ArrowSchema* schema, int64_t n_children); +void ArrowSchemaRelease(struct ArrowSchema* schema); + +// Array functions +int ArrowArrayInitFromSchema(struct ArrowArray* array, struct ArrowSchema* schema, struct ArrowError* error); +int ArrowArrayStartAppending(struct ArrowArray* array); +int ArrowArrayAppendNull(struct ArrowArray* array, int64_t n); +int ArrowArrayAppendInt(struct ArrowArray* array, int64_t value); +int ArrowArrayAppendDouble(struct ArrowArray* array, double value); +int ArrowArrayAppendFloat(struct ArrowArray* array, float value); +int ArrowArrayAppendBool(struct ArrowArray* array, uint8_t value); +int ArrowArrayAppendString(struct ArrowArray* array, struct ArrowStringView value); +int ArrowArrayAppendBytes(struct ArrowArray* array, struct ArrowBufferView value); +int ArrowArrayAppendDecimal128(struct ArrowArray* array, struct ArrowDecimal128 value); +int ArrowArrayAppendList(struct ArrowArray* array); +int ArrowArrayAppendStruct(struct ArrowArray* array); +int ArrowArrayFinishBuildingDefault(struct ArrowArray* array, struct ArrowError* error); +void ArrowArrayRelease(struct ArrowArray* array); + +#ifdef __cplusplus +} +#endif + +#endif // NANOARROW_H_INCLUDED diff --git a/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/py.typed b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/py.typed new file mode 100644 index 000000000000..1242d4327701 --- /dev/null +++ b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/py.typed @@ -0,0 +1 @@ +# Marker file for PEP 561. diff --git a/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/python.py b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/python.py new file mode 100644 index 000000000000..294f316e538d --- /dev/null +++ b/packages/google-cloud-spanner-arrow/src/google_cloud_spanner_arrow/python.py @@ -0,0 +1,234 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pure-Python fallback implementation of Spanner Arrow converter.""" + +import base64 +import decimal +import math +from typing import Any, List, Optional, Sequence + +try: + import pyarrow as pa + import pyarrow.compute as pc + _HAS_PYARROW = True +except ImportError: # pragma: NO COVER + _HAS_PYARROW = False + pa = None + pc = None + +# Spanner TypeCode constants (matching google.cloud.spanner_v1.types.TypeCode) +SPANNER_TYPE_UNSPECIFIED = 0 +SPANNER_TYPE_BOOL = 1 +SPANNER_TYPE_INT64 = 2 +SPANNER_TYPE_FLOAT64 = 3 +SPANNER_TYPE_TIMESTAMP = 4 +SPANNER_TYPE_DATE = 5 +SPANNER_TYPE_STRING = 6 +SPANNER_TYPE_BYTES = 7 +SPANNER_TYPE_ARRAY = 8 +SPANNER_TYPE_STRUCT = 9 +SPANNER_TYPE_NUMERIC = 10 +SPANNER_TYPE_JSON = 11 +SPANNER_TYPE_PROTO = 13 +SPANNER_TYPE_ENUM = 14 +SPANNER_TYPE_FLOAT32 = 15 +SPANNER_TYPE_INTERVAL = 16 +SPANNER_TYPE_UUID = 17 + + +def _check_pyarrow(): + if not _HAS_PYARROW: + raise ImportError( + "pyarrow is required to use Arrow features. " + "Install it with `pip install pyarrow`." + ) + + +def _get_type_code(type_obj: Any) -> int: + if isinstance(type_obj, int): + return type_obj + if hasattr(type_obj, "code"): + return type_obj.code if isinstance(type_obj.code, int) else int(type_obj.code) + return SPANNER_TYPE_STRING + + +def spanner_type_to_arrow_type(spanner_type: Any) -> "pa.DataType": + """Map a Spanner Type to a PyArrow DataType.""" + _check_pyarrow() + code = _get_type_code(spanner_type) + + if code == SPANNER_TYPE_BOOL: + return pa.bool_() + elif code in (SPANNER_TYPE_INT64, SPANNER_TYPE_ENUM): + return pa.int64() + elif code == SPANNER_TYPE_FLOAT32: + return pa.float32() + elif code == SPANNER_TYPE_FLOAT64: + return pa.float64() + elif code in (SPANNER_TYPE_STRING, SPANNER_TYPE_JSON, SPANNER_TYPE_INTERVAL, SPANNER_TYPE_UUID): + return pa.string() + elif code in (SPANNER_TYPE_BYTES, SPANNER_TYPE_PROTO): + return pa.binary() + elif code == SPANNER_TYPE_TIMESTAMP: + return pa.timestamp("us", tz="UTC") + elif code == SPANNER_TYPE_DATE: + return pa.date32() + elif code == SPANNER_TYPE_NUMERIC: + return pa.decimal128(38, 9) + elif code == SPANNER_TYPE_ARRAY: + elem_type = getattr(spanner_type, "array_element_type", SPANNER_TYPE_STRING) + return pa.list_(spanner_type_to_arrow_type(elem_type)) + elif code == SPANNER_TYPE_STRUCT: + struct_type = getattr(spanner_type, "struct_type", None) + fields = getattr(struct_type, "fields", ()) + arrow_fields = [ + pa.field(f.name, spanner_type_to_arrow_type(f.type_)) for f in fields + ] + return pa.struct(arrow_fields) + return pa.string() + + +def fields_to_arrow_schema(fields: Sequence[Any]) -> "pa.Schema": + """Convert sequence of Spanner Field descriptors to a PyArrow Schema.""" + _check_pyarrow() + arrow_fields = [] + for f in fields: + if isinstance(f, tuple): + name, type_code = f[0], f[1] + if len(f) >= 3 and _get_type_code(type_code) == SPANNER_TYPE_STRUCT: + sub_fields = [ + pa.field(sf[0], spanner_type_to_arrow_type(sf[1])) for sf in f[2] + ] + arrow_fields.append(pa.field(name, pa.struct(sub_fields))) + elif len(f) >= 3 and _get_type_code(type_code) == SPANNER_TYPE_ARRAY: + sub_elem = f[2] + sub_type = sub_elem[1] if isinstance(sub_elem, tuple) else sub_elem + arrow_fields.append( + pa.field(name, pa.list_(spanner_type_to_arrow_type(sub_type))) + ) + else: + arrow_fields.append(pa.field(name, spanner_type_to_arrow_type(type_code))) + else: + name = getattr(f, "name", "col") + type_obj = getattr(f, "type_", SPANNER_TYPE_STRING) + arrow_fields.append(pa.field(name, spanner_type_to_arrow_type(type_obj))) + return pa.schema(arrow_fields) + + +def _extract_cell_value(cell: Any, type_code: int) -> Any: + if cell is None: + return None + + if hasattr(cell, "WhichOneof"): + kind = cell.WhichOneof("kind") + if kind == "null_value" or kind is None: + return None + elif kind == "bool_value": + return cell.bool_value + elif kind == "number_value": + return cell.number_value + elif kind == "string_value": + val_str = cell.string_value + if type_code in (SPANNER_TYPE_BYTES, SPANNER_TYPE_PROTO): + return base64.b64decode(val_str) + elif type_code in (SPANNER_TYPE_FLOAT32, SPANNER_TYPE_FLOAT64): + if val_str == "NaN": + return float("nan") + elif val_str == "Infinity": + return float("inf") + elif val_str == "-Infinity": + return float("-inf") + return float(val_str) + return val_str + elif kind == "list_value": + return [ + _extract_cell_value(elem, SPANNER_TYPE_STRING) + for elem in cell.list_value.values + ] + elif kind == "struct_value": + return { + k: _extract_cell_value(v, SPANNER_TYPE_STRING) + for k, v in cell.struct_value.fields.items() + } + return None + + if type_code in (SPANNER_TYPE_BYTES, SPANNER_TYPE_PROTO) and isinstance(cell, str): + return base64.b64decode(cell) + return cell + + +def convert_column_to_arrow_array( + column_values: List[Any], + arrow_field: "pa.Field", + type_code: Optional[int] = None, +) -> "pa.Array": + """Convert a single column's raw values into a PyArrow Array with fast casting.""" + arrow_type = arrow_field.type + if not column_values: + return pa.array([], type=arrow_type) + + first_non_null = next((v for v in column_values if v is not None), None) + if isinstance(first_non_null, str): + if type_code in (SPANNER_TYPE_INT64, SPANNER_TYPE_ENUM): + return pc.cast(pa.array(column_values, type=pa.string()), pa.int64()) + elif type_code == SPANNER_TYPE_FLOAT32: + return pc.cast(pa.array(column_values, type=pa.string()), pa.float32()) + elif type_code == SPANNER_TYPE_FLOAT64: + return pc.cast(pa.array(column_values, type=pa.string()), pa.float64()) + elif type_code == SPANNER_TYPE_DATE: + return pc.cast(pa.array(column_values, type=pa.string()), pa.date32()) + elif type_code == SPANNER_TYPE_TIMESTAMP: + return pc.cast( + pa.array(column_values, type=pa.string()), + pa.timestamp("us", tz="UTC"), + ) + elif type_code == SPANNER_TYPE_NUMERIC: + return pc.cast(pa.array(column_values, type=pa.string()), arrow_type) + + return pa.array(column_values, type=arrow_type) + + +def rows_to_arrow_batch( + fields: Sequence[Any], rows: Sequence[Sequence[Any]], schema: Optional["pa.Schema"] = None +) -> "pa.RecordBatch": + """Convert sequence of rows to pyarrow.RecordBatch in pure Python.""" + _check_pyarrow() + if schema is None: + schema = fields_to_arrow_schema(fields) + + num_cols = len(fields) + type_codes = [ + _get_type_code(f[1] if isinstance(f, tuple) else getattr(f, "type_", SPANNER_TYPE_STRING)) + for f in fields + ] + + if not rows: + return pa.RecordBatch.from_arrays( + [pa.array([], type=f.type) for f in schema], schema=schema + ) + + columns_data: List[List[Any]] = [[] for _ in range(num_cols)] + for row in rows: + row_len = len(row) + for c in range(num_cols): + cell = row[c] if c < row_len else None + columns_data[c].append(_extract_cell_value(cell, type_codes[c])) + + arrays = [ + convert_column_to_arrow_array(col_vals, pa_field, type_code) + for col_vals, pa_field, type_code in zip(columns_data, schema, type_codes) + ] + + return pa.RecordBatch.from_arrays(arrays, schema=schema) diff --git a/packages/google-cloud-spanner-arrow/tests/unit/test_c_extension.py b/packages/google-cloud-spanner-arrow/tests/unit/test_c_extension.py new file mode 100644 index 000000000000..ec37e4f52f1d --- /dev/null +++ b/packages/google-cloud-spanner-arrow/tests/unit/test_c_extension.py @@ -0,0 +1,226 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import base64 +import datetime +import decimal +import math +import unittest +import zoneinfo + +import pyarrow as pa +from google.protobuf.struct_pb2 import ListValue, Struct, Value + +import google_cloud_spanner_arrow as sa +from google_cloud_spanner_arrow import cext, python + + +class TestSpannerArrowCExtension(unittest.TestCase): + def test_implementation_is_c(self): + self.assertEqual(sa.implementation, "c") + + def test_empty_rows(self): + fields = [("id", 2), ("name", 6)] + batch = sa.rows_to_arrow_batch(fields, []) + self.assertEqual(batch.num_rows, 0) + self.assertEqual(batch.num_columns, 2) + self.assertEqual(batch.schema.names, ["id", "name"]) + self.assertEqual(batch.schema.field("id").type, pa.int64()) + self.assertEqual(batch.schema.field("name").type, pa.string()) + + def test_basic_types(self): + fields = [ + ("id", 2), + ("name", 6), + ("active", 1), + ("score", 3), + ("f32", 15), + ] + rows = [ + ["1", "Alice", True, 95.5, 12.5], + ["2", "Bob", False, 82.0, -0.5], + ] + batch = sa.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 2) + self.assertEqual(batch.column("id").to_pylist(), [1, 2]) + self.assertEqual(batch.column("name").to_pylist(), ["Alice", "Bob"]) + self.assertEqual(batch.column("active").to_pylist(), [True, False]) + self.assertEqual(batch.column("score").to_pylist(), [95.5, 82.0]) + self.assertAlmostEqual(batch.column("f32").to_pylist()[0], 12.5, places=4) + self.assertAlmostEqual(batch.column("f32").to_pylist()[1], -0.5, places=4) + + def test_protobuf_value_objects(self): + fields = [ + ("id", 2), + ("name", 6), + ("active", 1), + ("score", 3), + ] + rows = [ + [ + Value(string_value="10"), + Value(string_value="Charlie"), + Value(bool_value=True), + Value(number_value=88.5), + ], + [ + Value(string_value="20"), + Value(string_value="David"), + Value(bool_value=False), + Value(number_value=99.0), + ], + ] + batch = sa.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 2) + self.assertEqual(batch.column("id").to_pylist(), [10, 20]) + self.assertEqual(batch.column("name").to_pylist(), ["Charlie", "David"]) + self.assertEqual(batch.column("active").to_pylist(), [True, False]) + self.assertEqual(batch.column("score").to_pylist(), [88.5, 99.0]) + + def test_null_values(self): + fields = [ + ("id", 2), + ("name", 6), + ("date_col", 5), + ("ts_col", 4), + ("num_col", 10), + ("bytes_col", 7), + ] + rows = [ + [None, None, None, None, None, None], + [ + Value(null_value=0), + Value(null_value=0), + Value(null_value=0), + Value(null_value=0), + Value(null_value=0), + Value(null_value=0), + ], + ] + batch = sa.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 2) + for col_name in ["id", "name", "date_col", "ts_col", "num_col", "bytes_col"]: + self.assertEqual(batch.column(col_name).to_pylist(), [None, None]) + + def test_advanced_types_date_timestamp_numeric_bytes(self): + fields = [ + ("date_col", 5), + ("ts_col", 4), + ("num_col", 10), + ("bytes_col", 7), + ] + raw_bytes = b"quantum-accelerator-bytes" + b64_bytes = base64.b64encode(raw_bytes).decode("ascii") + + rows = [ + [ + Value(string_value="2023-01-15"), + Value(string_value="2023-01-15T10:30:00.123456Z"), + Value(string_value="12345.678900000"), + Value(string_value=b64_bytes), + ], + [ + "2024-12-31", + "2024-12-31T23:59:59.999999Z", + "-0.000000001", + raw_bytes, + ], + ] + batch = sa.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 2) + + # Dates + self.assertEqual( + batch.column("date_col").to_pylist(), + [datetime.date(2023, 1, 15), datetime.date(2024, 12, 31)], + ) + + # Timestamps (UTC) + utc = zoneinfo.ZoneInfo("UTC") + self.assertEqual( + batch.column("ts_col").to_pylist(), + [ + datetime.datetime(2023, 1, 15, 10, 30, 0, 123456, tzinfo=utc), + datetime.datetime(2024, 12, 31, 23, 59, 59, 999999, tzinfo=utc), + ], + ) + + # Numerics + self.assertEqual( + batch.column("num_col").to_pylist(), + [decimal.Decimal("12345.678900000"), decimal.Decimal("-0.000000001")], + ) + + # Bytes + self.assertEqual(batch.column("bytes_col").to_pylist(), [raw_bytes, raw_bytes]) + + def test_float_nan_and_infinities(self): + fields = [("val_f64", 3), ("val_f32", 15)] + rows = [ + [Value(string_value="NaN"), Value(string_value="NaN")], + [Value(string_value="Infinity"), Value(string_value="Infinity")], + [Value(string_value="-Infinity"), Value(string_value="-Infinity")], + ] + batch = sa.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 3) + + f64_list = batch.column("val_f64").to_pylist() + self.assertTrue(math.isnan(f64_list[0])) + self.assertEqual(f64_list[1], float("inf")) + self.assertEqual(f64_list[2], float("-inf")) + + f32_list = batch.column("val_f32").to_pylist() + self.assertTrue(math.isnan(f32_list[0])) + self.assertEqual(f32_list[1], float("inf")) + self.assertEqual(f32_list[2], float("-inf")) + + def test_arrays_and_structs(self): + fields = [ + ("arr", 8, ("item", 6)), + ( + "st", + 9, + ( + ("f_int", 2), + ("f_str", 6), + ), + ), + ] + rows = [ + [ + Value( + list_value=ListValue( + values=[Value(string_value="x"), Value(string_value="y")] + ) + ), + Value( + struct_value=Struct( + fields={ + "f_int": Value(string_value="10"), + "f_str": Value(string_value="hello"), + } + ) + ), + ] + ] + batch = sa.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 1) + self.assertEqual(batch.column("arr").to_pylist(), [["x", "y"]]) + self.assertEqual( + batch.column("st").to_pylist(), [{"f_int": 10, "f_str": "hello"}] + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/packages/google-cloud-spanner-arrow/tests/unit/test_concurrency.py b/packages/google-cloud-spanner-arrow/tests/unit/test_concurrency.py new file mode 100644 index 000000000000..d843108dbdfc --- /dev/null +++ b/packages/google-cloud-spanner-arrow/tests/unit/test_concurrency.py @@ -0,0 +1,49 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import concurrent.futures +import unittest + +import google_cloud_spanner_arrow as sa + + +class TestSpannerArrowConcurrency(unittest.TestCase): + def test_multi_threaded_batch_creation(self): + fields = [("id", 2), ("name", 6), ("score", 3)] + num_threads = 8 + rows_per_thread = 5000 + + def worker(thread_id: int): + rows = [ + [f"{thread_id * 100000 + i}", f"user_{thread_id}_{i}", float(i)] + for i in range(rows_per_thread) + ] + batch = sa.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, rows_per_thread) + self.assertEqual(batch.column("id").to_pylist()[0], thread_id * 100000) + self.assertEqual( + batch.column("id").to_pylist()[-1], + thread_id * 100000 + rows_per_thread - 1, + ) + return batch.num_rows + + with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: + futures = [executor.submit(worker, tid) for tid in range(num_threads)] + results = [f.result() for f in concurrent.futures.as_completed(futures)] + + self.assertEqual(sum(results), num_threads * rows_per_thread) + + +if __name__ == "__main__": + unittest.main() diff --git a/packages/google-cloud-spanner-arrow/tests/unit/test_pure_python_fallback.py b/packages/google-cloud-spanner-arrow/tests/unit/test_pure_python_fallback.py new file mode 100644 index 000000000000..8e2374f845ce --- /dev/null +++ b/packages/google-cloud-spanner-arrow/tests/unit/test_pure_python_fallback.py @@ -0,0 +1,86 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +import decimal +import unittest +import zoneinfo + +import pyarrow as pa +from google.protobuf.struct_pb2 import Value + +from google_cloud_spanner_arrow import python + + +class TestSpannerArrowPurePythonFallback(unittest.TestCase): + def test_basic_types(self): + fields = [ + ("id", 2), + ("name", 6), + ("active", 1), + ("score", 3), + ] + rows = [ + ["1", "Alice", True, 95.5], + ["2", "Bob", False, 82.0], + ] + batch = python.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 2) + self.assertEqual(batch.column("id").to_pylist(), [1, 2]) + self.assertEqual(batch.column("name").to_pylist(), ["Alice", "Bob"]) + self.assertEqual(batch.column("active").to_pylist(), [True, False]) + self.assertEqual(batch.column("score").to_pylist(), [95.5, 82.0]) + + def test_nulls(self): + fields = [("id", 2), ("name", 6)] + rows = [[None, None], [Value(null_value=0), Value(null_value=0)]] + batch = python.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 2) + self.assertEqual(batch.column("id").to_pylist(), [None, None]) + self.assertEqual(batch.column("name").to_pylist(), [None, None]) + + def test_advanced_types(self): + fields = [ + ("date_col", 5), + ("ts_col", 4), + ("num_col", 10), + ("bytes_col", 7), + ] + raw_bytes = b"fallback-bytes" + rows = [ + [ + Value(string_value="2023-01-15"), + Value(string_value="2023-01-15T10:30:00.123456Z"), + Value(string_value="12345.678900000"), + raw_bytes, + ] + ] + batch = python.rows_to_arrow_batch(fields, rows) + self.assertEqual(batch.num_rows, 1) + self.assertEqual( + batch.column("date_col").to_pylist(), [datetime.date(2023, 1, 15)] + ) + utc = zoneinfo.ZoneInfo("UTC") + self.assertEqual( + batch.column("ts_col").to_pylist(), + [datetime.datetime(2023, 1, 15, 10, 30, 0, 123456, tzinfo=utc)], + ) + self.assertEqual( + batch.column("num_col").to_pylist(), [decimal.Decimal("12345.678900000")] + ) + self.assertEqual(batch.column("bytes_col").to_pylist(), [raw_bytes]) + + +if __name__ == "__main__": + unittest.main() diff --git a/packages/google-cloud-spanner/benchmark_comprehensive_arrow.py b/packages/google-cloud-spanner/benchmark_comprehensive_arrow.py new file mode 100644 index 000000000000..c7556f74b30b --- /dev/null +++ b/packages/google-cloud-spanner/benchmark_comprehensive_arrow.py @@ -0,0 +1,340 @@ +# Copyright 2026 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Comprehensive benchmark comparing: +1. Traditional Python row decoding +2. Customer PyArrow conversion (Rows -> dicts -> pa.Table) +3. Pure-Python PyArrow (_arrow.py) +4. C-Accelerated PyArrow (google-cloud-spanner-arrow) + +Measures Single-Threaded and Multi-Threaded setups for both Small and Large result sets. +""" + +import concurrent.futures +import gc +import statistics +import time +from typing import Tuple + +import pyarrow as pa +from google.cloud import spanner +from google.cloud.spanner_v1 import _arrow as pure_py_arrow +from google_cloud_spanner_arrow import cext as spanner_arrow_cext +from google_cloud_spanner_arrow import python as spanner_arrow_python + +PROJECT_ID = "appdev-soda-spanner-staging" +INSTANCE_ID = "knut-test-ycsb" +DATABASE_ID = "spring-data-jpa" + +# Query generating diverse column types dynamically +SQL = """SELECT + MOD(FARM_FINGERPRINT(GENERATE_UUID()), 2) = 0 AS random_bool, + CAST(GENERATE_UUID() AS BYTES) AS random_bytes, + DATE_FROM_UNIX_DATE(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 2932896))) AS random_date, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS FLOAT32) AS random_float32, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS FLOAT64) AS random_float64, + MAKE_INTERVAL(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 10)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 12)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 28)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 24)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 60)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 60))) AS random_interval, + TO_JSON('{"key": "' || GENERATE_UUID() || '"}') AS random_json, + FARM_FINGERPRINT(GENERATE_UUID()) AS random_int64, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS NUMERIC) AS random_numeric, + GENERATE_UUID() AS random_string, + TIMESTAMP_MICROS(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 1230219000000000))) AS random_timestamp, + NEW_UUID() AS random_uuid +FROM UNNEST(GENERATE_ARRAY(1, @num_rows)) AS n""" + + +# ------------------------------------------------------------- +# Read Operations (measuring rows 2...N to isolate client parsing) +# ------------------------------------------------------------- + +def read_traditional_rows(database, num_rows: int) -> Tuple[float, int]: + """Read rows using traditional Python client row decoding.""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + row_iter = iter(results) + try: + first = next(row_iter) + for _ in first: + pass + except StopIteration: + return 0.0, 0 + + start = time.perf_counter() + count = 1 + for row in row_iter: + for _ in row: + pass + count += 1 + elapsed = time.perf_counter() - start + return elapsed, count + + +def read_customer_to_arrow(database, num_rows: int) -> Tuple[float, int]: + """Customer approach: iterate rows into dicts, call pa.Table.from_pylist().""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + row_iter = iter(results) + try: + first = next(row_iter) + for _ in first: + pass + except StopIteration: + return 0.0, 0 + + col_names = [col.name for col in results.fields] + start = time.perf_counter() + rows_list = [] + for row in row_iter: + row_dict = {} + for col_name, cell in zip(col_names, row): + if hasattr(cell, "months"): + row_dict[col_name] = str(cell) + else: + row_dict[col_name] = cell + rows_list.append(row_dict) + table = pa.Table.from_pylist(rows_list) + elapsed = time.perf_counter() - start + return elapsed, len(table) + 1 + + +def read_pure_python_arrow(database, num_rows: int, max_chunk_size: int = 65536) -> Tuple[float, int]: + """Pure-Python PyArrow stream ingestion.""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + results._lazy_decode = True + results._consume_next() + + fields = results.fields + pa_schema = spanner_arrow_python.fields_to_arrow_schema(fields) + + start = time.perf_counter() + batches = [] + accumulated = [] + while True: + if results._rows: + accumulated.extend(results._rows) + results._rows = [] + while len(accumulated) >= max_chunk_size: + chunk = accumulated[:max_chunk_size] + accumulated = accumulated[max_chunk_size:] + batches.append(spanner_arrow_python.rows_to_arrow_batch(fields, chunk, schema=pa_schema)) + if results._done: + break + try: + results._consume_next() + except StopIteration: + break + if accumulated: + batches.append(spanner_arrow_python.rows_to_arrow_batch(fields, accumulated, schema=pa_schema)) + table = pa.Table.from_batches(batches, schema=pa_schema) + elapsed = time.perf_counter() - start + return elapsed, table.num_rows + + +def read_c_accelerated_arrow(database, num_rows: int, max_chunk_size: int = 65536) -> Tuple[float, int]: + """Native C-accelerated PyArrow stream ingestion.""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + results._lazy_decode = True + results._consume_next() + + fields = results.fields + pa_schema = spanner_arrow_python.fields_to_arrow_schema(fields) + + start = time.perf_counter() + batches = [] + accumulated = [] + while True: + if results._rows: + accumulated.extend(results._rows) + results._rows = [] + while len(accumulated) >= max_chunk_size: + chunk = accumulated[:max_chunk_size] + accumulated = accumulated[max_chunk_size:] + batches.append(spanner_arrow_cext.rows_to_arrow_batch(fields, chunk, schema=pa_schema)) + if results._done: + break + try: + results._consume_next() + except StopIteration: + break + if accumulated: + batches.append(spanner_arrow_cext.rows_to_arrow_batch(fields, accumulated, schema=pa_schema)) + table = pa.Table.from_batches(batches, schema=pa_schema) + elapsed = time.perf_counter() - start + return elapsed, table.num_rows + + +# ------------------------------------------------------------- +# Multi-Threaded Concurrent Workloads +# ------------------------------------------------------------- + +def run_concurrent_workload(database, num_rows_per_query: int, num_threads: int, read_fn): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: + futures = [executor.submit(read_fn, database, num_rows_per_query) for _ in range(num_threads)] + results = [f.result() for f in concurrent.futures.as_completed(futures)] + total_time = time.perf_counter() - start + total_rows = sum(r[1] for r in results) + return total_time, total_rows + + +# ------------------------------------------------------------- +# Benchmark Runner +# ------------------------------------------------------------- + +def run_all_benchmarks(): + print(f"Connecting to Cloud Spanner: {PROJECT_ID} / {INSTANCE_ID} / {DATABASE_ID}") + + client = spanner.Client(project=PROJECT_ID) + instance = client.instance(INSTANCE_ID) + database = instance.database(DATABASE_ID) + + # Warmup + print("Warming up database connection and JIT/gRPC channels...") + read_traditional_rows(database, 1000) + read_pure_python_arrow(database, 1000) + read_c_accelerated_arrow(database, 1000) + + # 1. Single-Threaded Benchmarks + print("\n" + "=" * 80) + print(" 1. SINGLE-THREADED BENCHMARKS (Small, Medium, Large Result Sets)") + print("=" * 80) + + single_sizes = [ + ("Small", 2000), + ("Medium", 10000), + ("Large", 50000), + ("Very Large", 100000), + ] + iterations = 3 + + for label, num_rows in single_sizes: + print(f"\n--- [{label}] Result Set: {num_rows:,} rows (12 diverse columns) ---") + + times_trad = [] + times_cust = [] + times_py_arrow = [] + times_c_arrow = [] + + for it in range(iterations): + gc.collect() + t, _ = read_traditional_rows(database, num_rows) + times_trad.append(t) + + gc.collect() + t, _ = read_customer_to_arrow(database, num_rows) + times_cust.append(t) + + gc.collect() + t, _ = read_pure_python_arrow(database, num_rows) + times_py_arrow.append(t) + + gc.collect() + t, _ = read_c_accelerated_arrow(database, num_rows) + times_c_arrow.append(t) + + avg_trad = statistics.mean(times_trad) + avg_cust = statistics.mean(times_cust) + avg_py_arrow = statistics.mean(times_py_arrow) + avg_c_arrow = statistics.mean(times_c_arrow) + + rps_trad = num_rows / avg_trad + rps_cust = num_rows / avg_cust + rps_py = num_rows / avg_py_arrow + rps_c = num_rows / avg_c_arrow + + print(f" 1. Traditional Rows: {avg_trad*1000:7.1f} ms | {rps_trad:9,.0f} rows/s") + print(f" 2. Customer (Rows -> PyArrow): {avg_cust*1000:7.1f} ms | {rps_cust:9,.0f} rows/s") + print(f" 3. Pure-Python PyArrow: {avg_py_arrow*1000:7.1f} ms | {rps_py:9,.0f} rows/s") + print(f" 4. C-Accelerated PyArrow: {avg_c_arrow*1000:7.1f} ms | {rps_c:9,.0f} rows/s") + print(f" ==> Speedup over Traditional: {avg_trad/avg_c_arrow:5.2f}x faster") + print(f" ==> Speedup over Customer: {avg_cust/avg_c_arrow:5.2f}x faster") + print(f" ==> Speedup over Pure-Python: {avg_py_arrow/avg_c_arrow:5.2f}x faster") + + # 2. Multi-Threaded Benchmarks + print("\n" + "=" * 80) + print(" 2. MULTI-THREADED CONCURRENT BENCHMARKS (4 & 8 Threads)") + print("=" * 80) + + multi_configs = [ + ("Small Concurrent (8 threads x 2,000 rows = 16,000 rows)", 2000, 8), + ("Medium Concurrent (4 threads x 10,000 rows = 40,000 rows)", 10000, 4), + ("Large Concurrent (4 threads x 50,000 rows = 200,000 rows)", 50000, 4), + ("Large Concurrent (8 threads x 25,000 rows = 200,000 rows)", 25000, 8), + ] + + for title, rows_per_query, num_threads in multi_configs: + total_rows = rows_per_query * num_threads + print(f"\n--- {title} ---") + + times_trad = [] + times_cust = [] + times_py_arrow = [] + times_c_arrow = [] + + for it in range(iterations): + gc.collect() + t, _ = run_concurrent_workload(database, rows_per_query, num_threads, read_traditional_rows) + times_trad.append(t) + + gc.collect() + t, _ = run_concurrent_workload(database, rows_per_query, num_threads, read_customer_to_arrow) + times_cust.append(t) + + gc.collect() + t, _ = run_concurrent_workload(database, rows_per_query, num_threads, read_pure_python_arrow) + times_py_arrow.append(t) + + gc.collect() + t, _ = run_concurrent_workload(database, rows_per_query, num_threads, read_c_accelerated_arrow) + times_c_arrow.append(t) + + avg_trad = statistics.mean(times_trad) + avg_cust = statistics.mean(times_cust) + avg_py = statistics.mean(times_py_arrow) + avg_c = statistics.mean(times_c_arrow) + + rps_trad = total_rows / avg_trad + rps_cust = total_rows / avg_cust + rps_py = total_rows / avg_py + rps_c = total_rows / avg_c + + print(f" 1. Traditional Rows: {avg_trad*1000:7.1f} ms | {rps_trad:9,.0f} rows/s (total throughput)") + print(f" 2. Customer (Rows -> PyArrow): {avg_cust*1000:7.1f} ms | {rps_cust:9,.0f} rows/s (total throughput)") + print(f" 3. Pure-Python PyArrow: {avg_py*1000:7.1f} ms | {rps_py:9,.0f} rows/s (total throughput)") + print(f" 4. C-Accelerated PyArrow: {avg_c*1000:7.1f} ms | {rps_c:9,.0f} rows/s (total throughput)") + print(f" ==> Speedup over Traditional: {avg_trad/avg_c:5.2f}x faster") + print(f" ==> Speedup over Customer: {avg_cust/avg_c:5.2f}x faster") + print(f" ==> Speedup over Pure-Python: {avg_py/avg_c:5.2f}x faster") + + +if __name__ == "__main__": + run_all_benchmarks() diff --git a/packages/google-cloud-spanner/benchmark_end_to_end_concurrency.py b/packages/google-cloud-spanner/benchmark_end_to_end_concurrency.py new file mode 100644 index 000000000000..0cc352ceb937 --- /dev/null +++ b/packages/google-cloud-spanner/benchmark_end_to_end_concurrency.py @@ -0,0 +1,207 @@ +# Copyright 2026 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-End concurrency benchmark measuring full query lifecycle: +- snapshot.execute_sql() RPC dispatch +- gRPC streaming network transport +- Stream ingestion and decoding +Across 1, 4, 8, 16, and 32 concurrent threads. +""" + +import concurrent.futures +import gc +import statistics +import time +from typing import Tuple + +import pyarrow as pa +from google.cloud import spanner +from google_cloud_spanner_arrow import cext as spanner_arrow_cext +from google_cloud_spanner_arrow import python as spanner_arrow_python + +PROJECT_ID = "appdev-soda-spanner-staging" +INSTANCE_ID = "knut-test-ycsb" +DATABASE_ID = "spring-data-jpa" + +SQL = """SELECT + MOD(FARM_FINGERPRINT(GENERATE_UUID()), 2) = 0 AS random_bool, + CAST(GENERATE_UUID() AS BYTES) AS random_bytes, + DATE_FROM_UNIX_DATE(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 2932896))) AS random_date, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS FLOAT32) AS random_float32, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS FLOAT64) AS random_float64, + MAKE_INTERVAL(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 10)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 12)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 28)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 24)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 60)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 60))) AS random_interval, + TO_JSON('{"key": "' || GENERATE_UUID() || '"}') AS random_json, + FARM_FINGERPRINT(GENERATE_UUID()) AS random_int64, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS NUMERIC) AS random_numeric, + GENERATE_UUID() AS random_string, + TIMESTAMP_MICROS(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 1230219000000000))) AS random_timestamp, + NEW_UUID() AS random_uuid +FROM UNNEST(GENERATE_ARRAY(1, @num_rows)) AS n""" + + +# ------------------------------------------------------------- +# Full End-to-End Query Functions (Measuring execute_sql -> done) +# ------------------------------------------------------------- + +def e2e_traditional_rows(database, num_rows: int) -> Tuple[float, int]: + """Full end-to-end traditional Python row iteration.""" + start = time.perf_counter() + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + count = 0 + for row in results: + for _ in row: + pass + count += 1 + elapsed = time.perf_counter() - start + return elapsed, count + + +def e2e_pure_python_arrow(database, num_rows: int) -> Tuple[float, int]: + """Full end-to-end pure-Python PyArrow streaming table.""" + start = time.perf_counter() + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + results._lazy_decode = True + batches = [] + accumulated = [] + fields = None + pa_schema = None + max_chunk_size = 65536 + + while True: + try: + results._consume_next() + except StopIteration: + break + if fields is None and results._metadata: + fields = results.fields + pa_schema = spanner_arrow_python.fields_to_arrow_schema(fields) + if results._rows: + accumulated.extend(results._rows) + results._rows = [] + while len(accumulated) >= max_chunk_size: + chunk = accumulated[:max_chunk_size] + accumulated = accumulated[max_chunk_size:] + batches.append(spanner_arrow_python.rows_to_arrow_batch(fields, chunk, schema=pa_schema)) + if results._done: + break + if accumulated: + batches.append(spanner_arrow_python.rows_to_arrow_batch(fields, accumulated, schema=pa_schema)) + table = pa.Table.from_batches(batches, schema=pa_schema) + count = table.num_rows + elapsed = time.perf_counter() - start + return elapsed, count + + +def e2e_direct_wire_c_arrow(database, num_rows: int) -> Tuple[float, int]: + """Full end-to-end Direct Wire C-extension Arrow streaming.""" + start = time.perf_counter() + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + chunks = [] + fields = None + pa_schema = None + for resp in results._response_iterator: + if fields is None and resp.metadata: + fields = resp.metadata.row_type.fields + pa_schema = spanner_arrow_python.fields_to_arrow_schema(fields) + chunks.append(resp._pb.SerializeToString()) + batch = spanner_arrow_cext.wire_prs_to_arrow_batch(fields, chunks, schema=pa_schema) + count = batch.num_rows + elapsed = time.perf_counter() - start + return elapsed, count + + +def run_concurrent_workload(database, num_rows_per_query: int, num_threads: int, fn): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: + futures = [executor.submit(fn, database, num_rows_per_query) for _ in range(num_threads)] + results = [f.result() for f in concurrent.futures.as_completed(futures)] + total_time = time.perf_counter() - start + total_rows = sum(r[1] for r in results) + return total_time, total_rows + + +def run_e2e_benchmarks(): + print(f"Connecting to Cloud Spanner: {PROJECT_ID} / {INSTANCE_ID} / {DATABASE_ID}") + client = spanner.Client(project=PROJECT_ID) + instance = client.instance(INSTANCE_ID) + database = instance.database(DATABASE_ID) + + # Warmup + print("Warming up gRPC channels and Spanner query caches...") + e2e_traditional_rows(database, 1000) + e2e_direct_wire_c_arrow(database, 1000) + + thread_counts = [1, 4, 8, 16, 32] + num_rows_per_query = 30000 # 30,000 rows x 12 cols per concurrent query + iterations = 3 + + print("\n" + "=" * 86) + print(" END-TO-END CONCURRENCY BENCHMARK (1, 4, 8, 16, 32 Concurrent Queries)") + print(f" Each thread runs ExecuteStreamingSql() for {num_rows_per_query:,} rows (12 diverse columns)") + print(" Measures full wall-clock time from execute_sql() RPC dispatch to final Table/Batch") + print("=" * 86) + + for num_threads in thread_counts: + total_rows = num_threads * num_rows_per_query + print(f"\n--- {num_threads} Concurrent Threads ({total_rows:,} total rows across {num_threads} queries) ---") + + times_trad = [] + times_py_arrow = [] + times_wire_c = [] + + for _ in range(iterations): + gc.collect() + t, r = run_concurrent_workload(database, num_rows_per_query, num_threads, e2e_traditional_rows) + times_trad.append(t) + + gc.collect() + t, r = run_concurrent_workload(database, num_rows_per_query, num_threads, e2e_pure_python_arrow) + times_py_arrow.append(t) + + gc.collect() + t, r = run_concurrent_workload(database, num_rows_per_query, num_threads, e2e_direct_wire_c_arrow) + times_wire_c.append(t) + + avg_trad = statistics.mean(times_trad) + avg_py = statistics.mean(times_py_arrow) + avg_wire = statistics.mean(times_wire_c) + + rps_trad = total_rows / avg_trad + rps_py = total_rows / avg_py + rps_wire = total_rows / avg_wire + + print(f" 1. Traditional Python Rows: {avg_trad*1000:7.1f} ms | {rps_trad:9,.0f} rows/s (total throughput)") + print(f" 2. Pure-Python PyArrow: {avg_py*1000:7.1f} ms | {rps_py:9,.0f} rows/s (total throughput)") + print(f" 3. Direct Wire C-Ext PyArrow: {avg_wire*1000:7.1f} ms | {rps_wire:9,.0f} rows/s (total throughput)") + print(f" ==> Wire C Speedup vs Traditional: {avg_trad/avg_wire:5.2f}x faster") + print(f" ==> Wire C Speedup vs Pure-Python: {avg_py/avg_wire:5.2f}x faster") + + +if __name__ == "__main__": + run_e2e_benchmarks() diff --git a/packages/google-cloud-spanner/benchmark_read_large_result_set.py b/packages/google-cloud-spanner/benchmark_read_large_result_set.py new file mode 100644 index 000000000000..2d95603469b4 --- /dev/null +++ b/packages/google-cloud-spanner/benchmark_read_large_result_set.py @@ -0,0 +1,206 @@ +# Copyright 2026 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark comparing traditional row reading vs direct PyArrow conversion. + +Reuses the read-large-result-set query from spanner-client-benchmarks. +Measures rows 2...N to exclude initial query execution latency on Spanner. +""" + +import gc +import statistics +import time +import pyarrow as pa +from google.cloud import spanner + +PROJECT_ID = "appdev-soda-spanner-staging" +INSTANCE_ID = "knut-test-ycsb" +DATABASE_ID = "spring-data-jpa" + +# Query generating random query results on the fly without inserting data +SQL = """SELECT + MOD(FARM_FINGERPRINT(GENERATE_UUID()), 2) = 0 AS random_bool, + CAST(GENERATE_UUID() AS BYTES) AS random_bytes, + DATE_FROM_UNIX_DATE(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 2932896))) AS random_date, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS FLOAT32) AS random_float32, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS FLOAT64) AS random_float64, + MAKE_INTERVAL(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 10)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 12)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 28)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 24)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 60)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 60))) AS random_interval, + TO_JSON('{"key": "' || GENERATE_UUID() || '"}') AS random_json, + FARM_FINGERPRINT(GENERATE_UUID()) AS random_int64, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS NUMERIC) AS random_numeric, + GENERATE_UUID() AS random_string, + TIMESTAMP_MICROS(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 1230219000000000))) AS random_timestamp, + NEW_UUID() AS random_uuid +FROM UNNEST(GENERATE_ARRAY(1, @num_rows)) AS n""" + + +def benchmark_traditional_rows(database, num_rows): + """Traditional row iteration measuring rows 2...N.""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + row_iterator = iter(results) + try: + first_row = next(row_iterator) + for cell in first_row: + pass + except StopIteration: + return 0.0 + + # Measure iteration and decoding of remaining rows (rows 2...N) + start_time = time.perf_counter() + count = 1 + for row in row_iterator: + for cell in row: + pass + count += 1 + end_time = time.perf_counter() + return (end_time - start_time), count + + +def benchmark_rows_to_pyarrow_customer_way(database, num_rows): + """Customer's current path: traditional rows -> PyArrow Table (measuring 2...N).""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + row_iterator = iter(results) + try: + first_row = next(row_iterator) + for cell in first_row: + pass + except StopIteration: + return 0.0 + + col_names = [col.name for col in results.fields] + start_time = time.perf_counter() + rows_list = [] + for row in row_iterator: + row_dict = {} + for col_name, cell in zip(col_names, row): + if hasattr(cell, "months"): # Spanner Interval + row_dict[col_name] = str(cell) + else: + row_dict[col_name] = cell + rows_list.append(row_dict) + table = pa.Table.from_pylist(rows_list) + end_time = time.perf_counter() + return (end_time - start_time), len(table) + 1 + + +def benchmark_direct_arrow_batches(database, num_rows, max_chunk_size=65536): + """Direct-to-Arrow streaming batches (measuring rows 2...N after initial stream initialization).""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + # Initialize stream with lazy decoding (fetches first PartialResultSet / executes query on Spanner) + results._lazy_decode = True + results._consume_next() + + start_time = time.perf_counter() + total_rows = 0 + for batch in results.to_arrow_batches(max_chunk_size=max_chunk_size): + total_rows += batch.num_rows + end_time = time.perf_counter() + return (end_time - start_time), total_rows + + +def benchmark_direct_arrow_table(database, num_rows, max_chunk_size=65536): + """Direct-to-Arrow complete Table materialization (measuring rows 2...N).""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + # Initialize stream with lazy decoding (fetches first PartialResultSet / executes query on Spanner) + results._lazy_decode = True + results._consume_next() + + start_time = time.perf_counter() + table = results.to_arrow(max_chunk_size=max_chunk_size) + end_time = time.perf_counter() + return (end_time - start_time), table.num_rows + + +def run_benchmarks(): + print(f"Connecting to Cloud Spanner: {PROJECT_ID} / {INSTANCE_ID} / {DATABASE_ID}") + client = spanner.Client(project=PROJECT_ID) + instance = client.instance(INSTANCE_ID) + database = instance.database(DATABASE_ID) + + # Warmup + print("\nWarming up connection and caches...") + benchmark_traditional_rows(database, 1000) + benchmark_direct_arrow_table(database, 1000) + + test_sizes = [10000, 50000, 100000] + iterations = 3 + + for num_rows in test_sizes: + print(f"\n========================================================") + print(f" BENCHMARK: {num_rows:,} ROWS (12 diverse columns)") + print(f" Measuring time for rows 2...N (excluding query compile)") + print(f"========================================================") + + results_trad = [] + results_cust = [] + results_arrow_batch = [] + results_arrow_table = [] + + for i in range(iterations): + gc.collect() + t_trad, count_trad = benchmark_traditional_rows(database, num_rows) + results_trad.append(t_trad) + + gc.collect() + t_cust, count_cust = benchmark_rows_to_pyarrow_customer_way(database, num_rows) + results_cust.append(t_cust) + + gc.collect() + t_batch, count_batch = benchmark_direct_arrow_batches(database, num_rows) + results_arrow_batch.append(t_batch) + + gc.collect() + t_table, count_table = benchmark_direct_arrow_table(database, num_rows) + results_arrow_table.append(t_table) + + print(f" Run {i+1}/{iterations}:") + print(f" - Traditional Rows: {t_trad*1000:.1f} ms ({num_rows/t_trad:,.0f} rows/s)") + print(f" - Customer (Rows -> PyArrow Table): {t_cust*1000:.1f} ms ({num_rows/t_cust:,.0f} rows/s)") + print(f" - Direct to_arrow_batches(): {t_batch*1000:.1f} ms ({num_rows/t_batch:,.0f} rows/s)") + print(f" - Direct to_arrow() Table: {t_table*1000:.1f} ms ({num_rows/t_table:,.0f} rows/s)") + + avg_trad = statistics.mean(results_trad) + avg_cust = statistics.mean(results_cust) + avg_batch = statistics.mean(results_arrow_batch) + avg_table = statistics.mean(results_arrow_table) + + print(f"\n --- Summary for {num_rows:,} rows (Average of {iterations} runs) ---") + print(f" 1. Traditional Python Rows: {avg_trad*1000:7.1f} ms | {num_rows/avg_trad:9,.0f} rows/s") + print(f" 2. Customer Path (Rows -> PyArrow):{avg_cust*1000:7.1f} ms | {num_rows/avg_cust:9,.0f} rows/s") + print(f" 3. Direct to_arrow_batches(): {avg_batch*1000:7.1f} ms | {num_rows/avg_batch:9,.0f} rows/s --> {avg_cust/avg_batch:.2f}x faster than customer") + print(f" 4. Direct to_arrow() Table: {avg_table*1000:7.1f} ms | {num_rows/avg_table:9,.0f} rows/s --> {avg_cust/avg_table:.2f}x faster than customer") + + +if __name__ == "__main__": + run_benchmarks() diff --git a/packages/google-cloud-spanner/benchmark_wire_vs_object_arrow.py b/packages/google-cloud-spanner/benchmark_wire_vs_object_arrow.py new file mode 100644 index 000000000000..512e045f3bb0 --- /dev/null +++ b/packages/google-cloud-spanner/benchmark_wire_vs_object_arrow.py @@ -0,0 +1,222 @@ +# Copyright 2026 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark comparing: +1. Python Protobuf Object Deserialization + Pure-Python Arrow +2. Python Protobuf Object Deserialization + C-Extension Arrow (rows_to_c_batch) +3. Direct Protobuf Wire Decoding in C -> Arrow (wire_prs_to_c_batch) + (Measuring pure parsing/decoding throughput and multi-threaded scaling). +""" + +import concurrent.futures +import gc +import statistics +import time +from typing import List + +import pyarrow as pa +from google.cloud import spanner +from google_cloud_spanner_arrow import cext as spanner_arrow_cext +from google_cloud_spanner_arrow import python as spanner_arrow_python + +PROJECT_ID = "appdev-soda-spanner-staging" +INSTANCE_ID = "knut-test-ycsb" +DATABASE_ID = "spring-data-jpa" + +SQL = """SELECT + MOD(FARM_FINGERPRINT(GENERATE_UUID()), 2) = 0 AS random_bool, + CAST(GENERATE_UUID() AS BYTES) AS random_bytes, + DATE_FROM_UNIX_DATE(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 2932896))) AS random_date, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS FLOAT32) AS random_float32, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS FLOAT64) AS random_float64, + MAKE_INTERVAL(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 10)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 12)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 28)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 24)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 60)), ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 60))) AS random_interval, + TO_JSON('{"key": "' || GENERATE_UUID() || '"}') AS random_json, + FARM_FINGERPRINT(GENERATE_UUID()) AS random_int64, + CAST(FARM_FINGERPRINT(GENERATE_UUID()) / FARM_FINGERPRINT(GENERATE_UUID()) AS NUMERIC) AS random_numeric, + GENERATE_UUID() AS random_string, + TIMESTAMP_MICROS(ABS(MOD(FARM_FINGERPRINT(GENERATE_UUID()), 1230219000000000))) AS random_timestamp, + NEW_UUID() AS random_uuid +FROM UNNEST(GENERATE_ARRAY(1, @num_rows)) AS n""" + + +def fetch_wire_chunks(database, num_rows: int): + """Fetch raw protobuf wire byte chunks and assembled rows from Spanner.""" + with database.snapshot() as snapshot: + results = snapshot.execute_sql( + SQL, + params={"num_rows": num_rows}, + param_types={"num_rows": spanner.param_types.INT64}, + ) + results._lazy_decode = True + chunks = [] + python_rows = [] + + raw_iter = results._response_iterator + + class InterceptingIterator: + def __init__(self, it): + self.it = it + def __iter__(self): + return self + def __next__(self): + resp = next(self.it) + chunks.append(resp._pb.SerializeToString()) + return resp + + results._response_iterator = InterceptingIterator(raw_iter) + + while True: + try: + results._consume_next() + except StopIteration: + break + if results._rows: + python_rows.extend(results._rows) + results._rows = [] + if results._done: + break + + fields = results.fields + return fields, chunks, python_rows + + +# ------------------------------------------------------------- +# Parsing Benchmarks (Pure parsing CPU speed excluding network) +# ------------------------------------------------------------- + +def bench_pure_python_parsing(fields, python_rows: List[List]): + start = time.perf_counter() + batch = spanner_arrow_python.rows_to_arrow_batch(fields, python_rows) + elapsed = time.perf_counter() - start + return elapsed, batch.num_rows + + +def bench_c_object_parsing(fields, python_rows: List[List]): + start = time.perf_counter() + batch = spanner_arrow_cext.rows_to_arrow_batch(fields, python_rows) + elapsed = time.perf_counter() - start + return elapsed, batch.num_rows + + +def bench_c_direct_wire_parsing(fields, wire_chunks: List[bytes]): + start = time.perf_counter() + batch = spanner_arrow_cext.wire_prs_to_arrow_batch(fields, wire_chunks) + elapsed = time.perf_counter() - start + return elapsed, batch.num_rows + + +def run_benchmark(): + print(f"Connecting to Cloud Spanner: {PROJECT_ID} / {INSTANCE_ID} / {DATABASE_ID}") + client = spanner.Client(project=PROJECT_ID) + instance = client.instance(INSTANCE_ID) + database = instance.database(DATABASE_ID) + + test_sizes = [10000, 50000, 100000] + iterations = 5 + + print("\n" + "=" * 84) + print(" 1. PARSING & CONVERSION CPU THROUGHPUT (Single-Threaded)") + print(" Measuring pure decoding & Arrow construction time") + print("=" * 84) + + for num_rows in test_sizes: + print(f"\nFetching test dataset of {num_rows:,} rows from Spanner...") + fields, wire_chunks, python_rows = fetch_wire_chunks(database, num_rows) + total_wire_bytes = sum(len(c) for c in wire_chunks) + print(f"Dataset: {len(python_rows):,} rows | 12 cols | {total_wire_bytes / (1024*1024):.2f} MB wire protobuf") + + actual_rows = len(python_rows) + times_py = [] + times_c_obj = [] + times_c_wire = [] + + for _ in range(iterations): + gc.collect() + t, r = bench_pure_python_parsing(fields, python_rows) + times_py.append(t) + + gc.collect() + t, r = bench_c_object_parsing(fields, python_rows) + times_c_obj.append(t) + + gc.collect() + t, r = bench_c_direct_wire_parsing(fields, wire_chunks) + times_c_wire.append(t) + + avg_py = statistics.mean(times_py) + avg_c_obj = statistics.mean(times_c_obj) + avg_c_wire = statistics.mean(times_c_wire) + + rps_py = actual_rows / avg_py + rps_c_obj = actual_rows / avg_c_obj + rps_c_wire = actual_rows / avg_c_wire + + print(f" A. Pure-Python (Values -> Arrow): {avg_py*1000:7.2f} ms | {rps_py:10,.0f} rows/s") + print(f" B. C-Ext Objects (Values -> Arrow): {avg_c_obj*1000:7.2f} ms | {rps_c_obj:10,.0f} rows/s") + print(f" C. Direct Wire C-Ext (Raw Proto -> Arrow):{avg_c_wire*1000:7.2f} ms | {rps_c_wire:10,.0f} rows/s") + print(f" ==> Direct Wire vs Pure-Python: {avg_py/avg_c_wire:6.1f}x FASTER!") + print(f" ==> Direct Wire vs C-Ext Object: {avg_c_obj/avg_c_wire:6.1f}x FASTER!") + + print("\n" + "=" * 84) + print(" 2. MULTI-THREADED CONCURRENT SCALING (8 Threads Parallel Parsing)") + print("=" * 84) + + num_rows = 100000 + num_threads = 8 + print(f"\nBenchmarking 8 parallel threads decoding 100k rows each (800,000 rows total)...") + fields, wire_chunks, python_rows = fetch_wire_chunks(database, num_rows) + actual_rows = len(python_rows) + + def run_multi(fn, data): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: + futures = [executor.submit(fn, fields, data) for _ in range(num_threads)] + results = [f.result() for f in concurrent.futures.as_completed(futures)] + return time.perf_counter() - start, sum(r[1] for r in results) + + times_mt_py = [] + times_mt_c_obj = [] + times_mt_c_wire = [] + + for _ in range(iterations): + gc.collect() + t, r = run_multi(bench_pure_python_parsing, python_rows) + times_mt_py.append(t) + + gc.collect() + t, r = run_multi(bench_c_object_parsing, python_rows) + times_mt_c_obj.append(t) + + gc.collect() + t, r = run_multi(bench_c_direct_wire_parsing, wire_chunks) + times_mt_c_wire.append(t) + + avg_mt_py = statistics.mean(times_mt_py) + avg_mt_obj = statistics.mean(times_mt_c_obj) + avg_mt_wire = statistics.mean(times_mt_c_wire) + + total_rows = actual_rows * num_threads + rps_mt_py = total_rows / avg_mt_py + rps_mt_obj = total_rows / avg_mt_obj + rps_mt_wire = total_rows / avg_mt_wire + + print(f" A. Pure-Python (8 threads): {avg_mt_py*1000:7.2f} ms | {rps_mt_py:10,.0f} rows/s (total)") + print(f" B. C-Ext Objects (8 threads): {avg_mt_obj*1000:7.2f} ms | {rps_mt_obj:10,.0f} rows/s (total)") + print(f" C. Direct Wire C-Ext (8 threads): {avg_mt_wire*1000:7.2f} ms | {rps_mt_wire:10,.0f} rows/s (total)") + print(f" ==> Multi-Threaded Speedup over Pure-Python: {avg_mt_py/avg_mt_wire:6.1f}x FASTER!") + print(f" ==> Multi-Threaded Speedup over C-Ext Object: {avg_mt_obj/avg_mt_wire:6.1f}x FASTER!") + + +if __name__ == "__main__": + run_benchmark() diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_arrow.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_arrow.py new file mode 100644 index 000000000000..14d0cb7e1dd7 --- /dev/null +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_arrow.py @@ -0,0 +1,246 @@ +# Copyright 2026 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Helper utilities for converting Cloud Spanner types and schemas to PyArrow.""" + +import base64 +from typing import Any, List, Optional, Sequence + +from google.cloud.spanner_v1.types.type import TypeCode + +try: + import pyarrow as pa + import pyarrow.compute as pc + + _HAS_PYARROW = True +except ImportError: # pragma: NO COVER + _HAS_PYARROW = False + pa = None + pc = None + +_NO_PYARROW_ERROR = ( + "pyarrow is required to use Arrow features. " + "Install it with `pip install google-cloud-spanner[pyarrow]` or `pip install pyarrow`." +) + + +def _check_pyarrow(): + """Verify pyarrow is installed.""" + if not _HAS_PYARROW: + raise ImportError(_NO_PYARROW_ERROR) + + +def spanner_type_to_arrow_type(spanner_type) -> "pa.DataType": + """Map a Spanner Type to a PyArrow DataType. + + :type spanner_type: :class:`~google.cloud.spanner_v1.types.Type` + :param spanner_type: Spanner column type. + + :rtype: :class:`pyarrow.DataType` + :returns: PyArrow DataType corresponding to the Spanner type. + """ + _check_pyarrow() + code = spanner_type.code + + if code == TypeCode.BOOL: + return pa.bool_() + elif code == TypeCode.INT64: + return pa.int64() + elif code == TypeCode.FLOAT32: + return pa.float32() + elif code == TypeCode.FLOAT64: + return pa.float64() + elif code == TypeCode.STRING: + return pa.string() + elif code == TypeCode.BYTES: + return pa.binary() + elif code == TypeCode.TIMESTAMP: + return pa.timestamp("us", tz="UTC") + elif code == TypeCode.DATE: + return pa.date32() + elif code == TypeCode.NUMERIC: + return pa.decimal128(38, 9) + elif code == TypeCode.JSON: + return pa.string() + elif code == TypeCode.PROTO: + return pa.binary() + elif code == TypeCode.ENUM: + return pa.int64() + elif code == TypeCode.INTERVAL: + return pa.string() + elif code == TypeCode.UUID: + return pa.string() + elif code == TypeCode.ARRAY: + element_type = spanner_type_to_arrow_type(spanner_type.array_element_type) + return pa.list_(element_type) + elif code == TypeCode.STRUCT: + fields = [ + pa.field(f.name, spanner_type_to_arrow_type(f.type_)) + for f in spanner_type.struct_type.fields + ] + return pa.struct(fields) + return pa.string() + + +def spanner_schema_to_arrow_schema(fields: Sequence[Any]) -> "pa.Schema": + """Convert Spanner row_type.fields to a PyArrow Schema. + + :type fields: Sequence of :class:`~google.cloud.spanner_v1.types.StructType.Field` + :param fields: List of Spanner fields describing column names and types. + + :rtype: :class:`pyarrow.Schema` + :returns: PyArrow Schema corresponding to the Spanner fields. + """ + _check_pyarrow() + arrow_fields = [ + pa.field(f.name, spanner_type_to_arrow_type(f.type_)) for f in fields + ] + return pa.schema(arrow_fields) + + +def _extract_cell_value(cell: Any, spanner_type: Any = None) -> Any: + """Extract raw value from protobuf Value or Python object for fast Arrow ingestion.""" + if cell is None: + return None + + type_code = spanner_type.code if hasattr(spanner_type, "code") else spanner_type + + # Check if cell is a google.protobuf.Value + if hasattr(cell, "WhichOneof"): + kind = cell.WhichOneof("kind") + if kind == "null_value" or kind is None: + return None + elif kind == "bool_value": + return cell.bool_value + elif kind == "number_value": + return cell.number_value + elif kind == "string_value": + val_str = cell.string_value + if type_code == TypeCode.BYTES: + return base64.b64decode(val_str) + elif type_code in (TypeCode.FLOAT32, TypeCode.FLOAT64): + return float(val_str) + return val_str + elif kind == "list_value": + if type_code == TypeCode.STRUCT and hasattr(spanner_type, "struct_type"): + struct_fields = spanner_type.struct_type.fields + return { + f.name: _extract_nested_element(elem, f.type_) + for f, elem in zip(struct_fields, cell.list_value.values) + } + element_type = ( + spanner_type.array_element_type + if hasattr(spanner_type, "array_element_type") + else None + ) + return [ + _extract_nested_element(elem, element_type) + for elem in cell.list_value.values + ] + elif kind == "struct_value": + struct_fields_dict = { + f.name: f.type_ + for f in getattr( + getattr(spanner_type, "struct_type", None), "fields", () + ) + } + return { + k: _extract_nested_element(v, struct_fields_dict.get(k)) + for k, v in cell.struct_value.fields.items() + } + return None + + # Cell is already a Python object + if type_code == TypeCode.BYTES and isinstance(cell, str): + return base64.b64decode(cell) + if type_code in ( + TypeCode.STRING, + TypeCode.INTERVAL, + TypeCode.JSON, + TypeCode.UUID, + ) and not isinstance(cell, str): + return str(cell) + return cell + + +def _extract_nested_element(elem: Any, spanner_type: Any = None) -> Any: + """Extract nested array/struct elements into python scalars for PyArrow nested builders.""" + if elem is None: + return None + type_code = spanner_type.code if hasattr(spanner_type, "code") else spanner_type + val = _extract_cell_value(elem, spanner_type) + if isinstance(val, str): + if type_code == TypeCode.INT64: + return int(val) + elif type_code == TypeCode.NUMERIC: + import decimal + + return decimal.Decimal(val) + return val + + +def extract_columns_from_rows( + rows: Sequence[Sequence[Any]], field_types: Sequence[Any] +) -> List[List[Any]]: + """Extract columnar data from a batch of rows using fast list comprehensions.""" + num_columns = len(field_types) + return [ + [_extract_cell_value(row[idx], field_types[idx]) for row in rows] + for idx in range(num_columns) + ] + + +def convert_column_to_arrow_array( + column_values: List[Any], + arrow_field: "pa.Field", + type_code: Optional[int] = None, +) -> "pa.Array": + """Convert a single column's raw values into a PyArrow Array using fast C++ parsing. + + :type column_values: List[Any] + :param column_values: Extracted raw values for the column. + + :type arrow_field: :class:`pyarrow.Field` + :param arrow_field: PyArrow target field. + + :type type_code: Optional[int] + :param type_code: Spanner TypeCode for the column. + + :rtype: :class:`pyarrow.Array` + :returns: PyArrow Array for the column. + """ + _check_pyarrow() + arrow_type = arrow_field.type + + # If column is empty, return empty array with appropriate type + if not column_values: + return pa.array([], type=arrow_type) + + # If column contains string-encoded primitives from protobuf, use fast C++ casting + first_non_null = next((v for v in column_values if v is not None), None) + if isinstance(first_non_null, str): + if type_code == TypeCode.INT64: + return pc.cast(pa.array(column_values, type=pa.string()), pa.int64()) + elif type_code == TypeCode.DATE: + return pc.cast(pa.array(column_values, type=pa.string()), pa.date32()) + elif type_code == TypeCode.TIMESTAMP: + return pc.cast( + pa.array(column_values, type=pa.string()), + pa.timestamp("us", tz="UTC"), + ) + elif type_code == TypeCode.NUMERIC: + return pc.cast(pa.array(column_values, type=pa.string()), arrow_type) + + # Standard array construction for BOOL, FLOAT, STRING, BYTES, JSON, ARRAY, STRUCT, and native objects + return pa.array(column_values, type=arrow_type) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py index d16955d88abb..fd6b874c81a8 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py @@ -298,6 +298,148 @@ def to_dict_list(self): ) return rows + @CrossSync.convert + async def to_arrow_batches(self, max_chunk_size: int = 65536): + """Yield query results as a sequence of PyArrow RecordBatches. + + :type max_chunk_size: int + :param max_chunk_size: Target maximum number of rows per :class:`pyarrow.RecordBatch`. + Defaults to 65,536. + + :rtype: Iterator[:class:`pyarrow.RecordBatch`] + :returns: An iterator yielding RecordBatch objects. + """ + try: + import google_cloud_spanner_arrow as spanner_arrow_c + _HAS_ARROW_ACCELERATOR = True + except ImportError: + spanner_arrow_c = None + _HAS_ARROW_ACCELERATOR = False + + from google.cloud.spanner_v1._arrow import ( + _check_pyarrow, + convert_column_to_arrow_array, + extract_columns_from_rows, + spanner_schema_to_arrow_schema, + ) + + _check_pyarrow() + import pyarrow as pa + + self._lazy_decode = True + + if self._metadata is None: + try: + await self._consume_next() + except CrossSync.StopIteration: + pass + if self._metadata is None: + return + + fields = self.fields + pa_schema = spanner_schema_to_arrow_schema(fields) + field_types = [f.type_ for f in fields] + type_codes = [f.type_.code for f in fields] + + accumulated_rows = [] + + while True: + if self._rows: + accumulated_rows.extend(self._rows) + self._rows = [] + + while len(accumulated_rows) >= max_chunk_size: + batch_rows = accumulated_rows[:max_chunk_size] + accumulated_rows = accumulated_rows[max_chunk_size:] + if _HAS_ARROW_ACCELERATOR: + yield spanner_arrow_c.rows_to_arrow_batch( + fields, batch_rows, schema=pa_schema + ) + else: + col_buffers = extract_columns_from_rows(batch_rows, field_types) + arrays = [ + convert_column_to_arrow_array(col_data, pa_field, type_code) + for col_data, pa_field, type_code in zip( + col_buffers, pa_schema, type_codes + ) + ] + yield pa.RecordBatch.from_arrays(arrays, schema=pa_schema) + + if self._done: + break + + try: + await self._consume_next() + except CrossSync.StopIteration: + break + + if accumulated_rows: + if _HAS_ARROW_ACCELERATOR: + yield spanner_arrow_c.rows_to_arrow_batch( + fields, accumulated_rows, schema=pa_schema + ) + else: + col_buffers = extract_columns_from_rows(accumulated_rows, field_types) + arrays = [ + convert_column_to_arrow_array(col_data, pa_field, type_code) + for col_data, pa_field, type_code in zip( + col_buffers, pa_schema, type_codes + ) + ] + yield pa.RecordBatch.from_arrays(arrays, schema=pa_schema) + + @CrossSync.convert + async def to_arrow(self, max_chunk_size: int = 65536): + """Return the result of a query as a PyArrow Table. + + :type max_chunk_size: int + :param max_chunk_size: Target maximum number of rows per chunk when reading. + Defaults to 65,536. + + :rtype: :class:`pyarrow.Table` + :returns: A PyArrow Table containing all rows from the query result. + """ + from google.cloud.spanner_v1._arrow import ( + _check_pyarrow, + spanner_schema_to_arrow_schema, + ) + + _check_pyarrow() + import pyarrow as pa + + batches = [] + if CrossSync.is_async: + async for batch in self.to_arrow_batches(max_chunk_size=max_chunk_size): + batches.append(batch) + else: + for batch in self.to_arrow_batches(max_chunk_size=max_chunk_size): + batches.append(batch) + + if not batches: + if self._metadata is not None: + schema = spanner_schema_to_arrow_schema(self.fields) + else: + schema = pa.schema([]) + return pa.Table.from_batches([], schema=schema) + return pa.Table.from_batches(batches) + + @CrossSync.convert + async def to_dataframe(self, max_chunk_size: int = 65536, **kwargs): + """Return the result of a query as a Pandas DataFrame. + + :type max_chunk_size: int + :param max_chunk_size: Target maximum number of rows per chunk when reading. + Defaults to 65,536. + + :type kwargs: dict + :param kwargs: Keyword arguments forwarded to :meth:`pyarrow.Table.to_pandas`. + + :rtype: :class:`pandas.DataFrame` + :returns: A Pandas DataFrame containing all rows from the query result. + """ + table = await self.to_arrow(max_chunk_size=max_chunk_size) + return table.to_pandas(**kwargs) + class Unmergeable(ValueError): """Unable to merge two values. diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py index 8facd015151d..32f513acc7c2 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py @@ -258,6 +258,138 @@ def to_dict_list(self): ) return rows + def to_arrow_batches(self, max_chunk_size: int = 65536): + """Yield query results as a sequence of PyArrow RecordBatches. + + :type max_chunk_size: int + :param max_chunk_size: Target maximum number of rows per :class:`pyarrow.RecordBatch`. + Defaults to 65,536. + + :rtype: Iterator[:class:`pyarrow.RecordBatch`] + :returns: An iterator yielding RecordBatch objects. + """ + try: + import google_cloud_spanner_arrow as spanner_arrow_c + _HAS_ARROW_ACCELERATOR = True + except ImportError: + spanner_arrow_c = None + _HAS_ARROW_ACCELERATOR = False + + from google.cloud.spanner_v1._arrow import ( + _check_pyarrow, + convert_column_to_arrow_array, + extract_columns_from_rows, + spanner_schema_to_arrow_schema, + ) + + _check_pyarrow() + import pyarrow as pa + + self._lazy_decode = True + + if self._metadata is None: + try: + self._consume_next() + except StopIteration: + pass + if self._metadata is None: + return + + fields = self.fields + pa_schema = spanner_schema_to_arrow_schema(fields) + field_types = [f.type_ for f in fields] + type_codes = [f.type_.code for f in fields] + + accumulated_rows = [] + + while True: + if self._rows: + accumulated_rows.extend(self._rows) + self._rows = [] + + while len(accumulated_rows) >= max_chunk_size: + batch_rows = accumulated_rows[:max_chunk_size] + accumulated_rows = accumulated_rows[max_chunk_size:] + if _HAS_ARROW_ACCELERATOR: + yield spanner_arrow_c.rows_to_arrow_batch( + fields, batch_rows, schema=pa_schema + ) + else: + col_buffers = extract_columns_from_rows(batch_rows, field_types) + arrays = [ + convert_column_to_arrow_array(col_data, pa_field, type_code) + for col_data, pa_field, type_code in zip( + col_buffers, pa_schema, type_codes + ) + ] + yield pa.RecordBatch.from_arrays(arrays, schema=pa_schema) + + if self._done: + break + + try: + self._consume_next() + except StopIteration: + break + + if accumulated_rows: + if _HAS_ARROW_ACCELERATOR: + yield spanner_arrow_c.rows_to_arrow_batch( + fields, accumulated_rows, schema=pa_schema + ) + else: + col_buffers = extract_columns_from_rows(accumulated_rows, field_types) + arrays = [ + convert_column_to_arrow_array(col_data, pa_field, type_code) + for col_data, pa_field, type_code in zip( + col_buffers, pa_schema, type_codes + ) + ] + yield pa.RecordBatch.from_arrays(arrays, schema=pa_schema) + + def to_arrow(self, max_chunk_size: int = 65536): + """Return the result of a query as a PyArrow Table. + + :type max_chunk_size: int + :param max_chunk_size: Target maximum number of rows per chunk when reading. + Defaults to 65,536. + + :rtype: :class:`pyarrow.Table` + :returns: A PyArrow Table containing all rows from the query result. + """ + from google.cloud.spanner_v1._arrow import ( + _check_pyarrow, + spanner_schema_to_arrow_schema, + ) + + _check_pyarrow() + import pyarrow as pa + + batches = list(self.to_arrow_batches(max_chunk_size=max_chunk_size)) + if not batches: + if self._metadata is not None: + schema = spanner_schema_to_arrow_schema(self.fields) + else: + schema = pa.schema([]) + return pa.Table.from_batches([], schema=schema) + return pa.Table.from_batches(batches) + + def to_dataframe(self, max_chunk_size: int = 65536, **kwargs): + """Return the result of a query as a Pandas DataFrame. + + :type max_chunk_size: int + :param max_chunk_size: Target maximum number of rows per chunk when reading. + Defaults to 65,536. + + :type kwargs: dict + :param kwargs: Keyword arguments forwarded to :meth:`pyarrow.Table.to_pandas`. + + :rtype: :class:`pandas.DataFrame` + :returns: A Pandas DataFrame containing all rows from the query result. + """ + table = self.to_arrow(max_chunk_size=max_chunk_size) + return table.to_pandas(**kwargs) + class Unmergeable(ValueError): """Unable to merge two values. diff --git a/packages/google-cloud-spanner/setup.py b/packages/google-cloud-spanner/setup.py index 7c5878bc7a6c..b95009871343 100644 --- a/packages/google-cloud-spanner/setup.py +++ b/packages/google-cloud-spanner/setup.py @@ -64,6 +64,7 @@ ] extras = { "libcst": "libcst >= 0.2.5", + "pyarrow": ["pyarrow >= 14.0.0"], "test": [ "pytest", "mock", @@ -71,6 +72,7 @@ "pytest-cov", "pytest-asyncio", "pytest-xdist", + "pyarrow >= 14.0.0", ], } diff --git a/packages/google-cloud-spanner/tests/unit/test_streamed_arrow.py b/packages/google-cloud-spanner/tests/unit/test_streamed_arrow.py new file mode 100644 index 000000000000..0e4461a80d44 --- /dev/null +++ b/packages/google-cloud-spanner/tests/unit/test_streamed_arrow.py @@ -0,0 +1,388 @@ +# Copyright 2026 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import base64 +import decimal +import unittest +from unittest import mock + +from google.protobuf.struct_pb2 import ListValue, Struct, Value + +from google.cloud.spanner_v1 import _arrow, StructType, Type +from google.cloud.spanner_v1.streamed import StreamedResultSet +from google.cloud.spanner_v1.types.result_set import ( + PartialResultSet, + ResultSetMetadata, +) +from google.cloud.spanner_v1.types.type import TypeCode + +try: + import pyarrow as pa + _HAS_PYARROW = True +except ImportError: + _HAS_PYARROW = False + + +class _MockIterator(object): + def __init__(self, *values): + self._values = list(values) + + def __iter__(self): + return self + + def __next__(self): + if not self._values: + raise StopIteration + return self._values.pop(0) + + +@unittest.skipUnless(_HAS_PYARROW, "pyarrow is required for these tests") +class TestArrowHelpers(unittest.TestCase): + def test_spanner_type_to_arrow_type(self): + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.BOOL)), pa.bool_()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.INT64)), pa.int64()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.FLOAT32)), pa.float32()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.FLOAT64)), pa.float64()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.STRING)), pa.string()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.BYTES)), pa.binary()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.TIMESTAMP)), pa.timestamp("us", tz="UTC")) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.DATE)), pa.date32()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.NUMERIC)), pa.decimal128(38, 9)) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.JSON)), pa.string()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.PROTO)), pa.binary()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.ENUM)), pa.int64()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.INTERVAL)), pa.string()) + self.assertEqual(_arrow.spanner_type_to_arrow_type(Type(code=TypeCode.UUID)), pa.string()) + + # Array of INT64 + array_type = Type(code=TypeCode.ARRAY, array_element_type=Type(code=TypeCode.INT64)) + self.assertEqual(_arrow.spanner_type_to_arrow_type(array_type), pa.list_(pa.int64())) + + # Struct + struct_type = Type( + code=TypeCode.STRUCT, + struct_type=StructType( + fields=[ + StructType.Field(name="f1", type_=Type(code=TypeCode.STRING)), + StructType.Field(name="f2", type_=Type(code=TypeCode.INT64)), + ] + ), + ) + expected_struct = pa.struct([pa.field("f1", pa.string()), pa.field("f2", pa.int64())]) + self.assertEqual(_arrow.spanner_type_to_arrow_type(struct_type), expected_struct) + + def test_spanner_schema_to_arrow_schema(self): + fields = [ + StructType.Field(name="col_id", type_=Type(code=TypeCode.INT64)), + StructType.Field(name="col_name", type_=Type(code=TypeCode.STRING)), + ] + schema = _arrow.spanner_schema_to_arrow_schema(fields) + self.assertEqual(len(schema), 2) + self.assertEqual(schema.field("col_id").type, pa.int64()) + self.assertEqual(schema.field("col_name").type, pa.string()) + + def test_extract_cell_value(self): + # Null + val_null = Value(null_value=0) + self.assertIsNone(_arrow._extract_cell_value(val_null, TypeCode.INT64)) + + # Bool + val_bool = Value(bool_value=True) + self.assertTrue(_arrow._extract_cell_value(val_bool, TypeCode.BOOL)) + + # Number + val_num = Value(number_value=3.14) + self.assertEqual(_arrow._extract_cell_value(val_num, TypeCode.FLOAT64), 3.14) + + # String int64 + val_int = Value(string_value="123456") + self.assertEqual(_arrow._extract_cell_value(val_int, TypeCode.INT64), "123456") + + # Bytes (base64) + raw_bytes = b"hello world" + b64_str = base64.b64encode(raw_bytes).decode("ascii") + val_bytes = Value(string_value=b64_str) + self.assertEqual(_arrow._extract_cell_value(val_bytes, TypeCode.BYTES), raw_bytes) + + # Float NaN + val_nan = Value(string_value="NaN") + import math + self.assertTrue(math.isnan(_arrow._extract_cell_value(val_nan, TypeCode.FLOAT64))) + + # List + val_list = Value(list_value=ListValue(values=[Value(string_value="a"), Value(string_value="b")])) + self.assertEqual(_arrow._extract_cell_value(val_list, TypeCode.ARRAY), ["a", "b"]) + + def test_check_pyarrow_missing(self): + with mock.patch("google.cloud.spanner_v1._arrow._HAS_PYARROW", False): + with self.assertRaises(ImportError): + _arrow._check_pyarrow() + + +@unittest.skipUnless(_HAS_PYARROW, "pyarrow is required for these tests") +class TestStreamedResultSetArrow(unittest.TestCase): + def _make_metadata(self, fields): + metadata = ResultSetMetadata( + row_type=StructType(fields=[]) + ) + for name, code in fields: + metadata.row_type.fields.append( + StructType.Field(name=name, type_=Type(code=code)) + ) + return metadata + + def _make_partial_result_set( + self, values=(), metadata=None, stats=None, chunked_value=False, last=False + ): + results = PartialResultSet( + metadata=metadata, stats=stats, chunked_value=chunked_value, last=last + ) + for v in values: + results.values.append(v) + return results + + def test_to_arrow_empty(self): + iterator = _MockIterator() + streamed = StreamedResultSet(iterator) + table = streamed.to_arrow() + self.assertEqual(table.num_rows, 0) + self.assertEqual(table.num_columns, 0) + + def test_to_arrow_basic_query(self): + metadata = self._make_metadata([ + ("id", TypeCode.INT64), + ("name", TypeCode.STRING), + ("active", TypeCode.BOOL), + ("score", TypeCode.FLOAT64), + ]) + + prs1 = self._make_partial_result_set( + metadata=metadata, + values=[ + Value(string_value="1"), + Value(string_value="Alice"), + Value(bool_value=True), + Value(number_value=95.5), + Value(string_value="2"), + Value(string_value="Bob"), + Value(bool_value=False), + Value(number_value=82.0), + ], + last=True, + ) + + streamed = StreamedResultSet(_MockIterator(prs1)) + table = streamed.to_arrow() + + self.assertEqual(table.num_rows, 2) + self.assertEqual(table.num_columns, 4) + self.assertEqual(table.column("id").to_pylist(), [1, 2]) + self.assertEqual(table.column("name").to_pylist(), ["Alice", "Bob"]) + self.assertEqual(table.column("active").to_pylist(), [True, False]) + self.assertEqual(table.column("score").to_pylist(), [95.5, 82.0]) + + def test_to_arrow_batches_chunk_size(self): + metadata = self._make_metadata([("id", TypeCode.INT64)]) + values = [Value(string_value=str(i)) for i in range(10)] + prs = self._make_partial_result_set(metadata=metadata, values=values, last=True) + + streamed = StreamedResultSet(_MockIterator(prs)) + batches = list(streamed.to_arrow_batches(max_chunk_size=3)) + + self.assertEqual(len(batches), 4) # 3, 3, 3, 1 + self.assertEqual(batches[0].num_rows, 3) + self.assertEqual(batches[1].num_rows, 3) + self.assertEqual(batches[2].num_rows, 3) + self.assertEqual(batches[3].num_rows, 1) + + # Verify combined + table = pa.Table.from_batches(batches) + self.assertEqual(table.column("id").to_pylist(), list(range(10))) + + def test_to_arrow_with_chunked_values(self): + metadata = self._make_metadata([ + ("id", TypeCode.INT64), + ("description", TypeCode.STRING), + ]) + + # Chunk 1: id=1, description='hello ' (chunked) + prs1 = self._make_partial_result_set( + metadata=metadata, + values=[Value(string_value="1"), Value(string_value="hello ")], + chunked_value=True, + ) + # Chunk 2: 'world' (continuation), id=2, description='test' + prs2 = self._make_partial_result_set( + values=[ + Value(string_value="world"), + Value(string_value="2"), + Value(string_value="test"), + ], + last=True, + ) + + streamed = StreamedResultSet(_MockIterator(prs1, prs2)) + table = streamed.to_arrow() + + self.assertEqual(table.num_rows, 2) + self.assertEqual(table.column("id").to_pylist(), [1, 2]) + self.assertEqual(table.column("description").to_pylist(), ["hello world", "test"]) + + def test_to_arrow_advanced_types(self): + metadata = self._make_metadata([ + ("date_col", TypeCode.DATE), + ("ts_col", TypeCode.TIMESTAMP), + ("num_col", TypeCode.NUMERIC), + ("bytes_col", TypeCode.BYTES), + ]) + + raw_bytes = b"sample_bytes" + prs = self._make_partial_result_set( + metadata=metadata, + values=[ + Value(string_value="2023-01-15"), + Value(string_value="2023-01-15T10:30:00.123456Z"), + Value(string_value="12345.678900000"), + Value(string_value=base64.b64encode(raw_bytes).decode("ascii")), + ], + last=True, + ) + + streamed = StreamedResultSet(_MockIterator(prs)) + table = streamed.to_arrow() + + self.assertEqual(table.num_rows, 1) + self.assertEqual(table.column("num_col").to_pylist(), [decimal.Decimal("12345.678900000")]) + self.assertEqual(table.column("bytes_col").to_pylist(), [raw_bytes]) + + def test_to_dataframe(self): + try: + import pandas as pd + except ImportError: + return # Skip if pandas is not installed + + metadata = self._make_metadata([("id", TypeCode.INT64), ("name", TypeCode.STRING)]) + prs = self._make_partial_result_set( + metadata=metadata, + values=[Value(string_value="42"), Value(string_value="Answer")], + last=True, + ) + + streamed = StreamedResultSet(_MockIterator(prs)) + df = streamed.to_dataframe() + + self.assertEqual(len(df), 1) + self.assertEqual(df["id"].iloc[0], 42) + self.assertEqual(df["name"].iloc[0], "Answer") + + def test_to_arrow_nulls(self): + metadata = self._make_metadata([ + ("id", TypeCode.INT64), + ("name", TypeCode.STRING), + ("date_col", TypeCode.DATE), + ("ts_col", TypeCode.TIMESTAMP), + ("num_col", TypeCode.NUMERIC), + ]) + prs = self._make_partial_result_set( + metadata=metadata, + values=[ + Value(null_value=0), + Value(null_value=0), + Value(null_value=0), + Value(null_value=0), + Value(null_value=0), + ], + last=True, + ) + streamed = StreamedResultSet(_MockIterator(prs)) + table = streamed.to_arrow() + self.assertEqual(table.num_rows, 1) + self.assertIsNone(table.column("id").to_pylist()[0]) + self.assertIsNone(table.column("name").to_pylist()[0]) + self.assertIsNone(table.column("date_col").to_pylist()[0]) + self.assertIsNone(table.column("ts_col").to_pylist()[0]) + self.assertIsNone(table.column("num_col").to_pylist()[0]) + + def test_to_arrow_arrays_and_structs(self): + metadata = ResultSetMetadata( + row_type=StructType( + fields=[ + StructType.Field( + name="arr", + type_=Type( + code=TypeCode.ARRAY, + array_element_type=Type(code=TypeCode.STRING), + ), + ), + StructType.Field( + name="st", + type_=Type( + code=TypeCode.STRUCT, + struct_type=StructType( + fields=[ + StructType.Field( + name="f_int", + type_=Type(code=TypeCode.INT64), + ), + StructType.Field( + name="f_str", + type_=Type(code=TypeCode.STRING), + ), + ] + ), + ), + ), + ] + ) + ) + + arr_val = Value( + list_value=ListValue(values=[Value(string_value="x"), Value(string_value="y")]) + ) + st_val = Value( + struct_value=Struct( + fields={"f_int": Value(string_value="10"), "f_str": Value(string_value="hello")} + ) + ) + + prs = self._make_partial_result_set( + metadata=metadata, + values=[arr_val, st_val], + last=True, + ) + + streamed = StreamedResultSet(_MockIterator(prs)) + table = streamed.to_arrow() + + self.assertEqual(table.num_rows, 1) + self.assertEqual(table.column("arr").to_pylist(), [["x", "y"]]) + + def test_to_arrow_missing_pyarrow_raises(self): + metadata = self._make_metadata([("id", TypeCode.INT64)]) + prs = self._make_partial_result_set(metadata=metadata, values=[Value(string_value="1")], last=True) + streamed = StreamedResultSet(_MockIterator(prs)) + + with mock.patch("google.cloud.spanner_v1._arrow._HAS_PYARROW", False): + with self.assertRaises(ImportError): + streamed.to_arrow() + + with self.assertRaises(ImportError): + list(streamed.to_arrow_batches()) + + with self.assertRaises(ImportError): + streamed.to_dataframe() + + +if __name__ == "__main__": + unittest.main()