🗂️ Extract, Combine, and Project Multiple Tables¶

from tabularix import group, regex, non_empty

# Pattern for the vertical header column (one-dimensional cell sequence).
header_pattern = group(regex(r"^(Date|Fiscal Year)$").repeat(2, 2))

# Pattern for the vertical data column (one-dimensional cell sequence).
data_pattern = group(non_empty().repeat(2, 2))

from typing import cast
import polars as pl
from tabularix import Sheet, extract_table_with_header_and_data

def extract_metadata(sheet: Sheet) -> pl.DataFrame:
    # Extract the metadata as a horizontal table.
    table = extract_table_with_header_and_data(
        sheet,
        header_pattern,
        data_pattern,
        main_direction="LR",
        inner_direction="TB",
        clean_names=False,  # Keep original capitalization ('Date', 'Fiscal Year').
    )

    # Cast to pl.DataFrame to satisfy static type checkers.
    df = cast(pl.DataFrame, pl.from_arrow(table.to_arrow()))

    return df

from tabularix import any, empty, grid, group, regex, value

# Locate the territory title (one-dimensional group pattern).
territory_pattern = group(regex(r"^(North|South|East|West)$"))
territory_matcher = territory_pattern.to_matcher(direction="LR")

# Locate the two-row merged header (two-dimensional grid pattern).
header_row1 = group(
    value("Product"),
    group(
        regex(r"\d{4}"),
        empty()  # Matches the merged cell next to the year.
    ).zero_or_more()
)
header_row2 = group(
    empty(),  # Matches the merged cell below "Product".
    group(
        value("Expected"),
        value("Actual")
    ).zero_or_more()
)
header_pattern = grid(header_row1, header_row2)
header_matcher = header_pattern.to_matcher(outer_direction="TB", inner_direction="LR")

# Locate the footer row.
footer_pattern = group(
    value("Total"),
    any().zero_or_more()
)
footer_matcher = footer_pattern.to_matcher(direction="LR")

def extract_territory_tables(sheet: Sheet) -> pl.DataFrame:
    dfs = []
    search_row = 0

    while search_row < sheet.shape[0]:
        # 1. Match the territory name (e.g. "North").
        territory_range = sheet.search_range(territory_matcher, start_row=search_row)
        if territory_range is None:
            break

        territory = str(sheet.get_cell_value(territory_range.start_row, territory_range.start_col))

        # 2. Match the 2-row header range below the title.
        header_range = sheet.search_range(header_matcher, start_row=territory_range.end_row + 1)
        if header_range is None:
            raise ValueError(f"Header not found for {territory}")

        # 3. Match the footer relative to the header.
        footer_range = sheet.search_range_relative(footer_matcher, below=header_range)
        if footer_range is None:
            raise ValueError(f"Footer not found for {territory}")

        # 4. Extract data rows between header and footer.
        data_range = sheet.get_range_between(header_range, footer_range)
        table = sheet.extract_table(
            data_range,
            header=header_range,
            clean_names=True,
            flatten_header=False,
        )

        # 5. Convert to Polars and insert territory as the first column.
        df = cast(pl.DataFrame, pl.from_arrow(table.to_arrow()))
        df = df.select([pl.lit(territory).alias("territory"), pl.all()])

        # 6. Unpack the product struct.
        df = df.with_columns(pl.col("product").struct.field("product"))

        # 7. Unpivot the year columns to rows (all columns not in index are unpivoted).
        df = df.unpivot(
            on=None,
            index=["territory", "product"],
            variable_name="year",
            value_name="metrics",
        )

        # 8. Unnest the metrics struct (containing expected and actual fields).
        df = df.unnest("metrics")

        dfs.append(df)

        # Move the cursor below this table's footer.
        search_row = footer_range.end_row + 1

    return pl.concat(dfs)

# Broadcast metadata to all territory rows.
projected_df = territories_df.join(metadata_df, how="cross")