DataFrame#

def __arrow_c_schema__(self) -> Any:
    return self.schema().to_pyarrow_schema().__arrow_c_schema__()

def __arrow_c_stream__(self, requested_schema: Any = None) -> Any:
    """Export as an Arrow C stream (PyCapsule).

    This triggers materialization of the DataFrame.
    Enables ``pa.table(daft_df)`` and other Arrow PyCapsule consumers.
    """
    self.collect()
    assert self._result is not None
    mp = self._result._get_merged_micropartition(self.schema())
    return mp._micropartition.__arrow_c_stream__(requested_schema)

def __contains__(self, col_name: str) -> bool:
    """Returns whether the column exists in the dataframe.

    Args:
        col_name (str): column name

    Returns:
        bool: whether the column exists in the dataframe.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
        >>> "x" in df
        True

    """
    return col_name in self.column_names

def __getitem__(self, item: int | str | slice | Iterable[str | int]) -> Union[Expression, "DataFrame"]:
    """Gets a column from the DataFrame as an Expression (``df["mycol"]``).

    Args:
        item (Union[int, str, slice, Iterable[Union[str, int]]]): The column to get. Can be an integer index, a string column name, a slice for multiple columns, or an iterable of column names or indices.

    Returns:
        Union[Expression, DataFrame]: If a single column is requested, returns an Expression representing that column.
        If multiple columns are requested (via a slice or iterable), returns a new DataFrame containing those columns.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})
        >>> df["a"]  # Get a single column
        col(a)
        >>> df["b"]  # Get another single column
        col(b)
        >>> df[0]  # Get the first column by index
        col(a)
        >>> df[1:3]  # Get a slice of columns
        ╭───────┬───────╮
        │ b     ┆ c     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╰───────┴───────╯
        <BLANKLINE>
        (No data to display: Dataframe not materialized, use .collect() to materialize)
        >>> df[["a", "c"]]  # Get multiple columns by name
        ╭───────┬───────╮
        │ a     ┆ c     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╰───────┴───────╯
        <BLANKLINE>
        (No data to display: Dataframe not materialized, use .collect() to materialize)
        >>> df[["a", 1]]  # Get multiple columns by name and index
        ╭───────┬───────╮
        │ a     ┆ b     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╰───────┴───────╯
        <BLANKLINE>
        (No data to display: Dataframe not materialized, use .collect() to materialize)
        >>> df[0:2]  # Get a slice of columns by index
        ╭───────┬───────╮
        │ a     ┆ b     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╰───────┴───────╯
        <BLANKLINE>
        (No data to display: Dataframe not materialized, use .collect() to materialize)
        >>> df[["a", "b", 2]]  # Get a mix of column names and indices
        ╭───────┬───────┬───────╮
        │ a     ┆ b     ┆ c     │
        │ ---   ┆ ---   ┆ ---   │
        │ Int64 ┆ Int64 ┆ Int64 │
        ╰───────┴───────┴───────╯
        <BLANKLINE>
        (No data to display: Dataframe not materialized, use .collect() to materialize)

    """
    result: Expression | None

    if isinstance(item, int):
        schema = self._builder.schema()
        if item < -len(schema) or item >= len(schema):
            raise ValueError(f"{item} out of bounds for {schema}")
        result = ExpressionsProjection.from_schema(schema)[item]
        assert result is not None
        return result
    elif isinstance(item, str):
        schema = self._builder.schema()
        if item not in schema.column_names() and item != "*":
            raise ValueError(f"{item} does not exist in schema {schema}")

        return col(item)
    elif isinstance(item, Iterable):
        schema = self._builder.schema()

        columns = []
        for it in item:
            if isinstance(it, str):
                result = col(schema[it].name)
                columns.append(result)
            elif isinstance(it, int):
                if it < -len(schema) or it >= len(schema):
                    raise ValueError(f"{it} out of bounds for {schema}")
                field = list(self._builder.schema())[it]
                columns.append(col(field.name))
            else:
                raise ValueError(f"unknown indexing type: {type(it)}")
        return self.select(*columns)
    elif isinstance(item, slice):
        schema = self._builder.schema()
        columns_exprs: ExpressionsProjection = ExpressionsProjection.from_schema(schema)
        selected_columns = columns_exprs[item]
        return self.select(*[typing.cast("ColumnInputType", c) for c in selected_columns])
    else:
        raise ValueError(f"unknown indexing type: {type(item)}")

@DataframePublicAPI
def __iter__(self) -> Iterator[dict[str, Any]]:
    """Alias of `self.iter_rows()` with default arguments for convenient access of data.

    Returns:
        Iterator[dict[str, Any]]: An iterator over the rows of the DataFrame, where each row is a dictionary
        mapping column names to values.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"foo": [1, 2, 3], "bar": ["a", "b", "c"]})
        >>> for row in df:
        ...     print(row)
        {'foo': 1, 'bar': 'a'}
        {'foo': 2, 'bar': 'b'}
        {'foo': 3, 'bar': 'c'}

    Tip:
        See also [`df.iter_rows()`][daft.DataFrame.iter_rows]: iterator over rows with more options
    """
    return self.iter_rows(results_buffer_size=None)

def __len__(self) -> int:
    """Returns the count of rows when dataframe is materialized.

    If dataframe is not materialized yet, raises a runtime error.

    Returns:
        int: count of rows.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"x": [1, 2, 3], "y": [4, 5, 6]})
        >>> df = df.collect()
        >>> len(df)
        3

    """
    if self._result is not None:
        return len(self._result)

    message = (
        "Cannot call len() on an unmaterialized dataframe:"
        " either materialize your dataframe with df.collect() first before calling len(),"
        " or use `df.count_rows()` instead which will calculate the total number of rows."
    )
    raise RuntimeError(message)

@DataframePublicAPI
def agg(self, *to_agg: Expression | Iterable[Expression]) -> "DataFrame":
    """Perform aggregations on this DataFrame.

    Allows for mixed aggregations for multiple columns and will return a single row that aggregated the entire DataFrame.

    Args:
        *to_agg (Expression): aggregation expressions

    Returns:
        DataFrame: DataFrame with aggregated results

    Examples:
        >>> import daft
        >>> from daft import col
        >>> df = daft.from_pydict(
        ...     {"student_id": [1, 2, 3, 4], "test1": [0.5, 0.4, 0.6, 0.7], "test2": [0.9, 0.8, 0.7, 1.0]}
        ... )
        >>> agg_df = df.agg(
        ...     df["test1"].mean(),
        ...     df["test2"].mean(),
        ...     ((df["test1"] + df["test2"]) / 2).min().alias("total_min"),
        ...     ((df["test1"] + df["test2"]) / 2).max().alias("total_max"),
        ... )
        >>> agg_df.show()
        ╭─────────┬────────────────────┬────────────────────┬───────────╮
        │ test1   ┆ test2              ┆ total_min          ┆ total_max │
        │ ---     ┆ ---                ┆ ---                ┆ ---       │
        │ Float64 ┆ Float64            ┆ Float64            ┆ Float64   │
        ╞═════════╪════════════════════╪════════════════════╪═══════════╡
        │ 0.55    ┆ 0.8500000000000001 ┆ 0.6000000000000001 ┆ 0.85      │
        ╰─────────┴────────────────────┴────────────────────┴───────────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)

    """
    to_agg_list = (
        list(to_agg[0])
        if (len(to_agg) == 1 and not isinstance(to_agg[0], Expression))
        else list(typing.cast("tuple[Expression]", to_agg))
    )

    for expr in to_agg_list:
        if not isinstance(expr, Expression):
            raise ValueError(f"DataFrame.agg() only accepts expression type, received: {type(expr)}")

    return self._agg(to_agg_list, group_by=None)

@DataframePublicAPI
def agg_concat(self, *cols: ColumnInputType, delimiter: str | None = None) -> "DataFrame":
    """Performs a global concatenation agg on the DataFrame.

    Args:
        *cols (Union[str, Expression]): columns that are lists or strings to concatenate
        delimiter: Optional delimiter to insert between concatenated string values. Only supported for string
            columns.

    Returns:
        DataFrame: Globally aggregated list or string. Should be a single row.

    Examples:
        >>> import daft
        >>> from daft import col
        >>> df = daft.from_pydict({"col_a": [[1, 2], [3, 4]]})
        >>> df = df.agg_concat("col_a")
        >>> df.show()
        ╭──────────────╮
        │ col_a        │
        │ ---          │
        │ List[Int64]  │
        ╞══════════════╡
        │ [1, 2, 3, 4] │
        ╰──────────────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)
    """
    return self._apply_agg_fn(lambda expr: Expression.string_agg(expr, delimiter=delimiter), cols)

@DataframePublicAPI
def agg_list(self, *cols: ColumnInputType) -> "DataFrame":
    """Performs a global list agg on the DataFrame.

    Args:
        *cols (Union[str, Expression]): columns to form into a list
    Returns:
        DataFrame: Globally aggregated list. Should be a single row.

    Examples:
        >>> import daft
        >>> from daft import col
        >>> df = daft.from_pydict({"col_a": [1, 2, 3]})
        >>> df = df.agg_list("col_a")
        >>> df.show()
        ╭─────────────╮
        │ col_a       │
        │ ---         │
        │ List[Int64] │
        ╞═════════════╡
        │ [1, 2, 3]   │
        ╰─────────────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)
    """
    return self._apply_agg_fn(Expression.list_agg, cols)

@DataframePublicAPI
def agg_set(self, *cols: ColumnInputType) -> "DataFrame":
    """Performs a global set agg on the DataFrame (ignoring nulls).

    Args:
        *cols (Union[str, Expression]): columns to form into a set

    Returns:
        DataFrame: Globally aggregated set. Should be a single row.

    Examples:
        >>> import daft
        >>> from daft import col
        >>> df = daft.from_pydict({"col_a": [1, 2, 2, 3]})
        >>> df = df.agg_set("col_a")
        >>> df.show()
        ╭─────────────╮
        │ col_a       │
        │ ---         │
        │ List[Int64] │
        ╞═════════════╡
        │ [1, 2, 3]   │
        ╰─────────────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)
    """
    return self._apply_agg_fn(Expression.list_agg_distinct, cols)

@DataframePublicAPI
def any_value(self, *cols: ColumnInputType) -> "DataFrame":
    """Returns an arbitrary value on this DataFrame.

    Values for each column are not guaranteed to be from the same row.

    Args:
        *cols (Union[str, Expression]): columns to get an arbitrary value from
    Returns:
        DataFrame: DataFrame with any values.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"col_a": [1, 2, 3]})
        >>> df = df.any_value("col_a")
        >>> df.show()
        ╭───────╮
        │ col_a │
        │ ---   │
        │ Int64 │
        ╞═══════╡
        │ 1     │
        ╰───────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)
    """
    return self._apply_agg_fn(Expression.any_value, cols)

@DataframePublicAPI
def collect(self, num_preview_rows: int | None = 8) -> "DataFrame":
    """Executes the entire DataFrame and materializes the results.

    Args:
        num_preview_rows: Number of rows to preview. Defaults to 8.

    Returns:
        DataFrame: DataFrame with materialized results.

    Note:
        This call is **blocking** and will execute the DataFrame when called

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"x": [1, 2, 3], "y": [4, 5, 6]})
        >>> df = df.collect()
        >>> df.show()
        ╭───────┬───────╮
        │ x     ┆ y     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╞═══════╪═══════╡
        │ 1     ┆ 4     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 2     ┆ 5     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 3     ┆ 6     │
        ╰───────┴───────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)
    """
    self._materialize_results()
    assert self._result is not None
    dataframe_len = len(self._result)
    if num_preview_rows is not None:
        self._num_preview_rows = num_preview_rows
    else:
        self._num_preview_rows = dataframe_len
    return self

@DataframePublicAPI
def concat(self, other: "DataFrame") -> "DataFrame":
    """Concatenates two DataFrames together in a "vertical" concatenation.

    The resulting DataFrame has number of rows equal to the sum of the number of rows of the input DataFrames.

    Args:
        other (DataFrame): other DataFrame to concatenate

    Returns:
        DataFrame: DataFrame with rows from `self` on top and rows from `other` at the bottom.

    Note:
        DataFrames being concatenated **must have exactly the same schema**. You may wish to use the
        [df.select()][daft.DataFrame.select] and [expr.cast()][daft.expressions.Expression.cast] methods
        to ensure schema compatibility before concatenation.

    Examples:
        >>> import daft
        >>> df1 = daft.from_pydict({"a": [1, 2], "b": [3, 4]})
        >>> df2 = daft.from_pydict({"a": [5, 6], "b": [7, 8]})
        >>> concatenated_df = df1.concat(df2)
        >>> concatenated_df.show()
        ╭───────┬───────╮
        │ a     ┆ b     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╞═══════╪═══════╡
        │ 1     ┆ 3     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 2     ┆ 4     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 5     ┆ 7     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 6     ┆ 8     │
        ╰───────┴───────╯
        <BLANKLINE>
        (Showing first 4 of 4 rows)
    """
    if self.schema() != other.schema():
        raise ValueError(
            f"DataFrames must have exactly the same schema for concatenation!\nExpected:\n{self.schema()}\n\nReceived:\n{other.schema()}"
        )
    builder = self._builder.concat(other._builder)
    return DataFrame(builder)

@DataframePublicAPI
def count(self, *cols: ColumnInputType | int) -> "DataFrame":
    """Performs a global count on the DataFrame.

    Args:
        *cols (Union[str, Expression, int]): columns to count
    Returns:
        DataFrame: Globally aggregated count. Should be a single row.


    Examples:
        If no columns are specified (i.e. in the case you call `df.count()`), or only the literal string "*",
        this functions very similarly to a COUNT(*) operation in SQL and will return a new dataframe with a
        single column with the name "count".

        >>> import daft
        >>> from daft import col
        >>> df = daft.from_pydict({"foo": [1, None, None], "bar": [None, 2, 2], "baz": [3, 4, 5]})
        >>> df.count().show()  # equivalent to df.count("*").show()
        ╭────────╮
        │ count  │
        │ ---    │
        │ UInt64 │
        ╞════════╡
        │ 3      │
        ╰────────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)

        However, specifying some column names would instead change the behavior to count all non-null values,
        similar to a SQL command for `SELECT COUNT(foo), COUNT(bar) FROM df`. Also, using `df.count(col("*"))`
        will expand out into count() for each column.

        >>> df.count("foo", "bar").show()
        ╭────────┬────────╮
        │ foo    ┆ bar    │
        │ ---    ┆ ---    │
        │ UInt64 ┆ UInt64 │
        ╞════════╪════════╡
        │ 1      ┆ 2      │
        ╰────────┴────────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)

        >>> df.count(df["*"]).show()
        ╭────────┬────────┬────────╮
        │ foo    ┆ bar    ┆ baz    │
        │ ---    ┆ ---    ┆ ---    │
        │ UInt64 ┆ UInt64 ┆ UInt64 │
        ╞════════╪════════╪════════╡
        │ 1      ┆ 2      ┆ 3      │
        ╰────────┴────────┴────────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)

    """
    # Special case: treat this as a COUNT(*) operation which is likely what most people would expect
    # If user passes in "*", also do this behavior (by default it would count each column individually)
    if (
        len(cols) == 0
        or (len(cols) == 1 and isinstance(cols[0], str) and cols[0] == "*")
        or (len(cols) == 1 and isinstance(cols[0], int))
    ):
        builder = self._builder.count()
        return DataFrame(builder)

    if any(isinstance(c, str) and c == "*" for c in cols):
        # we do not support hybrid count-all and count-nonnull
        raise ValueError("Cannot call count() with both * and column names")

    # Otherwise, perform a column-wise count on the specified columns
    return self._apply_agg_fn(Expression.count, typing.cast("tuple[ColumnInputType, ...]", cols))

@DataframePublicAPI
def count_distinct(self, *cols: ColumnInputType) -> "DataFrame":
    """Performs a global count of distinct values on the DataFrame.

    Args:
        *cols (Union[str, Expression]): columns to count distinct values
    Returns:
        DataFrame: Globally aggregated count of distinct values. Should be a single row.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"col_a": [1, 2, 2, 3, 3, 3]})
        >>> df = df.count_distinct("col_a")
        >>> df.show()
        ╭────────╮
        │ col_a  │
        │ ---    │
        │ UInt64 │
        ╞════════╡
        │ 3      │
        ╰────────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)
    """
    return self._apply_agg_fn(Expression.count_distinct, cols)

@DataframePublicAPI
def count_rows(self) -> int:
    """Executes the Dataframe to count the number of rows.

    Returns:
        int: count of the number of rows in this DataFrame.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
        >>> df.count_rows()
        3

    Note:
        This will execute the DataFrame and return the number of rows in it.

    """
    if self._result is not None:
        return len(self._result)
    builder = self._builder.count()
    count_df = DataFrame(builder)
    # Expects builder to produce a single-partition, single-row DataFrame containing
    # a "count" column, where the lone value represents the row count for the DataFrame.
    return count_df.to_pydict()["count"][0]

@DataframePublicAPI
def describe(self) -> "DataFrame":
    """Returns the Schema of the DataFrame, which provides information about each column, as a new DataFrame.

    Returns:
        DataFrame: A dataframe where each row is a column name and its corresponding type.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"a": [1, 2, 3], "b": ["x", "y", "z"]})
        >>> df.describe().show()
        ╭─────────────┬────────╮
        │ column_name ┆ type   │
        │ ---         ┆ ---    │
        │ String      ┆ String │
        ╞═════════════╪════════╡
        │ a           ┆ Int64  │
        ├╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
        │ b           ┆ String │
        ╰─────────────┴────────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)
    """
    builder = self.__builder.describe()
    return DataFrame(builder)

@DataframePublicAPI
def distinct(self, *on: ColumnInputType) -> "DataFrame":
    """Computes distinct rows, dropping duplicates.

    Optionally, specify a subset of columns to perform distinct on.

    Args:
        *on (Union[str, Expression]): columns to perform distinct on. Defaults to all columns.

    Returns:
        DataFrame: DataFrame that has only distinct rows.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"x": [1, 2, 2], "y": [4, 5, 5], "z": [7, 8, 8]})
        >>> distinct_df = df.distinct()
        >>> distinct_df = distinct_df.sort("x")
        >>> distinct_df.show()
        ╭───────┬───────┬───────╮
        │ x     ┆ y     ┆ z     │
        │ ---   ┆ ---   ┆ ---   │
        │ Int64 ┆ Int64 ┆ Int64 │
        ╞═══════╪═══════╪═══════╡
        │ 1     ┆ 4     ┆ 7     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 2     ┆ 5     ┆ 8     │
        ╰───────┴───────┴───────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)
        >>> # Pass a subset of columns to perform distinct on
        >>> # Note that output for z is non-deterministic. Both 8 and 9 are possible.
        >>> df = daft.from_pydict({"x": [1, 2, 2], "y": [4, 5, 5], "z": [7, 8, 9]})
        >>> df.distinct("x", daft.col("y")).sort("x").show()
        ╭───────┬───────┬───────╮
        │ x     ┆ y     ┆ z     │
        │ ---   ┆ ---   ┆ ---   │
        │ Int64 ┆ Int64 ┆ Int64 │
        ╞═══════╪═══════╪═══════╡
        │ 1     ┆ 4     ┆ 7     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 2     ┆ 5     ┆ 8     │
        ╰───────┴───────┴───────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)
    """
    builder = self._builder.distinct(column_inputs_to_expressions(on))
    return DataFrame(builder)

@DataframePublicAPI
def drop_duplicates(self, *subset: ColumnInputType) -> "DataFrame":
    """Computes distinct rows, dropping duplicates.

    Alias for [DataFrame.distinct][daft.DataFrame.distinct].

    Args:
        *subset (Union[str, Expression]): columns to perform distinct on. Defaults to all columns.

    Returns:
        DataFrame: DataFrame that has only distinct rows.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"x": [1, 2, 2], "y": [4, 5, 5], "z": [7, 8, 8]})
        >>> distinct_df = df.drop_duplicates()
        >>> distinct_df = distinct_df.sort("x")
        >>> distinct_df.show()
        ╭───────┬───────┬───────╮
        │ x     ┆ y     ┆ z     │
        │ ---   ┆ ---   ┆ ---   │
        │ Int64 ┆ Int64 ┆ Int64 │
        ╞═══════╪═══════╪═══════╡
        │ 1     ┆ 4     ┆ 7     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 2     ┆ 5     ┆ 8     │
        ╰───────┴───────┴───────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)
    """
    return self.distinct(*subset)

@DataframePublicAPI
def drop_nan(self, *cols: ColumnInputType) -> "DataFrame":
    """Drops rows that contains NaNs. If cols is None it will drop rows with any NaN value.

    If column names are supplied, it will drop only those rows that contains NaNs in one of these columns.

    Args:
        *cols (str): column names by which rows containing nans/NULLs should be filtered

    Returns:
        DataFrame: DataFrame without NaNs in specified/all columns

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"a": [1.0, 2.2, 3.5, float("nan")]})
        >>> df.drop_nan().collect()  # drops rows where any column contains NaN values
        ╭─────────╮
        │ a       │
        │ ---     │
        │ Float64 │
        ╞═════════╡
        │ 1       │
        ├╌╌╌╌╌╌╌╌╌┤
        │ 2.2     │
        ├╌╌╌╌╌╌╌╌╌┤
        │ 3.5     │
        ╰─────────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)

        >>> import daft
        >>> df = daft.from_pydict({"a": [1.6, 2.5, 3.3, float("nan")]})
        >>> df.drop_nan("a").collect()  # drops rows where column `a` contains NaN values
        ╭─────────╮
        │ a       │
        │ ---     │
        │ Float64 │
        ╞═════════╡
        │ 1.6     │
        ├╌╌╌╌╌╌╌╌╌┤
        │ 2.5     │
        ├╌╌╌╌╌╌╌╌╌┤
        │ 3.3     │
        ╰─────────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)

    """
    if len(cols) == 0:
        columns = column_inputs_to_expressions(self.column_names)
    else:
        columns = column_inputs_to_expressions(cols)
    float_columns = [
        column
        for column in columns
        if (
            column._to_field(self.schema()).dtype == DataType.float32()
            or column._to_field(self.schema()).dtype == DataType.float64()
        )
    ]

    # avoid superfluous .where with empty iterable when nothing to filter.
    if not float_columns:
        return self

    from daft.functions import is_nan, when

    return self.where(
        ~reduce(
            lambda x, y: when(x.is_null(), lit(False)).otherwise(x) | when(y.is_null(), lit(False)).otherwise(y),
            (is_nan(x) for x in float_columns),
        )
    )

@DataframePublicAPI
def drop_null(self, *cols: ColumnInputType) -> "DataFrame":
    """Drops rows that contains NaNs or NULLs. If cols is None it will drop rows with any NULL value.

    If column names are supplied, it will drop only those rows that contains NULLs in one of these columns.

    Args:
        *cols (str): column names by which rows containing nans should be filtered

    Returns:
        DataFrame: DataFrame without missing values in specified/all columns

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"a": [1.6, 2.5, None, float("NaN")]})
        >>> df.drop_null("a").collect()
        ╭─────────╮
        │ a       │
        │ ---     │
        │ Float64 │
        ╞═════════╡
        │ 1.6     │
        ├╌╌╌╌╌╌╌╌╌┤
        │ 2.5     │
        ├╌╌╌╌╌╌╌╌╌┤
        │ NaN     │
        ╰─────────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)


    """
    if len(cols) == 0:
        columns = column_inputs_to_expressions(self.column_names)
    else:
        columns = column_inputs_to_expressions(cols)
    return self.where(~reduce(lambda x, y: x | y, (x.is_null() for x in columns)))

@DataframePublicAPI
def except_all(self, other: "DataFrame") -> "DataFrame":
    """Returns the set difference of two DataFrames, considering duplicates.

    Args:
        other (DataFrame): DataFrame to except with

    Returns:
        DataFrame: DataFrame with the set difference of the two DataFrames, considering duplicates

    Examples:
        >>> import daft
        >>> df1 = daft.from_pydict({"a": [1, 1, 2, 2], "b": [4, 4, 6, 6]})
        >>> df2 = daft.from_pydict({"a": [1, 2, 2], "b": [4, 6, 6]})
        >>> df1.except_all(df2).collect()
        ╭───────┬───────╮
        │ a     ┆ b     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╞═══════╪═══════╡
        │ 1     ┆ 4     │
        ╰───────┴───────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)

    """
    builder = self._builder.except_all(other._builder)
    return DataFrame(builder)

@DataframePublicAPI
def except_distinct(self, other: "DataFrame") -> "DataFrame":
    """Returns the set difference of two DataFrames.

    Args:
        other (DataFrame): DataFrame to except with

    Returns:
        DataFrame: DataFrame with the set difference of the two DataFrames

    Examples:
        >>> import daft
        >>> df1 = daft.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
        >>> df2 = daft.from_pydict({"a": [1, 2, 3], "b": [4, 8, 6]})
        >>> df1.except_distinct(df2).collect()
        ╭───────┬───────╮
        │ a     ┆ b     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╞═══════╪═══════╡
        │ 2     ┆ 5     │
        ╰───────┴───────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)

    """
    builder = self._builder.except_distinct(other._builder)
    return DataFrame(builder)

@DataframePublicAPI
def exclude(self, *names: str) -> "DataFrame":
    """Drops columns from the current DataFrame by name.

    This is equivalent of performing a select with all the columns but the ones excluded.

    Args:
        *names (str): names to exclude

    Returns:
        DataFrame: DataFrame with some columns excluded.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
        >>> df_without_x = df.exclude("x")
        >>> df_without_x.show()
        ╭───────┬───────╮
        │ y     ┆ z     │
        │ ---   ┆ ---   │
        │ Int64 ┆ Int64 │
        ╞═══════╪═══════╡
        │ 4     ┆ 7     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 5     ┆ 8     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ 6     ┆ 9     │
        ╰───────┴───────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)
    """
    builder = self._builder.exclude(list(names))
    return DataFrame(builder)

@DataframePublicAPI
def explain(
    self, show_all: bool = False, format: str = "ascii", simple: bool = False, file: io.IOBase | None = None
) -> Any:
    r"""Prints the (logical and physical) plans that will be executed to produce this DataFrame.

    Defaults to showing the unoptimized logical plan. Use `show_all=True` to show the unoptimized logical plan,
    the optimized logical plan, and the physical plan.

    Args:
        show_all (bool): Whether to show the optimized logical plan and the physical plan in addition to the
            unoptimized logical plan.
        format (str): The format to print the plan in. one of 'ascii' or 'mermaid'
        simple (bool): Whether to only show the type of op for each node in the plan, rather than showing details
            of how each op is configured.

        file (Optional[io.IOBase]): Location to print the output to, or defaults to None which defaults to the default location for
            print (in Python, that should be sys.stdout)

    Returns:
        Union[None, str, MermaidFormatter]:
            - If `format="mermaid"` and running in a notebook, returns a `MermaidFormatter` instance for rich rendering.
            - If `format="mermaid"` and not in a notebook, returns a string representation of the plan.
            - Otherwise, prints the plan(s) to the specified file or stdout and returns `None`.

    Examples:
        >>> import daft
        >>>
        >>> df = daft.from_pydict({"x": [1, 2, 3]})
        >>>
        >>> def double(df, column: str):
        ...     return df.select((df[column] * df[column]).alias(column))
        >>>
        >>> df = df.pipe(double, "x")
        >>>
        >>> df.explain()
        == Unoptimized Logical Plan ==
        <BLANKLINE>
        * Project: col(x) * col(x) as x
        |
        * Source:
        |   Number of partitions = 1
        |   Output schema = x#Int64
        <BLANKLINE>
        <BLANKLINE>
        <BLANKLINE>
        Set `show_all=True` to also see the Optimized and Physical plans. This will run the query optimizer.

    """
    is_cached = self._result_cache is not None
    if format == "mermaid":
        from daft.dataframe.display import MermaidFormatter
        from daft.utils import in_notebook

        instance = MermaidFormatter(self.__builder, show_all, simple, is_cached)
        if file is not None:
            # if we are printing to a file, we print the markdown representation of the plan
            text = instance._repr_markdown_()
            print(text, file=file)
        if in_notebook():
            # if in a notebook, we return the class instance and let jupyter display it
            return instance
        else:
            # if we are not in a notebook, we return the raw markdown instead of the class instance
            return repr(instance)

    print_to_file = partial(print, file=file)

    if self._result_cache is not None:
        print_to_file("Result is cached and will skip computation\n")
        print_to_file(self._builder.pretty_print(simple, format=format))

        print_to_file("However here is the logical plan used to produce this result:\n", file=file)

    builder = self.__builder
    print_to_file("== Unoptimized Logical Plan ==\n")
    print_to_file(builder.pretty_print(simple, format=format))
    if show_all:
        print_to_file("\n== Optimized Logical Plan ==\n")
        execution_config = get_context().daft_execution_config
        builder = builder.optimize(execution_config)
        print_to_file(builder.pretty_print(simple))
        print_to_file("\n== Physical Plan ==\n")
        if get_or_create_runner().name != "native":
            from daft.daft import DistributedPhysicalPlan

            distributed_plan = DistributedPhysicalPlan.from_logical_plan_builder(
                builder._builder, "<tmp>", execution_config
            )
            if format == "ascii":
                print_to_file(distributed_plan.repr_ascii(simple))
            elif format == "mermaid":
                print_to_file(distributed_plan.repr_mermaid(MermaidOptions(simple)))
        else:
            native_executor = NativeExecutor()
            print_to_file(
                native_executor.pretty_print(builder, get_context().daft_execution_config, simple, format=format)
            )
    else:
        print_to_file(
            "\n \nSet `show_all=True` to also see the Optimized and Physical plans. This will run the query optimizer.",
        )
    return None