Google Git
Sign in
apache/spark/master/./python/pyspark/sql/connect/dataframe.py
blob: 95d6cddd50ba6bf604c4d0ccc8a8ca27451f469f [file]
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# mypy: disable-error-code="override"
import copy
import functools
import json
import os
import random
import sys
import warnings
from collections.abc import Iterable
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
Iterator,
List,
NoReturn,
Optional,
Sequence,
Tuple,
Type,
Union,
cast,
overload,
)
import pyarrow as pa
import pyspark.sql.connect.plan as plan
from pyspark import _NoValue
from pyspark._globals import _NoValueType
from pyspark.errors import (
PySparkAttributeError,
PySparkNotImplementedError,
PySparkRuntimeError,
PySparkTypeError,
PySparkValueError,
)
from pyspark.errors.exceptions.base import (
PySparkIndexError,
SessionNotSameException,
)
from pyspark.resource import ResourceProfile
from pyspark.serializers import CPickleSerializer
from pyspark.sql.column import Column
from pyspark.sql.connect.column import Column as ConnectColumn
from pyspark.sql.connect.expressions import (
ColumnReference,
DirectShufflePartitionID,
SortOrder,
SubqueryExpression,
UnresolvedRegex,
UnresolvedStar,
)
from pyspark.sql.connect.functions import builtin as F
from pyspark.sql.connect.group import GroupedData
from pyspark.sql.connect.logging import logger
from pyspark.sql.connect.merge import MergeIntoWriter
from pyspark.sql.connect.readwriter import DataFrameWriter, DataFrameWriterV2
from pyspark.sql.connect.streaming.readwriter import DataStreamWriter
from pyspark.sql.conversion import ArrowTableToRowsConversion
from pyspark.sql.dataframe import (
DataFrame as ParentDataFrame,
)
from pyspark.sql.dataframe import (
DataFrameNaFunctions as ParentDataFrameNaFunctions,
)
from pyspark.sql.dataframe import (
DataFrameStatFunctions as ParentDataFrameStatFunctions,
)
from pyspark.sql.pandas.functions import _validate_vectorized_udf # type: ignore[attr-defined]
from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema
from pyspark.sql.table_arg import TableArg
from pyspark.sql.types import Row, StructType, _create_row
from pyspark.storagelevel import StorageLevel
from pyspark.util import PythonEvalType, is_remote_only
if TYPE_CHECKING:
from pyspark.core.rdd import RDD
from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame
from pyspark.sql.connect._typing import (
ArrowMapIterFunction,
ColumnOrName,
ColumnOrNameOrOrdinal,
LiteralType,
OptionalPrimitiveType,
PandasMapIterFunction,
PrimitiveType,
)
from pyspark.sql.connect.observation import Observation
from pyspark.sql.connect.session import SparkSession
from pyspark.sql.metrics import ExecutionInfo
from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike
from pyspark.sql.plot import PySparkPlotAccessor
class DataFrame(ParentDataFrame):
def __new__(
cls,
*args: Any,
**kwargs: Any,
) -> "DataFrame":
# ParentDataFrame by default creates classic DataFrame for backward compatibility.
# We have to do an explicit object.__new__ to avoid that.
return object.__new__(cls)
def __init__(
self,
plan: plan.LogicalPlan,
session: "SparkSession",
):
"""Creates a new data frame"""
self._plan = plan
if self._plan is None:
raise PySparkRuntimeError(
errorClass="MISSING_VALID_PLAN",
messageParameters={"operator": "__init__"},
)
self._session: "SparkSession" = session # type: ignore[assignment]
if self._session is None:
raise PySparkRuntimeError(
errorClass="NO_ACTIVE_SESSION",
messageParameters={"operator": "__init__"},
)
# Check whether _repr_html is supported or not, we use it to avoid calling RPC twice
# by __repr__ and _repr_html_ while eager evaluation opens.
self._support_repr_html = False
self._cached_schema: Optional[StructType] = None
self._cached_schema_serialized: Optional[bytes] = None
self._execution_info: Optional["ExecutionInfo"] = None
def __reduce__(self) -> Tuple:
"""
Custom method for serializing the DataFrame object using Pickle. Since the DataFrame
overrides "__getattr__" method, the default serialization method does not work.
Returns
-------
The tuple containing the information needed to reconstruct the object.
"""
return (
DataFrame,
(
self._plan,
self._session,
),
{
"_support_repr_html": self._support_repr_html,
"_cached_schema": self._cached_schema,
},
)
def __repr__(self) -> str:
if not self._support_repr_html:
(
repl_eager_eval_enabled,
repl_eager_eval_max_num_rows,
repl_eager_eval_truncate,
) = self._session._client.get_configs(
"spark.sql.repl.eagerEval.enabled",
"spark.sql.repl.eagerEval.maxNumRows",
"spark.sql.repl.eagerEval.truncate",
)
if repl_eager_eval_enabled == "true":
return self._show_string(
n=int(cast(str, repl_eager_eval_max_num_rows)),
truncate=int(cast(str, repl_eager_eval_truncate)),
vertical=False,
)
return "DataFrame[%s]" % (", ".join("%s: %s" % c for c in self.dtypes))
def _repr_html_(self) -> Optional[str]:
if not self._support_repr_html:
self._support_repr_html = True
(
repl_eager_eval_enabled,
repl_eager_eval_max_num_rows,
repl_eager_eval_truncate,
) = self._session._client.get_configs(
"spark.sql.repl.eagerEval.enabled",
"spark.sql.repl.eagerEval.maxNumRows",
"spark.sql.repl.eagerEval.truncate",
)
if repl_eager_eval_enabled == "true":
table, _ = DataFrame(
plan.HtmlString(
child=self._plan,
num_rows=int(cast(str, repl_eager_eval_max_num_rows)),
truncate=int(cast(str, repl_eager_eval_truncate)),
),
session=self._session,
)._to_table()
return table[0][0].as_py()
else:
return None
def _to_cols(
self,
*cols: Union[Sequence["ColumnOrNameOrOrdinal"], "ColumnOrNameOrOrdinal"],
arg_name: str = "cols",
allow_ordinal: bool = False,
allow_sequence: bool = True,
sort: bool = False,
) -> List[Column]:
def _raise() -> NoReturn:
if allow_ordinal:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column, int or str",
"arg_name": arg_name,
"arg_type": type(c).__name__,
},
)
else:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column or str",
"arg_name": arg_name,
"arg_type": type(c).__name__,
},
)
if (
len(cols) == 1
and not isinstance(cols[0], (int, str, Column))
and isinstance(cols[0], Sequence)
):
if allow_sequence:
cols = tuple(cols[0])
else:
_raise()
_cols: List[Column] = []
for c in cols:
if isinstance(c, Column):
col = c
elif isinstance(c, str):
col = F.col(c)
elif isinstance(c, int) and not isinstance(c, bool):
if allow_ordinal:
if c < 1:
raise PySparkIndexError(
errorClass="INDEX_NOT_POSITIVE", messageParameters={"index": str(c)}
)
col = self[c - 1]
else:
_raise()
else:
_raise()
if sort:
if not isinstance(col._expr, SortOrder):
col = col.asc()
_cols.append(col)
return _cols
@property
def write(self) -> "DataFrameWriter":
def cb(qe: "ExecutionInfo") -> None:
self._execution_info = qe
return DataFrameWriter(self._plan, self._session, cb)
@functools.cache
def isEmpty(self) -> bool:
return len(self.select().take(1)) == 0
@overload
def select(self, *cols: "ColumnOrName") -> ParentDataFrame: ...
@overload
def select(
self, __cols: Union[Sequence["ColumnOrName"], "ColumnOrName"]
) -> ParentDataFrame: ...
def select(self, *cols: Union[Sequence["ColumnOrName"], "ColumnOrName"]) -> ParentDataFrame:
return DataFrame(
plan.Project(self._plan, self._to_cols(*cols)),
session=self._session,
)
def selectExpr(self, *expr: Union[str, List[str]]) -> ParentDataFrame:
sql_expr = []
if len(expr) == 1 and isinstance(expr[0], list):
expr = expr[0] # type: ignore[assignment]
for element in expr:
if isinstance(element, str):
sql_expr.append(F.expr(element))
else:
sql_expr.extend([F.expr(e) for e in element])
return DataFrame(plan.Project(self._plan, sql_expr), session=self._session)
def agg(self, *exprs: Union[Column, Dict[str, str]]) -> ParentDataFrame:
if not exprs:
raise PySparkValueError(
errorClass="CANNOT_BE_EMPTY",
messageParameters={"item": "exprs"},
)
if len(exprs) == 1 and isinstance(exprs[0], dict):
measures = [F._invoke_function(f, F.col(e)) for e, f in exprs[0].items()]
return self.groupBy().agg(*measures)
else:
# other expressions
assert all(isinstance(c, Column) for c in exprs), "all exprs should be Expression"
exprs = cast(Tuple[Column, ...], exprs)
return self.groupBy().agg(*exprs)
def alias(self, alias: str) -> ParentDataFrame:
res = DataFrame(plan.SubqueryAlias(self._plan, alias), session=self._session)
res._cached_schema = self._cached_schema
return res
def metadataColumn(self, colName: str) -> Column:
if not isinstance(colName, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "colName",
"expected_type": "str",
"arg_type": type(colName).__name__,
},
)
return self._col(colName, is_metadata_column=True)
def colRegex(self, colName: str) -> Column:
if not isinstance(colName, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "colName",
"expected_type": "str",
"arg_type": type(colName).__name__,
},
)
return ConnectColumn(UnresolvedRegex(colName, self._plan._plan_id))
@property
def dtypes(self) -> List[Tuple[str, str]]:
return [(str(f.name), f.dataType.simpleString()) for f in self._schema.fields]
@property
def columns(self) -> List[str]:
return [field.name for field in self._schema.fields]
@property
def sparkSession(self) -> "SparkSession":
return self._session
def count(self) -> int:
table, _ = self.agg(F._invoke_function("count", F.lit(1)))._to_table() # type: ignore[operator]
return table[0][0].as_py()
def crossJoin(self, other: ParentDataFrame) -> ParentDataFrame:
other = self._check_same_session(other)
return DataFrame(
plan.Join(left=self._plan, right=other._plan, on=None, how="cross"),
session=self._session,
)
def zip(self, other: ParentDataFrame) -> ParentDataFrame:
other = self._check_same_session(other)
return DataFrame(
plan.Zip(self._plan, other._plan),
session=self._session,
)
def _check_same_session(self, other: ParentDataFrame) -> "DataFrame":
if (
not isinstance(other, DataFrame)
or self._session.session_id != other._session.session_id
):
raise SessionNotSameException(
errorClass="SESSION_NOT_SAME",
messageParameters={},
)
return other
def coalesce(self, numPartitions: int) -> ParentDataFrame:
if not numPartitions > 0:
raise PySparkValueError(
errorClass="VALUE_NOT_POSITIVE",
messageParameters={"arg_name": "numPartitions", "arg_value": str(numPartitions)},
)
res = DataFrame(
plan.Repartition(self._plan, num_partitions=numPartitions, shuffle=False),
self._session,
)
res._cached_schema = self._cached_schema
return res
def repartition(
self, numPartitions: Union[int, "ColumnOrName"], *cols: "ColumnOrName"
) -> ParentDataFrame:
if isinstance(numPartitions, int):
if not numPartitions > 0:
raise PySparkValueError(
errorClass="VALUE_NOT_POSITIVE",
messageParameters={
"arg_name": "numPartitions",
"arg_value": str(numPartitions),
},
)
if len(cols) == 0:
res = DataFrame(
plan.Repartition(self._plan, numPartitions, shuffle=True),
self._session,
)
else:
res = DataFrame(
plan.RepartitionByExpression(self._plan, numPartitions, self._to_cols(cols)),
self.sparkSession,
)
elif isinstance(numPartitions, (str, Column)):
cols = (numPartitions,) + cols
res = DataFrame(
plan.RepartitionByExpression(self._plan, None, self._to_cols(cols)),
self.sparkSession,
)
else:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column or str",
"arg_name": "numPartitions",
"arg_type": type(numPartitions).__name__,
},
)
res._cached_schema = self._cached_schema
return res
def repartitionByRange(
self, numPartitions: Union[int, "ColumnOrName"], *cols: "ColumnOrName"
) -> ParentDataFrame:
if isinstance(numPartitions, int):
if not numPartitions > 0:
raise PySparkValueError(
errorClass="VALUE_NOT_POSITIVE",
messageParameters={
"arg_name": "numPartitions",
"arg_value": str(numPartitions),
},
)
if len(cols) == 0:
raise PySparkValueError(
errorClass="CANNOT_BE_EMPTY",
messageParameters={"item": "cols"},
)
else:
res = DataFrame(
plan.RepartitionByExpression(
self._plan, numPartitions, self._to_cols(cols, sort=True)
),
self.sparkSession,
)
elif isinstance(numPartitions, (str, Column)):
res = DataFrame(
plan.RepartitionByExpression(
self._plan, None, self._to_cols((numPartitions,) + cols, sort=True)
),
self.sparkSession,
)
else:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column, int or str",
"arg_name": "numPartitions",
"arg_type": type(numPartitions).__name__,
},
)
res._cached_schema = self._cached_schema
return res
def repartitionById(
self, numPartitions: int, partitionIdCol: "ColumnOrName"
) -> ParentDataFrame:
if not isinstance(numPartitions, int) or isinstance(numPartitions, bool):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "int",
"arg_name": "numPartitions",
"arg_type": type(numPartitions).__name__,
},
)
if numPartitions <= 0:
raise PySparkValueError(
errorClass="VALUE_NOT_POSITIVE",
messageParameters={
"arg_name": "numPartitions",
"arg_value": str(numPartitions),
},
)
partition_connect_col = cast(ConnectColumn, F._to_col(partitionIdCol))
direct_partition_expr = DirectShufflePartitionID(partition_connect_col._expr)
direct_partition_col = ConnectColumn(direct_partition_expr)
res = DataFrame(
plan.RepartitionByExpression(self._plan, numPartitions, [direct_partition_col]),
self._session,
)
res._cached_schema = self._cached_schema
return res
def dropDuplicates(self, subset: Optional[List[str]] = None) -> ParentDataFrame:
if subset is not None and not isinstance(subset, (list, tuple)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "subset",
"expected_type": "list or tuple",
"arg_type": type(subset).__name__,
},
)
if subset is None:
res = DataFrame(
plan.Deduplicate(child=self._plan, all_columns_as_keys=True), session=self._session
)
else:
for c in subset:
if not isinstance(c, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "subset",
"expected_type": "str",
"arg_type": type(c).__name__,
},
)
res = DataFrame(
plan.Deduplicate(child=self._plan, column_names=subset), session=self._session
)
res._cached_schema = self._cached_schema
return res
drop_duplicates = dropDuplicates
def dropDuplicatesWithinWatermark(self, subset: Optional[List[str]] = None) -> ParentDataFrame:
if subset is not None and not isinstance(subset, (list, tuple)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "subset",
"expected_type": "list or tuple",
"arg_type": type(subset).__name__,
},
)
if subset is None:
return DataFrame(
plan.Deduplicate(child=self._plan, all_columns_as_keys=True, within_watermark=True),
session=self._session,
)
else:
for c in subset:
if not isinstance(c, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "subset",
"expected_type": "str",
"arg_type": type(c).__name__,
},
)
return DataFrame(
plan.Deduplicate(child=self._plan, column_names=subset, within_watermark=True),
session=self._session,
)
def distinct(self) -> ParentDataFrame:
res = DataFrame(
plan.Deduplicate(child=self._plan, all_columns_as_keys=True), session=self._session
)
res._cached_schema = self._cached_schema
return res
def drop(self, *cols: "ColumnOrName") -> ParentDataFrame:
# We can't convert names to columns here because drop has different behavior
# for names and columns.
_cols = list(cols)
if any(not isinstance(c, (str, Column)) for c in _cols):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column or str",
"arg_name": "cols",
"arg_type": type(cols).__name__,
},
)
return DataFrame(
plan.Drop(
child=self._plan,
columns=_cols,
),
session=self._session,
)
def filter(self, condition: Union[Column, str]) -> ParentDataFrame:
if isinstance(condition, str):
expr = F.expr(condition)
else:
expr = condition
res = DataFrame(plan.Filter(child=self._plan, filter=expr), session=self._session)
res._cached_schema = self._cached_schema
return res
def first(self) -> Optional[Row]:
return self.head()
@overload
def groupBy(self, *cols: "ColumnOrNameOrOrdinal") -> "GroupedData": ...
@overload
def groupBy(self, __cols: Sequence["ColumnOrNameOrOrdinal"]) -> "GroupedData": ...
def groupBy(
self, *cols: Union[Sequence["ColumnOrNameOrOrdinal"], "ColumnOrNameOrOrdinal"]
) -> "GroupedData":
return GroupedData(
df=self, group_type="groupby", grouping_cols=self._to_cols(*cols, allow_ordinal=True)
)
groupby = groupBy # type: ignore[assignment]
@overload
def rollup(self, *cols: "ColumnOrNameOrOrdinal") -> "GroupedData": ...
@overload
def rollup(self, __cols: Sequence["ColumnOrNameOrOrdinal"]) -> "GroupedData": ...
def rollup(
self, *cols: Union[Sequence["ColumnOrNameOrOrdinal"], "ColumnOrNameOrOrdinal"]
) -> "GroupedData":
return GroupedData(
df=self, group_type="rollup", grouping_cols=self._to_cols(*cols, allow_ordinal=True)
)
@overload
def cube(self, *cols: "ColumnOrName") -> "GroupedData": ...
@overload
def cube(self, __cols: Sequence["ColumnOrNameOrOrdinal"]) -> "GroupedData": ...
def cube(
self, *cols: Union[Sequence["ColumnOrNameOrOrdinal"], "ColumnOrNameOrOrdinal"]
) -> "GroupedData":
return GroupedData(
df=self, group_type="cube", grouping_cols=self._to_cols(*cols, allow_ordinal=True)
)
def groupingSets(
self,
groupingSets: Sequence[Sequence["ColumnOrNameOrOrdinal"]],
*cols: "ColumnOrNameOrOrdinal",
) -> "GroupedData":
gsets: List[List[Column]] = []
for grouping_set in groupingSets:
gsets.append(self._to_cols(grouping_set, arg_name="groupingSets", allow_ordinal=True))
return GroupedData(
df=self,
group_type="grouping_sets",
grouping_cols=self._to_cols(cols, allow_ordinal=True),
grouping_sets=gsets,
)
@overload
def head(self) -> Optional[Row]: ...
@overload
def head(self, n: int) -> List[Row]: ...
def head(self, n: Optional[int] = None) -> Union[Optional[Row], List[Row]]:
if n is None:
rs = self.head(1)
return rs[0] if rs else None
return self.take(n)
def take(self, num: int) -> List[Row]:
return self.limit(num).collect()
def join(
self,
other: ParentDataFrame,
on: Optional[Union[str, List[str], Column, List[Column]]] = None,
how: Optional[str] = None,
) -> ParentDataFrame:
other = self._check_same_session(other)
if how is not None and isinstance(how, str):
how = how.lower().replace("_", "")
return DataFrame(
plan.Join(left=self._plan, right=other._plan, on=on, how=how),
session=self._session,
)
def lateralJoin(
self,
other: ParentDataFrame,
on: Optional[Column] = None,
how: Optional[str] = None,
) -> ParentDataFrame:
other = self._check_same_session(other)
if how is not None and isinstance(how, str):
how = how.lower().replace("_", "")
return DataFrame(
plan.LateralJoin(left=self._plan, right=other._plan, on=on, how=how),
session=self._session,
)
def nearestByJoin(
self,
other: ParentDataFrame,
rankingExpression: Column,
numResults: int,
mode: str,
direction: str,
*,
joinType: str = "inner",
) -> ParentDataFrame:
other = self._check_same_session(other)
return DataFrame(
plan.NearestByJoin(
left=self._plan,
right=other._plan,
ranking_expression=rankingExpression,
num_results=int(numResults),
join_type=joinType,
mode=mode,
direction=direction,
),
session=self._session,
)
def _joinAsOf(
self,
other: ParentDataFrame,
leftAsOfColumn: Union[str, Column],
rightAsOfColumn: Union[str, Column],
on: Optional[Union[str, List[str], Column, List[Column]]] = None,
how: Optional[str] = None,
*,
tolerance: Optional[Column] = None,
allowExactMatches: bool = True,
direction: str = "backward",
) -> ParentDataFrame:
other = self._check_same_session(other)
if how is None:
how = "inner"
assert isinstance(how, str), "how should be a string"
if tolerance is not None:
assert isinstance(tolerance, Column), "tolerance should be Column"
def _convert_col(df: ParentDataFrame, col: "ColumnOrName") -> Column:
if isinstance(col, Column):
return col
else:
return df._col(col) # type: ignore[operator]
return DataFrame(
plan.AsOfJoin(
left=self._plan,
right=other._plan,
left_as_of=_convert_col(self, leftAsOfColumn),
right_as_of=_convert_col(other, rightAsOfColumn),
on=on,
how=how,
tolerance=tolerance,
allow_exact_matches=allowExactMatches,
direction=direction,
),
session=self._session,
)
def limit(self, num: int) -> ParentDataFrame:
res = DataFrame(plan.Limit(child=self._plan, limit=num), session=self._session)
res._cached_schema = self._cached_schema
return res
def tail(self, num: int) -> List[Row]:
return DataFrame(plan.Tail(child=self._plan, limit=num), session=self._session).collect()
def sort(
self,
*cols: Union[Sequence["ColumnOrNameOrOrdinal"], "ColumnOrNameOrOrdinal"],
**kwargs: Any,
) -> ParentDataFrame:
_cols = self._preapare_cols_for_sort(F.col, cols, kwargs)
res = DataFrame(
plan.Sort(
self._plan,
columns=[F._sort_col(c) for c in _cols],
is_global=True,
),
session=self._session,
)
res._cached_schema = self._cached_schema
return res
orderBy = sort
def sortWithinPartitions(
self,
*cols: Union[Sequence["ColumnOrNameOrOrdinal"], "ColumnOrNameOrOrdinal"],
**kwargs: Any,
) -> ParentDataFrame:
_cols = self._preapare_cols_for_sort(F.col, cols, kwargs)
res = DataFrame(
plan.Sort(
self._plan,
columns=[F._sort_col(c) for c in _cols],
is_global=False,
),
session=self._session,
)
res._cached_schema = self._cached_schema
return res
def sample(
self,
withReplacement: Optional[Union[float, bool]] = None,
fraction: Optional[Union[int, float]] = None,
seed: Optional[int] = None,
) -> ParentDataFrame:
_w, _f, _s = self._preapare_args_for_sample(withReplacement, fraction, seed)
res = DataFrame(
plan.Sample(
child=self._plan,
lower_bound=0.0,
upper_bound=_f,
with_replacement=_w,
seed=_s,
),
session=self._session,
)
res._cached_schema = self._cached_schema
return res
def withColumnRenamed(self, existing: str, new: str) -> ParentDataFrame:
return self.withColumnsRenamed({existing: new})
def withColumnsRenamed(self, colsMap: Dict[str, str]) -> ParentDataFrame:
if not isinstance(colsMap, dict):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "dict",
"arg_name": "colsMap",
"arg_type": type(colsMap).__name__,
},
)
return DataFrame(plan.WithColumnsRenamed(self._plan, colsMap), self._session)
def _show_string(
self, n: int = 20, truncate: Union[bool, int] = True, vertical: bool = False
) -> str:
if not isinstance(n, int) or isinstance(n, bool):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "int",
"arg_name": "n",
"arg_type": type(n).__name__,
},
)
if not isinstance(vertical, bool):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "bool",
"arg_name": "vertical",
"arg_type": type(vertical).__name__,
},
)
_truncate: int = -1
if isinstance(truncate, bool) and truncate:
_truncate = 20
else:
try:
_truncate = int(truncate)
except ValueError:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "bool",
"arg_name": "truncate",
"arg_type": type(truncate).__name__,
},
)
table, _ = DataFrame(
plan.ShowString(
child=self._plan,
num_rows=n,
truncate=_truncate,
vertical=vertical,
),
session=self._session,
)._to_table()
return table[0][0].as_py()
def withColumns(self, *colsMap: Dict[str, Column]) -> ParentDataFrame:
# Below code is to help enable kwargs in future.
assert len(colsMap) == 1
colsMap = colsMap[0] # type: ignore[assignment]
if not isinstance(colsMap, dict):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "dict",
"arg_name": "colsMap",
"arg_type": type(colsMap).__name__,
},
)
names: List[str] = []
columns: List[Column] = []
for columnName, column in colsMap.items():
names.append(columnName)
columns.append(column)
return DataFrame(
plan.WithColumns(
self._plan,
columnNames=names,
columns=columns,
),
session=self._session,
)
def withColumn(self, colName: str, col: Column) -> ParentDataFrame:
if not isinstance(col, Column):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column",
"arg_name": "col",
"arg_type": type(col).__name__,
},
)
return DataFrame(
plan.WithColumns(
self._plan,
columnNames=[colName],
columns=[col],
),
session=self._session,
)
def withMetadata(self, columnName: str, metadata: Dict[str, Any]) -> ParentDataFrame:
if not isinstance(metadata, dict):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "dict",
"arg_name": "metadata",
"arg_type": type(metadata).__name__,
},
)
return DataFrame(
plan.WithColumns(
self._plan,
columnNames=[columnName],
columns=[self[columnName]],
metadata=[json.dumps(metadata)],
),
session=self._session,
)
def unpivot(
self,
ids: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]],
values: Optional[Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]]],
variableColumnName: str,
valueColumnName: str,
) -> ParentDataFrame:
assert ids is not None, "ids must not be None"
def _convert_cols(
cols: Optional[Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]]],
) -> List[Column]:
if cols is None:
return []
elif isinstance(cols, (tuple, list)):
return [F._to_col(c) for c in cols]
else:
return [F._to_col(cols)]
return DataFrame(
plan.Unpivot(
self._plan,
_convert_cols(ids),
_convert_cols(values) if values is not None else None,
variableColumnName,
valueColumnName,
),
self._session,
)
melt = unpivot
def withWatermark(self, eventTime: str, delayThreshold: str) -> ParentDataFrame:
# TODO: reuse error handling code in sql.DataFrame.withWatermark()
if not eventTime or type(eventTime) is not str:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "eventTime",
"expected_type": "str",
"arg_type": type(eventTime).__name__,
},
)
if not delayThreshold or type(delayThreshold) is not str:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "delayThreshold",
"expected_type": "str",
"arg_type": type(delayThreshold).__name__,
},
)
return DataFrame(
plan.WithWatermark(
self._plan,
event_time=eventTime,
delay_threshold=delayThreshold,
),
session=self._session,
)
def hint(
self, name: str, *parameters: Union["PrimitiveType", "Column", List["PrimitiveType"]]
) -> ParentDataFrame:
if len(parameters) == 1 and isinstance(parameters[0], list):
parameters = parameters[0] # type: ignore[assignment]
if not isinstance(name, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "name",
"expected_type": "str",
"arg_type": type(name).__name__,
},
)
allowed_types = (str, float, int, Column, list)
allowed_primitive_types = (str, float, int)
allowed_types_repr = ", ".join(
[t.__name__ for t in allowed_types[:-1]]
+ ["list[" + t.__name__ + "]" for t in allowed_primitive_types]
)
for p in parameters:
if not isinstance(p, allowed_types):
raise PySparkTypeError(
errorClass="INVALID_ITEM_FOR_CONTAINER",
messageParameters={
"arg_name": "parameters",
"allowed_types": allowed_types_repr,
"item_type": type(p).__name__,
},
)
if isinstance(p, list):
if not all(isinstance(e, allowed_primitive_types) for e in p):
raise PySparkTypeError(
errorClass="INVALID_ITEM_FOR_CONTAINER",
messageParameters={
"arg_name": "parameters",
"allowed_types": allowed_types_repr,
"item_type": type(p).__name__ + "[" + type(p[0]).__name__ + "]",
},
)
res = DataFrame(
plan.Hint(self._plan, name, [F.lit(p) for p in list(parameters)]),
session=self._session,
)
res._cached_schema = self._cached_schema
return res
def randomSplit(
self,
weights: List[float],
seed: Optional[int] = None,
) -> List[ParentDataFrame]:
for w in weights:
if w < 0.0:
raise PySparkValueError(
errorClass="VALUE_NOT_POSITIVE",
messageParameters={"arg_name": "weights", "arg_value": str(w)},
)
seed = seed if seed is not None else random.randint(0, sys.maxsize)
total = sum(weights)
if total <= 0:
raise PySparkValueError(
errorClass="VALUE_NOT_POSITIVE",
messageParameters={"arg_name": "sum(weights)", "arg_value": str(total)},
)
proportions = list(map(lambda x: x / total, weights))
normalizedCumWeights = [0.0]
for v in proportions:
normalizedCumWeights.append(normalizedCumWeights[-1] + v)
j = 1
length = len(normalizedCumWeights)
splits = []
while j < length:
lowerBound = normalizedCumWeights[j - 1]
upperBound = normalizedCumWeights[j]
samplePlan = DataFrame(
plan.Sample(
child=self._plan,
lower_bound=lowerBound,
upper_bound=upperBound,
with_replacement=False,
seed=int(seed),
deterministic_order=True,
),
session=self._session,
)
samplePlan._cached_schema = self._cached_schema
splits.append(samplePlan)
j += 1
return splits # type: ignore[return-value]
def observe(
self,
observation: Union["Observation", str],
*exprs: Column,
) -> ParentDataFrame:
from pyspark.sql.connect.observation import Observation
if len(exprs) == 0:
raise PySparkValueError(
errorClass="CANNOT_BE_EMPTY",
messageParameters={"item": "exprs"},
)
if not all(isinstance(c, Column) for c in exprs):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "exprs",
"expected_type": "list[Column]",
"arg_type": type(exprs).__name__,
},
)
if isinstance(observation, Observation):
res = observation._on(self, *exprs)
elif isinstance(observation, str):
res = DataFrame(
plan.CollectMetrics(self._plan, observation, list(exprs)),
self._session,
)
else:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "observation",
"expected_type": "Observation or str",
"arg_type": type(observation).__name__,
},
)
res._cached_schema = self._cached_schema
return res
def show(self, n: int = 20, truncate: Union[bool, int] = True, vertical: bool = False) -> None:
print(self._show_string(n, truncate, vertical))
def _merge_cached_schema(self, other: ParentDataFrame) -> Optional[StructType]:
# to avoid type coercion, only propagate the schema
# when the cached schemas are exactly the same
if self._cached_schema is not None and self._cached_schema == other._cached_schema:
return self.schema
return None
def union(self, other: ParentDataFrame) -> ParentDataFrame:
other = self._check_same_session(other)
return self.unionAll(other)
def unionAll(self, other: ParentDataFrame) -> ParentDataFrame:
other = self._check_same_session(other)
res = DataFrame(
plan.SetOperation(self._plan, other._plan, "union", is_all=True),
session=self._session,
)
res._cached_schema = self._merge_cached_schema(other)
return res
def unionByName(
self, other: ParentDataFrame, allowMissingColumns: bool = False
) -> ParentDataFrame:
other = self._check_same_session(other)
res = DataFrame(
plan.SetOperation(
self._plan,
other._plan,
"union",
by_name=True,
allow_missing_columns=allowMissingColumns,
),
session=self._session,
)
res._cached_schema = self._merge_cached_schema(other)
return res
def subtract(self, other: ParentDataFrame) -> ParentDataFrame:
other = self._check_same_session(other)
res = DataFrame(
plan.SetOperation(self._plan, other._plan, "except", is_all=False),
session=self._session,
)
res._cached_schema = self._merge_cached_schema(other)
return res
def exceptAll(self, other: ParentDataFrame) -> ParentDataFrame:
other = self._check_same_session(other)
res = DataFrame(
plan.SetOperation(self._plan, other._plan, "except", is_all=True),
session=self._session,
)
res._cached_schema = self._merge_cached_schema(other)
return res
def zipWithIndex(self, indexColName: str = "index") -> ParentDataFrame:
return self.select(
F.col("*"), F._invoke_function("distributed_sequence_id").alias(indexColName)
)
def intersect(self, other: ParentDataFrame) -> ParentDataFrame:
other = self._check_same_session(other)
res = DataFrame(
plan.SetOperation(self._plan, other._plan, "intersect", is_all=False),
session=self._session,
)
res._cached_schema = self._merge_cached_schema(other)
return res
def intersectAll(self, other: ParentDataFrame) -> ParentDataFrame:
other = self._check_same_session(other)
res = DataFrame(
plan.SetOperation(self._plan, other._plan, "intersect", is_all=True),
session=self._session,
)
res._cached_schema = self._merge_cached_schema(other)
return res
def where(self, condition: Union[Column, str]) -> ParentDataFrame:
if not isinstance(condition, (str, Column)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column or str",
"arg_name": "condition",
"arg_type": type(condition).__name__,
},
)
return self.filter(condition)
@property
def na(self) -> ParentDataFrameNaFunctions:
return DataFrameNaFunctions(self)
def fillna(
self,
value: Union["LiteralType", Dict[str, "LiteralType"]],
subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None,
) -> ParentDataFrame:
if not isinstance(value, (float, int, str, bool, dict)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "bool, dict, float, int or str",
"arg_name": "value",
"arg_type": type(value).__name__,
},
)
if isinstance(value, dict):
if len(value) == 0:
raise PySparkValueError(
errorClass="CANNOT_BE_EMPTY",
messageParameters={"item": "value"},
)
for c, v in value.items():
if not isinstance(c, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "key type of dict",
"expected_type": "str",
"arg_type": type(c).__name__,
},
)
if not isinstance(v, (bool, int, float, str)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "bool, float, int or str",
"arg_name": "value type of dict",
"arg_type": type(v).__name__,
},
)
_cols: List[str] = []
if subset is not None:
if isinstance(subset, str):
_cols = [subset]
elif isinstance(subset, (tuple, list)):
for c in subset:
if not isinstance(c, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "cols",
"expected_type": "list, str or tuple",
"arg_type": type(c).__name__,
},
)
_cols = list(subset)
else:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "subset",
"expected_type": "list or tuple",
"arg_type": type(subset).__name__,
},
)
if isinstance(value, dict):
_cols = list(value.keys())
_values = [value[c] for c in _cols]
else:
_values = [value]
return DataFrame(
plan.NAFill(child=self._plan, cols=_cols, values=_values),
session=self._session,
)
def dropna(
self,
how: str = "any",
thresh: Optional[int] = None,
subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None,
) -> ParentDataFrame:
min_non_nulls: Optional[int] = None
if how is not None:
if not isinstance(how, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "how",
"expected_type": "str",
"arg_type": type(how).__name__,
},
)
if how == "all":
min_non_nulls = 1
elif how == "any":
min_non_nulls = None
else:
raise PySparkValueError(
errorClass="VALUE_NOT_ALLOWED",
messageParameters={"arg_name": "how", "allowed_values": "['any', 'all']"},
)
if thresh is not None:
if not isinstance(thresh, int):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "int",
"arg_name": "thresh",
"arg_type": type(thresh).__name__,
},
)
# 'thresh' overwrites 'how'
min_non_nulls = thresh
_cols: List[str] = []
if subset is not None:
if isinstance(subset, str):
_cols = [subset]
elif isinstance(subset, (tuple, list)):
for c in subset:
if not isinstance(c, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "cols",
"expected_type": "list, str or tuple",
"arg_type": type(c).__name__,
},
)
_cols = list(subset)
else:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "subset",
"expected_type": "list, str or tuple",
"arg_type": type(subset).__name__,
},
)
return DataFrame(
plan.NADrop(child=self._plan, cols=_cols, min_non_nulls=min_non_nulls),
session=self._session,
)
def replace(
self,
to_replace: Union[
"LiteralType", List["LiteralType"], Dict["LiteralType", "OptionalPrimitiveType"]
],
value: Optional[
Union["OptionalPrimitiveType", List["OptionalPrimitiveType"], _NoValueType]
] = _NoValue,
subset: Optional[List[str]] = None,
) -> ParentDataFrame:
if value is _NoValue:
if isinstance(to_replace, dict):
value = None
else:
raise PySparkTypeError(
errorClass="ARGUMENT_REQUIRED",
messageParameters={"arg_name": "value", "condition": "`to_replace` is dict"},
)
# Helper functions
def all_of(types: Union[Type, Tuple[Type, ...]]) -> Callable[[Iterable], bool]:
"""Given a type or tuple of types and a sequence of xs
check if each x is instance of type(s)
>>> all_of(bool)([True, False])
True
>>> all_of(str)(["a", 1])
False
"""
def all_of_(xs: Iterable) -> bool:
return all(isinstance(x, types) for x in xs)
return all_of_
all_of_bool = all_of(bool)
all_of_str = all_of(str)
all_of_numeric = all_of((float, int))
# Validate input types
valid_types = (bool, float, int, str, list, tuple)
if not isinstance(to_replace, valid_types + (dict,)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "bool, dict, float, int, str or tuple",
"arg_name": "to_replace",
"arg_type": type(to_replace).__name__,
},
)
if (
not isinstance(value, valid_types)
and value is not None
and not isinstance(to_replace, dict)
):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "bool, float, int, list, None, str or tuple",
"arg_name": "value",
"arg_type": type(value).__name__,
},
)
if isinstance(to_replace, (list, tuple)) and isinstance(value, (list, tuple)):
if len(to_replace) != len(value):
raise PySparkValueError(
errorClass="LENGTH_SHOULD_BE_THE_SAME",
messageParameters={
"arg1": "to_replace",
"arg2": "value",
"arg1_length": str(len(to_replace)),
"arg2_length": str(len(value)),
},
)
if not (subset is None or isinstance(subset, (list, tuple, str))):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "subset",
"expected_type": "list, str or tuple",
"arg_type": type(subset).__name__,
},
)
# Reshape input arguments if necessary
if isinstance(to_replace, (float, int, str)):
to_replace = [to_replace]
if isinstance(to_replace, dict):
rep_dict = to_replace
if value is not None:
warnings.warn("to_replace is a dict and value is not None. value will be ignored.")
else:
if isinstance(value, (float, int, str)) or value is None:
value = [value for _ in range(len(to_replace))]
rep_dict = dict(zip(to_replace, cast("Iterable[Optional[Union[float, str]]]", value)))
if isinstance(subset, str):
subset = [subset]
# Verify we were not passed in mixed type generics.
if not any(
all_of_type(rep_dict.keys())
and all_of_type(x for x in rep_dict.values() if x is not None)
for all_of_type in [all_of_bool, all_of_str, all_of_numeric]
):
raise PySparkValueError(
errorClass="MIXED_TYPE_REPLACEMENT",
messageParameters={},
)
def _convert_int_to_float(v: Any) -> Any:
# a bool is also an int
if v is not None and not isinstance(v, bool) and isinstance(v, int):
return float(v)
else:
return v
_replacements = []
for k, v in rep_dict.items():
_k = _convert_int_to_float(k)
_v = _convert_int_to_float(v)
_replacements.append((F.lit(_k), F.lit(_v)))
return DataFrame(
plan.NAReplace(
child=self._plan,
cols=subset,
replacements=_replacements,
),
session=self._session,
)
@property
def stat(self) -> ParentDataFrameStatFunctions:
return DataFrameStatFunctions(self)
def summary(self, *statistics: str) -> ParentDataFrame:
_statistics: List[str] = list(statistics)
for s in _statistics:
if not isinstance(s, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "statistics",
"expected_type": "list[str]",
"arg_type": type(s).__name__,
},
)
return DataFrame(
plan.StatSummary(child=self._plan, statistics=_statistics),
session=self._session,
)
def describe(self, *cols: Union[str, List[str]]) -> ParentDataFrame:
if len(cols) == 1 and isinstance(cols[0], list):
cols = cols[0] # type: ignore[assignment]
_cols = []
for column in cols:
if isinstance(column, str):
_cols.append(column)
else:
_cols.extend([s for s in column])
return DataFrame(
plan.StatDescribe(child=self._plan, cols=_cols),
session=self._session,
)
def cov(self, col1: str, col2: str) -> float:
if not isinstance(col1, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "col1",
"expected_type": "str",
"arg_type": type(col1).__name__,
},
)
if not isinstance(col2, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "col2",
"expected_type": "str",
"arg_type": type(col2).__name__,
},
)
table, _ = DataFrame(
plan.StatCov(child=self._plan, col1=col1, col2=col2),
session=self._session,
)._to_table()
return table[0][0].as_py()
def corr(self, col1: str, col2: str, method: Optional[str] = None) -> float:
if not isinstance(col1, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "col1",
"expected_type": "str",
"arg_type": type(col1).__name__,
},
)
if not isinstance(col2, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "col2",
"expected_type": "str",
"arg_type": type(col2).__name__,
},
)
if not method:
method = "pearson"
if not method == "pearson":
raise PySparkValueError(
errorClass="VALUE_NOT_ALLOWED",
messageParameters={"arg_name": "method", "allowed_values": "['pearson']"},
)
table, _ = DataFrame(
plan.StatCorr(child=self._plan, col1=col1, col2=col2, method=method),
session=self._session,
)._to_table()
return table[0][0].as_py()
def approxQuantile(
self,
col: Union[str, List[str], Tuple[str]],
probabilities: Union[List[float], Tuple[float]],
relativeError: float,
) -> Union[List[float], List[List[float]]]:
if not isinstance(col, (str, list, tuple)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "col",
"expected_type": "list, str or tuple",
"arg_type": type(col).__name__,
},
)
isStr = isinstance(col, str)
if isinstance(col, tuple):
col = list(col)
elif isStr:
col = [cast(str, col)]
for c in col:
if not isinstance(c, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "columns",
"expected_type": "list[str]",
"arg_type": type(c).__name__,
},
)
if not isinstance(probabilities, (list, tuple)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "probabilities",
"expected_type": "list or tuple",
"arg_type": type(probabilities).__name__,
},
)
if isinstance(probabilities, tuple):
probabilities = list(probabilities)
for p in probabilities:
if not isinstance(p, (float, int)) or p < 0 or p > 1:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "probabilities",
"expected_type": "list[float, int]",
"arg_type": type(p).__name__,
},
)
if not isinstance(relativeError, (float, int)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "float or int",
"arg_name": "relativeError",
"arg_type": type(relativeError).__name__,
},
)
if relativeError < 0:
raise PySparkValueError(
errorClass="NEGATIVE_VALUE",
messageParameters={
"arg_name": "relativeError",
"arg_value": str(relativeError),
},
)
relativeError = float(relativeError)
table, _ = DataFrame(
plan.StatApproxQuantile(
child=self._plan,
cols=list(col),
probabilities=probabilities,
relativeError=relativeError,
),
session=self._session,
)._to_table()
jaq = [q.as_py() for q in table[0][0]]
jaq_list = [list(j) for j in jaq]
return jaq_list[0] if isStr else jaq_list
def crosstab(self, col1: str, col2: str) -> ParentDataFrame:
if not isinstance(col1, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "col1",
"expected_type": "str",
"arg_type": type(col1).__name__,
},
)
if not isinstance(col2, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "col2",
"expected_type": "str",
"arg_type": type(col2).__name__,
},
)
return DataFrame(
plan.StatCrosstab(child=self._plan, col1=col1, col2=col2),
session=self._session,
)
def freqItems(
self, cols: Union[List[str], Tuple[str]], support: Optional[float] = None
) -> ParentDataFrame:
if isinstance(cols, tuple):
cols = list(cols)
if not isinstance(cols, list):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "cols",
"expected_type": "list or tuple",
"arg_type": type(cols).__name__,
},
)
if not support:
support = 0.01
return DataFrame(
plan.StatFreqItems(child=self._plan, cols=cols, support=support),
session=self._session,
)
def sampleBy(
self, col: "ColumnOrName", fractions: Dict[Any, float], seed: Optional[int] = None
) -> ParentDataFrame:
if not isinstance(col, (str, Column)):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column or str",
"arg_name": "col",
"arg_type": type(col).__name__,
},
)
if not isinstance(fractions, dict):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "dict",
"arg_name": "fractions",
"arg_type": type(fractions).__name__,
},
)
_fractions = []
for k, v in fractions.items():
if not isinstance(k, (float, int, str)):
raise PySparkTypeError(
errorClass="DISALLOWED_TYPE_FOR_CONTAINER",
messageParameters={
"arg_name": "fractions",
"arg_type": type(fractions).__name__,
"allowed_types": "float, int, str",
"item_type": type(k).__name__,
},
)
_fractions.append((F.lit(k), float(v)))
seed = seed if seed is not None else random.randint(0, sys.maxsize)
return DataFrame(
plan.StatSampleBy(
child=self._plan,
col=F._to_col(col),
fractions=_fractions,
seed=seed,
),
session=self._session,
)
def _ipython_key_completions_(self) -> List[str]:
"""Returns the names of columns in this :class:`DataFrame`.
Examples
--------
>>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], ["age", "name"])
>>> df._ipython_key_completions_()
['age', 'name']
Would return illegal identifiers.
>>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], ["age 1", "name?1"])
>>> df._ipython_key_completions_()
['age 1', 'name?1']
"""
return self.columns
def __getattr__(self, name: str) -> "Column":
if name in ["_jseq", "_jdf", "_jmap", "_jcols", "rdd"]:
raise PySparkAttributeError(
errorClass="JVM_ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": name}
)
# Only eagerly validate the column name when:
# 1, PYSPARK_VALIDATE_COLUMN_NAME_LEGACY is set 1; or
# 2, the name starts with '__', because it is likely a python internal method and
# an AttributeError might be expected to check whether the attribute exists.
# For example:
# pickle/cloudpickle need to check whether method '__setstate__' is defined or not,
# and it internally invokes __getattr__("__setstate__").
# Returning a dataframe column self._col("__setstate__") in this case will break
# the serialization of connect dataframe and features built atop it (e.g. FEB).
if os.environ.get("PYSPARK_VALIDATE_COLUMN_NAME_LEGACY") == "1" or name.startswith("__"):
if name not in self.columns:
raise PySparkAttributeError(
errorClass="ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": name}
)
return self._col(name)
@overload
def __getitem__(self, item: Union[int, str]) -> Column: ...
@overload
def __getitem__(self, item: Union[Column, List, Tuple]) -> ParentDataFrame: ...
def __getitem__(
self, item: Union[int, str, Column, List, Tuple]
) -> Union[Column, ParentDataFrame]:
if isinstance(item, str):
if item == "*":
return ConnectColumn(
UnresolvedStar(
unparsed_target=None,
plan_id=self._plan._plan_id,
)
)
else:
# TODO: revisit classic Spark's Dataset.col
# if (sparkSession.sessionState.conf.supportQuotedRegexColumnName) {
# colRegex(colName)
# } else {
# ConnectColumn(addDataFrameIdToCol(resolve(colName)))
# }
# validate the column name
if os.environ.get("PYSPARK_VALIDATE_COLUMN_NAME_LEGACY") == "1" and not hasattr(
self._session, "is_mock_session"
):
from pyspark.sql.connect.types import verify_col_name
# Try best to verify the column name with cached schema
# If fails, fall back to the server side validation
if not verify_col_name(item, self._schema):
self.select(item).isLocal()
return self._col(item)
elif isinstance(item, Column):
return self.filter(item)
elif isinstance(item, (list, tuple)):
return self.select(*item)
elif isinstance(item, int):
return F.col(self.columns[item])
else:
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "Column, int, list, str or tuple",
"arg_name": "item",
"arg_type": type(item).__name__,
},
)
def _col(self, name: str, is_metadata_column: bool = False) -> Column:
return ConnectColumn(
ColumnReference(
unparsed_identifier=name,
plan_id=self._plan._plan_id,
is_metadata_column=is_metadata_column,
)
)
def __dir__(self) -> List[str]:
attrs = set(dir(DataFrame))
attrs.update(self.columns)
return sorted(attrs)
def collect(self) -> List[Row]:
table, schema = self._to_table()
# not all datatypes are supported in arrow based collect
# here always verify the schema by from_arrow_schema
schema2 = from_arrow_schema(table.schema, prefer_timestamp_ntz=True)
schema = schema or schema2
assert schema is not None and isinstance(schema, StructType)
return ArrowTableToRowsConversion.convert(
table, schema, binary_as_bytes=self._get_binary_as_bytes()
)
def _get_binary_as_bytes(self) -> bool:
"""Get the binary_as_bytes configuration value from Spark session."""
conf_value = self._session.conf.get("spark.sql.execution.pyspark.binaryAsBytes", "true")
return conf_value is not None and conf_value.lower() == "true"
def _to_table(self) -> Tuple["pa.Table", Optional[StructType]]:
query = self._plan.to_proto(self._session.client)
table, schema, self._execution_info = self._session.client.to_table(
query, self._plan.observations
)
assert table is not None
return (table, schema)
def toArrow(self) -> "pa.Table":
schema = to_arrow_schema(
self.schema,
error_on_duplicated_field_names_in_struct=True,
timezone="UTC",
)
table, _ = self._to_table()
return table.cast(schema)
def toPandas(self) -> "PandasDataFrameLike":
return self._to_pandas()
def _to_pandas(self, **kwargs: Any) -> "PandasDataFrameLike":
query = self._plan.to_proto(self._session.client)
pdf, ei = self._session.client.to_pandas(query, self._plan.observations, **kwargs)
self._execution_info = ei
return pdf
def transpose(self, indexColumn: Optional["ColumnOrName"] = None) -> ParentDataFrame:
return DataFrame(
plan.Transpose(self._plan, [F._to_col(indexColumn)] if indexColumn is not None else []),
self._session,
)
def asTable(self) -> TableArg:
from pyspark.sql.connect.table_arg import TableArg as ConnectTableArg
return ConnectTableArg(SubqueryExpression(self._plan, subquery_type="table_arg"))
def scalar(self) -> Column:
return ConnectColumn(SubqueryExpression(self._plan, subquery_type="scalar"))
def exists(self) -> Column:
return ConnectColumn(SubqueryExpression(self._plan, subquery_type="exists"))
@property
def _schema(self) -> StructType:
# Schema caching is correct in most cases. Connect is lazy by nature. This means that
# we only resolve the plan when it is submitted for execution or analysis. We do not
# cache intermediate resolved plan. If the input (changes table, view redefinition,
# etc...) of the plan changes between the schema() call, and a subsequent action, the
# cached schema might be inconsistent with the end schema.
if self._cached_schema is None:
query = self._plan.to_proto(self._session.client)
self._cached_schema = self._session.client.schema(query)
try:
self._cached_schema_serialized = CPickleSerializer().dumps(self._schema)
except Exception as e:
logger.warning(f"DataFrame schema pickle dumps failed with exception: {e}.")
self._cached_schema_serialized = None
return self._cached_schema
@property
def schema(self) -> StructType:
# self._schema call will cache the schema and serialize it if it is not cached yet.
_schema = self._schema
if self._cached_schema_serialized is not None:
try:
return CPickleSerializer().loads(self._cached_schema_serialized)
except Exception as e:
logger.warning(f"DataFrame schema pickle loads failed with exception: {e}.")
# In case of pickle ser/de failure, fallback to deepcopy approach.
return copy.deepcopy(_schema)
@functools.cache
def isLocal(self) -> bool:
query = self._plan.to_proto(self._session.client)
result = self._session.client._analyze(method="is_local", plan=query).is_local
assert result is not None
return result
@functools.cached_property
def isStreaming(self) -> bool:
query = self._plan.to_proto(self._session.client)
result = self._session.client._analyze(method="is_streaming", plan=query).is_streaming
assert result is not None
return result
def printSchema(self, level: Optional[int] = None) -> None:
if level:
print(self.schema.treeString(level))
else:
print(self.schema.treeString())
@functools.cache
def inputFiles(self) -> List[str]:
query = self._plan.to_proto(self._session.client)
result = self._session.client._analyze(method="input_files", plan=query).input_files
assert result is not None
return result
def to(self, schema: StructType) -> ParentDataFrame:
assert schema is not None
res = DataFrame(
plan.ToSchema(child=self._plan, schema=schema),
session=self._session,
)
res._cached_schema = schema
return res
def toDF(self, *cols: str) -> ParentDataFrame:
for col_ in cols:
if not isinstance(col_, str):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"arg_name": "cols",
"expected_type": "list[str]",
"arg_type": type(col_).__name__,
},
)
return DataFrame(plan.ToDF(self._plan, list(cols)), self._session)
def transform(
self, func: Callable[..., ParentDataFrame], *args: Any, **kwargs: Any
) -> ParentDataFrame:
result = func(self, *args, **kwargs)
assert isinstance(result, DataFrame), (
"Func returned an instance of type [%s], should have been DataFrame." % type(result)
)
return result
def _explain_string(
self, extended: Optional[Union[bool, str]] = None, mode: Optional[str] = None
) -> str:
if extended is not None and mode is not None:
raise PySparkValueError(
errorClass="CANNOT_SET_TOGETHER",
messageParameters={"arg_list": "extended and mode"},
)
# For the no argument case: df.explain()
is_no_argument = extended is None and mode is None
# For the cases below:
# explain(True)
# explain(extended=False)
is_extended_case = isinstance(extended, bool) and mode is None
# For the case when extended is mode:
# df.explain("formatted")
is_extended_as_mode = isinstance(extended, str) and mode is None
# For the mode specified:
# df.explain(mode="formatted")
is_mode_case = extended is None and isinstance(mode, str)
if not (is_no_argument or is_extended_case or is_extended_as_mode or is_mode_case):
argtypes = [str(type(arg)) for arg in [extended, mode] if arg is not None]
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "bool or str",
"arg_name": "extended (optional) and mode (optional)",
"arg_type": ", ".join(argtypes),
},
)
# Sets an explain mode depending on a given argument
if is_no_argument:
explain_mode = "simple"
elif is_extended_case:
explain_mode = "extended" if extended else "simple"
elif is_mode_case:
explain_mode = cast(str, mode)
elif is_extended_as_mode:
explain_mode = cast(str, extended)
query = self._plan.to_proto(self._session.client)
return self._session.client.explain_string(query, explain_mode)
def explain(
self, extended: Optional[Union[bool, str]] = None, mode: Optional[str] = None
) -> None:
print(self._explain_string(extended=extended, mode=mode))
def createTempView(self, name: str) -> None:
command = plan.CreateView(
child=self._plan, name=name, is_global=False, replace=False
).command(session=self._session.client)
_, _, ei = self._session.client.execute_command(command, self._plan.observations)
self._execution_info = ei
def createOrReplaceTempView(self, name: str) -> None:
command = plan.CreateView(
child=self._plan, name=name, is_global=False, replace=True
).command(session=self._session.client)
_, _, ei = self._session.client.execute_command(command, self._plan.observations)
self._execution_info = ei
def createGlobalTempView(self, name: str) -> None:
command = plan.CreateView(
child=self._plan, name=name, is_global=True, replace=False
).command(session=self._session.client)
_, _, ei = self._session.client.execute_command(command, self._plan.observations)
self._execution_info = ei
def createOrReplaceGlobalTempView(self, name: str) -> None:
command = plan.CreateView(
child=self._plan, name=name, is_global=True, replace=True
).command(session=self._session.client)
_, _, ei = self._session.client.execute_command(command, self._plan.observations)
self._execution_info = ei
def cache(self) -> ParentDataFrame:
return self.persist()
def persist(
self,
storageLevel: StorageLevel = (StorageLevel.MEMORY_AND_DISK_DESER),
) -> ParentDataFrame:
relation = self._plan.plan(self._session.client)
self._session.client._analyze(
method="persist", relation=relation, storage_level=storageLevel
)
return self
@property
def storageLevel(self) -> StorageLevel:
relation = self._plan.plan(self._session.client)
storage_level = self._session.client._analyze(
method="get_storage_level", relation=relation
).storage_level
assert storage_level is not None
return storage_level
def unpersist(self, blocking: bool = False) -> ParentDataFrame:
relation = self._plan.plan(self._session.client)
self._session.client._analyze(method="unpersist", relation=relation, blocking=blocking)
return self
@property
def is_cached(self) -> bool:
return self.storageLevel != StorageLevel.NONE
def toLocalIterator(self, prefetchPartitions: bool = False) -> Iterator[Row]:
query = self._plan.to_proto(self._session.client)
binary_as_bytes = self._get_binary_as_bytes()
schema: Optional[StructType] = None
for schema_or_table in self._session.client.to_table_as_iterator(
query, self._plan.observations
):
if isinstance(schema_or_table, StructType):
assert schema is None
schema = schema_or_table
else:
assert isinstance(schema_or_table, pa.Table)
table = schema_or_table
if schema is None:
schema = from_arrow_schema(table.schema, prefer_timestamp_ntz=True)
yield from ArrowTableToRowsConversion.convert(
table, schema, binary_as_bytes=binary_as_bytes
)
def pandas_api(
self, index_col: Optional[Union[str, List[str]]] = None
) -> "PandasOnSparkDataFrame":
from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame
from pyspark.pandas.internal import InternalFrame
from pyspark.pandas.namespace import _get_index_map
index_spark_columns, index_names = _get_index_map(self, index_col)
internal = InternalFrame(
spark_frame=self,
index_spark_columns=index_spark_columns,
index_names=index_names, # type: ignore[arg-type]
)
return PandasOnSparkDataFrame(internal)
def registerTempTable(self, name: str) -> None:
warnings.warn("Deprecated in 2.0, use createOrReplaceTempView instead.", FutureWarning)
self.createOrReplaceTempView(name)
def _map_partitions(
self,
func: "PandasMapIterFunction",
schema: Union[StructType, str],
evalType: int,
barrier: bool,
profile: Optional[ResourceProfile],
) -> ParentDataFrame:
from pyspark.sql.connect.udf import UserDefinedFunction
_validate_vectorized_udf(func, evalType)
if isinstance(schema, str):
schema = cast(StructType, self._session._parse_ddl(schema))
udf_obj = UserDefinedFunction(
func,
returnType=schema,
evalType=evalType,
)
res = DataFrame(
plan.MapPartitions(
child=self._plan,
function=udf_obj,
cols=self.columns,
is_barrier=barrier,
profile=profile,
),
session=self._session,
)
if isinstance(schema, StructType):
res._cached_schema = schema
return res
def mapInPandas(
self,
func: "PandasMapIterFunction",
schema: Union[StructType, str],
barrier: bool = False,
profile: Optional[ResourceProfile] = None,
) -> ParentDataFrame:
return self._map_partitions(
func, schema, PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, barrier, profile
)
def mapInArrow(
self,
func: "ArrowMapIterFunction",
schema: Union[StructType, str],
barrier: bool = False,
profile: Optional[ResourceProfile] = None,
) -> ParentDataFrame:
return self._map_partitions(
func, schema, PythonEvalType.SQL_MAP_ARROW_ITER_UDF, barrier, profile
)
def foreach(self, f: Callable[[Row], None]) -> None:
def foreach_func(row: Any) -> None:
f(row)
self.select(F.struct(*self._schema.fieldNames()).alias("row")).select(
F.udf(foreach_func, StructType())("row") # type: ignore[arg-type]
).collect()
def foreachPartition(self, f: Callable[[Iterator[Row]], None]) -> None:
schema = self._schema
binary_as_bytes = self._get_binary_as_bytes()
field_converters = [
ArrowTableToRowsConversion._create_converter(
f.dataType, none_on_identity=False, binary_as_bytes=binary_as_bytes
)
for f in schema.fields
]
def foreach_partition_func(itr: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
def flatten() -> Iterator[Row]:
for table in itr:
columnar_data = [
ArrowTableToRowsConversion._to_pylist(column) for column in table.columns
]
for i in range(0, table.num_rows):
values = [
field_converters[j](columnar_data[j][i]) # type: ignore[misc]
for j in range(table.num_columns)
]
yield _create_row(fields=schema.fieldNames(), values=values)
f(flatten())
return iter([])
self.mapInArrow(foreach_partition_func, schema=StructType()).collect()
@property
def writeStream(self) -> DataStreamWriter:
return DataStreamWriter(plan=self._plan, session=self._session)
def sameSemantics(self, other: ParentDataFrame) -> bool:
if not isinstance(other, DataFrame):
raise PySparkTypeError(
errorClass="NOT_EXPECTED_TYPE",
messageParameters={
"expected_type": "DataFrame",
"arg_name": "other",
"arg_type": type(other).__name__,
},
)
other = self._check_same_session(other)
return self._session.client.same_semantics(
plan=self._plan.to_proto(self._session.client),
other=other._plan.to_proto(other._session.client),
)
@functools.cache
def semanticHash(self) -> int:
return self._session.client.semantic_hash(
plan=self._plan.to_proto(self._session.client),
)
def writeTo(self, table: str) -> "DataFrameWriterV2":
def cb(ei: "ExecutionInfo") -> None:
self._execution_info = ei
return DataFrameWriterV2(self._plan, self._session, table, cb)
def mergeInto(self, table: str, condition: Column) -> "MergeIntoWriter":
def cb(ei: "ExecutionInfo") -> None:
self._execution_info = ei
return MergeIntoWriter(
self._plan,
self._session,
table,
condition, # type: ignore[arg-type]
cb,
)
def offset(self, num: int) -> ParentDataFrame:
return DataFrame(plan.Offset(child=self._plan, offset=num), session=self._session)
def checkpoint(self, eager: bool = True) -> ParentDataFrame:
cmd = plan.Checkpoint(child=self._plan, local=False, eager=eager)
_, properties, self._execution_info = self._session.client.execute_command(
cmd.command(self._session.client)
)
assert "checkpoint_command_result" in properties
checkpointed = properties["checkpoint_command_result"]
assert isinstance(checkpointed._plan, plan.CachedRemoteRelation)
return checkpointed
def localCheckpoint(
self, eager: bool = True, storageLevel: Optional[StorageLevel] = None
) -> ParentDataFrame:
cmd = plan.Checkpoint(child=self._plan, local=True, eager=eager, storage_level=storageLevel)
_, properties, self._execution_info = self._session.client.execute_command(
cmd.command(self._session.client)
)
assert "checkpoint_command_result" in properties
checkpointed = properties["checkpoint_command_result"]
assert isinstance(checkpointed._plan, plan.CachedRemoteRelation)
return checkpointed
def toJSON(self) -> ParentDataFrame:
return self.select(F.to_json(F.struct(F.col("*"))).alias("value"))
if not is_remote_only():
@property
def rdd(self) -> "RDD[Row]":
raise PySparkNotImplementedError(
errorClass="NOT_IMPLEMENTED",
messageParameters={"feature": "rdd"},
)
@property
def executionInfo(self) -> Optional["ExecutionInfo"]:
return self._execution_info
@property
def plot(self) -> "PySparkPlotAccessor":
from pyspark.sql.plot import PySparkPlotAccessor
return PySparkPlotAccessor(self)
class DataFrameNaFunctions(ParentDataFrameNaFunctions):
def __init__(self, df: ParentDataFrame):
self.df = df
def fill(
self,
value: Union["LiteralType", Dict[str, "LiteralType"]],
subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None,
) -> ParentDataFrame:
return self.df.fillna(value=value, subset=subset)
def drop(
self,
how: str = "any",
thresh: Optional[int] = None,
subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None,
) -> ParentDataFrame:
return self.df.dropna(how=how, thresh=thresh, subset=subset)
def replace(
self,
to_replace: Union[List["LiteralType"], Dict["LiteralType", "OptionalPrimitiveType"]],
value: Optional[
Union["OptionalPrimitiveType", List["OptionalPrimitiveType"], _NoValueType]
] = _NoValue,
subset: Optional[List[str]] = None,
) -> ParentDataFrame:
return self.df.replace(to_replace, value, subset) # type: ignore[arg-type]
class DataFrameStatFunctions(ParentDataFrameStatFunctions):
def __init__(self, df: ParentDataFrame):
self.df = df
def cov(self, col1: str, col2: str) -> float:
return self.df.cov(col1, col2)
def corr(self, col1: str, col2: str, method: Optional[str] = None) -> float:
return self.df.corr(col1, col2, method)
def approxQuantile(
self,
col: Union[str, List[str], Tuple[str]],
probabilities: Union[List[float], Tuple[float]],
relativeError: float,
) -> Union[List[float], List[List[float]]]:
return self.df.approxQuantile(col, probabilities, relativeError)
def crosstab(self, col1: str, col2: str) -> ParentDataFrame:
return self.df.crosstab(col1, col2)
def freqItems(
self, cols: Union[List[str], Tuple[str]], support: Optional[float] = None
) -> ParentDataFrame:
return self.df.freqItems(cols, support)
def sampleBy(
self, col: str, fractions: Dict[Any, float], seed: Optional[int] = None
) -> ParentDataFrame:
return self.df.sampleBy(col, fractions, seed)
def _test() -> None:
import doctest
import os
import sys
import pyspark.sql.dataframe
from pyspark.sql import SparkSession as PySparkSession
from pyspark.testing.utils import have_pandas, have_pyarrow
from pyspark.util import is_remote_only
# It inherits docstrings but doctests cannot detect them so we run
# the parent classe's doctests here directly.
os.chdir(os.environ["SPARK_HOME"])
globs = pyspark.sql.dataframe.__dict__.copy()
if not is_remote_only():
del pyspark.sql.dataframe.DataFrame.toJSON.__doc__
del pyspark.sql.dataframe.DataFrame.rdd.__doc__
if not have_pandas or not have_pyarrow:
del pyspark.sql.dataframe.DataFrame.toPandas.__doc__
del pyspark.sql.dataframe.DataFrame.mapInPandas.__doc__
del pyspark.sql.dataframe.DataFrame.pandas_api.__doc__
if not have_pyarrow:
del pyspark.sql.dataframe.DataFrame.toArrow.__doc__
del pyspark.sql.dataframe.DataFrame.mapInArrow.__doc__
else:
import pyarrow as pa
from pyspark.loose_version import LooseVersion
if LooseVersion(pa.__version__) < LooseVersion("21.0.0"):
del pyspark.sql.dataframe.DataFrame.mapInArrow.__doc__
globs["spark"] = (
PySparkSession.builder.appName("sql.connect.dataframe tests")
.remote(os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[4]"))
.getOrCreate()
)
failure_count, test_count = doctest.testmod(
pyspark.sql.dataframe,
globs=globs,
optionflags=doctest.ELLIPSIS
| doctest.NORMALIZE_WHITESPACE
| doctest.IGNORE_EXCEPTION_DETAIL,
)
globs["spark"].stop()
if failure_count:
sys.exit(-1)
if __name__ == "__main__":
_test()