summaryrefslogtreecommitdiffstats
path: root/sqlglot/dataframe
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2023-01-30 17:08:37 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2023-01-30 17:08:37 +0000
commitbe1cb18ea28222fca384a5459a024b7e9af5cadb (patch)
tree4698c9069380a7c30ceb51129f93f6c8662315e4 /sqlglot/dataframe
parentReleasing debian version 10.5.6-1. (diff)
downloadsqlglot-be1cb18ea28222fca384a5459a024b7e9af5cadb.tar.xz
sqlglot-be1cb18ea28222fca384a5459a024b7e9af5cadb.zip
Merging upstream version 10.5.10.
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'sqlglot/dataframe')
-rw-r--r--sqlglot/dataframe/README.md32
-rw-r--r--sqlglot/dataframe/sql/dataframe.py9
-rw-r--r--sqlglot/dataframe/sql/session.py4
3 files changed, 21 insertions, 24 deletions
diff --git a/sqlglot/dataframe/README.md b/sqlglot/dataframe/README.md
index 54d3856..02179f4 100644
--- a/sqlglot/dataframe/README.md
+++ b/sqlglot/dataframe/README.md
@@ -1,29 +1,29 @@
# PySpark DataFrame SQL Generator
-This is a drop-in replacement for the PysPark DataFrame API that will generate SQL instead of executing DataFrame operations directly. This, when combined with the transpiling support in SQLGlot, allows one to write PySpark DataFrame code and execute it on other engines like [DuckDB](https://duckdb.org/), [Presto](https://prestodb.io/), [Spark](https://spark.apache.org/), [Snowflake](https://www.snowflake.com/en/), and [BigQuery](https://cloud.google.com/bigquery/).
+This is a drop-in replacement for the PySpark DataFrame API that will generate SQL instead of executing DataFrame operations directly. This, when combined with the transpiling support in SQLGlot, allows one to write PySpark DataFrame code and execute it on other engines like [DuckDB](https://duckdb.org/), [Presto](https://prestodb.io/), [Spark](https://spark.apache.org/), [Snowflake](https://www.snowflake.com/en/), and [BigQuery](https://cloud.google.com/bigquery/).
Currently many of the common operations are covered and more functionality will be added over time. Please [open an issue](https://github.com/tobymao/sqlglot/issues) or [PR](https://github.com/tobymao/sqlglot/pulls) with your feedback or contribution to help influence what should be prioritized next and make sure your use case is properly supported.
# How to use
## Instructions
-* [Install SQLGlot](https://github.com/tobymao/sqlglot/blob/main/README.md#install) and that is all that is required to just generate SQL. [The examples](#examples) show generating SQL and then executing that SQL on a specific engine and that will require that engine's client library
-* Find/replace all `from pyspark.sql` with `from sqlglot.dataframe`
-* Prior to any `spark.read.table` or `spark.table` run `sqlglot.schema.add_table('<table_name>', <column_structure>)`
+* [Install SQLGlot](https://github.com/tobymao/sqlglot/blob/main/README.md#install) and that is all that is required to just generate SQL. [The examples](#examples) show generating SQL and then executing that SQL on a specific engine and that will require that engine's client library.
+* Find/replace all `from pyspark.sql` with `from sqlglot.dataframe`.
+* Prior to any `spark.read.table` or `spark.table` run `sqlglot.schema.add_table('<table_name>', <column_structure>)`.
* The column structure can be defined the following ways:
- * Dictionary where the keys are column names and values are string of the Spark SQL type name
- * Ex: {'cola': 'string', 'colb': 'int'}
- * PySpark DataFrame `StructType` similar to when using `createDataFrame`
+ * Dictionary where the keys are column names and values are string of the Spark SQL type name.
+ * Ex: `{'cola': 'string', 'colb': 'int'}`
+ * PySpark DataFrame `StructType` similar to when using `createDataFrame`.
* Ex: `StructType([StructField('cola', StringType()), StructField('colb', IntegerType())])`
- * A string of names and types similar to what is supported in `createDataFrame`
+ * A string of names and types similar to what is supported in `createDataFrame`.
* Ex: `cola: STRING, colb: INT`
- * [Not Recommended] A list of string column names without type
- * Ex: ['cola', 'colb']
- * The lack of types may limit functionality in future releases
- * See [Registering Custom Schema](#registering-custom-schema-class) for information on how to skip this step if the information is stored externally
-* Add `.sql(pretty=True)` to your final DataFrame command to return a list of sql statements to run that command
+ * [Not Recommended] A list of string column names without type.
+ * Ex: `['cola', 'colb']`
+ * The lack of types may limit functionality in future releases.
+ * See [Registering Custom Schema](#registering-custom-schema-class) for information on how to skip this step if the information is stored externally.
+* Add `.sql(pretty=True)` to your final DataFrame command to return a list of sql statements to run that command.
* In most cases a single SQL statement is returned. Currently the only exception is when caching DataFrames which isn't supported in other dialects.
- * Spark is the default output dialect. See [dialects](https://github.com/tobymao/sqlglot/tree/main/sqlglot/dialects) for a full list of dialects
+ * Spark is the default output dialect. See [dialects](https://github.com/tobymao/sqlglot/tree/main/sqlglot/dialects) for a full list of dialects.
* Ex: `.sql(pretty=True, dialect='bigquery')`
## Examples
@@ -51,7 +51,7 @@ df = (
print(df.sql(pretty=True)) # Spark will be the dialect used by default
```
-Output:
+
```sparksql
SELECT
`employee`.`age` AS `age`,
@@ -206,7 +206,7 @@ sql_statements = (
.createDataFrame(data, schema)
.groupBy(F.col("age"))
.agg(F.countDistinct(F.col("employee_id")).alias("num_employees"))
- .sql(dialect="bigquery")
+ .sql(dialect="spark")
)
pyspark = PySparkSession.builder.master("local[*]").getOrCreate()
diff --git a/sqlglot/dataframe/sql/dataframe.py b/sqlglot/dataframe/sql/dataframe.py
index a17bb9d..65a37f5 100644
--- a/sqlglot/dataframe/sql/dataframe.py
+++ b/sqlglot/dataframe/sql/dataframe.py
@@ -111,16 +111,13 @@ class DataFrame:
return DataFrameNaFunctions(self)
def _replace_cte_names_with_hashes(self, expression: exp.Select):
- expression = expression.copy()
- ctes = expression.ctes
replacement_mapping = {}
- for cte in ctes:
+ for cte in expression.ctes:
old_name_id = cte.args["alias"].this
new_hashed_id = exp.to_identifier(
self._create_hash_from_expression(cte.this), quoted=old_name_id.args["quoted"]
)
replacement_mapping[old_name_id] = new_hashed_id
- cte.set("alias", exp.TableAlias(this=new_hashed_id))
expression = expression.transform(replace_id_value, replacement_mapping)
return expression
@@ -183,7 +180,7 @@ class DataFrame:
expression = df.expression
hint_expression = expression.args.get("hint") or exp.Hint(expressions=[])
for hint in df.pending_partition_hints:
- hint_expression.args.get("expressions").append(hint)
+ hint_expression.append("expressions", hint)
df.pending_hints.remove(hint)
join_aliases = {
@@ -209,7 +206,7 @@ class DataFrame:
sequence_id_expression.set("this", matching_cte.args["alias"].this)
df.pending_hints.remove(hint)
break
- hint_expression.args.get("expressions").append(hint)
+ hint_expression.append("expressions", hint)
if hint_expression.expressions:
expression.set("hint", hint_expression)
return df
diff --git a/sqlglot/dataframe/sql/session.py b/sqlglot/dataframe/sql/session.py
index c4a22c6..af589b0 100644
--- a/sqlglot/dataframe/sql/session.py
+++ b/sqlglot/dataframe/sql/session.py
@@ -129,7 +129,7 @@ class SparkSession:
@property
def _random_name(self) -> str:
- return f"a{str(uuid.uuid4())[:8]}"
+ return "r" + uuid.uuid4().hex
@property
def _random_branch_id(self) -> str:
@@ -145,7 +145,7 @@ class SparkSession:
@property
def _random_id(self) -> str:
- id = f"a{str(uuid.uuid4())[:8]}"
+ id = self._random_name
self.known_ids.add(id)
return id