<jupyter_start><jupyter_text>Convert Spark Talbes and Functions to sql-language-server Schema FileThis notebooks creates a simple table with nested columns. It then exports the table and spark functions to a json file which sql-language-server can use to code-complete SQL statement for spark sql.This notebook serves as an illustration of how to do this conversion. This code will eventually be packaged into a pip install so that it can be reused more easily.This code could also be leveraged by the %%sparksql (sparksql-magic) which would make it transparent to the user. Since generation of the schema takes a few seconds, %%sparksql could detect how up-to-date the schema file is and re-generated it only when it passes a certain threshold (say 15min).<jupyter_code>import findspark
findspark.init()

import pyspark
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
%load_ext sparksql_magic
%%sparksql
DROP TABLE IF EXISTS student<jupyter_output><empty_output><jupyter_text>Alias `%%sparksql` to be `%%sql` which provides syntax highligthing<jupyter_code>%alias_magic --cell sql sparksql
%%sql
CREATE TABLE student (
    id INT, 
    name STRING,
    `nestedwithspaces` STRUCT<`sub field`:STRUCT<`sub field2`:STRING>>,
    age INT,
    books ARRAY<STRUCT<`title`:STRING, `chapters`:ARRAY<STRUCT<`paragraph`:STRING>>>>,
    struct_col STRUCT<`address`:STRUCT<`streetName`:STRING, `streetNumber`:BIGINT>>,
    map_col MAP<STRING, MAP<STRING, STRUCT<`start`:BIGINT,`end`:BIGINT>>>

    ) USING PARQUET
%%sql
/*
    We should have a student table in our database.    
*/
SHOW TABLES
%%sql
/*
    We can see the column nestedwithspaces contains sub-fields with spaces.
    Also books is an array of structs and map_col is of MapType.
*/
DESCRIBE TABLE student<jupyter_output><empty_output><jupyter_text>Generate Schema File for sql-language-serverNormally the spark catalog would point to a data warehouse full of schemas and tables. In this demo we only have a single student table which we just created.<jupyter_code>import findspark
findspark.init()
import json
import pyspark
from pyspark.sql import SparkSession
from pyspark.sql.types import *

spark = SparkSession.builder.getOrCreate()
def getTypeName(t):
    if type(t) == LongType:
        return 'long'
    if type(t) == IntegerType:
        return 'integer'
    if type(t) == StringType:
        return 'string'
    if type(t) == ArrayType:
        return 'array'
    if type(t) == StructType:
        return 'struct'
    if type(t) == MapType:
        return 'map'

def getPath(path, name):
    if ' ' in name:
        name = '`' + name + '`'
    if len(path) > 0:
        return f'{path}.{name}'
    return name

def getChildren(field, path, fields):
    if type(field) == StructField:
        getChildren(field.dataType, getPath(path, field.name), fields)
    elif type(field) == MapType:
        getChildren(field.valueType, getPath(path, 'key'), fields)
    elif type(field) == ArrayType:
        getChildren(field.elementType, path, fields)
    elif type(field) == StructType:
        for name in field.fieldNames():
            child = field[name]
            fields.append({
                'columnName': getPath(path, name),
                'metadata': child.metadata, 
                'type': getTypeName(child.dataType),
                'description': getTypeName(child.dataType)
            })
            getChildren(child, path, fields)

def getColumns(name):
    fields = []
    getChildren(spark.table(name).schema, '', fields)
    return fields

def getTables(database):
    rows = spark.sql(f'SHOW TABLES IN {database}').collect()
    return list(map(lambda r: {
        "tableName": r.tableName,
        "columns": getColumns(r.tableName), 
        "database": None
    }, rows))

def getDescription(name):
    rows = spark.sql(f'DESCRIBE FUNCTION EXTENDED {name}').collect()
    textLines = list(map(lambda r: r.function_desc, rows))
    return "\n".join(textLines)

def getFunctions():
    rows = spark.sql('SHOW FUNCTIONS').collect()
    return list(map(lambda f: {
        "name": f.function, 
        "description": getDescription(f.function)
    }, rows))

def getSparkDatabaseSchema():
    return {
        "tables": getTables('default'), 
        "functions": getFunctions()
    }

sparkdb_schema = getSparkDatabaseSchema()
# Save schema to disk. sql-language-server will pickup any changes to this file.
outputFileName = '/tmp/sparkdb.schema.json'
with open(outputFileName, 'w') as fout:
    json.dump(sparkdb_schema, fout, sort_keys=True, indent=2)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1076555023923445, 'fraction_numerical': 0.001087429317094389, 'mean_word_length': 3.6454545454545455, 'pattern_counts': {'":': 7, '<': 24, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '424590', 'license': 'MIT', 'n_tokens_mistral': 1418, 'n_tokens_neox': 1369, 'n_words': 432, 'path': 'example/jupyterlab/ExportSparkSchema.ipynb', 'repo_name': 'Leko/sql-language-server', 'size': '4598'}