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                expressions=[
 92                    exp.Values(
 93                        expressions=data_expressions,
 94                        alias=exp.TableAlias(
 95                            this=exp.to_identifier(self._auto_incrementing_name),
 96                            columns=[exp.to_identifier(col_name) for col_name in column_mapping],
 97                        ),
 98                    ),
 99                ],
100            ),
101        }
102
103        sel_expression = exp.Select(**select_kwargs)
104        return DataFrame(self, sel_expression)
105
106    def sql(self, sqlQuery: str) -> DataFrame:
107        expression = sqlglot.parse_one(sqlQuery, read="spark")
108        if isinstance(expression, exp.Select):
109            df = DataFrame(self, expression)
110            df = df._convert_leaf_to_cte()
111        elif isinstance(expression, (exp.Create, exp.Insert)):
112            select_expression = expression.expression.copy()
113            if isinstance(expression, exp.Insert):
114                select_expression.set("with", expression.args.get("with"))
115                expression.set("with", None)
116            del expression.args["expression"]
117            df = DataFrame(self, select_expression, output_expression_container=expression)  # type: ignore
118            df = df._convert_leaf_to_cte()
119        else:
120            raise ValueError(
121                "Unknown expression type provided in the SQL. Please create an issue with the SQL."
122            )
123        return df
124
125    @property
126    def _auto_incrementing_name(self) -> str:
127        name = f"a{self.incrementing_id}"
128        self.incrementing_id += 1
129        return name
130
131    @property
132    def _random_name(self) -> str:
133        return "r" + uuid.uuid4().hex
134
135    @property
136    def _random_branch_id(self) -> str:
137        id = self._random_id
138        self.known_branch_ids.add(id)
139        return id
140
141    @property
142    def _random_sequence_id(self):
143        id = self._random_id
144        self.known_sequence_ids.add(id)
145        return id
146
147    @property
148    def _random_id(self) -> str:
149        id = self._random_name
150        self.known_ids.add(id)
151        return id
152
153    @property
154    def _join_hint_names(self) -> t.Set[str]:
155        return {"BROADCAST", "MERGE", "SHUFFLE_HASH", "SHUFFLE_REPLICATE_NL"}
156
157    def _add_alias_to_mapping(self, name: str, sequence_id: str):
158        self.name_to_sequence_id_mapping[name].append(sequence_id)

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