Skip to content

mecfs_bio.build_system.task.dataframe_output

How a task writes a dataframe asset: the output format types, and the single writer that interprets them.

Classes:

Functions:

Attributes:

OutFormat module-attribute

OutFormat = ParquetOutFormat | CSVOutFormat

ParquetCompression module-attribute

ParquetCompression = Literal[
    "snappy", "zstd", "gzip", "brotli", "lz4", "none"
]

CSVOutFormat

Attributes:

sep instance-attribute

sep: str

ParquetOutFormat

Attributes:

write_options class-attribute instance-attribute

write_options: ParquetWriteOptions | None = None

ParquetWriteOptions

How to encode a parquet output, when the defaults are not good enough.

Attributes:

byte_stream_split_floats class-attribute instance-attribute

byte_stream_split_floats: bool = False

compression class-attribute instance-attribute

compression: ParquetCompression = 'zstd'

compression_level class-attribute instance-attribute

compression_level: int | None = None

float_column_names

float_column_names(table: Table) -> list[str]

Names of the floating-point columns, the ones worth byte-stream-splitting.

Source code in mecfs_bio/build_system/task/dataframe_output.py
def float_column_names(table: pyarrow.Table) -> list[str]:
    """Names of the floating-point columns, the ones worth byte-stream-splitting."""
    return [
        field.name for field in table.schema if pyarrow.types.is_floating(field.type)
    ]

get_extension_and_read_spec_from_format

get_extension_and_read_spec_from_format(
    out_format: OutFormat,
) -> tuple[str, DataFrameReadSpec]
Source code in mecfs_bio/build_system/task/dataframe_output.py
def get_extension_and_read_spec_from_format(
    out_format: OutFormat,
) -> tuple[str, DataFrameReadSpec]:
    if isinstance(out_format, CSVOutFormat):
        read_spec = DataFrameReadSpec(DataFrameTextFormat(separator=out_format.sep))
        if out_format.sep == "\t":
            extension = ".tsv"
        elif out_format.sep == ",":
            extension = ".csv"
        else:
            raise ValueError("Unknown sep")
    elif isinstance(out_format, ParquetOutFormat):
        read_spec = DataFrameReadSpec(DataFrameParquetFormat())
        extension = ".parquet"
    else:
        raise ValueError(f"Unknown format {out_format}")
    return extension, read_spec

write_df_according_to_format

write_df_according_to_format(
    df: LazyFrame, out_path: Path, out_format: OutFormat
) -> None

Write a dataframe to out_path in the requested format. The frame's backend is preserved: narwhals dispatches to the underlying library's own writer

Source code in mecfs_bio/build_system/task/dataframe_output.py
def write_df_according_to_format(
    df: narwhals.LazyFrame, out_path: Path, out_format: OutFormat
) -> None:
    """Write a dataframe to out_path in the requested format.
    The frame's backend is preserved: narwhals dispatches to the underlying
    library's own writer
    """
    if isinstance(out_format, CSVOutFormat):
        df.collect().to_pandas().to_csv(out_path, index=False, sep=out_format.sep)
    elif isinstance(out_format, ParquetOutFormat):
        if out_format.write_options is None:
            df.sink_parquet(out_path)
        else:
            options = out_format.write_options
            table = df.collect().to_arrow()
            write_parquet_table(
                table=table,
                out_path=out_path,
                compression=options.compression,
                compression_level=options.compression_level,
                byte_stream_split_columns=(
                    float_column_names(table)
                    if options.byte_stream_split_floats
                    else []
                ),
            )
    else:
        raise ValueError(f"Unknown format {out_format}")

write_parquet_table

write_parquet_table(
    table: Table,
    out_path: Path,
    compression: ParquetCompression,
    compression_level: int | None,
    byte_stream_split_columns: Sequence[str],
) -> None

Write an arrow table to parquet with explicit encoding control.

Dictionary encoding takes precedence over BYTE_STREAM_SPLIT in the parquet writer: a column left dictionary-enabled is written as RLE_DICTIONARY and the requested split is silently dropped, producing a file byte-identical to one written without it. Dictionary encoding is therefore disabled on exactly the split columns and left on for the rest, where it is what makes low-cardinality string columns small.

Pass an empty byte_stream_split_columns to disable the split entirely.

Source code in mecfs_bio/build_system/task/dataframe_output.py
def write_parquet_table(
    table: pyarrow.Table,
    out_path: Path,
    compression: ParquetCompression,
    compression_level: int | None,
    byte_stream_split_columns: Sequence[str],
) -> None:
    """Write an arrow table to parquet with explicit encoding control.

    Dictionary encoding takes precedence over BYTE_STREAM_SPLIT in the parquet
    writer: a column left dictionary-enabled is written as RLE_DICTIONARY and
    the requested split is silently dropped, producing a file byte-identical to
    one written without it. Dictionary encoding is therefore disabled on exactly
    the split columns and left on for the rest, where it is what makes
    low-cardinality string columns small.

    Pass an empty byte_stream_split_columns to disable the split entirely.
    """
    split_columns = list(byte_stream_split_columns)
    missing = set(split_columns) - set(table.schema.names)
    assert not missing, f"byte_stream_split_columns not in frame: {missing}"
    other_columns = [n for n in table.schema.names if n not in set(split_columns)]
    pyarrow.parquet.write_table(
        table,
        out_path,
        compression=compression,
        compression_level=compression_level,
        use_byte_stream_split=split_columns if split_columns else False,
        use_dictionary=other_columns if split_columns else True,
    )