from __future__ import annotations

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

from polars.utils import 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, 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


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, 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 {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


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: str = "utf8",
    low_memory: bool = False,
    rechunk: bool = True,
    use_pyarrow: bool = False,
    storage_options: dict | None = None,
    skip_rows_after_header: int = 0,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    sample_size: int = 1024,
    **kwargs: Any,
) -> 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]``: A null value per column.
        - ``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``.
    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
        Allowed encodings: ``utf8`` or ``utf8-lossy``.
        Lossy means that invalid utf8 values are replaced with ``�``
        characters.
    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 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.

    Returns
    -------
    DataFrame

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

    """

    # Map legacy arguments to current ones and remove them from kwargs.
    has_header = kwargs.pop("has_headers", has_header)
    dtypes = kwargs.pop("dtype", dtypes)
    n_rows = kwargs.pop("stop_after_n_rows", n_rows)

    if columns is None:
        columns = kwargs.pop("projection", None)

    _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)

    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 isinstance(column, str) and 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 encoding == "utf8"
        and not low_memory
        and null_values is None
        and parse_dates
    ):
        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, **storage_options) as data:
            tbl = pa.csv.read_csv(
                data,
                pa.csv.ReadOptions(
                    skip_rows=skip_rows, autogenerate_column_names=not has_header
                ),
                pa.csv.ParseOptions(delimiter=sep),
                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 = cast(DataFrame, 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, **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,
            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,
        )

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


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, type[DataType]] | 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: str = "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,
    **kwargs: Any,
) -> 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]``: A null value per column.
        - ``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
        Allowed encodings: ``utf8`` or ``utf8-lossy``.
        Lossy means that invalid utf8 values are replaced with ``�``
        characters.
    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``.

    Returns
    -------
    LazyFrame

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

    Examples
    --------
    >>> (
    ...     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"]}
    ... )
    >>> df.to_csv("mydf.csv")
    >>> pl.scan_csv(
    ...     "mydf.csv", 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     │
    └─────────┴──────────┘

    """

    # Map legacy arguments to current ones and remove them from kwargs.
    has_header = kwargs.pop("has_headers", has_header)
    dtypes = kwargs.pop("dtype", dtypes)
    n_rows = kwargs.pop("stop_after_n_rows", n_rows)

    _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,
    )


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 | None = None,
    **kwargs: Any,
) -> 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.
    """

    # Map legacy arguments to current ones and remove them from kwargs.
    n_rows = kwargs.pop("stop_after_n_rows", n_rows)

    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,
    )


def scan_parquet(
    file: str | Path,
    n_rows: int | None = None,
    cache: bool = True,
    parallel: str = "auto",
    rechunk: bool = True,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    storage_options: dict | None = None,
    low_memory: bool = False,
    **kwargs: Any,
) -> 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
        Any of { '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.
    """

    # Map legacy arguments to current ones and remove them from kwargs.
    n_rows = kwargs.pop("stop_after_n_rows", n_rows)

    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 read_avro(
    file: str | Path | BytesIO | BinaryIO,
    columns: list[int] | list[str] | None = None,
    n_rows: int | None = None,
    **kwargs: Any,
) -> 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)
    if columns is None:
        columns = kwargs.pop("projection", None)

    return DataFrame._read_avro(file, n_rows=n_rows, columns=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 | None = None,
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    rechunk: bool = True,
    **kwargs: Any,
) -> 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
        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.
    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
    """

    # Map legacy arguments to current ones and remove them from kwargs.
    n_rows = kwargs.pop("stop_after_n_rows", n_rows)

    if columns is None:
        columns = kwargs.pop("projection", None)

    if use_pyarrow:
        if row_count_name is not None:
            raise ValueError(
                "``row_count_name`` cannot be used with ``use_pyarrow=True``."
            )
        if n_rows:
            raise ValueError("``n_rows`` cannot be used with ``use_pyarrow=True``.")

    storage_options = storage_options or {}
    with _prepare_file_arg(file, **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)
            return DataFrame._from_arrow(tbl, rechunk=rechunk)

        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,
        )


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 | None = None,
    parallel: str = "auto",
    row_count_name: str | None = None,
    row_count_offset: int = 0,
    low_memory: bool = False,
    **kwargs: Any,
) -> 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
        Any of { '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: bool
        Reduce memory pressure at the expense of performance.
    **kwargs
        kwargs for `pyarrow.parquet.read_table <https://arrow.apache.org/docs/python/generated/pyarrow.parquet.read_table.html>`_.

    Returns
    -------
    DataFrame
    """

    # Map legacy arguments to current ones and remove them from kwargs.
    n_rows = kwargs.pop("stop_after_n_rows", n_rows)

    if columns is None:
        columns = kwargs.pop("projection", None)

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

    storage_options = storage_options or {}
    with _prepare_file_arg(source, **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(  # type: ignore[return-value]
                pa.parquet.read_table(
                    source_prep,
                    memory_map=memory_map,
                    columns=columns,
                    **kwargs,
                )
            )

        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 cast(DataFrame, 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 | None = None,
    read_csv_options: dict | None = None,
) -> DataFrame:
    """
    Read Excel (XLSX) sheet into a DataFrame by converting an Excel
    sheet with ``xlsx2csv.Xlsx2csv().convert()`` to CSV and parsing
    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

    Alternative
    -----------
    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
    except ImportError:
        raise ImportError(
            "xlsx2csv is not installed. Please run `pip install xlsx2csv`."
        )

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

    if not xlsx2csv_options:
        xlsx2csv_options = {}

    if not read_csv_options:
        read_csv_options = {}

    # Override xlsx2csv eprint function so in case an error occurs
    # it raises an exception instead of writing to stderr.
    def _eprint(*args: Any, **kwargs: Any) -> None:
        raise xlsx2csv.XlsxException(format(*args))

    xlsx2csv.eprint = _eprint

    # Create Xlsx2csv instance.
    xlsx2csv_instance = xlsx2csv.Xlsx2csv(file, **xlsx2csv_options)

    if sheet_name:
        sheet_id = xlsx2csv_instance.getSheetIdByName(sheet_name)

        if not sheet_id:
            raise xlsx2csv.XlsxException(f"Sheet '{sheet_name}' not found.")

    csv_buffer = StringIO()

    # Convert sheet from XSLX document to CSV.
    xlsx2csv_instance.convert(outfile=csv_buffer, sheetid=sheet_id)

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

    # Parse CSV output.
    return read_csv(csv_buffer, **read_csv_options)


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)
