sqlframe 3.16.0__py3-none-any.whl → 3.17.0__py3-none-any.whl

sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.16.0'
-__version_tuple__ = version_tuple = (3, 16, 0)
+__version__ = version = '3.17.0'
+__version_tuple__ = version_tuple = (3, 17, 0)

sqlframe/base/column.py

@@ -291,6 +291,7 @@ class Column:
             this=self.column_expression,
             alias=alias.this if isinstance(alias, exp.Column) else alias,
         )
+        new_expression._meta = {"display_name": name, **(new_expression._meta or {})}
         return Column(new_expression)
 
     def asc(self) -> Column:

sqlframe/base/dataframe.py

@@ -233,6 +233,7 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         last_op: Operation = Operation.INIT,
         pending_hints: t.Optional[t.List[exp.Expression]] = None,
         output_expression_container: t.Optional[OutputExpressionContainer] = None,
+        display_name_mapping: t.Optional[t.Dict[str, str]] = None,
         **kwargs,
     ):
         self.session = session
@@ -246,6 +247,7 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         self.pending_hints = pending_hints or []
         self.output_expression_container = output_expression_container or exp.Select()
         self.temp_views: t.List[exp.Select] = []
+        self.display_name_mapping = display_name_mapping or {}
 
     def __getattr__(self, column_name: str) -> Column:
         return self[column_name]
@@ -385,13 +387,14 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         return Column.ensure_cols(ensure_list(cols))  # type: ignore
 
     def _ensure_and_normalize_cols(
-        self, cols, expression: t.Optional[exp.Select] = None
+        self, cols, expression: t.Optional[exp.Select] = None, skip_star_expansion: bool = False
     ) -> t.List[Column]:
         from sqlframe.base.normalize import normalize
 
         cols = self._ensure_list_of_columns(cols)
         normalize(self.session, expression or self.expression, cols)
-        cols = list(flatten([self._expand_star(col) for col in cols]))
+        if not skip_star_expansion:
+            cols = list(flatten([self._expand_star(col) for col in cols]))
         self._resolve_ambiguous_columns(cols)
         return cols
 
@@ -592,6 +595,23 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             )
         return [col]
 
+    def _update_display_name_mapping(
+        self, normalized_columns: t.List[Column], user_input: t.Iterable[ColumnOrName]
+    ) -> None:
+        from sqlframe.base.column import Column
+
+        normalized_aliases = [x.alias_or_name for x in normalized_columns]
+        user_display_names = [
+            x.expression.meta.get("display_name") if isinstance(x, Column) else x
+            for x in user_input
+        ]
+        zipped = {
+            k: v
+            for k, v in dict(zip(normalized_aliases, user_display_names)).items()
+            if v is not None
+        }
+        self.display_name_mapping.update(zipped)
+
     def _get_expressions(
         self,
         optimize: bool = True,
@@ -611,6 +631,16 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             select_expression = select_expression.transform(
                 replace_id_value, replacement_mapping
             ).assert_is(exp.Select)
+            for index, column in enumerate(select_expression.expressions):
+                column_name = quote_preserving_alias_or_name(column)
+                if column_name in self.display_name_mapping:
+                    display_name_identifier = exp.to_identifier(
+                        self.display_name_mapping[column_name], quoted=True
+                    )
+                    display_name_identifier._meta = {"case_sensitive": True, **(column._meta or {})}
+                    select_expression.expressions[index] = exp.alias_(
+                        column.unalias(), display_name_identifier, quoted=True
+                    )
             if optimize:
                 select_expression = t.cast(
                     exp.Select,
@@ -803,6 +833,17 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         if isinstance(cols[0], list):
             cols = cols[0]  # type: ignore
         columns = self._ensure_and_normalize_cols(cols)
+        if "skip_update_display_name_mapping" not in kwargs:
+            unexpanded_columns = self._ensure_and_normalize_cols(cols, skip_star_expansion=True)
+            user_cols = list(cols)
+            star_columns = []
+            for index, user_col in enumerate(cols):
+                if "*" in (user_col if isinstance(user_col, str) else user_col.alias_or_name):
+                    star_columns.append(index)
+            for index in star_columns:
+                unexpanded_columns.pop(index)
+                user_cols.pop(index)
+            self._update_display_name_mapping(unexpanded_columns, user_cols)
         kwargs["append"] = kwargs.get("append", False)
         # If an expression is `CAST(x AS DATETYPE)` then we want to alias so that `x` is the result column name
         columns = [
@@ -852,6 +893,7 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
     @operation(Operation.SELECT)
     def agg(self, *exprs, **kwargs) -> Self:
         cols = self._ensure_and_normalize_cols(exprs)
+        self._update_display_name_mapping(cols, exprs)
         return self.groupBy().agg(*cols)
 
     @operation(Operation.FROM)
@@ -1051,7 +1093,9 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         new_df = self.copy(expression=join_expression)
         new_df.pending_join_hints.extend(self.pending_join_hints)
         new_df.pending_hints.extend(other_df.pending_hints)
-        new_df = new_df.select.__wrapped__(new_df, *select_column_names)  # type: ignore
+        new_df = new_df.select.__wrapped__(  # type: ignore
+            new_df, *select_column_names, skip_update_display_name_mapping=True
+        )
         return new_df
 
     @operation(Operation.ORDER_BY)
@@ -1441,20 +1485,18 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
     def withColumnRenamed(self, existing: str, new: str) -> Self:
         expression = self.expression.copy()
         existing = self.session._normalize_string(existing)
-        new = self.session._normalize_string(new)
-        existing_columns = [
-            expression
-            for expression in expression.expressions
-            if expression.alias_or_name == existing
-        ]
-        if not existing_columns:
+        columns = self._get_outer_select_columns(expression)
+        results = []
+        found_match = False
+        for column in columns:
+            if column.alias_or_name == existing:
+                column = column.alias(new)
+                self._update_display_name_mapping([column], [new])
+                found_match = True
+            results.append(column)
+        if not found_match:
             raise ValueError("Tried to rename a column that doesn't exist")
-        for existing_column in existing_columns:
-            if isinstance(existing_column, exp.Column):
-                existing_column.replace(exp.alias_(existing_column, new))
-            else:
-                existing_column.set("alias", exp.to_identifier(new))
-        return self.copy(expression=expression)
+        return self.select.__wrapped__(self, *results, skip_update_display_name_mapping=True)  # type: ignore
 
     @operation(Operation.SELECT)
     def withColumns(self, *colsMap: t.Dict[str, Column]) -> Self:
@@ -1495,23 +1537,27 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         if len(colsMap) != 1:
             raise ValueError("Only a single map is supported")
         col_map = {
-            self._ensure_and_normalize_col(k).alias_or_name: self._ensure_and_normalize_col(v)
+            self._ensure_and_normalize_col(k): (self._ensure_and_normalize_col(v), k)
             for k, v in colsMap[0].items()
         }
         existing_cols = self._get_outer_select_columns(self.expression)
         existing_col_names = [x.alias_or_name for x in existing_cols]
         select_columns = existing_cols
-        for column_name, col_value in col_map.items():
+        for col, (col_value, display_name) in col_map.items():
+            column_name = col.alias_or_name
             existing_col_index = (
                 existing_col_names.index(column_name) if column_name in existing_col_names else None
             )
             if existing_col_index is not None:
                 select_columns[existing_col_index] = col_value.alias(  # type: ignore
-                    column_name
-                ).expression
+                    display_name
+                )
             else:
-                select_columns.append(col_value.alias(column_name))
-        return self.select.__wrapped__(self, *select_columns)  # type: ignore
+                select_columns.append(col_value.alias(display_name))
+        self._update_display_name_mapping(
+            [col for col in col_map], [name for _, name in col_map.values()]
+        )
+        return self.select.__wrapped__(self, *select_columns, skip_update_display_name_mapping=True)  # type: ignore
 
     @operation(Operation.SELECT)
     def drop(self, *cols: t.Union[str, Column]) -> Self:

sqlframe/base/functions.py

@@ -39,11 +39,19 @@ def col(column_name: t.Union[ColumnOrName, t.Any]) -> Column:
 
     dialect = _BaseSession().input_dialect
     if isinstance(column_name, str):
-        return Column(
-            expression.to_column(column_name, dialect=dialect).transform(
-                dialect.normalize_identifier
-            )
+        col_expression = expression.to_column(column_name, dialect=dialect).transform(
+            dialect.normalize_identifier
         )
+        case_sensitive_expression = expression.to_column(column_name, dialect=dialect)
+        if not isinstance(
+            case_sensitive_expression, (expression.Star, expression.Literal, expression.Null)
+        ):
+            col_expression._meta = {
+                "display_name": case_sensitive_expression.this.this,
+                **(col_expression._meta or {}),
+            }
+
+        return Column(col_expression)
     return Column(column_name)
 
 

sqlframe/base/session.py

@@ -507,9 +507,14 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, TABLE, CONN, UDF_REGIS
         result = self._cur.fetchall()
         if not self._cur.description:
             return []
+        case_sensitive_cols = []
+        for col in self._cur.description:
+            col_id = exp.parse_identifier(col[0], dialect=self.execution_dialect)
+            col_id._meta = {"case_sensitive": True, **(col_id._meta or {})}
+            case_sensitive_cols.append(col_id)
         columns = [
-            normalize_string(x[0], from_dialect="execution", to_dialect="output", is_column=True)
-            for x in self._cur.description
+            normalize_string(x, from_dialect="execution", to_dialect="output")
+            for x in case_sensitive_cols
         ]
         return [self._to_row(columns, row) for row in result]
 

sqlframe/spark/session.py

@@ -79,17 +79,18 @@ class SparkSession(
         if skip_rows:
             return []
         assert self._last_df is not None
-        return [
-            Row(
-                **{
-                    normalize_string(
-                        k, from_dialect="execution", to_dialect="output", is_column=True
-                    ): v
-                    for k, v in row.asDict().items()
-                }
-            )
-            for row in self._last_df.collect()
-        ]
+        results = []
+        for row in self._last_df.collect():
+            rows_normalized = {}
+            for k, v in row.asDict().items():
+                col_id = exp.parse_identifier(k, dialect=self.execution_dialect)
+                col_id._meta = {"case_sensitive": True, **(col_id._meta or {})}
+                col_name = normalize_string(
+                    col_id, from_dialect="execution", to_dialect="output", is_column=True
+                )
+                rows_normalized[col_name] = v
+            results.append(Row(**rows_normalized))
+        return results
 
     def _execute(self, sql: str) -> None:
         self._last_df = self.spark_session.sql(sql)

{sqlframe-3.16.0.dist-info → sqlframe-3.17.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.16.0
+Version: 3.17.0
 Summary: Turning PySpark Into a Universal DataFrame API
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman

{sqlframe-3.16.0.dist-info → sqlframe-3.17.0.dist-info}/RECORD

@@ -1,19 +1,19 @@
 sqlframe/__init__.py,sha256=wfqm98eLoLid9oV_FzzpG5loKC6LxOhj2lXpfN7SARo,3138
-sqlframe/_version.py,sha256=CtTis8a_OeN0EsLFoVgtqX-ARqHjuin2ATomgRROY1Y,413
+sqlframe/_version.py,sha256=KdbrTz1mygb-tPODYZu2E4Sk2KYmeTUCHVpQLRpXAXo,413
 sqlframe/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlframe/base/_typing.py,sha256=b2clI5HI1zEZKB_3Msx3FeAJQyft44ubUifJwQRVXyQ,1298
 sqlframe/base/catalog.py,sha256=SzFQalTWdhWzxUY-4ut1f9TfOECp_JmJEgNPfrRKCe0,38457
-sqlframe/base/column.py,sha256=wRghgieYAA51aw4WuFQWOvl0TFOToZbBhBuIamEzxx4,18011
-sqlframe/base/dataframe.py,sha256=KKBwtn73xNGt2gRwUB8Vri7Ee6_ivP5a_qij4Eq96zE,76622
+sqlframe/base/column.py,sha256=oHVwkSWABO3ZlAbgBShsxSSlgbI06BOup5XJrRhgqJI,18097
+sqlframe/base/dataframe.py,sha256=SQtwoQKpq-12WXuplOPN21fXQPvjF_D9WLcPPFA12Zs,78973
 sqlframe/base/decorators.py,sha256=ms-CvDOIW3T8IVB9VqDmLwAiaEsqXLYRXEqVQaxktiM,1890
 sqlframe/base/exceptions.py,sha256=9Uwvqn2eAkDpqm4BrRgbL61qM-GMCbJEMAW8otxO46s,370
 sqlframe/base/function_alternatives.py,sha256=NV31IaEhVYmfUSWetAEFISAvLzs2DxQ7bp-iMNgj0hQ,53786
-sqlframe/base/functions.py,sha256=o8zwbS8zCsyNe5arcb6dbAGBL8a1tH99rGyRimwzzUk,220614
+sqlframe/base/functions.py,sha256=1LHxazgC9tZ_GzyWNsjU945SRnAsQjUH2easMJLU3h4,221012
 sqlframe/base/group.py,sha256=fsyG5990_Pd7gFPjTFrH9IEoAquL_wEkVpIlBAIkZJU,4091
 sqlframe/base/normalize.py,sha256=nXAJ5CwxVf4DV0GsH-q1w0p8gmjSMlv96k_ez1eVul8,3880
 sqlframe/base/operations.py,sha256=xSPw74e59wYvNd6U1AlwziNCTG6Aftrbl4SybN9u9VE,3450
 sqlframe/base/readerwriter.py,sha256=w8926cqIrXF7NGHiINw5UHzP_3xpjsqbijTBTzycBRM,26605
-sqlframe/base/session.py,sha256=s9M9_nbtOQQgLyEBZs-ijkMeHkYkILHfBc8JsU2SLmU,26369
+sqlframe/base/session.py,sha256=0eBE_HYEb3npyyOGM7zS_VR8WgzvfgVI-PFLCK9Hy0M,26628
 sqlframe/base/table.py,sha256=rCeh1W5SWbtEVfkLAUiexzrZwNgmZeptLEmLcM1ABkE,6961
 sqlframe/base/transforms.py,sha256=y0j3SGDz3XCmNGrvassk1S-owllUWfkHyMgZlY6SFO4,467
 sqlframe/base/types.py,sha256=iBNk9bpFtb2NBIogYS8i7OlQZMRvpR6XxqzBebsjQDU,12280
@@ -110,7 +110,7 @@ sqlframe/spark/functions.py,sha256=MYCgHsjRQWylT-rezWRBuLV6BivcaVarbaQtP4T0toQ,3
 sqlframe/spark/functions.pyi,sha256=GyOdUzv2Z7Qt99JAKEPKgV2t2Rn274OuqwAfcoAXlN0,24259
 sqlframe/spark/group.py,sha256=MrvV_v-YkBc6T1zz882WrEqtWjlooWIyHBCmTQg3fCA,379
 sqlframe/spark/readwriter.py,sha256=zXZcCPWpQMMN90wdIx8AD4Y5tWBcpRSL4-yKX2aZyik,874
-sqlframe/spark/session.py,sha256=1kgi69uztJxJ6bJpgkpRxllOYgVrizKXA5iT88-jWKA,5421
+sqlframe/spark/session.py,sha256=9qG-J5L8gmiy384GZFSBT2tHF8akqqJNij23Y3pheMs,5651
 sqlframe/spark/table.py,sha256=puWV8h_CqA64zwpzq0ydY9LoygMAvprkODyxyzZeF9M,186
 sqlframe/spark/types.py,sha256=KwNyuXIo-2xVVd4bZED3YrQOobKCtemlxGrJL7DrTC8,34
 sqlframe/spark/udf.py,sha256=owB8NDaGVkUQ0WGm7SZt2t9zfvLFCfi0W48QiPfgjck,1153
@@ -129,8 +129,8 @@ sqlframe/standalone/udf.py,sha256=azmgtUjHNIPs0WMVNId05SHwiYn41MKVBhKXsQJ5dmY,27
 sqlframe/standalone/window.py,sha256=6GKPzuxeSapJakBaKBeT9VpED1ACdjggDv9JRILDyV0,35
 sqlframe/testing/__init__.py,sha256=VVCosQhitU74A3NnE52O4mNtGZONapuEXcc20QmSlnQ,132
 sqlframe/testing/utils.py,sha256=PFsGZpwNUE_4-g_f43_vstTqsK0AQ2lBneb5Eb6NkFo,13008
-sqlframe-3.16.0.dist-info/LICENSE,sha256=VZu79YgW780qxaFJMr0t5ZgbOYEh04xWoxaWOaqIGWk,1068
-sqlframe-3.16.0.dist-info/METADATA,sha256=SMpgyXmxbVMqeeRuByF19qKm9iLDYubcniTCYBUmyNo,8970
-sqlframe-3.16.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-sqlframe-3.16.0.dist-info/top_level.txt,sha256=T0_RpoygaZSF6heeWwIDQgaP0varUdSK1pzjeJZRjM8,9
-sqlframe-3.16.0.dist-info/RECORD,,
+sqlframe-3.17.0.dist-info/LICENSE,sha256=VZu79YgW780qxaFJMr0t5ZgbOYEh04xWoxaWOaqIGWk,1068
+sqlframe-3.17.0.dist-info/METADATA,sha256=K8kfOT5t6cEBs4YsIK76QCFBPW2NEcDcsPMkEhWCLUI,8970
+sqlframe-3.17.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+sqlframe-3.17.0.dist-info/top_level.txt,sha256=T0_RpoygaZSF6heeWwIDQgaP0varUdSK1pzjeJZRjM8,9
+sqlframe-3.17.0.dist-info/RECORD,,