floe-python 0.8.0__py3-none-any.whl

floe/read/json.py

@@ -0,0 +1,236 @@
+"""JSON reader: every selected column as a raw string (Rust: `io/read/json.rs`, `json_selector.rs`)."""
+
+import json
+import math
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, NoReturn, Self
+
+import polars as pl
+
+from floe.checks.cast import decimal_type
+from floe.config.entity import WHITE_SPACE
+from floe.config.types import ColumnConfig, SourceOptions
+from floe.errors import FloeIoError
+
+Selector = list[str | int]  # a field name or an array index
+
+
+@dataclass(frozen=True)
+class JsonChunk:
+    """`offset` is the whole-file index of the chunk's first row (add it to the row index)."""
+
+    offset: int
+    df: pl.DataFrame
+
+
+def parse_selector(value: str) -> Selector:
+    """A path of fields and array indexes, e.g. `a.b[0].c` (a `.` right after `]` is optional)."""
+    trimmed = value.strip()
+    if not trimmed:
+        raise ValueError("selector is empty")
+    tokens: Selector = []
+    buf = ""
+    after_index = False
+    chars = iter(trimmed)
+    for ch in chars:
+        if ch == ".":
+            if buf:
+                tokens.append(buf)
+                buf = ""
+            elif not after_index:
+                raise ValueError("selector contains empty field")
+            after_index = False
+        elif ch == "[":
+            if buf:
+                tokens.append(buf)
+                buf = ""
+            digits = ""
+            for nxt in chars:
+                if nxt == "]":
+                    break
+                digits += nxt
+            if not digits:
+                raise ValueError("selector index is empty")
+            if not (digits.isascii() and digits.isdigit()):
+                raise ValueError(f"selector index is not numeric: {digits}")
+            if int(digits) >= 1 << 64:
+                raise ValueError(f"selector index is invalid: {digits}")
+            tokens.append(int(digits))
+            after_index = True
+        elif ch == "]":
+            raise ValueError("selector contains unexpected ']'")
+        else:
+            buf += ch
+            after_index = False
+    if buf:
+        tokens.append(buf)
+    return tokens
+
+
+def _is(value: Any, kind: type) -> bool:
+    """isinstance without narrowing: pyright strict rejects the Unknown containers it infers."""
+    return isinstance(value, kind)
+
+
+def _compact(value: Any) -> str:
+    return json.dumps(
+        value, separators=(",", ":"), ensure_ascii=False, sort_keys=True
+    )  # serde_json sorts keys
+
+
+class _Number(float):
+    """A JSON float that keeps its token, so a `decimal(p,s)` column reads it without rounding."""
+
+    text: str
+
+    def __new__(cls, text: str) -> Self:
+        self = super().__new__(cls, text)
+        self.text = text
+        return self
+
+
+def _cell(
+    value: Any,
+    tokens: Selector,
+    source: str,
+    cast_mode: str,
+    where: str,
+    exact: bool = False,
+) -> str | None:
+    """Missing paths and JSON null are null; objects and arrays are an error unless `coerce`."""
+    row: Any = value
+    for token in tokens:
+        if isinstance(token, str):
+            if not _is(row, dict) or token not in row:
+                return None
+        elif not _is(row, list) or token >= len(row):
+            return None
+        row = row[token]
+    if row is None:
+        return None
+    if _is(row, dict) or _is(row, list):
+        if cast_mode != "coerce":
+            raise FloeIoError(f"non-scalar value for selector {source} at {where}")
+        return _compact(row)
+    if isinstance(row, bool):
+        return "true" if row else "false"
+    if exact and isinstance(row, _Number):
+        return row.text
+    return str(row)
+
+
+def _reject_constant(name: str) -> NoReturn:
+    raise ValueError(f"invalid number: {name}")
+
+
+def _float(text: str) -> float:
+    value = _Number(text)
+    if math.isinf(value):
+        raise ValueError("number out of range")
+    return value
+
+
+def _int(text: str) -> float:
+    """serde reads integers beyond i64/u64 as f64."""
+    value = int(text)
+    return value if -(1 << 63) <= value < 1 << 64 else _float(text)
+
+
+def json_loads(text: str) -> Any:
+    """`json.loads` with serde's rules: no NaN/Infinity, big numbers kept, unpaired surrogates rejected."""
+    value = json.loads(
+        text, parse_constant=_reject_constant, parse_float=_float, parse_int=_int
+    )
+    # serde rejects unpaired surrogate escapes; Python keeps them and only fails later
+    # ponytail: a second pass over each value; a parse_string hook would avoid it
+    json.dumps(value, ensure_ascii=False).encode("utf-8")
+    return value
+
+
+def json_records(path: Path, mode: str) -> Iterator[tuple[str, Any]]:
+    """Yield (location, object) per row; `where` is `line N` for ndjson, `index N` for arrays."""
+    try:
+        if mode == "ndjson":
+            with path.open("rb") as handle:
+                for number, raw in enumerate(handle, 1):
+                    try:
+                        line = raw.decode("utf-8").strip(WHITE_SPACE)  # Rust str::trim
+                        if not line:
+                            continue
+                        value = json_loads(line)
+                    except ValueError as err:  # includes UnicodeDecodeError
+                        raise FloeIoError(
+                            f"json parse error at line {number}: {err}"
+                        ) from err
+                    if not _is(value, dict):
+                        raise FloeIoError(f"expected json object at line {number}")
+                    yield f"line {number}", value
+            return
+        # ponytail: array mode parses the whole file in memory, as Rust; stream if arrays outgrow it
+        try:
+            array: Any = json_loads(path.read_text(encoding="utf-8"))
+        except ValueError as err:
+            raise FloeIoError(f"json parse error: {err}") from err
+    except OSError as err:
+        raise FloeIoError(f"failed to read json at {path}: {err}") from err
+    if not _is(array, list):
+        raise FloeIoError("expected json array at root")
+    for index, value in enumerate(array):
+        if not _is(value, dict):
+            raise FloeIoError(f"expected json object at index {index}")
+        yield f"index {index}", value
+
+
+def read_json(
+    path: Path,
+    options: SourceOptions,
+    columns: Sequence[ColumnConfig],
+    cast_mode: str = "strict",
+    chunk_rows: int = 1_000_000,
+) -> Iterator[JsonChunk]:
+    """Yield the rows in chunks of at most `chunk_rows`, columns named by selector; always at least one.
+
+    `options.json_mode` is `array` (default) or `ndjson`. A selector that finds nothing is null.
+    Objects and arrays are an error unless `cast_mode` is `coerce`, which keeps their compact JSON.
+    Scalars are stringified with Python's `str`; Polars would render `1.0` as `1`, so it is not used.
+    """
+    mode = options.json_mode or "array"
+    if mode not in ("array", "ndjson"):
+        raise FloeIoError(
+            f"unsupported source.options.json_mode={mode} (allowed: array, ndjson)"
+        )
+    names = [c.source if c.source is not None else c.name for c in columns]
+    if len(set(names)) != len(names):
+        dup = next(n for n in names if names.count(n) > 1)
+        raise FloeIoError(f"duplicate json selector source: {dup}")
+    plans: list[Selector] = []
+    for name in names:
+        try:
+            plans.append(parse_selector(name))
+        except ValueError as err:
+            raise FloeIoError(f"invalid selector {name}: {err}") from err
+
+    def frame(rows: list[list[str | None]]) -> pl.DataFrame:
+        cols = zip(*rows, strict=True) if rows else [() for _ in names]
+        return pl.DataFrame(
+            {n: pl.Series(n, list(c), dtype=pl.String) for n, c in zip(names, cols)}
+        )
+
+    exact = [decimal_type(c.column_type) is not None for c in columns]
+    rows: list[list[str | None]] = []
+    done = 0
+    for where, value in json_records(path, mode):
+        rows.append(
+            [
+                _cell(value, p, n, cast_mode, where, e)
+                for p, n, e in zip(plans, names, exact)
+            ]
+        )
+        if len(rows) == chunk_rows:
+            yield JsonChunk(done, frame(rows))
+            done += len(rows)
+            rows = []
+    if rows or done == 0:
+        yield JsonChunk(done, frame(rows))

floe/read/orc.py

@@ -0,0 +1,42 @@
+"""ORC reader: native column types, one chunk per stripe (Rust: `io/read/orc.rs`)."""
+# pyarrow ships no type stubs, so strict pyright cannot see its types.
+# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false, reportUnknownArgumentType=false
+
+from collections.abc import Collection, Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import polars as pl
+import pyarrow as pa
+import pyarrow.orc as paorc
+
+from floe.errors import FloeIoError
+
+
+@dataclass(frozen=True)
+class OrcChunk:
+    """`offset` is the whole-file index of the chunk's first row."""
+
+    offset: int
+    df: pl.DataFrame
+
+
+def read_orc(path: Path, columns: Collection[str] = ()) -> Iterator[OrcChunk]:
+    """Yield one chunk per stripe; always at least one (maybe empty).
+
+    Columns keep their native types (casting them is #758, D15). `columns` projects to the
+    declared columns that exist in the file; none of them present reads everything, as Rust.
+    """
+    try:
+        orc = paorc.ORCFile(path)
+        picked = [name for name in orc.schema.names if name in columns] or None
+        done = 0
+        for i in range(orc.nstripes):
+            df = pl.DataFrame(orc.read_stripe(i, columns=picked))
+            yield OrcChunk(done, df)
+            done += df.height
+        if done == 0:
+            yield OrcChunk(0, pl.DataFrame(orc.read(columns=picked)))
+    except (pl.exceptions.PolarsError, pa.ArrowException, OSError) as err:
+        raise FloeIoError(f"orc read failed: {err}") from err

floe/read/parquet.py

@@ -0,0 +1,43 @@
+"""Parquet reader: native column types, in bounded chunks (Rust: `io/read/parquet.rs`)."""
+
+from collections.abc import Collection, Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import polars as pl
+
+from floe.errors import FloeIoError
+
+
+@dataclass(frozen=True)
+class ParquetChunk:
+    """`offset` is the whole-file index of the chunk's first row."""
+
+    offset: int
+    df: pl.DataFrame
+
+
+def read_parquet(
+    path: Path, columns: Collection[str] = (), chunk_rows: int = 1_000_000
+) -> Iterator[ParquetChunk]:
+    """Yield the file in chunks of at most `chunk_rows` rows; always at least one (maybe empty).
+
+    Columns keep their native types (casting them is #758, D15). `columns` projects to the
+    declared columns that exist in the file; none of them present reads everything, as Rust.
+    """
+    try:
+        lf = pl.scan_parquet(path)
+        picked = [name for name in lf.collect_schema().names() if name in columns]
+    except (pl.exceptions.PolarsError, OSError) as err:
+        raise FloeIoError(f"failed to scan parquet at {path}: {err}") from err
+    if picked:
+        lf = lf.select(picked)
+    try:
+        done = 0
+        for df in lf.collect_batches(chunk_size=chunk_rows):
+            yield ParquetChunk(done, df)
+            done += df.height
+        if done == 0:
+            yield ParquetChunk(0, lf.head(0).collect())
+    except (pl.exceptions.PolarsError, OSError) as err:
+        raise FloeIoError(f"parquet read failed: {err}") from err

floe/read/xlsx.py

@@ -0,0 +1,75 @@
+"""XLSX reader: every cell as a raw string (Rust: `io/read/xlsx.rs`), via `fastexcel` (calamine)."""
+
+from pathlib import Path
+
+import fastexcel
+import polars as pl
+
+from floe.config.types import SourceOptions
+from floe.errors import FloeIoError
+
+
+def _fail(rule: str, message: str) -> FloeIoError:
+    return FloeIoError(f"{rule}: {message}")
+
+
+def _header(cells: tuple[str | None, ...]) -> list[str]:
+    """Trimmed names; blank cells become `column_N`; repeats become `<name>_dup_<count>` (as Rust)."""
+    counts: dict[str, int] = {}
+    used: set[str] = set()
+    names: list[str] = []
+    for index, cell in enumerate(cells):
+        name = (cell or "").strip() or f"column_{index + 1}"
+        count = counts.get(name, 0) + 1
+        candidate = name if count == 1 else f"{name}_dup_{count}"
+        while candidate in used:
+            count += 1
+            candidate = f"{name}_dup_{count}"
+        counts[name] = count
+        used.add(candidate)
+        names.append(candidate)
+    return names
+
+
+def read_xlsx(path: Path, options: SourceOptions) -> pl.DataFrame:
+    """Read the whole sheet; rows and columns count from the sheet's used range, as in Rust.
+
+    The data rows start at `data_row` (default `header_row + 1`). Numbers, booleans and dates are
+    rendered by calamine (see the tests for the one rendering that differs from Rust: dates).
+    """
+    # ponytail: calamine loads the whole sheet in memory, so there is one frame, not chunks.
+    header_row = 1 if options.header_row is None else options.header_row
+    data_row = header_row + 1 if options.data_row is None else options.data_row
+    try:
+        book = fastexcel.read_excel(path)
+    except fastexcel.FastExcelError as err:
+        raise _fail("xlsx_read_error", f"failed to open xlsx at {path}: {err}") from err
+    names = book.sheet_names
+    if not names:
+        raise _fail("xlsx_read_error", f"no sheets found in {path}")
+    sheet = names[0] if options.sheet is None else options.sheet
+    if sheet not in names:
+        raise _fail(
+            "xlsx_sheet_error",
+            f"sheet {sheet} not found in {path} (available: {', '.join(names)})",
+        )
+    try:
+        # the eager=True overload of fastexcel's stubs returns an untyped pyarrow table
+        loaded = book.load_sheet(  # pyright: ignore[reportUnknownMemberType]
+            sheet, header_row=None, dtypes="string"
+        )
+        raw = loaded.to_polars()
+    except fastexcel.FastExcelError as err:
+        raise _fail(
+            "xlsx_read_error", f"failed to read sheet {sheet} in {path}: {err}"
+        ) from err
+    if raw.width == 0:
+        raise _fail("xlsx_read_error", f"sheet is empty in {path}")
+    header_idx = max(header_row - 1, 0)
+    if header_idx >= raw.height:
+        raise _fail(
+            "xlsx_read_error",
+            f"header_row={header_row} is outside sheet bounds in {path}",
+        )
+    names = _header(raw.row(header_idx))
+    return raw.slice(max(data_row - 1, 0)).rename(dict(zip(raw.columns, names)))

floe/read/xml.py

@@ -0,0 +1,227 @@
+"""XML reader: every selected column as a raw string, streamed in bounded chunks (Rust: `io/read/xml.rs`)."""
+
+import re
+from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from xml.etree import ElementTree as ET
+
+import polars as pl
+
+from floe.config.types import ColumnConfig, SourceOptions
+from floe.errors import FloeConfigError, FloeIoError
+
+_READ_SIZE = 1 << 16
+Selector = list[tuple[bool, str]]  # (is_attribute, name)
+
+
+@dataclass(frozen=True)
+class XmlChunk:
+    """`offset` is the whole-file index of the chunk's first row (add it to the row index)."""
+
+    offset: int
+    df: pl.DataFrame
+
+
+def parse_selector(value: str) -> Selector:
+    """A `.` or `/` separated path of element names, optionally ending in `@attribute`."""
+    trimmed = value.strip()
+    if not trimmed:
+        raise ValueError("selector is empty")
+    parts = [p.strip() for p in re.split(r"[./]", trimmed)]
+    if "" in parts:
+        raise ValueError("selector contains empty token")
+    tokens: Selector = []
+    for index, part in enumerate(parts):
+        if part.startswith("@"):
+            if index + 1 != len(parts):
+                raise ValueError("attribute selector must be final token")
+            attribute = part.lstrip("@").strip()
+            if not attribute:
+                raise ValueError("attribute selector is empty")
+            tokens.append((True, attribute))
+        elif "@" in part:
+            raise ValueError(f"invalid selector token: {part}")
+        else:
+            tokens.append((False, part))
+    return tokens
+
+
+def _matches(element: ET.Element, tag: str, namespace: str | None) -> bool:
+    prefix, _, local = tag.rpartition(":")
+    ns, _, name = (
+        element.tag[1:].partition("}")
+        if element.tag.startswith("{")
+        else ("", "", element.tag)
+    )
+    if name != local:
+        return False
+    if prefix and namespace is None:
+        return True
+    return namespace is None or ns == namespace
+
+
+def _text(element: ET.Element) -> str | None:
+    return "".join(element.itertext()).strip() or None
+
+
+def _evaluate(
+    row: ET.Element, tokens: Selector, namespace: str | None, value_tag: str | None
+) -> str | None:
+    node = row
+    for is_attribute, name in tokens:
+        if is_attribute:
+            return node.get(name)
+        child = next((c for c in node if _matches(c, name, namespace)), None)
+        if child is None:
+            return None
+        node = child
+    if value_tag:
+        node = next(
+            (
+                d
+                for d in node.iter()
+                if d is not node and _matches(d, value_tag, namespace)
+            ),
+            node,
+        )
+    return _text(node)
+
+
+class _Rows:
+    """Parser target: evaluates each finished row element, then clears it (bounded memory)."""
+
+    def __init__(
+        self,
+        evaluate: Callable[[ET.Element], list[str | None]],
+        row_tag: str,
+        namespace: str | None,
+    ) -> None:
+        self.builder = ET.TreeBuilder()
+        self.evaluate, self.row_tag, self.namespace = evaluate, row_tag, namespace
+        self.rows: list[list[str | None]] = []
+        self.depth = 0  # ponytail: row_tag elements nested in a row are content, not rows (Rust yields both)
+
+    def start(self, tag: str, attrib: dict[str, str]) -> ET.Element:
+        self.depth += self.depth > 0
+        element = self.builder.start(tag, attrib)
+        if self.depth == 0 and _matches(element, self.row_tag, self.namespace):
+            self.depth = 1
+        return element
+
+    def end(self, tag: str) -> ET.Element:
+        element = self.builder.end(tag)
+        if self.depth == 1 and _matches(element, self.row_tag, self.namespace):
+            self.rows.append(self.evaluate(element))
+            element.clear()
+            # ponytail: the emptied row stays in its parent; drop it there if huge files need it
+        self.depth -= self.depth > 0
+        return element
+
+    def data(self, data: str) -> None:
+        self.builder.data(data)
+
+    def close(self) -> ET.Element:
+        return self.builder.close()
+
+    def doctype(self, name: str, pubid: str | None, system: str | None) -> None:
+        raise ET.ParseError("DTDs and external entities are not supported")
+
+
+def _options(options: SourceOptions) -> tuple[str, str | None, str | None]:
+    row_tag = (options.row_tag or "").strip()
+    if not row_tag:
+        raise FloeConfigError("source.options.row_tag is required for xml input")
+    return (
+        row_tag,
+        (options.namespace or "").strip() or None,
+        (options.value_tag or "").strip() or None,
+    )
+
+
+def read_xml(
+    path: Path,
+    options: SourceOptions,
+    columns: Sequence[ColumnConfig],
+    chunk_rows: int = 1_000_000,
+) -> Iterator[XmlChunk]:
+    """Yield the rows under `row_tag` in chunks of at most `chunk_rows`, columns named by selector.
+
+    Missing elements and attributes, and empty text, are null. A file without rows is an error.
+    """
+    row_tag, namespace, value_tag = _options(options)
+    names = [c.source if c.source is not None else c.name for c in columns]
+    if len(set(names)) != len(names):
+        dup = next(n for n in names if names.count(n) > 1)
+        raise FloeIoError(f"duplicate xml selector source: {dup}")
+    plans: list[Selector] = []
+    for name in names:
+        try:
+            plans.append(parse_selector(name))
+        except ValueError as err:
+            raise FloeIoError(f"invalid selector {name}: {err}") from err
+
+    target = _Rows(
+        lambda row: [_evaluate(row, p, namespace, value_tag) for p in plans],
+        row_tag,
+        namespace,
+    )
+    parser = ET.XMLParser(target=target)
+    done = 0
+
+    def frame() -> pl.DataFrame:
+        cols = zip(*target.rows, strict=True) if target.rows else [() for _ in names]
+        return pl.DataFrame(
+            {n: pl.Series(n, list(c), dtype=pl.String) for n, c in zip(names, cols)}
+        )
+
+    try:
+        with path.open("rb") as handle:
+            while block := handle.read(_READ_SIZE):
+                parser.feed(block)
+                while len(target.rows) >= chunk_rows:
+                    rest = target.rows[chunk_rows:]
+                    target.rows = target.rows[:chunk_rows]
+                    yield XmlChunk(done, frame())
+                    done += chunk_rows
+                    target.rows = rest
+        parser.close()
+    except ET.ParseError as err:
+        raise FloeIoError(f"xml parse error: {err}") from err
+    except OSError as err:
+        raise FloeIoError(f"failed to read xml at {path}: {err}") from err
+    if target.rows:
+        yield XmlChunk(done, frame())
+    elif done == 0:
+        raise FloeIoError(f"row_tag={row_tag} produced no rows")
+
+
+def _in_namespace(tag: str, namespace: str | None) -> bool:
+    return namespace is None or tag[1:].partition("}")[0] == namespace
+
+
+def xml_columns(path: Path, options: SourceOptions) -> list[str]:
+    """Sorted child element names of the first row: the probe Rust's `read_input_columns` runs."""
+    row_tag, namespace, _ = _options(options)
+    try:
+        row = None  # the first matching element in document order, read once it closes
+        for event, element in ET.iterparse(path, events=("start", "end")):
+            if (
+                row is None
+                and event == "start"
+                and _matches(element, row_tag, namespace)
+            ):
+                row = element
+            elif element is row and event == "end":
+                return sorted(
+                    {
+                        c.tag.rpartition("}")[2]
+                        for c in element
+                        if _in_namespace(c.tag, namespace)
+                    }
+                )
+    except ET.ParseError as err:
+        raise FloeIoError(f"xml parse error: {err}") from err
+    except OSError as err:
+        raise FloeIoError(f"failed to read xml at {path}: {err}") from err
+    raise FloeIoError(f"row_tag={row_tag} not found in xml")