Extract

aef_bng.extract

Vectorised BNG reference generation and Arrow table extraction.

Converts reprojected raster arrays into tabular rows with BNG references.

Performance-critical: ~1M pixels per 10km chunk.

extract_pixels(data, chunk, year, mask_wkb=None)

Extract valid pixels from a reprojected chunk as an Arrow table.

Each embedding band is a separate int8 column (A00..A63). Includes easting/northing for GeoParquet geometry construction.

Parameters:

Name	Type	Description	Default
data	ndarray	Reprojected array of shape (64, rows, cols) int8.	required
chunk	ChunkSpec	The BNG chunk specification.	required
year	int	Year of the AEF embeddings.	required
mask_wkb	bytes | None	Optional boundary geometry (WKB) restricting output pixels.	None

Returns:

Type	Description
Table	Arrow table with columns: bng_ref, year, grid_10km_ref, grid_1km_ref,
Table	A00..A63, easting, northing.

Source code in aef_bng/extract.py

def extract_pixels(
    data: np.ndarray,
    chunk: ChunkSpec,
    year: int,
    mask_wkb: bytes | None = None,
) -> pa.Table:
    """Extract valid pixels from a reprojected chunk as an Arrow table.

    Each embedding band is a separate int8 column (A00..A63).
    Includes easting/northing for GeoParquet geometry construction.

    Args:
        data: Reprojected array of shape (64, rows, cols) int8.
        chunk: The BNG chunk specification.
        year: Year of the AEF embeddings.
        mask_wkb: Optional boundary geometry (WKB) restricting output pixels.

    Returns:
        Arrow table with columns: bng_ref, year, grid_10km_ref, grid_1km_ref,
        A00..A63, easting, northing.
    """
    result = _compute_pixel_data(data, chunk, mask_wkb)
    if result is None:
        return _empty_table()

    valid_eastings, valid_northings, valid_rows, valid_cols, ref_codes = result
    n_valid = len(ref_codes)

    columns = _key_columns(ref_codes, year, n_valid)
    _add_band_columns(columns, data, valid_rows, valid_cols, n_valid)

    columns["easting"] = pa.array(valid_eastings, type=pa.int32())
    columns["northing"] = pa.array(valid_northings, type=pa.int32())

    return pa.table(columns)

extract_pixels_spark(data, chunk, year, mask_wkb=None)

Extract valid pixels for the Spark/Unity Catalog path.

Each embedding band is a separate int8 column (A00..A63). A geometry_wkb binary column is always included, containing standard WKB polygons for each 10m cell (EPSG:27700).

Parameters:

Name	Type	Description	Default
data	ndarray	Reprojected array of shape (64, rows, cols) int8.	required
chunk	ChunkSpec	The BNG chunk specification.	required
year	int	Year of the AEF embeddings.	required
mask_wkb	bytes | None	Optional boundary geometry (WKB) restricting output pixels.	None

Returns:

Type	Description
Table	Arrow table with columns: bng_ref, year, grid_10km_ref, grid_1km_ref,
Table	A00..A63, geometry_wkb.

Source code in aef_bng/extract.py

def extract_pixels_spark(
    data: np.ndarray,
    chunk: ChunkSpec,
    year: int,
    mask_wkb: bytes | None = None,
) -> pa.Table:
    """Extract valid pixels for the Spark/Unity Catalog path.

    Each embedding band is a separate int8 column (A00..A63). A
    ``geometry_wkb`` binary column is always included, containing standard
    WKB polygons for each 10m cell (EPSG:27700).

    Args:
        data: Reprojected array of shape (64, rows, cols) int8.
        chunk: The BNG chunk specification.
        year: Year of the AEF embeddings.
        mask_wkb: Optional boundary geometry (WKB) restricting output pixels.

    Returns:
        Arrow table with columns: bng_ref, year, grid_10km_ref, grid_1km_ref,
        A00..A63, geometry_wkb.
    """
    result = _compute_pixel_data(data, chunk, mask_wkb)
    if result is None:
        return _empty_table_spark()

    valid_eastings, valid_northings, valid_rows, valid_cols, ref_codes = result
    n_valid = len(ref_codes)

    columns = _key_columns(ref_codes, year, n_valid)
    _add_band_columns(columns, data, valid_rows, valid_cols, n_valid)

    columns["geometry_wkb"] = _wkb_boxes(valid_eastings, valid_northings)

    return pa.table(columns)