Add typed Arrow execution backend boundary
This commit is contained in:
@@ -0,0 +1,222 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Iterable, Mapping
|
||||
|
||||
from govoplan_dataflow.backend.ir import (
|
||||
IrField,
|
||||
IrSchema,
|
||||
schema_from_fields,
|
||||
schema_from_rows,
|
||||
)
|
||||
|
||||
|
||||
class ArrowDependencyError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TypedBatch:
|
||||
schema: IrSchema
|
||||
columns: Mapping[str, tuple[Any, ...]]
|
||||
row_count: int
|
||||
byte_count: int
|
||||
|
||||
@classmethod
|
||||
def from_rows(
|
||||
cls,
|
||||
rows: Iterable[Mapping[str, Any]],
|
||||
*,
|
||||
schema: IrSchema | None = None,
|
||||
) -> "TypedBatch":
|
||||
normalized = [dict(row) for row in rows]
|
||||
resolved_schema = schema or schema_from_rows(normalized)
|
||||
names = _column_names(normalized, resolved_schema)
|
||||
columns = {
|
||||
name: tuple(row.get(name) for row in normalized)
|
||||
for name in names
|
||||
}
|
||||
return cls(
|
||||
schema=resolved_schema,
|
||||
columns=columns,
|
||||
row_count=len(normalized),
|
||||
byte_count=_estimated_bytes(normalized),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_arrow_ipc(cls, payload: bytes) -> "TypedBatch":
|
||||
pa = _pyarrow()
|
||||
with pa.ipc.open_stream(payload) as reader:
|
||||
table = reader.read_all()
|
||||
return cls.from_arrow_table(table)
|
||||
|
||||
@classmethod
|
||||
def from_arrow_table(cls, table: object) -> "TypedBatch":
|
||||
pa = _pyarrow()
|
||||
if not isinstance(table, pa.Table):
|
||||
raise TypeError("Expected a pyarrow.Table.")
|
||||
schema = schema_from_fields(
|
||||
tuple(
|
||||
IrField(
|
||||
name=field.name,
|
||||
type=_arrow_data_type(field.type),
|
||||
nullable=field.nullable,
|
||||
)
|
||||
for field in table.schema
|
||||
),
|
||||
)
|
||||
columns = {
|
||||
name: tuple(table.column(name).to_pylist())
|
||||
for name in table.column_names
|
||||
}
|
||||
return cls(
|
||||
schema=schema,
|
||||
columns=columns,
|
||||
row_count=table.num_rows,
|
||||
byte_count=table.nbytes,
|
||||
)
|
||||
|
||||
def to_rows(self) -> list[dict[str, Any]]:
|
||||
names = list(self.columns)
|
||||
return [
|
||||
{
|
||||
name: self.columns[name][index]
|
||||
for name in names
|
||||
}
|
||||
for index in range(self.row_count)
|
||||
]
|
||||
|
||||
def to_arrow_table(self) -> object:
|
||||
pa = _pyarrow()
|
||||
if self.columns:
|
||||
return pa.table(
|
||||
{
|
||||
name: list(values)
|
||||
for name, values in self.columns.items()
|
||||
}
|
||||
)
|
||||
fields = [
|
||||
pa.field(
|
||||
field.name,
|
||||
_pyarrow_type(field.type),
|
||||
nullable=field.nullable,
|
||||
)
|
||||
for field in self.schema.fields
|
||||
]
|
||||
schema = pa.schema(fields)
|
||||
return pa.Table.from_arrays(
|
||||
[pa.array([], type=field.type) for field in fields],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
def to_arrow_ipc(self) -> bytes:
|
||||
pa = _pyarrow()
|
||||
table = self.to_arrow_table()
|
||||
sink = pa.BufferOutputStream()
|
||||
with pa.ipc.new_stream(sink, table.schema) as writer:
|
||||
writer.write_table(table)
|
||||
return sink.getvalue().to_pybytes()
|
||||
|
||||
def ensure_within(
|
||||
self,
|
||||
*,
|
||||
max_rows: int,
|
||||
max_bytes: int,
|
||||
) -> None:
|
||||
if self.row_count > max_rows:
|
||||
raise ValueError(
|
||||
f"Typed batch has {self.row_count} rows; limit is {max_rows}."
|
||||
)
|
||||
if self.byte_count > max_bytes:
|
||||
raise ValueError(
|
||||
f"Typed batch uses {self.byte_count} bytes; limit is {max_bytes}."
|
||||
)
|
||||
|
||||
|
||||
def arrow_available() -> bool:
|
||||
try:
|
||||
_pyarrow()
|
||||
except ArrowDependencyError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _column_names(
|
||||
rows: list[dict[str, Any]],
|
||||
schema: IrSchema,
|
||||
) -> list[str]:
|
||||
names = [field.name for field in schema.fields]
|
||||
seen = set(names)
|
||||
for row in rows:
|
||||
for name in row:
|
||||
if name not in seen:
|
||||
names.append(name)
|
||||
seen.add(name)
|
||||
return names
|
||||
|
||||
|
||||
def _estimated_bytes(rows: list[dict[str, Any]]) -> int:
|
||||
encoded = json.dumps(
|
||||
rows,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
ensure_ascii=True,
|
||||
default=str,
|
||||
)
|
||||
return len(encoded.encode("utf-8"))
|
||||
|
||||
|
||||
def _pyarrow() -> Any:
|
||||
try:
|
||||
import pyarrow as pa
|
||||
except ImportError as exc:
|
||||
raise ArrowDependencyError(
|
||||
"Arrow execution requires the govoplan-dataflow analytics extra."
|
||||
) from exc
|
||||
return pa
|
||||
|
||||
|
||||
def _arrow_data_type(data_type: object) -> str:
|
||||
pa = _pyarrow()
|
||||
if pa.types.is_boolean(data_type):
|
||||
return "boolean"
|
||||
if pa.types.is_integer(data_type):
|
||||
return "integer"
|
||||
if pa.types.is_floating(data_type) or pa.types.is_decimal(data_type):
|
||||
return "number"
|
||||
if pa.types.is_date(data_type):
|
||||
return "date"
|
||||
if pa.types.is_timestamp(data_type):
|
||||
return "datetime"
|
||||
if pa.types.is_binary(data_type):
|
||||
return "binary"
|
||||
if pa.types.is_string(data_type) or pa.types.is_large_string(data_type):
|
||||
return "string"
|
||||
if pa.types.is_null(data_type):
|
||||
return "null"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _pyarrow_type(data_type: str) -> object:
|
||||
pa = _pyarrow()
|
||||
mapping = {
|
||||
"boolean": pa.bool_(),
|
||||
"integer": pa.int64(),
|
||||
"number": pa.float64(),
|
||||
"date": pa.date32(),
|
||||
"datetime": pa.timestamp("us", tz="UTC"),
|
||||
"binary": pa.binary(),
|
||||
"string": pa.string(),
|
||||
"null": pa.null(),
|
||||
"json": pa.string(),
|
||||
"unknown": pa.null(),
|
||||
}
|
||||
return mapping.get(data_type, pa.null())
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ArrowDependencyError",
|
||||
"TypedBatch",
|
||||
"arrow_available",
|
||||
]
|
||||
Reference in New Issue
Block a user