From ae1810aea22fdc2d1051b51db83161d748fdb7f2 Mon Sep 17 00:00:00 2001 From: v-chen_data Date: Thu, 26 Sep 2024 22:39:57 -0700 Subject: [PATCH] test --- .../command_utils/data_prep/convert_delta_to_json.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/llmfoundry/command_utils/data_prep/convert_delta_to_json.py b/llmfoundry/command_utils/data_prep/convert_delta_to_json.py index 808a18a8b6..67f73c8b06 100644 --- a/llmfoundry/command_utils/data_prep/convert_delta_to_json.py +++ b/llmfoundry/command_utils/data_prep/convert_delta_to_json.py @@ -216,7 +216,14 @@ def run_query( spark: Optional['SparkSession'] = None, collect: bool = True, ) -> Optional[Union[list['Row'], 'DataFrame', 'SparkDataFrame']]: - """Run SQL query via databricks-connect or databricks-sql.""" + """Run SQL query via databricks-connect or databricks-sql. + Args: + query (str): sql query + method (str): select from dbsql and dbconnect + cursor (Optional[Cursor]): connection.cursor + spark (Optional[SparkSession]): spark session + collect (bool): whether to get the underlying data from spark dataframe + """ if method == 'dbsql': if cursor is None: raise ValueError(f'cursor cannot be None if using method dbsql')