sqlglot.dataframe.sql

 20class SparkSession:
 21    known_ids: t.ClassVar[t.Set[str]] = set()
 22    known_branch_ids: t.ClassVar[t.Set[str]] = set()
 23    known_sequence_ids: t.ClassVar[t.Set[str]] = set()
 24    name_to_sequence_id_mapping: t.ClassVar[t.Dict[str, t.List[str]]] = defaultdict(list)
 25
 26    def __init__(self):
 27        self.incrementing_id = 1
 28
 29    def __getattr__(self, name: str) -> SparkSession:
 30        return self
 31
 32    def __call__(self, *args, **kwargs) -> SparkSession:
 33        return self
 34
 35    @property
 36    def read(self) -> DataFrameReader:
 37        return DataFrameReader(self)
 38
 39    def table(self, tableName: str) -> DataFrame:
 40        return self.read.table(tableName)
 41
 42    def createDataFrame(
 43        self,
 44        data: t.Sequence[t.Union[t.Dict[str, ColumnLiterals], t.List[ColumnLiterals], t.Tuple]],
 45        schema: t.Optional[SchemaInput] = None,
 46        samplingRatio: t.Optional[float] = None,
 47        verifySchema: bool = False,
 48    ) -> DataFrame:
 49        from sqlglot.dataframe.sql.dataframe import DataFrame
 50
 51        if samplingRatio is not None or verifySchema:
 52            raise NotImplementedError("Sampling Ratio and Verify Schema are not supported")
 53        if schema is not None and (
 54            not isinstance(schema, (StructType, str, list))
 55            or (isinstance(schema, list) and not isinstance(schema[0], str))
 56        ):
 57            raise NotImplementedError("Only schema of either list or string of list supported")
 58        if not data:
 59            raise ValueError("Must provide data to create into a DataFrame")
 60
 61        column_mapping: t.Dict[str, t.Optional[str]]
 62        if schema is not None:
 63            column_mapping = get_column_mapping_from_schema_input(schema)
 64        elif isinstance(data[0], dict):
 65            column_mapping = {col_name.strip(): None for col_name in data[0]}
 66        else:
 67            column_mapping = {f"_{i}": None for i in range(1, len(data[0]) + 1)}
 68
 69        data_expressions = [
 70            exp.Tuple(
 71                expressions=list(
 72                    map(
 73                        lambda x: F.lit(x).expression,
 74                        row if not isinstance(row, dict) else row.values(),
 75                    )
 76                )
 77            )
 78            for row in data
 79        ]
 80
 81        sel_columns = [
 82            F.col(name).cast(data_type).alias(name).expression
 83            if data_type is not None
 84            else F.col(name).expression
 85            for name, data_type in column_mapping.items()
 86        ]
 87
 88        select_kwargs = {
 89            "expressions": sel_columns,
 90            "from": exp.From(
 91                this=exp.Values(
 92                    expressions=data_expressions,
 93                    alias=exp.TableAlias(
 94                        this=exp.to_identifier(self._auto_incrementing_name),
 95                        columns=[exp.to_identifier(col_name) for col_name in column_mapping],
 96                    ),
 97                ),
 98            ),
 99        }
100
101        sel_expression = exp.Select(**select_kwargs)
102        return DataFrame(self, sel_expression)
103
104    def sql(self, sqlQuery: str) -> DataFrame:
105        expression = sqlglot.parse_one(sqlQuery, read="spark")
106        if isinstance(expression, exp.Select):
107            df = DataFrame(self, expression)
108            df = df._convert_leaf_to_cte()
109        elif isinstance(expression, (exp.Create, exp.Insert)):
110            select_expression = expression.expression.copy()
111            if isinstance(expression, exp.Insert):
112                select_expression.set("with", expression.args.get("with"))
113                expression.set("with", None)
114            del expression.args["expression"]
115            df = DataFrame(self, select_expression, output_expression_container=expression)  # type: ignore
116            df = df._convert_leaf_to_cte()
117        else:
118            raise ValueError(
119                "Unknown expression type provided in the SQL. Please create an issue with the SQL."
120            )
121        return df
122
123    @property
124    def _auto_incrementing_name(self) -> str:
125        name = f"a{self.incrementing_id}"
126        self.incrementing_id += 1
127        return name
128
129    @property
130    def _random_branch_id(self) -> str:
131        id = self._random_id
132        self.known_branch_ids.add(id)
133        return id
134
135    @property
136    def _random_sequence_id(self):
137        id = self._random_id
138        self.known_sequence_ids.add(id)
139        return id
140
141    @property
142    def _random_id(self) -> str:
143        id = "r" + uuid.uuid4().hex
144        self.known_ids.add(id)
145        return id
146
147    @property
148    def _join_hint_names(self) -> t.Set[str]:
149        return {"BROADCAST", "MERGE", "SHUFFLE_HASH", "SHUFFLE_REPLICATE_NL"}
150
151    def _add_alias_to_mapping(self, name: str, sequence_id: str):
152        self.name_to_sequence_id_mapping[name].append(sequence_id)

 45class DataFrame:
 46    def __init__(
 47        self,
 48        spark: SparkSession,
 49        expression: exp.Select,
 50        branch_id: t.Optional[str] = None,
 51        sequence_id: t.Optional[str] = None,
 52        last_op: Operation = Operation.INIT,
 53        pending_hints: t.Optional[t.List[exp.Expression]] = None,
 54        output_expression_container: t.Optional[OutputExpressionContainer] = None,
 55        **kwargs,
 56    ):
 57        self.spark = spark
 58        self.expression = expression
 59        self.branch_id = branch_id or self.spark._random_branch_id
 60        self.sequence_id = sequence_id or self.spark._random_sequence_id
 61        self.last_op = last_op
 62        self.pending_hints = pending_hints or []
 63        self.output_expression_container = output_expression_container or exp.Select()
 64
 65    def __getattr__(self, column_name: str) -> Column:
 66        return self[column_name]
 67
 68    def __getitem__(self, column_name: str) -> Column:
 69        column_name = f"{self.branch_id}.{column_name}"
 70        return Column(column_name)
 71
 72    def __copy__(self):
 73        return self.copy()
 74
 75    @property
 76    def sparkSession(self):
 77        return self.spark
 78
 79    @property
 80    def write(self):
 81        return DataFrameWriter(self)
 82
 83    @property
 84    def latest_cte_name(self) -> str:
 85        if not self.expression.ctes:
 86            from_exp = self.expression.args["from"]
 87            if from_exp.alias_or_name:
 88                return from_exp.alias_or_name
 89            table_alias = from_exp.find(exp.TableAlias)
 90            if not table_alias:
 91                raise RuntimeError(
 92                    f"Could not find an alias name for this expression: {self.expression}"
 93                )
 94            return table_alias.alias_or_name
 95        return self.expression.ctes[-1].alias
 96
 97    @property
 98    def pending_join_hints(self):
 99        return [hint for hint in self.pending_hints if isinstance(hint, exp.JoinHint)]
100
101    @property
102    def pending_partition_hints(self):
103        return [hint for hint in self.pending_hints if isinstance(hint, exp.Anonymous)]
104
105    @property
106    def columns(self) -> t.List[str]:
107        return self.expression.named_selects
108
109    @property
110    def na(self) -> DataFrameNaFunctions:
111        return DataFrameNaFunctions(self)
112
113    def _replace_cte_names_with_hashes(self, expression: exp.Select):
114        replacement_mapping = {}
115        for cte in expression.ctes:
116            old_name_id = cte.args["alias"].this
117            new_hashed_id = exp.to_identifier(
118                self._create_hash_from_expression(cte.this), quoted=old_name_id.args["quoted"]
119            )
120            replacement_mapping[old_name_id] = new_hashed_id
121            expression = expression.transform(replace_id_value, replacement_mapping)
122        return expression
123
124    def _create_cte_from_expression(
125        self,
126        expression: exp.Expression,
127        branch_id: t.Optional[str] = None,
128        sequence_id: t.Optional[str] = None,
129        **kwargs,
130    ) -> t.Tuple[exp.CTE, str]:
131        name = self._create_hash_from_expression(expression)
132        expression_to_cte = expression.copy()
133        expression_to_cte.set("with", None)
134        cte = exp.Select().with_(name, as_=expression_to_cte, **kwargs).ctes[0]
135        cte.set("branch_id", branch_id or self.branch_id)
136        cte.set("sequence_id", sequence_id or self.sequence_id)
137        return cte, name
138
139    @t.overload
140    def _ensure_list_of_columns(self, cols: t.Collection[ColumnOrLiteral]) -> t.List[Column]:
141        ...
142
143    @t.overload
144    def _ensure_list_of_columns(self, cols: ColumnOrLiteral) -> t.List[Column]:
145        ...
146
147    def _ensure_list_of_columns(self, cols):
148        return Column.ensure_cols(ensure_list(cols))
149
150    def _ensure_and_normalize_cols(self, cols, expression: t.Optional[exp.Select] = None):
151        cols = self._ensure_list_of_columns(cols)
152        normalize(self.spark, expression or self.expression, cols)
153        return cols
154
155    def _ensure_and_normalize_col(self, col):
156        col = Column.ensure_col(col)
157        normalize(self.spark, self.expression, col)
158        return col
159
160    def _convert_leaf_to_cte(self, sequence_id: t.Optional[str] = None) -> DataFrame:
161        df = self._resolve_pending_hints()
162        sequence_id = sequence_id or df.sequence_id
163        expression = df.expression.copy()
164        cte_expression, cte_name = df._create_cte_from_expression(
165            expression=expression, sequence_id=sequence_id
166        )
167        new_expression = df._add_ctes_to_expression(
168            exp.Select(), expression.ctes + [cte_expression]
169        )
170        sel_columns = df._get_outer_select_columns(cte_expression)
171        new_expression = new_expression.from_(cte_name).select(
172            *[x.alias_or_name for x in sel_columns]
173        )
174        return df.copy(expression=new_expression, sequence_id=sequence_id)
175
176    def _resolve_pending_hints(self) -> DataFrame:
177        df = self.copy()
178        if not self.pending_hints:
179            return df
180        expression = df.expression
181        hint_expression = expression.args.get("hint") or exp.Hint(expressions=[])
182        for hint in df.pending_partition_hints:
183            hint_expression.append("expressions", hint)
184            df.pending_hints.remove(hint)
185
186        join_aliases = {
187            join_table.alias_or_name
188            for join_table in get_tables_from_expression_with_join(expression)
189        }
190        if join_aliases:
191            for hint in df.pending_join_hints:
192                for sequence_id_expression in hint.expressions:
193                    sequence_id_or_name = sequence_id_expression.alias_or_name
194                    sequence_ids_to_match = [sequence_id_or_name]
195                    if sequence_id_or_name in df.spark.name_to_sequence_id_mapping:
196                        sequence_ids_to_match = df.spark.name_to_sequence_id_mapping[
197                            sequence_id_or_name
198                        ]
199                    matching_ctes = [
200                        cte
201                        for cte in reversed(expression.ctes)
202                        if cte.args["sequence_id"] in sequence_ids_to_match
203                    ]
204                    for matching_cte in matching_ctes:
205                        if matching_cte.alias_or_name in join_aliases:
206                            sequence_id_expression.set("this", matching_cte.args["alias"].this)
207                            df.pending_hints.remove(hint)
208                            break
209                hint_expression.append("expressions", hint)
210        if hint_expression.expressions:
211            expression.set("hint", hint_expression)
212        return df
213
214    def _hint(self, hint_name: str, args: t.List[Column]) -> DataFrame:
215        hint_name = hint_name.upper()
216        hint_expression = (
217            exp.JoinHint(
218                this=hint_name,
219                expressions=[exp.to_table(parameter.alias_or_name) for parameter in args],
220            )
221            if hint_name in JOIN_HINTS
222            else exp.Anonymous(
223                this=hint_name, expressions=[parameter.expression for parameter in args]
224            )
225        )
226        new_df = self.copy()
227        new_df.pending_hints.append(hint_expression)
228        return new_df
229
230    def _set_operation(self, klass: t.Callable, other: DataFrame, distinct: bool):
231        other_df = other._convert_leaf_to_cte()
232        base_expression = self.expression.copy()
233        base_expression = self._add_ctes_to_expression(base_expression, other_df.expression.ctes)
234        all_ctes = base_expression.ctes
235        other_df.expression.set("with", None)
236        base_expression.set("with", None)
237        operation = klass(this=base_expression, distinct=distinct, expression=other_df.expression)
238        operation.set("with", exp.With(expressions=all_ctes))
239        return self.copy(expression=operation)._convert_leaf_to_cte()
240
241    def _cache(self, storage_level: str):
242        df = self._convert_leaf_to_cte()
243        df.expression.ctes[-1].set("cache_storage_level", storage_level)
244        return df
245
246    @classmethod
247    def _add_ctes_to_expression(cls, expression: exp.Select, ctes: t.List[exp.CTE]) -> exp.Select:
248        expression = expression.copy()
249        with_expression = expression.args.get("with")
250        if with_expression:
251            existing_ctes = with_expression.expressions
252            existsing_cte_names = {x.alias_or_name for x in existing_ctes}
253            for cte in ctes:
254                if cte.alias_or_name not in existsing_cte_names:
255                    existing_ctes.append(cte)
256        else:
257            existing_ctes = ctes
258        expression.set("with", exp.With(expressions=existing_ctes))
259        return expression
260
261    @classmethod
262    def _get_outer_select_columns(cls, item: t.Union[exp.Expression, DataFrame]) -> t.List[Column]:
263        expression = item.expression if isinstance(item, DataFrame) else item
264        return [Column(x) for x in (expression.find(exp.Select) or exp.Select()).expressions]
265
266    @classmethod
267    def _create_hash_from_expression(cls, expression: exp.Expression) -> str:
268        value = expression.sql(dialect="spark").encode("utf-8")
269        return f"t{zlib.crc32(value)}"[:6]
270
271    def _get_select_expressions(
272        self,
273    ) -> t.List[t.Tuple[t.Union[t.Type[exp.Cache], OutputExpressionContainer], exp.Select]]:
274        select_expressions: t.List[
275            t.Tuple[t.Union[t.Type[exp.Cache], OutputExpressionContainer], exp.Select]
276        ] = []
277        main_select_ctes: t.List[exp.CTE] = []
278        for cte in self.expression.ctes:
279            cache_storage_level = cte.args.get("cache_storage_level")
280            if cache_storage_level:
281                select_expression = cte.this.copy()
282                select_expression.set("with", exp.With(expressions=copy(main_select_ctes)))
283                select_expression.set("cte_alias_name", cte.alias_or_name)
284                select_expression.set("cache_storage_level", cache_storage_level)
285                select_expressions.append((exp.Cache, select_expression))
286            else:
287                main_select_ctes.append(cte)
288        main_select = self.expression.copy()
289        if main_select_ctes:
290            main_select.set("with", exp.With(expressions=main_select_ctes))
291        expression_select_pair = (type(self.output_expression_container), main_select)
292        select_expressions.append(expression_select_pair)  # type: ignore
293        return select_expressions
294
295    def sql(self, dialect="spark", optimize=True, **kwargs) -> t.List[str]:
296        df = self._resolve_pending_hints()
297        select_expressions = df._get_select_expressions()
298        output_expressions: t.List[t.Union[exp.Select, exp.Cache, exp.Drop]] = []
299        replacement_mapping: t.Dict[exp.Identifier, exp.Identifier] = {}
300        for expression_type, select_expression in select_expressions:
301            select_expression = select_expression.transform(replace_id_value, replacement_mapping)
302            if optimize:
303                select_expression = t.cast(exp.Select, optimize_func(select_expression))
304            select_expression = df._replace_cte_names_with_hashes(select_expression)
305            expression: t.Union[exp.Select, exp.Cache, exp.Drop]
306            if expression_type == exp.Cache:
307                cache_table_name = df._create_hash_from_expression(select_expression)
308                cache_table = exp.to_table(cache_table_name)
309                original_alias_name = select_expression.args["cte_alias_name"]
310
311                replacement_mapping[exp.to_identifier(original_alias_name)] = exp.to_identifier(  # type: ignore
312                    cache_table_name
313                )
314                sqlglot.schema.add_table(
315                    cache_table_name,
316                    {
317                        expression.alias_or_name: expression.type.sql("spark")
318                        for expression in select_expression.expressions
319                    },
320                    dialect="spark",
321                )
322                cache_storage_level = select_expression.args["cache_storage_level"]
323                options = [
324                    exp.Literal.string("storageLevel"),
325                    exp.Literal.string(cache_storage_level),
326                ]
327                expression = exp.Cache(
328                    this=cache_table, expression=select_expression, lazy=True, options=options
329                )
330                # We will drop the "view" if it exists before running the cache table
331                output_expressions.append(exp.Drop(this=cache_table, exists=True, kind="VIEW"))
332            elif expression_type == exp.Create:
333                expression = df.output_expression_container.copy()
334                expression.set("expression", select_expression)
335            elif expression_type == exp.Insert:
336                expression = df.output_expression_container.copy()
337                select_without_ctes = select_expression.copy()
338                select_without_ctes.set("with", None)
339                expression.set("expression", select_without_ctes)
340                if select_expression.ctes:
341                    expression.set("with", exp.With(expressions=select_expression.ctes))
342            elif expression_type == exp.Select:
343                expression = select_expression
344            else:
345                raise ValueError(f"Invalid expression type: {expression_type}")
346            output_expressions.append(expression)
347
348        return [
349            expression.sql(**{"dialect": dialect, **kwargs}) for expression in output_expressions
350        ]
351
352    def copy(self, **kwargs) -> DataFrame:
353        return DataFrame(**object_to_dict(self, **kwargs))
354
355    @operation(Operation.SELECT)
356    def select(self, *cols, **kwargs) -> DataFrame:
357        cols = self._ensure_and_normalize_cols(cols)
358        kwargs["append"] = kwargs.get("append", False)
359        if self.expression.args.get("joins"):
360            ambiguous_cols = [
361                col
362                for col in cols
363                if isinstance(col.column_expression, exp.Column) and not col.column_expression.table
364            ]
365            if ambiguous_cols:
366                join_table_identifiers = [
367                    x.this for x in get_tables_from_expression_with_join(self.expression)
368                ]
369                cte_names_in_join = [x.this for x in join_table_identifiers]
370                # If we have columns that resolve to multiple CTE expressions then we want to use each CTE left-to-right
371                # and therefore we allow multiple columns with the same name in the result. This matches the behavior
372                # of Spark.
373                resolved_column_position: t.Dict[Column, int] = {col: -1 for col in ambiguous_cols}
374                for ambiguous_col in ambiguous_cols:
375                    ctes_with_column = [
376                        cte
377                        for cte in self.expression.ctes
378                        if cte.alias_or_name in cte_names_in_join
379                        and ambiguous_col.alias_or_name in cte.this.named_selects
380                    ]
381                    # Check if there is a CTE with this column that we haven't used before. If so, use it. Otherwise,
382                    # use the same CTE we used before
383                    cte = seq_get(ctes_with_column, resolved_column_position[ambiguous_col] + 1)
384                    if cte:
385                        resolved_column_position[ambiguous_col] += 1
386                    else:
387                        cte = ctes_with_column[resolved_column_position[ambiguous_col]]
388                    ambiguous_col.expression.set("table", cte.alias_or_name)
389        return self.copy(
390            expression=self.expression.select(*[x.expression for x in cols], **kwargs), **kwargs
391        )
392
393    @operation(Operation.NO_OP)
394    def alias(self, name: str, **kwargs) -> DataFrame:
395        new_sequence_id = self.spark._random_sequence_id
396        df = self.copy()
397        for join_hint in df.pending_join_hints:
398            for expression in join_hint.expressions:
399                if expression.alias_or_name == self.sequence_id:
400                    expression.set("this", Column.ensure_col(new_sequence_id).expression)
401        df.spark._add_alias_to_mapping(name, new_sequence_id)
402        return df._convert_leaf_to_cte(sequence_id=new_sequence_id)
403
404    @operation(Operation.WHERE)
405    def where(self, column: t.Union[Column, bool], **kwargs) -> DataFrame:
406        col = self._ensure_and_normalize_col(column)
407        return self.copy(expression=self.expression.where(col.expression))
408
409    filter = where
410
411    @operation(Operation.GROUP_BY)
412    def groupBy(self, *cols, **kwargs) -> GroupedData:
413        columns = self._ensure_and_normalize_cols(cols)
414        return GroupedData(self, columns, self.last_op)
415
416    @operation(Operation.SELECT)
417    def agg(self, *exprs, **kwargs) -> DataFrame:
418        cols = self._ensure_and_normalize_cols(exprs)
419        return self.groupBy().agg(*cols)
420
421    @operation(Operation.FROM)
422    def join(
423        self,
424        other_df: DataFrame,
425        on: t.Union[str, t.List[str], Column, t.List[Column]],
426        how: str = "inner",
427        **kwargs,
428    ) -> DataFrame:
429        other_df = other_df._convert_leaf_to_cte()
430        join_columns = self._ensure_list_of_columns(on)
431        # We will determine actual "join on" expression later so we don't provide it at first
432        join_expression = self.expression.join(
433            other_df.latest_cte_name, join_type=how.replace("_", " ")
434        )
435        join_expression = self._add_ctes_to_expression(join_expression, other_df.expression.ctes)
436        self_columns = self._get_outer_select_columns(join_expression)
437        other_columns = self._get_outer_select_columns(other_df)
438        # Determines the join clause and select columns to be used passed on what type of columns were provided for
439        # the join. The columns returned changes based on how the on expression is provided.
440        if isinstance(join_columns[0].expression, exp.Column):
441            """
442            Unique characteristics of join on column names only:
443            * The column names are put at the front of the select list
444            * The column names are deduplicated across the entire select list and only the column names (other dups are allowed)
445            """
446            table_names = [
447                table.alias_or_name
448                for table in get_tables_from_expression_with_join(join_expression)
449            ]
450            potential_ctes = [
451                cte
452                for cte in join_expression.ctes
453                if cte.alias_or_name in table_names
454                and cte.alias_or_name != other_df.latest_cte_name
455            ]
456            # Determine the table to reference for the left side of the join by checking each of the left side
457            # tables and see if they have the column being referenced.
458            join_column_pairs = []
459            for join_column in join_columns:
460                num_matching_ctes = 0
461                for cte in potential_ctes:
462                    if join_column.alias_or_name in cte.this.named_selects:
463                        left_column = join_column.copy().set_table_name(cte.alias_or_name)
464                        right_column = join_column.copy().set_table_name(other_df.latest_cte_name)
465                        join_column_pairs.append((left_column, right_column))
466                        num_matching_ctes += 1
467                if num_matching_ctes > 1:
468                    raise ValueError(
469                        f"Column {join_column.alias_or_name} is ambiguous. Please specify the table name."
470                    )
471                elif num_matching_ctes == 0:
472                    raise ValueError(
473                        f"Column {join_column.alias_or_name} does not exist in any of the tables."
474                    )
475            join_clause = functools.reduce(
476                lambda x, y: x & y,
477                [left_column == right_column for left_column, right_column in join_column_pairs],
478            )
479            join_column_names = [left_col.alias_or_name for left_col, _ in join_column_pairs]
480            # To match spark behavior only the join clause gets deduplicated and it gets put in the front of the column list
481            select_column_names = [
482                column.alias_or_name
483                if not isinstance(column.expression.this, exp.Star)
484                else column.sql()
485                for column in self_columns + other_columns
486            ]
487            select_column_names = [
488                column_name
489                for column_name in select_column_names
490                if column_name not in join_column_names
491            ]
492            select_column_names = join_column_names + select_column_names
493        else:
494            """
495            Unique characteristics of join on expressions:
496            * There is no deduplication of the results.
497            * The left join dataframe columns go first and right come after. No sort preference is given to join columns
498            """
499            join_columns = self._ensure_and_normalize_cols(join_columns, join_expression)
500            if len(join_columns) > 1:
501                join_columns = [functools.reduce(lambda x, y: x & y, join_columns)]
502            join_clause = join_columns[0]
503            select_column_names = [column.alias_or_name for column in self_columns + other_columns]
504
505        # Update the on expression with the actual join clause to replace the dummy one from before
506        join_expression.args["joins"][-1].set("on", join_clause.expression)
507        new_df = self.copy(expression=join_expression)
508        new_df.pending_join_hints.extend(self.pending_join_hints)
509        new_df.pending_hints.extend(other_df.pending_hints)
510        new_df = new_df.select.__wrapped__(new_df, *select_column_names)
511        return new_df
512
513    @operation(Operation.ORDER_BY)
514    def orderBy(
515        self,
516        *cols: t.Union[str, Column],
517        ascending: t.Optional[t.Union[t.Any, t.List[t.Any]]] = None,
518    ) -> DataFrame:
519        """
520        This implementation lets any ordered columns take priority over whatever is provided in `ascending`. Spark
521        has irregular behavior and can result in runtime errors. Users shouldn't be mixing the two anyways so this
522        is unlikely to come up.
523        """
524        columns = self._ensure_and_normalize_cols(cols)
525        pre_ordered_col_indexes = [
526            x
527            for x in [
528                i if isinstance(col.expression, exp.Ordered) else None
529                for i, col in enumerate(columns)
530            ]
531            if x is not None
532        ]
533        if ascending is None:
534            ascending = [True] * len(columns)
535        elif not isinstance(ascending, list):
536            ascending = [ascending] * len(columns)
537        ascending = [bool(x) for i, x in enumerate(ascending)]
538        assert len(columns) == len(
539            ascending
540        ), "The length of items in ascending must equal the number of columns provided"
541        col_and_ascending = list(zip(columns, ascending))
542        order_by_columns = [
543            exp.Ordered(this=col.expression, desc=not asc)
544            if i not in pre_ordered_col_indexes
545            else columns[i].column_expression
546            for i, (col, asc) in enumerate(col_and_ascending)
547        ]
548        return self.copy(expression=self.expression.order_by(*order_by_columns))
549
550    sort = orderBy
551
552    @operation(Operation.FROM)
553    def union(self, other: DataFrame) -> DataFrame:
554        return self._set_operation(exp.Union, other, False)
555
556    unionAll = union
557
558    @operation(Operation.FROM)
559    def unionByName(self, other: DataFrame, allowMissingColumns: bool = False):
560        l_columns = self.columns
561        r_columns = other.columns
562        if not allowMissingColumns:
563            l_expressions = l_columns
564            r_expressions = l_columns
565        else:
566            l_expressions = []
567            r_expressions = []
568            r_columns_unused = copy(r_columns)
569            for l_column in l_columns:
570                l_expressions.append(l_column)
571                if l_column in r_columns:
572                    r_expressions.append(l_column)
573                    r_columns_unused.remove(l_column)
574                else:
575                    r_expressions.append(exp.alias_(exp.Null(), l_column, copy=False))
576            for r_column in r_columns_unused:
577                l_expressions.append(exp.alias_(exp.Null(), r_column, copy=False))
578                r_expressions.append(r_column)
579        r_df = (
580            other.copy()._convert_leaf_to_cte().select(*self._ensure_list_of_columns(r_expressions))
581        )
582        l_df = self.copy()
583        if allowMissingColumns:
584            l_df = l_df._convert_leaf_to_cte().select(*self._ensure_list_of_columns(l_expressions))
585        return l_df._set_operation(exp.Union, r_df, False)
586
587    @operation(Operation.FROM)
588    def intersect(self, other: DataFrame) -> DataFrame:
589        return self._set_operation(exp.Intersect, other, True)
590
591    @operation(Operation.FROM)
592    def intersectAll(self, other: DataFrame) -> DataFrame:
593        return self._set_operation(exp.Intersect, other, False)
594
595    @operation(Operation.FROM)
596    def exceptAll(self, other: DataFrame) -> DataFrame:
597        return self._set_operation(exp.Except, other, False)
598
599    @operation(Operation.SELECT)
600    def distinct(self) -> DataFrame:
601        return self.copy(expression=self.expression.distinct())
602
603    @operation(Operation.SELECT)
604    def dropDuplicates(self, subset: t.Optional[t.List[str]] = None):
605        if not subset:
606            return self.distinct()
607        column_names = ensure_list(subset)
608        window = Window.partitionBy(*column_names).orderBy(*column_names)
609        return (
610            self.copy()
611            .withColumn("row_num", F.row_number().over(window))
612            .where(F.col("row_num") == F.lit(1))
613            .drop("row_num")
614        )
615
616    @operation(Operation.FROM)
617    def dropna(
618        self,
619        how: str = "any",
620        thresh: t.Optional[int] = None,
621        subset: t.Optional[t.Union[str, t.Tuple[str, ...], t.List[str]]] = None,
622    ) -> DataFrame:
623        minimum_non_null = thresh or 0  # will be determined later if thresh is null
624        new_df = self.copy()
625        all_columns = self._get_outer_select_columns(new_df.expression)
626        if subset:
627            null_check_columns = self._ensure_and_normalize_cols(subset)
628        else:
629            null_check_columns = all_columns
630        if thresh is None:
631            minimum_num_nulls = 1 if how == "any" else len(null_check_columns)
632        else:
633            minimum_num_nulls = len(null_check_columns) - minimum_non_null + 1
634        if minimum_num_nulls > len(null_check_columns):
635            raise RuntimeError(
636                f"The minimum num nulls for dropna must be less than or equal to the number of columns. "
637                f"Minimum num nulls: {minimum_num_nulls}, Num Columns: {len(null_check_columns)}"
638            )
639        if_null_checks = [
640            F.when(column.isNull(), F.lit(1)).otherwise(F.lit(0)) for column in null_check_columns
641        ]
642        nulls_added_together = functools.reduce(lambda x, y: x + y, if_null_checks)
643        num_nulls = nulls_added_together.alias("num_nulls")
644        new_df = new_df.select(num_nulls, append=True)
645        filtered_df = new_df.where(F.col("num_nulls") < F.lit(minimum_num_nulls))
646        final_df = filtered_df.select(*all_columns)
647        return final_df
648
649    @operation(Operation.FROM)
650    def fillna(
651        self,
652        value: t.Union[ColumnLiterals],
653        subset: t.Optional[t.Union[str, t.Tuple[str, ...], t.List[str]]] = None,
654    ) -> DataFrame:
655        """
656        Functionality Difference: If you provide a value to replace a null and that type conflicts
657        with the type of the column then PySpark will just ignore your replacement.
658        This will try to cast them to be the same in some cases. So they won't always match.
659        Best to not mix types so make sure replacement is the same type as the column
660
661        Possibility for improvement: Use `typeof` function to get the type of the column
662        and check if it matches the type of the value provided. If not then make it null.
663        """
664        from sqlglot.dataframe.sql.functions import lit
665
666        values = None
667        columns = None
668        new_df = self.copy()
669        all_columns = self._get_outer_select_columns(new_df.expression)
670        all_column_mapping = {column.alias_or_name: column for column in all_columns}
671        if isinstance(value, dict):
672            values = list(value.values())
673            columns = self._ensure_and_normalize_cols(list(value))
674        if not columns:
675            columns = self._ensure_and_normalize_cols(subset) if subset else all_columns
676        if not values:
677            values = [value] * len(columns)
678        value_columns = [lit(value) for value in values]
679
680        null_replacement_mapping = {
681            column.alias_or_name: (
682                F.when(column.isNull(), value).otherwise(column).alias(column.alias_or_name)
683            )
684            for column, value in zip(columns, value_columns)
685        }
686        null_replacement_mapping = {**all_column_mapping, **null_replacement_mapping}
687        null_replacement_columns = [
688            null_replacement_mapping[column.alias_or_name] for column in all_columns
689        ]
690        new_df = new_df.select(*null_replacement_columns)
691        return new_df
692
693    @operation(Operation.FROM)
694    def replace(
695        self,
696        to_replace: t.Union[bool, int, float, str, t.List, t.Dict],
697        value: t.Optional[t.Union[bool, int, float, str, t.List]] = None,
698        subset: t.Optional[t.Collection[ColumnOrName] | ColumnOrName] = None,
699    ) -> DataFrame:
700        from sqlglot.dataframe.sql.functions import lit
701
702        old_values = None
703        new_df = self.copy()
704        all_columns = self._get_outer_select_columns(new_df.expression)
705        all_column_mapping = {column.alias_or_name: column for column in all_columns}
706
707        columns = self._ensure_and_normalize_cols(subset) if subset else all_columns
708        if isinstance(to_replace, dict):
709            old_values = list(to_replace)
710            new_values = list(to_replace.values())
711        elif not old_values and isinstance(to_replace, list):
712            assert isinstance(value, list), "value must be a list since the replacements are a list"
713            assert len(to_replace) == len(
714                value
715            ), "the replacements and values must be the same length"
716            old_values = to_replace
717            new_values = value
718        else:
719            old_values = [to_replace] * len(columns)
720            new_values = [value] * len(columns)
721        old_values = [lit(value) for value in old_values]
722        new_values = [lit(value) for value in new_values]
723
724        replacement_mapping = {}
725        for column in columns:
726            expression = Column(None)
727            for i, (old_value, new_value) in enumerate(zip(old_values, new_values)):
728                if i == 0:
729                    expression = F.when(column == old_value, new_value)
730                else:
731                    expression = expression.when(column == old_value, new_value)  # type: ignore
732            replacement_mapping[column.alias_or_name] = expression.otherwise(column).alias(
733                column.expression.alias_or_name
734            )
735
736        replacement_mapping = {**all_column_mapping, **replacement_mapping}
737        replacement_columns = [replacement_mapping[column.alias_or_name] for column in all_columns]
738        new_df = new_df.select(*replacement_columns)
739        return new_df
740
741    @operation(Operation.SELECT)
742    def withColumn(self, colName: str, col: Column) -> DataFrame:
743        col = self._ensure_and_normalize_col(col)
744        existing_col_names = self.expression.named_selects
745        existing_col_index = (
746            existing_col_names.index(colName) if colName in existing_col_names else None
747        )
748        if existing_col_index:
749            expression = self.expression.copy()
750            expression.expressions[existing_col_index] = col.expression
751            return self.copy(expression=expression)
752        return self.copy().select(col.alias(colName), append=True)
753
754    @operation(Operation.SELECT)
755    def withColumnRenamed(self, existing: str, new: str):
756        expression = self.expression.copy()
757        existing_columns = [
758            expression
759            for expression in expression.expressions
760            if expression.alias_or_name == existing
761        ]
762        if not existing_columns:
763            raise ValueError("Tried to rename a column that doesn't exist")
764        for existing_column in existing_columns:
765            if isinstance(existing_column, exp.Column):
766                existing_column.replace(exp.alias_(existing_column, new))
767            else:
768                existing_column.set("alias", exp.to_identifier(new))
769        return self.copy(expression=expression)
770
771    @operation(Operation.SELECT)
772    def drop(self, *cols: t.Union[str, Column]) -> DataFrame:
773        all_columns = self._get_outer_select_columns(self.expression)
774        drop_cols = self._ensure_and_normalize_cols(cols)
775        new_columns = [
776            col
777            for col in all_columns
778            if col.alias_or_name not in [drop_column.alias_or_name for drop_column in drop_cols]
779        ]
780        return self.copy().select(*new_columns, append=False)
781
782    @operation(Operation.LIMIT)
783    def limit(self, num: int) -> DataFrame:
784        return self.copy(expression=self.expression.limit(num))
785
786    @operation(Operation.NO_OP)
787    def hint(self, name: str, *parameters: t.Optional[t.Union[str, int]]) -> DataFrame:
788        parameter_list = ensure_list(parameters)
789        parameter_columns = (
790            self._ensure_list_of_columns(parameter_list)
791            if parameters
792            else Column.ensure_cols([self.sequence_id])
793        )
794        return self._hint(name, parameter_columns)
795
796    @operation(Operation.NO_OP)
797    def repartition(
798        self, numPartitions: t.Union[int, ColumnOrName], *cols: ColumnOrName
799    ) -> DataFrame:
800        num_partition_cols = self._ensure_list_of_columns(numPartitions)
801        columns = self._ensure_and_normalize_cols(cols)
802        args = num_partition_cols + columns
803        return self._hint("repartition", args)
804
805    @operation(Operation.NO_OP)
806    def coalesce(self, numPartitions: int) -> DataFrame:
807        num_partitions = Column.ensure_cols([numPartitions])
808        return self._hint("coalesce", num_partitions)
809
810    @operation(Operation.NO_OP)
811    def cache(self) -> DataFrame:
812        return self._cache(storage_level="MEMORY_AND_DISK")
813
814    @operation(Operation.NO_OP)
815    def persist(self, storageLevel: str = "MEMORY_AND_DISK_SER") -> DataFrame:
816        """
817        Storage Level Options: https://spark.apache.org/docs/3.0.0-preview/sql-ref-syntax-aux-cache-cache-table.html
818        """
819        return self._cache(storageLevel)