cdm_data_loaders.coercion.pydantic ¶

def coerce_record_for_pyarrow(
    record: BaseModel | Mapping[str, Any],
    schema: pa.Schema,
    *,
    model: type[BaseModel] | None = None,
    by_alias: bool = False,
) -> dict[str, Any]:
    """Coerce a Pydantic ``BaseModel`` instance or record mapping to match a PyArrow schema.

    When ``record`` is a ``BaseModel`` instance (or ``model`` is supplied to validate a mapping),
    non-null default values and ``@computed_field`` properties are populated on the output dictionary.
    Only fields declared in ``schema`` are retained and coerced recursively so
    ``pyarrow.Table.from_pylist`` can construct a table matching ``schema``.

    :param record: Pydantic ``BaseModel`` instance or record dictionary.
    :type record: BaseModel | Mapping[str, Any]
    :param schema: Target PyArrow schema for the record's table.
    :type schema: pa.Schema
    :param model: Optional ``BaseModel`` class used to validate ``record`` when ``record`` is a mapping.
    :type model: type[BaseModel] | None
    :param by_alias: When ``True``, extract model fields using ``serialization_alias`` or ``alias``.
    :type by_alias: bool
    :return: Coerced record dictionary ready for ``DictBuffer`` or ``pa.Table.from_pylist``.
    :rtype: dict[str, Any]
    """
    if isinstance(record, BaseModel):
        source_map: Mapping[str, Any] = _extract_model_dict(record, by_alias=by_alias)
    elif model is not None:
        validated = model.model_validate(dict(record))
        source_map = _extract_model_dict(validated, by_alias=by_alias)
    else:
        source_map = record

    return {
        field.name: _coerce_value_for_pyarrow(
            source_map[field.name],
            field.type,
            field.metadata,
            by_alias=by_alias,
        )
        for field in schema
        if field.name in source_map
    }