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')