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)

 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.spark._random_name
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):
151        cols = self._ensure_list_of_columns(cols)
152        normalize(self.spark, 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.Select):
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 = optimize_func(select_expression, identify="always")
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                )
321                cache_storage_level = select_expression.args["cache_storage_level"]
322                options = [
323                    exp.Literal.string("storageLevel"),
324                    exp.Literal.string(cache_storage_level),
325                ]
326                expression = exp.Cache(
327                    this=cache_table, expression=select_expression, lazy=True, options=options
328                )
329                # We will drop the "view" if it exists before running the cache table
330                output_expressions.append(exp.Drop(this=cache_table, exists=True, kind="VIEW"))
331            elif expression_type == exp.Create:
332                expression = df.output_expression_container.copy()
333                expression.set("expression", select_expression)
334            elif expression_type == exp.Insert:
335                expression = df.output_expression_container.copy()
336                select_without_ctes = select_expression.copy()
337                select_without_ctes.set("with", None)
338                expression.set("expression", select_without_ctes)
339                if select_expression.ctes:
340                    expression.set("with", exp.With(expressions=select_expression.ctes))
341            elif expression_type == exp.Select:
342                expression = select_expression
343            else:
344                raise ValueError(f"Invalid expression type: {expression_type}")
345            output_expressions.append(expression)
346
347        return [
348            expression.sql(**{"dialect": dialect, **kwargs}) for expression in output_expressions
349        ]
350
351    def copy(self, **kwargs) -> DataFrame:
352        return DataFrame(**object_to_dict(self, **kwargs))
353
354    @operation(Operation.SELECT)
355    def select(self, *cols, **kwargs) -> DataFrame:
356        cols = self._ensure_and_normalize_cols(cols)
357        kwargs["append"] = kwargs.get("append", False)
358        if self.expression.args.get("joins"):
359            ambiguous_cols = [col for col in cols if not col.column_expression.table]
360            if ambiguous_cols:
361                join_table_identifiers = [
362                    x.this for x in get_tables_from_expression_with_join(self.expression)
363                ]
364                cte_names_in_join = [x.this for x in join_table_identifiers]
365                for ambiguous_col in ambiguous_cols:
366                    ctes_with_column = [
367                        cte
368                        for cte in self.expression.ctes
369                        if cte.alias_or_name in cte_names_in_join
370                        and ambiguous_col.alias_or_name in cte.this.named_selects
371                    ]
372                    # If the select column does not specify a table and there is a join
373                    # then we assume they are referring to the left table
374                    if len(ctes_with_column) > 1:
375                        table_identifier = self.expression.args["from"].args["expressions"][0].this
376                    else:
377                        table_identifier = ctes_with_column[0].args["alias"].this
378                    ambiguous_col.expression.set("table", table_identifier)
379        return self.copy(
380            expression=self.expression.select(*[x.expression for x in cols], **kwargs), **kwargs
381        )
382
383    @operation(Operation.NO_OP)
384    def alias(self, name: str, **kwargs) -> DataFrame:
385        new_sequence_id = self.spark._random_sequence_id
386        df = self.copy()
387        for join_hint in df.pending_join_hints:
388            for expression in join_hint.expressions:
389                if expression.alias_or_name == self.sequence_id:
390                    expression.set("this", Column.ensure_col(new_sequence_id).expression)
391        df.spark._add_alias_to_mapping(name, new_sequence_id)
392        return df._convert_leaf_to_cte(sequence_id=new_sequence_id)
393
394    @operation(Operation.WHERE)
395    def where(self, column: t.Union[Column, bool], **kwargs) -> DataFrame:
396        col = self._ensure_and_normalize_col(column)
397        return self.copy(expression=self.expression.where(col.expression))
398
399    filter = where
400
401    @operation(Operation.GROUP_BY)
402    def groupBy(self, *cols, **kwargs) -> GroupedData:
403        columns = self._ensure_and_normalize_cols(cols)
404        return GroupedData(self, columns, self.last_op)
405
406    @operation(Operation.SELECT)
407    def agg(self, *exprs, **kwargs) -> DataFrame:
408        cols = self._ensure_and_normalize_cols(exprs)
409        return self.groupBy().agg(*cols)
410
411    @operation(Operation.FROM)
412    def join(
413        self,
414        other_df: DataFrame,
415        on: t.Union[str, t.List[str], Column, t.List[Column]],
416        how: str = "inner",
417        **kwargs,
418    ) -> DataFrame:
419        other_df = other_df._convert_leaf_to_cte()
420        pre_join_self_latest_cte_name = self.latest_cte_name
421        columns = self._ensure_and_normalize_cols(on)
422        join_type = how.replace("_", " ")
423        if isinstance(columns[0].expression, exp.Column):
424            join_columns = [
425                Column(x).set_table_name(pre_join_self_latest_cte_name) for x in columns
426            ]
427            join_clause = functools.reduce(
428                lambda x, y: x & y,
429                [
430                    col.copy().set_table_name(pre_join_self_latest_cte_name)
431                    == col.copy().set_table_name(other_df.latest_cte_name)
432                    for col in columns
433                ],
434            )
435        else:
436            if len(columns) > 1:
437                columns = [functools.reduce(lambda x, y: x & y, columns)]
438            join_clause = columns[0]
439            join_columns = [
440                Column(x).set_table_name(pre_join_self_latest_cte_name)
441                if i % 2 == 0
442                else Column(x).set_table_name(other_df.latest_cte_name)
443                for i, x in enumerate(join_clause.expression.find_all(exp.Column))
444            ]
445        self_columns = [
446            column.set_table_name(pre_join_self_latest_cte_name, copy=True)
447            for column in self._get_outer_select_columns(self)
448        ]
449        other_columns = [
450            column.set_table_name(other_df.latest_cte_name, copy=True)
451            for column in self._get_outer_select_columns(other_df)
452        ]
453        column_value_mapping = {
454            column.alias_or_name
455            if not isinstance(column.expression.this, exp.Star)
456            else column.sql(): column
457            for column in other_columns + self_columns + join_columns
458        }
459        all_columns = [
460            column_value_mapping[name]
461            for name in {x.alias_or_name: None for x in join_columns + self_columns + other_columns}
462        ]
463        new_df = self.copy(
464            expression=self.expression.join(
465                other_df.latest_cte_name, on=join_clause.expression, join_type=join_type
466            )
467        )
468        new_df.expression = new_df._add_ctes_to_expression(
469            new_df.expression, other_df.expression.ctes
470        )
471        new_df.pending_hints.extend(other_df.pending_hints)
472        new_df = new_df.select.__wrapped__(new_df, *all_columns)
473        return new_df
474
475    @operation(Operation.ORDER_BY)
476    def orderBy(
477        self,
478        *cols: t.Union[str, Column],
479        ascending: t.Optional[t.Union[t.Any, t.List[t.Any]]] = None,
480    ) -> DataFrame:
481        """
482        This implementation lets any ordered columns take priority over whatever is provided in `ascending`. Spark
483        has irregular behavior and can result in runtime errors. Users shouldn't be mixing the two anyways so this
484        is unlikely to come up.
485        """
486        columns = self._ensure_and_normalize_cols(cols)
487        pre_ordered_col_indexes = [
488            x
489            for x in [
490                i if isinstance(col.expression, exp.Ordered) else None
491                for i, col in enumerate(columns)
492            ]
493            if x is not None
494        ]
495        if ascending is None:
496            ascending = [True] * len(columns)
497        elif not isinstance(ascending, list):
498            ascending = [ascending] * len(columns)
499        ascending = [bool(x) for i, x in enumerate(ascending)]
500        assert len(columns) == len(
501            ascending
502        ), "The length of items in ascending must equal the number of columns provided"
503        col_and_ascending = list(zip(columns, ascending))
504        order_by_columns = [
505            exp.Ordered(this=col.expression, desc=not asc)
506            if i not in pre_ordered_col_indexes
507            else columns[i].column_expression
508            for i, (col, asc) in enumerate(col_and_ascending)
509        ]
510        return self.copy(expression=self.expression.order_by(*order_by_columns))
511
512    sort = orderBy
513
514    @operation(Operation.FROM)
515    def union(self, other: DataFrame) -> DataFrame:
516        return self._set_operation(exp.Union, other, False)
517
518    unionAll = union
519
520    @operation(Operation.FROM)
521    def unionByName(self, other: DataFrame, allowMissingColumns: bool = False):
522        l_columns = self.columns
523        r_columns = other.columns
524        if not allowMissingColumns:
525            l_expressions = l_columns
526            r_expressions = l_columns
527        else:
528            l_expressions = []
529            r_expressions = []
530            r_columns_unused = copy(r_columns)
531            for l_column in l_columns:
532                l_expressions.append(l_column)
533                if l_column in r_columns:
534                    r_expressions.append(l_column)
535                    r_columns_unused.remove(l_column)
536                else:
537                    r_expressions.append(exp.alias_(exp.Null(), l_column))
538            for r_column in r_columns_unused:
539                l_expressions.append(exp.alias_(exp.Null(), r_column))
540                r_expressions.append(r_column)
541        r_df = (
542            other.copy()._convert_leaf_to_cte().select(*self._ensure_list_of_columns(r_expressions))
543        )
544        l_df = self.copy()
545        if allowMissingColumns:
546            l_df = l_df._convert_leaf_to_cte().select(*self._ensure_list_of_columns(l_expressions))
547        return l_df._set_operation(exp.Union, r_df, False)
548
549    @operation(Operation.FROM)
550    def intersect(self, other: DataFrame) -> DataFrame:
551        return self._set_operation(exp.Intersect, other, True)
552
553    @operation(Operation.FROM)
554    def intersectAll(self, other: DataFrame) -> DataFrame:
555        return self._set_operation(exp.Intersect, other, False)
556
557    @operation(Operation.FROM)
558    def exceptAll(self, other: DataFrame) -> DataFrame:
559        return self._set_operation(exp.Except, other, False)
560
561    @operation(Operation.SELECT)
562    def distinct(self) -> DataFrame:
563        return self.copy(expression=self.expression.distinct())
564
565    @operation(Operation.SELECT)
566    def dropDuplicates(self, subset: t.Optional[t.List[str]] = None):
567        if not subset:
568            return self.distinct()
569        column_names = ensure_list(subset)
570        window = Window.partitionBy(*column_names).orderBy(*column_names)
571        return (
572            self.copy()
573            .withColumn("row_num", F.row_number().over(window))
574            .where(F.col("row_num") == F.lit(1))
575            .drop("row_num")
576        )
577
578    @operation(Operation.FROM)
579    def dropna(
580        self,
581        how: str = "any",
582        thresh: t.Optional[int] = None,
583        subset: t.Optional[t.Union[str, t.Tuple[str, ...], t.List[str]]] = None,
584    ) -> DataFrame:
585        minimum_non_null = thresh or 0  # will be determined later if thresh is null
586        new_df = self.copy()
587        all_columns = self._get_outer_select_columns(new_df.expression)
588        if subset:
589            null_check_columns = self._ensure_and_normalize_cols(subset)
590        else:
591            null_check_columns = all_columns
592        if thresh is None:
593            minimum_num_nulls = 1 if how == "any" else len(null_check_columns)
594        else:
595            minimum_num_nulls = len(null_check_columns) - minimum_non_null + 1
596        if minimum_num_nulls > len(null_check_columns):
597            raise RuntimeError(
598                f"The minimum num nulls for dropna must be less than or equal to the number of columns. "
599                f"Minimum num nulls: {minimum_num_nulls}, Num Columns: {len(null_check_columns)}"
600            )
601        if_null_checks = [
602            F.when(column.isNull(), F.lit(1)).otherwise(F.lit(0)) for column in null_check_columns
603        ]
604        nulls_added_together = functools.reduce(lambda x, y: x + y, if_null_checks)
605        num_nulls = nulls_added_together.alias("num_nulls")
606        new_df = new_df.select(num_nulls, append=True)
607        filtered_df = new_df.where(F.col("num_nulls") < F.lit(minimum_num_nulls))
608        final_df = filtered_df.select(*all_columns)
609        return final_df
610
611    @operation(Operation.FROM)
612    def fillna(
613        self,
614        value: t.Union[ColumnLiterals],
615        subset: t.Optional[t.Union[str, t.Tuple[str, ...], t.List[str]]] = None,
616    ) -> DataFrame:
617        """
618        Functionality Difference: If you provide a value to replace a null and that type conflicts
619        with the type of the column then PySpark will just ignore your replacement.
620        This will try to cast them to be the same in some cases. So they won't always match.
621        Best to not mix types so make sure replacement is the same type as the column
622
623        Possibility for improvement: Use `typeof` function to get the type of the column
624        and check if it matches the type of the value provided. If not then make it null.
625        """
626        from sqlglot.dataframe.sql.functions import lit
627
628        values = None
629        columns = None
630        new_df = self.copy()
631        all_columns = self._get_outer_select_columns(new_df.expression)
632        all_column_mapping = {column.alias_or_name: column for column in all_columns}
633        if isinstance(value, dict):
634            values = list(value.values())
635            columns = self._ensure_and_normalize_cols(list(value))
636        if not columns:
637            columns = self._ensure_and_normalize_cols(subset) if subset else all_columns
638        if not values:
639            values = [value] * len(columns)
640        value_columns = [lit(value) for value in values]
641
642        null_replacement_mapping = {
643            column.alias_or_name: (
644                F.when(column.isNull(), value).otherwise(column).alias(column.alias_or_name)
645            )
646            for column, value in zip(columns, value_columns)
647        }
648        null_replacement_mapping = {**all_column_mapping, **null_replacement_mapping}
649        null_replacement_columns = [
650            null_replacement_mapping[column.alias_or_name] for column in all_columns
651        ]
652        new_df = new_df.select(*null_replacement_columns)
653        return new_df
654
655    @operation(Operation.FROM)
656    def replace(
657        self,
658        to_replace: t.Union[bool, int, float, str, t.List, t.Dict],
659        value: t.Optional[t.Union[bool, int, float, str, t.List]] = None,
660        subset: t.Optional[t.Collection[ColumnOrName] | ColumnOrName] = None,
661    ) -> DataFrame:
662        from sqlglot.dataframe.sql.functions import lit
663
664        old_values = None
665        new_df = self.copy()
666        all_columns = self._get_outer_select_columns(new_df.expression)
667        all_column_mapping = {column.alias_or_name: column for column in all_columns}
668
669        columns = self._ensure_and_normalize_cols(subset) if subset else all_columns
670        if isinstance(to_replace, dict):
671            old_values = list(to_replace)
672            new_values = list(to_replace.values())
673        elif not old_values and isinstance(to_replace, list):
674            assert isinstance(value, list), "value must be a list since the replacements are a list"
675            assert len(to_replace) == len(
676                value
677            ), "the replacements and values must be the same length"
678            old_values = to_replace
679            new_values = value
680        else:
681            old_values = [to_replace] * len(columns)
682            new_values = [value] * len(columns)
683        old_values = [lit(value) for value in old_values]
684        new_values = [lit(value) for value in new_values]
685
686        replacement_mapping = {}
687        for column in columns:
688            expression = Column(None)
689            for i, (old_value, new_value) in enumerate(zip(old_values, new_values)):
690                if i == 0:
691                    expression = F.when(column == old_value, new_value)
692                else:
693                    expression = expression.when(column == old_value, new_value)  # type: ignore
694            replacement_mapping[column.alias_or_name] = expression.otherwise(column).alias(
695                column.expression.alias_or_name
696            )
697
698        replacement_mapping = {**all_column_mapping, **replacement_mapping}
699        replacement_columns = [replacement_mapping[column.alias_or_name] for column in all_columns]
700        new_df = new_df.select(*replacement_columns)
701        return new_df
702
703    @operation(Operation.SELECT)
704    def withColumn(self, colName: str, col: Column) -> DataFrame:
705        col = self._ensure_and_normalize_col(col)
706        existing_col_names = self.expression.named_selects
707        existing_col_index = (
708            existing_col_names.index(colName) if colName in existing_col_names else None
709        )
710        if existing_col_index:
711            expression = self.expression.copy()
712            expression.expressions[existing_col_index] = col.expression
713            return self.copy(expression=expression)
714        return self.copy().select(col.alias(colName), append=True)
715
716    @operation(Operation.SELECT)
717    def withColumnRenamed(self, existing: str, new: str):
718        expression = self.expression.copy()
719        existing_columns = [
720            expression
721            for expression in expression.expressions
722            if expression.alias_or_name == existing
723        ]
724        if not existing_columns:
725            raise ValueError("Tried to rename a column that doesn't exist")
726        for existing_column in existing_columns:
727            if isinstance(existing_column, exp.Column):
728                existing_column.replace(exp.alias_(existing_column.copy(), new))
729            else:
730                existing_column.set("alias", exp.to_identifier(new))
731        return self.copy(expression=expression)
732
733    @operation(Operation.SELECT)
734    def drop(self, *cols: t.Union[str, Column]) -> DataFrame:
735        all_columns = self._get_outer_select_columns(self.expression)
736        drop_cols = self._ensure_and_normalize_cols(cols)
737        new_columns = [
738            col
739            for col in all_columns
740            if col.alias_or_name not in [drop_column.alias_or_name for drop_column in drop_cols]
741        ]
742        return self.copy().select(*new_columns, append=False)
743
744    @operation(Operation.LIMIT)
745    def limit(self, num: int) -> DataFrame:
746        return self.copy(expression=self.expression.limit(num))
747
748    @operation(Operation.NO_OP)
749    def hint(self, name: str, *parameters: t.Optional[t.Union[str, int]]) -> DataFrame:
750        parameter_list = ensure_list(parameters)
751        parameter_columns = (
752            self._ensure_list_of_columns(parameter_list)
753            if parameters
754            else Column.ensure_cols([self.sequence_id])
755        )
756        return self._hint(name, parameter_columns)
757
758    @operation(Operation.NO_OP)
759    def repartition(
760        self, numPartitions: t.Union[int, ColumnOrName], *cols: ColumnOrName
761    ) -> DataFrame:
762        num_partition_cols = self._ensure_list_of_columns(numPartitions)
763        columns = self._ensure_and_normalize_cols(cols)
764        args = num_partition_cols + columns
765        return self._hint("repartition", args)
766
767    @operation(Operation.NO_OP)
768    def coalesce(self, numPartitions: int) -> DataFrame:
769        num_partitions = Column.ensure_cols([numPartitions])
770        return self._hint("coalesce", num_partitions)
771
772    @operation(Operation.NO_OP)
773    def cache(self) -> DataFrame:
774        return self._cache(storage_level="MEMORY_AND_DISK")
775
776    @operation(Operation.NO_OP)
777    def persist(self, storageLevel: str = "MEMORY_AND_DISK_SER") -> DataFrame:
778        """
779        Storage Level Options: https://spark.apache.org/docs/3.0.0-preview/sql-ref-syntax-aux-cache-cache-table.html
780        """
781        return self._cache(storageLevel)