chore(spanner): decode Spanner query results to Arrow

[olavloite]

Adds a prototype for decoding Spanner query results directly to Arrow in an accelerator library written in C.

[gemini-code-assist]

Code Review

This pull request introduces google-cloud-spanner-arrow, a high-performance Apache Arrow accelerator package for Google Cloud Spanner. It includes a native C extension leveraging nanoarrow for GIL-free parsing of Spanner rows and raw protobuf wire bytes directly into Arrow RecordBatches, alongside a pure-Python fallback. The reviewer feedback highlights several critical issues in the C extension, including a Use-After-Free (UAF) vulnerability due to premature reference decrementing before releasing the GIL, memory leaks of heap-allocated Arrow structures, and ignored return values in row parsing. Additionally, the fallback detection logic in the main Spanner client needs to verify that the C extension compiled successfully, and the pure-Python fallback requires fixes to correctly handle nested schemas.

[gemini-code-assist]

There is a critical Use-After-Free (UAF) vulnerability here. The reference count of chunk_obj is decremented via Py_XDECREF(chunk_obj) immediately inside the loop. If chunk_obj is a temporary object (e.g., if py_wire_chunks is a generator or custom sequence), it will be destroyed and its internal buffer freed. When the GIL is subsequently released with Py_BEGIN_ALLOW_THREADS, raw_buffers[i].ptr becomes a dangling pointer pointing to freed memory, leading to crashes or memory corruption.

To fix this, keep the references to chunk_obj alive until after the GIL is re-acquired and parsing is complete.

    PyObject** chunk_objects = (PyObject**)calloc(num_chunks, sizeof(PyObject*));
    if (chunk_objects == 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);
        chunk_objects[i] = chunk_obj;
        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_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,
                &current_col_idx
            );
        }
    }
    Py_END_ALLOW_THREADS

    for (Py_ssize_t i = 0; i < num_chunks; i++) {
        Py_XDECREF(chunk_objects[i]);
    }
    free(chunk_objects);

[gemini-code-assist]

The out_schema and out_array structures are allocated on the heap via calloc but are never freed after being imported by PyArrow. Since PyArrow copies them by value, the original heap-allocated structures are leaked.

To fix this, we can add a simple helper function to free these pointers from Python right after _import_from_c returns.

[gemini-code-assist]

spanner_type_to_arrow_type and fields_to_arrow_schema do not correctly handle nested tuples (e.g., representing arrays or structs inside structs), defaulting them to pa.string() and causing schema mismatches.

To fix this, modify spanner_type_to_arrow_type to support normalized tuples recursively.

def spanner_type_to_arrow_type(spanner_type: Any) -> "pa.DataType":
    """Map a Spanner Type to a PyArrow DataType."""
    _check_pyarrow()
    if isinstance(spanner_type, tuple):
        code = _get_type_code(spanner_type[1])
        if code == SPANNER_TYPE_ARRAY:
            return pa.list_(spanner_type_to_arrow_type(spanner_type[2]))
        elif code == SPANNER_TYPE_STRUCT:
            return pa.struct([
                pa.field(sf[0], spanner_type_to_arrow_type(sf)) for sf in spanner_type[2]
            ])
        return spanner_type_to_arrow_type(spanner_type[1])

    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):
            arrow_fields.append(pa.field(f[0], spanner_type_to_arrow_type(f)))
        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)

[gemini-code-assist]

If google_cloud_spanner_arrow is installed but the C extension failed to compile, the code falls back to google_cloud_spanner_arrow.python instead of the correct, fully-featured built-in _arrow.py fallback.

To prevent this, check if the C extension is actually available by verifying spanner_arrow_c.implementation == 'c'.

Suggested change

	try:
	import google_cloud_spanner_arrow as spanner_arrow_c
	_HAS_ARROW_ACCELERATOR = True
	except ImportError:
	spanner_arrow_c = None
	_HAS_ARROW_ACCELERATOR = False
	try:
	import google_cloud_spanner_arrow as spanner_arrow_c
	_HAS_ARROW_ACCELERATOR = getattr(spanner_arrow_c, "implementation", None) == "c"
	except ImportError:
	spanner_arrow_c = None
	_HAS_ARROW_ACCELERATOR = False

[gemini-code-assist]

If google_cloud_spanner_arrow is installed but the C extension failed to compile, the code falls back to google_cloud_spanner_arrow.python instead of the correct, fully-featured built-in _arrow.py fallback.

To prevent this, check if the C extension is actually available by verifying spanner_arrow_c.implementation == 'c'.

Suggested change

	try:
	import google_cloud_spanner_arrow as spanner_arrow_c
	_HAS_ARROW_ACCELERATOR = True
	except ImportError:
	spanner_arrow_c = None
	_HAS_ARROW_ACCELERATOR = False
	try:
	import google_cloud_spanner_arrow as spanner_arrow_c
	_HAS_ARROW_ACCELERATOR = getattr(spanner_arrow_c, "implementation", None) == "c"
	except ImportError:
	spanner_arrow_c = None
	_HAS_ARROW_ACCELERATOR = False

[gemini-code-assist]

The return value of append_python_cell is completely ignored in py_rows_to_c_batch. If an allocation fails or an error occurs inside append_python_cell, the function will silently continue, leaving the array in an inconsistent or corrupted state, and any Python exceptions set during these calls will remain set and propagate randomly.

To fix this, check the return value of append_python_cell and abort if it is non-zero.

    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;
                if (append_python_cell(out_array->children[c], cell, f_info) != 0) {
                    Py_XDECREF(cell);
                    Py_XDECREF(f_info);
                    Py_XDECREF(row);
                    ArrowArrayRelease(out_array);
                    ArrowSchemaRelease(out_schema);
                    free(out_array);
                    free(out_schema);
                    return NULL;
                }
                Py_XDECREF(cell);
                Py_XDECREF(f_info);
            }
            out_array->length++;
        }
        Py_XDECREF(row);
    }