polars_lance

  1from collections.abc import Iterator
  2from pathlib import Path
  3from typing import Literal
  4
  5import polars as pl
  6from polars.io.plugins import register_io_source
  7
  8from polars_lance import _polars_lance
  9
 10__all__ = ["scan_lance", "write_lance"]
 11
 12
 13def scan_lance(
 14    source: str | Path,
 15    *,
 16    storage_options: dict[str, str] | None = None,
 17) -> pl.LazyFrame:
 18    """
 19    Lazily read from a Lance dataset.
 20
 21    Parameters
 22    ----------
 23    source
 24        Path or URI to a Lance dataset.
 25    storage_options
 26        Cloud storage configuration to read remote datasets on AWS S3,
 27        Azure Blob Storage, or Google Cloud Storage. Supported keys:
 28        - [aws](https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html)
 29        - [azure](https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html)
 30        - [gcp](https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html)
 31
 32    Returns
 33    -------
 34    LazyFrame
 35
 36    Examples
 37    --------
 38    Scan a local Lance dataset.
 39
 40    >>> scan_lance("example.lance")
 41
 42    Scan a remote Lance dataset on AWS S3.
 43
 44    >>> source = "s3://bucket/example.lance"
 45    >>> storage_options = {
 46    ...     "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
 47    ...     "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
 48    ...     "aws_region": "us-east-1",
 49    ... }
 50    >>> scan_lance(source, storage_options=storage_options)
 51    """
 52    source_str = str(source)
 53
 54    def io_source(
 55        with_columns: list[str] | None,
 56        predicate: pl.Expr | None,
 57        n_rows: int | None,
 58        batch_size: int | None,
 59    ) -> Iterator[pl.DataFrame]:
 60        lance_scanner = _polars_lance.LanceScanner(
 61            uri=source_str,
 62            with_columns=with_columns,
 63            predicate=predicate,
 64            n_rows=n_rows,
 65            batch_size=batch_size,
 66            storage_options=storage_options,
 67        )
 68
 69        while (df := lance_scanner.next()) is not None:
 70            yield df
 71
 72    return register_io_source(
 73        io_source=io_source,
 74        schema=_polars_lance.LanceScanner.schema_for_uri(
 75            uri=source_str,
 76            storage_options=storage_options,
 77        ),
 78    )
 79
 80
 81def write_lance(
 82    df: pl.DataFrame,
 83    target: str | Path,
 84    *,
 85    mode: Literal["error", "append", "overwrite"] = "error",
 86    storage_options: dict[str, str] | None = None,
 87    max_rows_per_file: int | None = None,
 88    max_bytes_per_file: int | None = None,
 89) -> None:
 90    """
 91    Write dataframe to a Lance dataset.
 92
 93    Parameters
 94    ----------
 95    df
 96        Dataframe to write.
 97    target
 98        Path or URI to the Lance dataset.
 99    mode : {'error', 'append', 'overwrite'}
100        How to behave if the target dataset already exists.
101        - `error`: raise an error
102        - `append`: append to the existing dataset
103        - `overwrite`: replace the existing dataset
104    storage_options
105        Cloud storage configuration to write remote datasets on AWS S3,
106        Azure Blob Storage, or Google Cloud Storage. Supported keys:
107        - [aws](https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html)
108        - [azure](https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html)
109        - [gcp](https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html)
110    max_rows_per_file
111        Maximum number of rows to write before starting a new data file.
112    max_bytes_per_file
113        Maximum number of bytes to write before starting a new data file. This is a soft
114        limit that is checked after a group is written, meaning that the actual file
115        size may exceed this limit.
116
117    Examples
118    --------
119    Write a local Lance dataset.
120
121    >>> df = pl.DataFrame({"id": [1, 2], "val": ["a", "b"]})
122    >>> write_lance(df, "example.lance")
123
124    Write a remote Lance dataset on AWS S3.
125
126    >>> df = pl.DataFrame({"id": [1, 2], "val": ["a", "b"]})
127    >>> target = "s3://bucket/example.lance"
128    >>> storage_options = {
129    ...     "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
130    ...     "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
131    ...     "aws_region": "us-east-1",
132    ... }
133    >>> write_lance(df, target, storage_options=storage_options)
134    """
135    _polars_lance.write_lance(
136        df,
137        target=str(target),
138        mode=mode,
139        storage_options=storage_options,
140        max_rows_per_file=max_rows_per_file,
141        max_bytes_per_file=max_bytes_per_file,
142    )