Question: How to self-define the transforms

[vince-zh]

Hi guys,
Really appreciated for developing the package, it is really useful! But I got one question here:

Previously I was using the following codes to convert date_add & date_diff function in presto to Spark SQL, and it works well, but since July 13th, I got error msg:__init__() got an unexpected keyword argument 'transforms' for Generator and __init__() got an unexpected keyword argument 'functions' for Parser. Would you mind helping to advise the root cause and maybe if there is any eaiser way to do so?

Really Appreciated!!
Thanks in advance!

Regards,
Vince Zheng

class date_add(Func): # put the presto func name here, e.g., date_add
    arg_types = {'unit': True, 'duration': True, 'time': True} 
    
class date_diff(Func):
    arg_types = {'unit': True, 'start_time': True, 'end_time': True}

class sqlTranslator:
  '''
  use this class to interpret and convert Presto SQL to Spark SQL
  '''
  
  param_dict = {'[[inc_start_date]]': "'{inc_start_date}'", '[[inc_end_date]]': "'{inc_end_date}'",'[[inc_start_datetime]]': "'{inc_start_date}'",'[[inc_end_datetime]]': "'{inc_end_date}'" }
  
  def __init__(self, query):
    for s,t in self.param_dict.items():
      query = query.replace(s,t)
    self.query = query
  
  
  def get_structure(self):
    '''
    return the query structures
    '''
    tokens = Tokenizer().tokenize(self.query)
    tree = Parser(functions={
    **date_add.default_parser_mappings(), 
    **date_diff.default_parser_mappings(),}).parse(tokens)[0] 
    return tree
  
  def get_spark_sql(self, pretty = True):
    '''
    return the spark query 
    '''
    
    output = Generator(TRANSFORMS={
    # spark format for date_add
    date_add: lambda self, e: f"({self.sql(e, 'time')} + interval '{self.sql(e, 'duration')}' {literal_eval(self.sql(e, 'unit'))}) "  if int(self.sql(e, 'duration')) > 0 
    else f"({self.sql(e, 'time')} - interval '{abs(int(self.sql(e, 'duration')))}' {literal_eval(self.sql(e, 'unit'))})" , 
    
    # spark format for date_diff
    date_diff: lambda self, e: f"floor((months_between({self.sql(e, 'end_time')},{self.sql(e, 'start_time')}))/12)" if literal_eval(self.sql(e, 'unit')) == 'year'
                                else f"floor((months_between({self.sql(e, 'end_time')},{self.sql(e, 'start_time')}))/3)" if literal_eval(self.sql(e, 'unit')) == 'quarter'
                                else f"floor(months_between({self.sql(e, 'end_time')},{self.sql(e, 'start_time')}))" if literal_eval(self.sql(e, 'unit')) == 'month'
                                else f"floor(datediff({self.sql(e, 'end_time')},{self.sql(e, 'start_time')})/7)" if literal_eval(self.sql(e, 'unit')) == 'week'
                                else f"datediff({self.sql(e, 'end_time')},{self.sql(e, 'start_time')})" if literal_eval(self.sql(e, 'unit')) == 'day'
                                else f"floor(((unix_timestamp({self.sql(e, 'end_time')}) - unix_timestamp({self.sql(e, 'start_time')}))/(3600)))" if literal_eval(self.sql(e, 'unit')) == 'hour'
                                else f"floor((unix_timestamp({self.sql(e, 'end_time')}) - unix_timestamp({self.sql(e, 'start_time')})))"  if literal_eval(self.sql(e, 'unit')) == 'second'
                                else None 
              }).generate(self.get_structure())
    
    spark_query = transpile(output,read = 'presto', write = 'spark', pretty = pretty)[0] 
    
    return spark_query

[tobymao]

you need to subclass now in sqlglot 4.0 look at dialects/ presto/hive for examples

https://github.com/tobymao/sqlglot/blob/main/sqlglot/dialects/presto.py

[vince-zh]

Hey @tobymao, it works after I subclass the spark dialect
May I have a follow-up question regarding the parsing of the Presto Functions whose arguments are not split by ',':

1. If I wanna transfer TRY(ARRAY[INDEX]) in Presto --> ELEMENT_AT(ARRAY,INDEX) in Spark how should I define the expression arguments?
2. Another case is: DATE(timestamp AT TIME ZONE 'Asia/Singapore') in Presto --> from_utc_timestamp(created, 'Asia/Singapore') in Spark

Any idea would be appreciated, thanks in advance!!

Regards,
Vince

[tobymao]

you can just make a mapping of ARRAY to be ELEMENT_AT and pass it into transforms, you could ignore TRY.

same with #2, look at the various dialects and how functions are overriden

[vince-zh]

Hi @tobymao

Really really appreciate your reply!!
I am so sorry that I think I might still have problems when taking your advice, would you mind helping to suggest more specifically again~

1. I think I might not make my case clear enough here. The "ARRAY" refers to the column in which the data type is ARRAY.
   e.g., I have a column/field called param which is a column of ARRAY, and user_id is one of the keys under the ARRAY
   in Presto, we use try(param['user_id']) to extract user_id (if_any)
   then in Spark, we wanna translate to element_at(param, 'user_id')
   how should I write the Parser/Generator for this case?

2. for 2, would you mind referring me to a similar example, where the args in the function are not split by a comma?

Thanks in advance!

Regards,
Vince