github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-08-02 00:28:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix:pyspark lazy load

Browse Source
This commit is contained in:
aries_ckt 2023-09-27 12:40:30 +08:00
parent 1d3ea2acd5
commit 92c25fe395
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
pilot/connections/conn_spark.py
View File

@ -17,12 +17,11 @@ class SparkConnect(BaseConnect):
driver: str = "spark" driver: str = "spark"
"""db dialect""" """db dialect"""
dialect: str = "sparksql" dialect: str = "sparksql"
from pyspark.sql import SparkSession, DataFrame
def __init__( def __init__(
self, self,
file_path: str, file_path: str,
spark_session: Optional[SparkSession] = None, spark_session: Optional = None,
engine_args: Optional[dict] = None, engine_args: Optional[dict] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
@ -30,6 +29,7 @@ class SparkConnect(BaseConnect):
return: Spark DataFrame return: Spark DataFrame
""" """
from pyspark.sql import SparkSession from pyspark.sql import SparkSession
self.spark_session = ( self.spark_session = (
spark_session or SparkSession.builder.appName("dbgpt_spark").getOrCreate() spark_session or SparkSession.builder.appName("dbgpt_spark").getOrCreate()
) )
@ -47,7 +47,7 @@ class SparkConnect(BaseConnect):
except Exception as e: except Exception as e:
print("load spark datasource error" + str(e)) print("load spark datasource error" + str(e))
def create_df(self, path) -> DataFrame: def create_df(self, path):
"""Create a Spark DataFrame from Datasource path(now support parquet, jdbc, orc, libsvm, csv, text, json.). """Create a Spark DataFrame from Datasource path(now support parquet, jdbc, orc, libsvm, csv, text, json.).
return: Spark DataFrame return: Spark DataFrame
reference:https://spark.apache.org/docs/latest/sql-data-sources-load-save-functions.html reference:https://spark.apache.org/docs/latest/sql-data-sources-load-save-functions.html