"""Functions for reading and writing data."""
from __future__ import annotations

from io import BytesIO, IOBase, StringIO
from pathlib import Path
from typing import TYPE_CHECKING, BinaryIO, Callable, Mapping, TextIO

from polars.utils import deprecated_alias, format_path, handle_projection_columns

try:
    import pyarrow as pa

    # do not remove these
    import pyarrow.csv
    import pyarrow.feather
    import pyarrow.parquet

    _PYARROW_AVAILABLE = True
except ImportError:
    _PYARROW_AVAILABLE = False

from polars.convert import from_arrow
from polars.datatypes import DataType, PolarsDataType, Utf8
from polars.internals import DataFrame, LazyFrame, _scan_ds
from polars.internals.io import _prepare_file_arg

try:
    import connectorx as cx

    _WITH_CX = True
except ImportError:
    _WITH_CX = False

if TYPE_CHECKING:
    from polars.internals.type_aliases import CsvEncoding, ParallelStrategy


def _check_arg_is_1byte(
    arg_name: str, arg: str | None, can_be_empty: bool = False
) -> None:
    if isinstance(arg, str):
        arg_byte_length = len(arg.encode("utf-8"))
        if can_be_empty:
            if arg_byte_length > 1:
                raise ValueError(
                    f'{arg_name}="{arg}" should be a single byte character or empty,'
                    f" but is {arg_byte_length} bytes long."
                )
        elif arg_byte_length != 1:
            raise ValueError(
                f'{arg_name}="{arg}" should be a single byte character, but is'
                f" {arg_byte_length} bytes long."
            )


def _update_columns(df: DataFrame, new_columns: list[str]) -> DataFrame:
    if df.width > len(new_columns):
        cols = df.columns
        for i, name in enumerate(new_columns):
            cols[i] = name
        new_columns = cols
    df.columns = new_columns
    return df


@deprecated_alias(
    has_headers="has_header",
    dtype="dtypes",
    stop_after_n_rows="n_rows",
    projection="columns",
)
def read_csv(
    file: str | TextIO | BytesIO | Path | BinaryIO | bytes,
    has_header: bool = True,
    columns: list[int] | list[str] | None = None,
    new_columns: list[str] | None = None,
    sep: str = ",",
    comment_char: str | None = None,
    quote_char: str | None = r'"',
    skip_rows: int = 0,
    dtypes: Mapping[str, type[DataType]] | list[type[DataType]] | None = None,
    null_values: str | list[str] | dict[str, str] | None = None,
    ignore_errors: bool = False,
    parse_dates: bool = False,
    n_threads: int | None = None,
    infer_schema_length: int | None = 100,
    batch_size: int = 8192,
    n_rows: int | None = None,
    encoding: CsvEncoding | str = "utf8",
    low_memory: bool = False,
    rechunk: bool = True,
    use_pyarrow: bool = False,
    storage_options: dict[str, object] | None = None,
    skip_rows_after_header: int = 0,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    sample_size: int = 1024,
    eol_char: str = "\n",
) -> DataFrame:
    """
    Read a CSV file into a DataFrame.

    Parameters
    ----------
    file
        Path to a file or a file-like object.
        By file-like object, we refer to objects with a ``read()``
        method, such as a file handler (e.g. via builtin ``open``
        function) or ``StringIO`` or ``BytesIO``.
        If ``fsspec`` is installed, it will be used to open remote
        files.
    has_header
        Indicate if the first row of dataset is a header or not.
        If set to False, column names will be autogenerated in the
        following format: ``column_x``, with ``x`` being an
        enumeration over every column in the dataset starting at 1.
    columns
        Columns to select. Accepts a list of column indices (starting
        at zero) or a list of column names.
    new_columns
        Rename columns right after parsing the CSV file. If the given
        list is shorter than the width of the DataFrame the remaining
        columns will have their original name.
    sep
        Single byte character to use as delimiter in the file.
    comment_char
        Single byte character that indicates the start of a comment line,
        for instance ``#``.
    quote_char
        Single byte character used for csv quoting, default = ``"``.
        Set to None to turn off special handling and escaping of quotes.
    skip_rows
        Start reading after ``skip_rows`` lines.
    dtypes
        Overwrite dtypes during inference.
    null_values
        Values to interpret as null values. You can provide a:

        - ``str``: All values equal to this string will be null.
        - ``List[str]``: All values equal to any string in this list will be null.
        - ``Dict[str, str]``: A dictionary that maps column name to a
          null value string.
    ignore_errors
        Try to keep reading lines if some lines yield errors.
        First try ``infer_schema_length=0`` to read all columns as
        ``pl.Utf8`` to check which values might cause an issue.
    parse_dates
        Try to automatically parse dates. If this does not succeed,
        the column remains of data type ``pl.Utf8``.
        If ``use_pyarrow=True``, dates will always be parsed.
    n_threads
        Number of threads to use in csv parsing.
        Defaults to the number of physical cpu's of your system.
    infer_schema_length
        Maximum number of lines to read to infer schema.
        If set to 0, all columns will be read as ``pl.Utf8``.
        If set to ``None``, a full table scan will be done (slow).
    batch_size
        Number of lines to read into the buffer at once.
        Modify this to change performance.
    n_rows
        Stop reading from CSV file after reading ``n_rows``.
        During multi-threaded parsing, an upper bound of ``n_rows``
        rows cannot be guaranteed.
    encoding : {'utf8', 'utf8-lossy', ...}
        Lossy means that invalid utf8 values are replaced with ``�``
        characters. When using other encodings than ``utf8`` or
        ``utf8-lossy``, the input is first decoded im memory with
        python. Defaults to ``utf8``.
    low_memory
        Reduce memory usage at expense of performance.
    rechunk
        Make sure that all columns are contiguous in memory by
        aggregating the chunks into a single array.
    use_pyarrow
        Try to use pyarrow's native CSV parser. This will always
        parse dates, even if ``parse_dates=False``.
        This is not always possible. The set of arguments given to
        this function determines if it is possible to use pyarrow's
        native parser. Note that pyarrow and polars may have a
        different strategy regarding type inference.
    storage_options
        Extra options that make sense for ``fsspec.open()`` or a
        particular storage connection.
        e.g. host, port, username, password, etc.
    skip_rows_after_header
        Skip this number of rows when the header is parsed.
    row_count_name
        If not None, this will insert a row count column with the given name into
        the DataFrame.
    row_count_offset
        Offset to start the row_count column (only used if the name is set).
    sample_size
        Set the sample size. This is used to sample statistics to estimate the
        allocation needed.
    eol_char
        Single byte end of line character.

    Returns
    -------
    DataFrame

    See Also
    --------
    scan_csv : Lazily read from a CSV file or multiple files via glob patterns.

    """
    _check_arg_is_1byte("sep", sep, False)
    _check_arg_is_1byte("comment_char", comment_char, False)
    _check_arg_is_1byte("quote_char", quote_char, True)
    _check_arg_is_1byte("eol_char", eol_char, False)

    projection, columns = handle_projection_columns(columns)

    if isinstance(file, bytes) and len(file) == 0:
        raise ValueError("Empty bytes data provided.")

    storage_options = storage_options or {}

    if columns and not has_header:
        for column in columns:
            if not column.startswith("column_"):
                raise ValueError(
                    'Specified column names do not start with "column_", '
                    "but autogenerated header names were requested."
                )

    if use_pyarrow and not _PYARROW_AVAILABLE:
        raise ImportError(
            "'pyarrow' is required when using 'read_csv(..., use_pyarrow=True)'."
        )

    if (
        use_pyarrow
        and dtypes is None
        and n_rows is None
        and n_threads is None
        and not low_memory
        and null_values is None
    ):
        include_columns = None

        if columns:
            if not has_header:
                # Convert 'column_1', 'column_2', ... column names to 'f0', 'f1', ...
                # column names for pyarrow, if CSV file does not contain a header.
                include_columns = [f"f{int(column[7:]) - 1}" for column in columns]
            else:
                include_columns = columns

        if not columns and projection:
            # Convert column indices from projection to 'f0', 'f1', ... column names
            # for pyarrow.
            include_columns = [f"f{column_idx}" for column_idx in projection]

        with _prepare_file_arg(
            file, encoding=None, use_pyarrow=True, **storage_options
        ) as data:
            tbl = pa.csv.read_csv(
                data,
                pa.csv.ReadOptions(
                    skip_rows=skip_rows,
                    autogenerate_column_names=not has_header,
                    encoding=encoding,
                ),
                pa.csv.ParseOptions(
                    delimiter=sep,
                    quote_char=quote_char if quote_char else False,
                    double_quote=quote_char is not None and quote_char == '"',
                ),
                pa.csv.ConvertOptions(
                    column_types=None,
                    include_columns=include_columns,
                    include_missing_columns=ignore_errors,
                ),
            )

        if not has_header:
            # Rename 'f0', 'f1', ... columns names autogenated by pyarrow
            # to 'column_1', 'column_2', ...
            tbl = tbl.rename_columns(
                [f"column_{int(column[1:]) + 1}" for column in tbl.column_names]
            )

        df = from_arrow(tbl, rechunk)
        if new_columns:
            return _update_columns(df, new_columns)
        return df

    if projection and dtypes and isinstance(dtypes, list):
        if len(projection) < len(dtypes):
            raise ValueError(
                "More dtypes overrides are specified than there are selected columns."
            )

        # Fix list of dtypes when used together with projection as polars CSV reader
        # wants a list of dtypes for the x first columns before it does the projection.
        dtypes_list: list[type[DataType]] = [Utf8] * (max(projection) + 1)

        for idx, column_idx in enumerate(projection):
            if idx < len(dtypes):
                dtypes_list[column_idx] = dtypes[idx]

        dtypes = dtypes_list

    if columns and dtypes and isinstance(dtypes, list):
        if len(columns) < len(dtypes):
            raise ValueError(
                "More dtypes overrides are specified than there are selected columns."
            )

        # Map list of dtypes when used together with selected columns as a dtypes dict
        # so the dtypes are applied to the correct column instead of the first x
        # columns.
        dtypes = {column: dtype for column, dtype in zip(columns, dtypes)}

    if new_columns and dtypes and isinstance(dtypes, dict):
        current_columns = None

        # As new column names are not available yet while parsing the CSV file, rename
        # column names in dtypes to old names (if possible) so they can be used during
        # CSV parsing.
        if columns:
            if len(columns) < len(new_columns):
                raise ValueError(
                    "More new column names are specified than there are selected"
                    " columns."
                )

            # Get column names of requested columns.
            current_columns = columns[0 : len(new_columns)]
        elif not has_header:
            # When there are no header, column names are autogenerated (and known).

            if projection:
                if columns and len(columns) < len(new_columns):
                    raise ValueError(
                        "More new column names are specified than there are selected"
                        " columns."
                    )
                # Convert column indices from projection to 'column_1', 'column_2', ...
                # column names.
                current_columns = [
                    f"column_{column_idx + 1}" for column_idx in projection
                ]
            else:
                # Generate autogenerated 'column_1', 'column_2', ... column names for
                # new column names.
                current_columns = [
                    f"column_{column_idx}"
                    for column_idx in range(1, len(new_columns) + 1)
                ]
        else:
            # When a header is present, column names are not known yet.

            if len(dtypes) <= len(new_columns):
                # If dtypes dictionary contains less or same amount of values than new
                # column names a list of dtypes can be created if all listed column
                # names in dtypes dictionary appear in the first consecutive new column
                # names.
                dtype_list = [
                    dtypes[new_column_name]
                    for new_column_name in new_columns[0 : len(dtypes)]
                    if new_column_name in dtypes
                ]

                if len(dtype_list) == len(dtypes):
                    dtypes = dtype_list

        if current_columns and isinstance(dtypes, dict):
            new_to_current = {
                new_column: current_column
                for new_column, current_column in zip(new_columns, current_columns)
            }
            # Change new column names to current column names in dtype.
            dtypes = {
                new_to_current.get(column_name, column_name): column_dtype
                for column_name, column_dtype in dtypes.items()
            }

    with _prepare_file_arg(
        file, encoding=encoding, use_pyarrow=False, **storage_options
    ) as data:
        df = DataFrame._read_csv(
            file=data,
            has_header=has_header,
            columns=columns if columns else projection,
            sep=sep,
            comment_char=comment_char,
            quote_char=quote_char,
            skip_rows=skip_rows,
            dtypes=dtypes,
            null_values=null_values,
            ignore_errors=ignore_errors,
            parse_dates=parse_dates,
            n_threads=n_threads,
            infer_schema_length=infer_schema_length,
            batch_size=batch_size,
            n_rows=n_rows,
            encoding=encoding if encoding == "utf8-lossy" else "utf8",
            low_memory=low_memory,
            rechunk=rechunk,
            skip_rows_after_header=skip_rows_after_header,
            row_count_name=row_count_name,
            row_count_offset=row_count_offset,
            sample_size=sample_size,
            eol_char=eol_char,
        )

    if new_columns:
        return _update_columns(df, new_columns)
    return df


@deprecated_alias(has_headers="has_header", dtype="dtypes", stop_after_n_rows="n_rows")
def scan_csv(
    file: str | Path,
    has_header: bool = True,
    sep: str = ",",
    comment_char: str | None = None,
    quote_char: str | None = r'"',
    skip_rows: int = 0,
    dtypes: dict[str, PolarsDataType] | None = None,
    null_values: str | list[str] | dict[str, str] | None = None,
    ignore_errors: bool = False,
    cache: bool = True,
    with_column_names: Callable[[list[str]], list[str]] | None = None,
    infer_schema_length: int | None = 100,
    n_rows: int | None = None,
    encoding: CsvEncoding = "utf8",
    low_memory: bool = False,
    rechunk: bool = True,
    skip_rows_after_header: int = 0,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    parse_dates: bool = False,
    eol_char: str = "\n",
) -> LazyFrame:
    """
    Lazily read from a CSV file or multiple files via glob patterns.

    This allows the query optimizer to push down predicates and
    projections to the scan level, thereby potentially reducing
    memory overhead.

    Parameters
    ----------
    file
        Path to a file.
    has_header
        Indicate if the first row of dataset is a header or not.
        If set to False, column names will be autogenerated in the
        following format: ``column_x``, with ``x`` being an
        enumeration over every column in the dataset starting at 1.
    sep
        Single byte character to use as delimiter in the file.
    comment_char
        Single byte character that indicates the start of a comment line,
        for instance ``#``.
    quote_char
        Single byte character used for csv quoting, default = ``"``.
        Set to None to turn off special handling and escaping of quotes.
    skip_rows
        Start reading after ``skip_rows`` lines. The header will be parsed at this
        offset.
    dtypes
        Overwrite dtypes during inference.
    null_values
        Values to interpret as null values. You can provide a:

        - ``str``: All values equal to this string will be null.
        - ``List[str]``: All values equal to any string in this list will be null.
        - ``Dict[str, str]``: A dictionary that maps column name to a
          null value string.
    ignore_errors
        Try to keep reading lines if some lines yield errors.
        First try ``infer_schema_length=0`` to read all columns as
        ``pl.Utf8`` to check which values might cause an issue.
    cache
        Cache the result after reading.
    with_column_names
        Apply a function over the column names.
        This can be used to update a schema just in time, thus before
        scanning.
    infer_schema_length
        Maximum number of lines to read to infer schema.
        If set to 0, all columns will be read as ``pl.Utf8``.
        If set to ``None``, a full table scan will be done (slow).
    n_rows
        Stop reading from CSV file after reading ``n_rows``.
    encoding : {'utf8', 'utf8-lossy'}
        Lossy means that invalid utf8 values are replaced with ``�``
        characters. Defaults to "utf8".
    low_memory
        Reduce memory usage in expense of performance.
    rechunk
        Reallocate to contiguous memory when all chunks/ files are parsed.
    skip_rows_after_header
        Skip this number of rows when the header is parsed.
    row_count_name
        If not None, this will insert a row count column with the given name into
        the DataFrame.
    row_count_offset
        Offset to start the row_count column (only used if the name is set).
    parse_dates
        Try to automatically parse dates. If this does not succeed,
        the column remains of data type ``pl.Utf8``.
    eol_char
        Single byte end of line character

    Returns
    -------
    LazyFrame

    See Also
    --------
    read_csv : Read a CSV file into a DataFrame.

    Examples
    --------
    >>> import pathlib
    >>>
    >>> (
    ...     pl.scan_csv("my_long_file.csv")  # lazy, doesn't do a thing
    ...     .select(
    ...         ["a", "c"]
    ...     )  # select only 2 columns (other columns will not be read)
    ...     .filter(
    ...         pl.col("a") > 10
    ...     )  # the filter is pushed down the scan, so less data is read into memory
    ...     .fetch(100)  # pushed a limit of 100 rows to the scan level
    ... )  # doctest: +SKIP

    We can use ``with_column_names`` to modify the header before scanning:

    >>> df = pl.DataFrame(
    ...     {"BrEeZaH": [1, 2, 3, 4], "LaNgUaGe": ["is", "terrible", "to", "read"]}
    ... )
    >>> path: pathlib.Path = dirpath / "mydf.csv"
    >>> df.write_csv(path)
    >>> pl.scan_csv(
    ...     path, with_column_names=lambda cols: [col.lower() for col in cols]
    ... ).fetch()
    shape: (4, 2)
    ┌─────────┬──────────┐
    │ breezah ┆ language │
    │ ---     ┆ ---      │
    │ i64     ┆ str      │
    ╞═════════╪══════════╡
    │ 1       ┆ is       │
    ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
    │ 2       ┆ terrible │
    ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
    │ 3       ┆ to       │
    ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
    │ 4       ┆ read     │
    └─────────┴──────────┘

    """
    _check_arg_is_1byte("sep", sep, False)
    _check_arg_is_1byte("comment_char", comment_char, False)
    _check_arg_is_1byte("quote_char", quote_char, True)

    if isinstance(file, (str, Path)):
        file = format_path(file)

    return LazyFrame.scan_csv(
        file=file,
        has_header=has_header,
        sep=sep,
        comment_char=comment_char,
        quote_char=quote_char,
        skip_rows=skip_rows,
        dtypes=dtypes,
        null_values=null_values,
        ignore_errors=ignore_errors,
        cache=cache,
        with_column_names=with_column_names,
        infer_schema_length=infer_schema_length,
        n_rows=n_rows,
        low_memory=low_memory,
        rechunk=rechunk,
        skip_rows_after_header=skip_rows_after_header,
        encoding=encoding,
        row_count_name=row_count_name,
        row_count_offset=row_count_offset,
        parse_dates=parse_dates,
        eol_char=eol_char,
    )


@deprecated_alias(stop_after_n_rows="n_rows")
def scan_ipc(
    file: str | Path,
    n_rows: int | None = None,
    cache: bool = True,
    rechunk: bool = True,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    storage_options: dict[str, object] | None = None,
    memory_map: bool = True,
) -> LazyFrame:
    """
    Lazily read from an Arrow IPC (Feather v2) file or multiple files via glob patterns.

    This allows the query optimizer to push down predicates and projections to the scan
    level, thereby potentially reducing memory overhead.

    Parameters
    ----------
    file
        Path to a IPC file.
    n_rows
        Stop reading from IPC file after reading ``n_rows``.
    cache
        Cache the result after reading.
    rechunk
        Reallocate to contiguous memory when all chunks/ files are parsed.
    row_count_name
        If not None, this will insert a row count column with give name into the
        DataFrame
    row_count_offset
        Offset to start the row_count column (only use if the name is set)
    storage_options
        Extra options that make sense for ``fsspec.open()`` or a
        particular storage connection.
        e.g. host, port, username, password, etc.
    memory_map
        Try to memory map the file. This can greatly improve performance on repeated
        queries as the OS may cache pages.
        Only uncompressed IPC files can be memory mapped.

    """
    return LazyFrame.scan_ipc(
        file=file,
        n_rows=n_rows,
        cache=cache,
        rechunk=rechunk,
        row_count_name=row_count_name,
        row_count_offset=row_count_offset,
        storage_options=storage_options,
        memory_map=memory_map,
    )


@deprecated_alias(stop_after_n_rows="n_rows")
def scan_parquet(
    file: str | Path,
    n_rows: int | None = None,
    cache: bool = True,
    parallel: ParallelStrategy = "auto",
    rechunk: bool = True,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    storage_options: dict[str, object] | None = None,
    low_memory: bool = False,
) -> LazyFrame:
    """
    Lazily read from a parquet file or multiple files via glob patterns.

    This allows the query optimizer to push down predicates and projections to the scan
    level, thereby potentially reducing memory overhead.

    Parameters
    ----------
    file
        Path to a file.
    n_rows
        Stop reading from parquet file after reading ``n_rows``.
    cache
        Cache the result after reading.
    parallel : {'auto', 'columns', 'row_groups', 'none'}
        This determines the direction of parallelism. 'auto' will try to determine the
        optimal direction.
    rechunk
        In case of reading multiple files via a glob pattern rechunk the final DataFrame
        into contiguous memory chunks.
    row_count_name
        If not None, this will insert a row count column with give name into the
        DataFrame
    row_count_offset
        Offset to start the row_count column (only use if the name is set)
    storage_options
        Extra options that make sense for ``fsspec.open()`` or a
        particular storage connection.
        e.g. host, port, username, password, etc.
    low_memory: bool
        Reduce memory pressure at the expense of performance.

    """
    if isinstance(file, (str, Path)):
        file = format_path(file)

    return LazyFrame.scan_parquet(
        file=file,
        n_rows=n_rows,
        cache=cache,
        parallel=parallel,
        rechunk=rechunk,
        row_count_name=row_count_name,
        row_count_offset=row_count_offset,
        storage_options=storage_options,
        low_memory=low_memory,
    )


def scan_ndjson(
    file: str | Path,
    infer_schema_length: int | None = 100,
    batch_size: int | None = 1024,
    n_rows: int | None = None,
    low_memory: bool = False,
    rechunk: bool = True,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
) -> LazyFrame:
    """
    Lazily read from a newline delimited JSON file.

    This allows the query optimizer to push down predicates and projections to the scan
    level, thereby potentially reducing memory overhead.

    Parameters
    ----------
    file
        Path to a file.
    infer_schema_length
        Infer the schema length from the first ``infer_schema_length`` rows.
    batch_size
        Number of rows to read in each batch.
    n_rows
        Stop reading from JSON file after reading ``n_rows``.
    low_memory
        Reduce memory pressure at the expense of performance.
    rechunk
        Reallocate to contiguous memory when all chunks/ files are parsed.
    row_count_name
        If not None, this will insert a row count column with give name into the
        DataFrame
    row_count_offset
        Offset to start the row_count column (only use if the name is set)

    """
    if isinstance(file, (str, Path)):
        file = format_path(file)

    return LazyFrame.scan_ndjson(
        file=file,
        infer_schema_length=infer_schema_length,
        batch_size=batch_size,
        n_rows=n_rows,
        low_memory=low_memory,
        rechunk=rechunk,
        row_count_name=row_count_name,
        row_count_offset=row_count_offset,
    )


@deprecated_alias(projection="columns")
def read_avro(
    file: str | Path | BytesIO | BinaryIO,
    columns: list[int] | list[str] | None = None,
    n_rows: int | None = None,
) -> DataFrame:
    """
    Read into a DataFrame from Apache Avro format.

    Parameters
    ----------
    file
        Path to a file or a file-like object.
    columns
        Columns to select. Accepts a list of column indices (starting at zero) or a list
        of column names.
    n_rows
        Stop reading from Apache Avro file after reading ``n_rows``.

    Returns
    -------
    DataFrame

    """
    if isinstance(file, (str, Path)):
        file = format_path(file)

    return DataFrame._read_avro(file, n_rows=n_rows, columns=columns)


@deprecated_alias(stop_after_n_rows="n_rows", projection="columns")
def read_ipc(
    file: str | BinaryIO | BytesIO | Path | bytes,
    columns: list[int] | list[str] | None = None,
    n_rows: int | None = None,
    use_pyarrow: bool = False,
    memory_map: bool = True,
    storage_options: dict[str, object] | None = None,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    rechunk: bool = True,
) -> DataFrame:
    """
    Read into a DataFrame from Arrow IPC (Feather v2) file.

    Parameters
    ----------
    file
        Path to a file or a file-like object.
        If ``fsspec`` is installed, it will be used to open remote files.
    columns
        Columns to select. Accepts a list of column indices (starting at zero) or a list
        of column names.
    n_rows
        Stop reading from IPC file after reading ``n_rows``.
        Only valid when `use_pyarrow=False`.
    use_pyarrow
        Use pyarrow or the native rust reader.
    memory_map
        Try to memory map the file. This can greatly improve performance on repeated
        queries as the OS may cache pages.
        Only uncompressed IPC files can be memory mapped.
    storage_options
        Extra options that make sense for ``fsspec.open()`` or a particular storage
        connection, e.g. host, port, username, password, etc.
    row_count_name
        If not None, this will insert a row count column with give name into the
        DataFrame
    row_count_offset
        Offset to start the row_count column (only use if the name is set)
    rechunk
        Make sure that all data is contiguous.

    Returns
    -------
    DataFrame

    """
    if use_pyarrow:
        if n_rows and not memory_map:
            raise ValueError(
                "``n_rows`` cannot be used with ``use_pyarrow=True` "
                "and memory_map=False`."
            )

    storage_options = storage_options or {}
    with _prepare_file_arg(file, use_pyarrow=use_pyarrow, **storage_options) as data:
        if use_pyarrow:
            if not _PYARROW_AVAILABLE:
                raise ImportError(
                    "'pyarrow' is required when using"
                    " 'read_ipc(..., use_pyarrow=True)'."
                )

            tbl = pa.feather.read_table(data, memory_map=memory_map, columns=columns)
            df = DataFrame._from_arrow(tbl, rechunk=rechunk)
            if row_count_name is not None:
                df = df.with_row_count(row_count_name, row_count_offset)
            if n_rows is not None:
                df = df.slice(0, n_rows)
            return df

        return DataFrame._read_ipc(
            data,
            columns=columns,
            n_rows=n_rows,
            row_count_name=row_count_name,
            row_count_offset=row_count_offset,
            rechunk=rechunk,
            memory_map=memory_map,
        )


@deprecated_alias(stop_after_n_rows="n_rows", projection="columns")
def read_parquet(
    source: str | Path | BinaryIO | BytesIO | bytes,
    columns: list[int] | list[str] | None = None,
    n_rows: int | None = None,
    use_pyarrow: bool = False,
    memory_map: bool = True,
    storage_options: dict[str, object] | None = None,
    parallel: ParallelStrategy = "auto",
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    low_memory: bool = False,
    pyarrow_options: dict[str, object] | None = None,
) -> DataFrame:
    """
    Read into a DataFrame from a parquet file.

    Parameters
    ----------
    source
        Path to a file, or a file-like object. If the path is a directory, that
        directory will be used as partition aware scan.
        If ``fsspec`` is installed, it will be used to open remote files.
    columns
        Columns to select. Accepts a list of column indices (starting at zero) or a list
        of column names.
    n_rows
        Stop reading from parquet file after reading ``n_rows``.
        Only valid when `use_pyarrow=False`.
    use_pyarrow
        Use pyarrow instead of the rust native parquet reader. The pyarrow reader is
        more stable.
    memory_map
        Memory map underlying file. This will likely increase performance.
        Only used when ``use_pyarrow=True``.
    storage_options
        Extra options that make sense for ``fsspec.open()`` or a particular storage
        connection, e.g. host, port, username, password, etc.
    parallel : {'auto', 'columns', 'row_groups', 'none'}
        This determines the direction of parallelism. 'auto' will try to determine the
        optimal direction.
    row_count_name
        If not None, this will insert a row count column with give name into the
        DataFrame.
    row_count_offset
        Offset to start the row_count column (only use if the name is set).
    low_memory
        Reduce memory pressure at the expense of performance.
    pyarrow_options
        Keyword arguments for `pyarrow.parquet.read_table
        <https://arrow.apache.org/docs/python/generated/pyarrow.parquet.read_table.html>`_.

    Returns
    -------
    DataFrame

    """  # noqa: E501
    if use_pyarrow and n_rows:
        raise ValueError("``n_rows`` cannot be used with ``use_pyarrow=True``.")

    storage_options = storage_options or {}
    pyarrow_options = pyarrow_options or {}

    with _prepare_file_arg(
        source, use_pyarrow=use_pyarrow, **storage_options
    ) as source_prep:
        if use_pyarrow:
            if not _PYARROW_AVAILABLE:
                raise ImportError(
                    "'pyarrow' is required when using"
                    " 'read_parquet(..., use_pyarrow=True)'."
                )

            return from_arrow(
                pa.parquet.read_table(
                    source_prep,
                    memory_map=memory_map,
                    columns=columns,
                    **pyarrow_options,
                )
            )

        return DataFrame._read_parquet(
            source_prep,
            columns=columns,
            n_rows=n_rows,
            parallel=parallel,
            row_count_name=row_count_name,
            row_count_offset=row_count_offset,
            low_memory=low_memory,
        )


def read_json(source: str | IOBase, json_lines: bool = False) -> DataFrame:
    """
    Read into a DataFrame from JSON format.

    Parameters
    ----------
    source
        Path to a file or a file-like object.
    json_lines
        Toggle between "JSON" and "NDJSON" format

    """
    return DataFrame._read_json(source, json_lines)


def read_sql(
    sql: list[str] | str,
    connection_uri: str,
    partition_on: str | None = None,
    partition_range: tuple[int, int] | None = None,
    partition_num: int | None = None,
    protocol: str | None = None,
) -> DataFrame:
    """
    Read a SQL query into a DataFrame.

    .. note::
        Make sure to install connectorx>=0.2.2. Read the documentation
        `here <https://sfu-db.github.io/connector-x/intro.html>`_.

    Reading a SQL query from the following data sources are supported:

        * Postgres
        * Mysql
        * Sqlite
        * Redshift (through postgres protocol)
        * Clickhouse (through mysql protocol)

    If a database source is not supported, an alternative solution is to first use
    pandas to load the SQL query, then converting the result into a polars DataFrame:

    >>> import pandas as pd
    >>> df = pl.from_pandas(pd.read_sql(sql, engine))  # doctest: +SKIP

    Parameters
    ----------
    sql
        Raw SQL query / queries.
    connection_uri
        Connectorx connection uri, for example

        * "postgresql://username:password@server:port/database"
    partition_on
        The column on which to partition the result.
    partition_range
        The value range of the partition column.
    partition_num
        How many partitions to generate.
    protocol
        Backend-specific transfer protocol directive; see connectorx documentation for
        details.

    Examples
    --------
    Read a DataFrame from a SQL query using a single thread:

    >>> uri = "postgresql://username:password@server:port/database"
    >>> query = "SELECT * FROM lineitem"
    >>> pl.read_sql(query, uri)  # doctest: +SKIP

    Read a DataFrame in parallel using 10 threads by automatically partitioning the
    provided SQL on the partition column:

    >>> uri = "postgresql://username:password@server:port/database"
    >>> query = "SELECT * FROM lineitem"
    >>> pl.read_sql(
    ...     query, uri, partition_on="partition_col", partition_num=10
    ... )  # doctest: +SKIP

    Read a DataFrame in parallel using 2 threads by explicitly providing two SQL
    queries:

    >>> uri = "postgresql://username:password@server:port/database"
    >>> queries = [
    ...     "SELECT * FROM lineitem WHERE partition_col <= 10",
    ...     "SELECT * FROM lineitem WHERE partition_col > 10",
    ... ]
    >>> pl.read_sql(queries, uri)  # doctest: +SKIP

    """
    if _WITH_CX:
        tbl = cx.read_sql(
            conn=connection_uri,
            query=sql,
            return_type="arrow2",
            partition_on=partition_on,
            partition_range=partition_range,
            partition_num=partition_num,
            protocol=protocol,
        )
        return from_arrow(tbl)
    else:
        raise ImportError(
            "connectorx is not installed. Please run `pip install connectorx>=0.2.2`."
        )


def read_excel(
    file: str | BytesIO | Path | BinaryIO | bytes,
    sheet_id: int | None = 1,
    sheet_name: str | None = None,
    xlsx2csv_options: dict[str, object] | None = None,
    read_csv_options: dict[str, object] | None = None,
) -> DataFrame:
    """
    Read Excel (XLSX) sheet into a DataFrame.

    Converts an Excel sheet with ``xlsx2csv.Xlsx2csv().convert()`` to CSV and parses the
    CSV output with :func:`read_csv`.

    Parameters
    ----------
    file
        Path to a file or a file-like object.
        By file-like object, we refer to objects with a ``read()`` method, such as a
        file handler (e.g. via builtin ``open`` function) or ``BytesIO``.
    sheet_id
        Sheet number to convert (0 for all sheets).
    sheet_name
        Sheet name to convert.
    xlsx2csv_options
        Extra options passed to ``xlsx2csv.Xlsx2csv()``.
        e.g.: ``{"skip_empty_lines": True}``
    read_csv_options
        Extra options passed to :func:`read_csv` for parsing the CSV file returned by
        ``xlsx2csv.Xlsx2csv().convert()``
        e.g.: ``{"has_header": False, "new_columns": ["a", "b", "c"],
        infer_schema_length=None}``

    Returns
    -------
    DataFrame

    Examples
    --------
    Read "My Datasheet" sheet from Excel sheet file to a DataFrame.

    >>> excel_file = "test.xlsx"
    >>> sheet_name = "My Datasheet"
    >>> pl.read_excel(
    ...     file=excel_file,
    ...     sheet_name=sheet_name,
    ... )  # doctest: +SKIP

    Read sheet 3 from Excel sheet file to a DataFrame while skipping empty lines in the
    sheet. As sheet 3 does not have header row, pass the needed settings to
    :func:`read_csv`.

    >>> excel_file = "test.xlsx"
    >>> pl.read_excel(
    ...     file=excel_file,
    ...     sheet_id=3,
    ...     xlsx2csv_options={"skip_empty_lines": True},
    ...     read_csv_options={"has_header": False, "new_columns": ["a", "b", "c"]},
    ... )  # doctest: +SKIP

    If the correct datatypes can't be determined by polars, look at :func:`read_csv`
    documentation to see which options you can pass to fix this issue. For example
    ``"infer_schema_length": None`` can be used to read the whole data twice, once to
    infer the correct output types and once to actually convert the input to the correct
    types. With `"infer_schema_length": 1000``, only the first 1000 lines are read
    twice.

    >>> excel_file = "test.xlsx"
    >>> pl.read_excel(
    ...     file=excel_file,
    ...     read_csv_options={"infer_schema_length": None},
    ... )  # doctest: +SKIP

    If :func:`read_excel` does not work or you need to read other types of spreadsheet
    files, you can try pandas ``pd.read_excel()``
    (supports `xls`, `xlsx`, `xlsm`, `xlsb`, `odf`, `ods` and `odt`).

    >>> excel_file = "test.xlsx"
    >>> pl.from_pandas(pd.read_excel(excel_file))  # doctest: +SKIP

    """
    try:
        import xlsx2csv  # type: ignore[import]
    except ImportError:
        raise ImportError(
            "xlsx2csv is not installed. Please run `pip install xlsx2csv`."
        ) from None

    if isinstance(file, (str, Path)):
        file = format_path(file)

    if not xlsx2csv_options:
        xlsx2csv_options = {}

    if not read_csv_options:
        read_csv_options = {}

    csv_buffer = StringIO()

    # Convert sheet from XSLX document to CSV.
    xlsx2csv.Xlsx2csv(file, **xlsx2csv_options).convert(
        outfile=csv_buffer, sheetid=sheet_id, sheetname=sheet_name
    )

    # Rewind buffer to start.
    csv_buffer.seek(0)

    # Parse CSV output.
    return read_csv(csv_buffer, **read_csv_options)  # type: ignore[arg-type]


def scan_ds(ds: pa.dataset.dataset) -> LazyFrame:
    """
    Scan a pyarrow dataset.

    This can be useful to connect to cloud or partitioned datasets.

    .. warning::
        This API is experimental and may change without it being considered a breaking
        change.

    Parameters
    ----------
    ds
        Pyarrow dataset to scan.

    Examples
    --------
    >>> import pyarrow.dataset as ds
    >>> dset = ds.dataset("s3://my-partitioned-folder/", format="ipc")  # doctest: +SKIP
    >>> out = (
    ...     pl.scan_ds(dset)
    ...     .filter("bools")
    ...     .select(["bools", "floats", "date"])
    ...     .collect()
    ... )  # doctest: +SKIP
    shape: (1, 3)
    ┌───────┬────────┬────────────┐
    │ bools ┆ floats ┆ date       │
    │ ---   ┆ ---    ┆ ---        │
    │ bool  ┆ f64    ┆ date       │
    ╞═══════╪════════╪════════════╡
    │ true  ┆ 2.0    ┆ 1970-05-04 │
    └───────┴────────┴────────────┘

    """
    return _scan_ds(ds)
