numbarrow.core.mapinarrow_factory

Overview

Factory for PySpark mapInArrow UDF functions. Bridges PySpark’s Arrow-based batch processing with Numba JIT-compiled functions by converting each pyarrow.RecordBatch column through the adapter layer before passing data to a user-supplied computation function.

Usage:

from numbarrow.core.mapinarrow_factory import Nullable, make_mapinarrow_func

def my_func(data_dict, bitmap_dict, broadcasts):
    # data_dict:   {name: np.ndarray} for a uniform column, or
    #              {name: {field: np.ndarray}} for a struct column
    # bitmap_dict: the same shape, each leaf a uint8 bitmap or None where
    #              the column carries no validity buffer; for a struct
    #              column the struct-level validity is folded into each
    #              field's bitmap
    # broadcasts:  {key: value}
    result = data_dict["input_col"] * broadcasts["scale"]
    # result is null wherever input_col is, so input_col's bitmap goes out
    # with it; a bare array would carry no nulls out, and a result with
    # nulls of its own packs them: np.packbits(valid, bitorder="little")
    return {"output_col": Nullable(result, bitmap_dict["input_col"])}

udf = make_mapinarrow_func(my_func, broadcasts={"scale": 1.5})
df_out = df_in.mapInArrow(udf, output_schema)

Every name in data_dict is also a key of bitmap_dict, so a batch that happens to contain no nulls is indexable exactly like one that does.

On the way out a bare array carries no nulls: a row that came in null goes out valid, holding whatever the UDF computed from the placeholder under the null, which is 0, 0.0 or '' in a batch from Spark. Nullable(data, bitmap) carries nulls out, the bitmap being in the layout bitmap_dict hands out. A result that is null exactly where one input column is, as in the example, passes that column’s bitmap through, Nullable(result, bitmap_dict[column]), or bitmap_dict[column][field] for a struct field; any other result builds its own, for instance np.packbits(valid, bitorder="little") from a boolean array valid. A packed bitmap carries no row count, so a bitmap the batch handed out is accepted only on a column of the length it covers, the batch’s rows for a column’s own bitmap and the flattened elements for a struct field’s; one that resizes the column needs a bitmap of its own.

For a StructArray column the struct-level validity is folded into each field’s bitmap, so a row that is null as a whole is visible to one is_null call per field. For a ListArray of structs the fold covers the flattened struct elements, NOT the outer list rows: a null outer row can be reported nowhere and also shifts the element-to-row mapping, so a list column whose null_count is non-zero raises NotImplementedError.

A struct or list-of-struct column’s fields sit under the column’s own name, data_dict[column][field], so a field never shares a namespace with another column or with another struct’s fields.

See test_mapinarrow_spark.py in the test suite for a complete runnable example.

Module

Factory for PySpark mapInArrow UDF functions.

Bridges PySpark’s Arrow-based batch processing with Numba JIT-compiled functions by converting each pyarrow.RecordBatch column through arrow_array_adapter() before passing the data to a user-supplied computation function.

class numbarrow.core.mapinarrow_factory.Nullable(data: object, bitmap: ndarray | None)[source]

Bases: NamedTuple

An output column with its validity, Nullable(data, bitmap).

data is anything a column may be: an ndarray, a list, a pyarrow.Array or a numpy record array. bitmap is a packed uint8 validity bitmap in the layout bitmap_dict hands out, (rows + 7) // 8 bytes with a set bit for a valid row, or None, which is what bitmap_dict holds for a column with no validity buffer. A bare array carries no nulls out of a UDF; this does.

bitmap: ndarray | None: Alias for field number 1

data: object: Alias for field number 0

numbarrow.core.mapinarrow_factory.make_mapinarrow_func(main_func: Callable, input_columns: list[str] | None = None, broadcasts: dict | None = None, output_schema: Schema | None = None)[source]

Creates a function that can be given as an argument to mapInArrow

Parameters:

main_func –
called once per pyarrow.RecordBatch as main_func(data_dict, bitmap_dict, broadcasts), returning a dict that maps each output column’s name to an ndarray, a list or a pyarrow.Array, from which a PyArrow RecordBatch is built. A numpy record array becomes a struct column, one child per field. A null comes out of a list, a tuple or an object array holding None, a pyarrow.Array, a numpy masked array, and a Nullable, Nullable(data, bitmap), whose bitmap is a packed uint8 validity bitmap in the layout bitmap_dict hands out, (rows + 7) // 8 bytes with a set bit for a valid row, or None. A bare array carries no nulls out: a row that came in null goes out valid, holding whatever the UDF computed from the placeholder under the null. A result that is null exactly where one input column is passes that column’s bitmap through, {"out": Nullable(result, bitmap_dict["value"])}, or bitmap_dict["column"]["field"] for a struct field, and any other result builds its own in that layout, the one is_null() reads, for instance np.packbits(valid, bitorder="little") from a boolean array. A bitmap that is not an ndarray, or is one of another length or dtype, raises naming the column, and so does a bitmap the batch handed out on a column whose row count is not the count that bitmap covers, the batch’s rows for a column’s own bitmap and the flattened elements for a struct field’s, since a packed bitmap cannot tell row counts apart inside one byte.

Spark binds the columns of that batch to the declared output schema by POSITION, not by name, and checks nothing about their names: it reads each Arrow vector through the accessor its declared type expects, so an int64 column declared as a timestamp reads as a timestamp, and two columns whose types share an accessor family swap silently when the dict is built in the other order. A column read through the accessor of another family fails in the JVM with java.lang.UnsupportedOperationException whatever its width, as float64 under LongType and int64 under DoubleType do, both 64 bits wide, and so does a width mismatch inside one family, such as int32 under LongType. Build the returned dict in the order the output schema declares, or pass output_schema and let Arrow bind it by name instead.

data_dict maps each selected column’s name to its data. A column of a uniform type maps to one array. A struct or list-of-struct column maps to a dict of its fields, data_dict[column][field], so a field never shares a namespace with another column or with another struct’s fields.

bitmap_dict has the same shape: a uint8 aligned array of bitmap data, or None where the column carries no validity buffer, and for a struct column a dict of those keyed by field. Every key of data_dict is present, so a null-free batch is indexable exactly like a batch containing nulls.

For a StructArray column the struct-level validity is folded into each field’s bitmap, so one is_null() call per field sees both a null field and a row that is null as a whole.

For a ListArray of structs the fold covers the flattened struct elements, NOT the outer list rows. A null outer row can be reported nowhere, and because the adapter returns the flattened elements with no offsets, such a row also shifts the element-to-row mapping, so a list column whose null_count is non-zero raises NotImplementedError.
input_columns – optional list of column names that will be expected to be needed for in data_dict for the calculation done by main_func. When not given, all columns in the iterated over PySpark DataFrame will be used. Names are matched exactly; a name the batch does not have raises KeyError listing the batch’s columns, since Spark’s case-insensitive projection may have spelled it differently, and a name the batch carries more than once, as an unaliased join produces, raises ValueError.
broadcasts – optional dictionary of broadcast values
output_schema –
optional pyarrow.Schema for the batch that is yielded. When given, the dict returned by main_func is bound to it BY NAME, so insertion order stops deciding, and every column is built with its declared type rather than inferred and cast: a string column declared large_string is built as one, a list of dicts declared struct is built field by field, and a list of dicts or of pairs declared map becomes a map, which no inferred type can be cast to. A name the schema declares but the dict omits raises KeyError, a key the schema does not name raises ValueError, and a dict key that no declared struct field has raises ValueError too, since Arrow matches struct fields by exact name and would otherwise fill the column with nulls.

What the declared type refuses is what pa.array refuses, and that depends on the shape the column arrives in. For an ndarray of a numeric or datetime dtype, or a pyarrow.Array, an integer out of the declared type’s range, a float with a fraction into an integer type and a timestamp unit change that drops digits all raise pyarrow.ArrowInvalid. A Python list, and any other sequence of Python objects, an object-dtype ndarray included, goes through pa.array’s sequence converter instead: an integer out of the declared type’s range still raises pyarrow.ArrowInvalid and one beyond int64 altogether raises OverflowError, but a float’s fraction and a timestamp’s extra digits are dropped silently. So the lossy conversions that pass without a word are a timestamp into date32 or date64, which floors to the day, float64 into float32, which overflows to inf, and, from a list alone, a fraction into an integer type and a timestamp unit change that drops digits.

Left as None the batch is built from the dict alone: insertion order decides, and every type is inferred from the value, so a unicode or bytes array comes back string or binary whatever type went in, a datetime64 array comes back a naive timestamp of its unit, except a day-unit one, which comes back date32, and an object array holding only None comes back null.