Is there a way to "set *" on a Spark/Databricks merge query when all columns are not in the source?

user_2670571 0
Viewed 333

I'm merging data from one table into another in Spark/Databricks. I can do and update set * if all columns are selected but this fails if all columns are not selected (e.g. if there is a col9 in the dst table in the query below). Is there a way to do this merge without repeating the long list of columns in the when matched (i.e. match by column name in the source query) and when not matched segments?

merge into demo dst
using (
select distinct
  col1,
  col2,
  col3,
  col4,
  col5,
  col6,
  col7,
  col8
from
  demo_raw
where 1=1
  and col1 = 'ac'
  and col2 is not null
) src
on src.col1 = dst.col1
when matched then
  update set *
when not matched then
  insert *
;
delta-lake databricks apache-spark
1 Answers

Here is a python solution that I'm somewhat content with:

Function Definition:

%python
def merge(srcTableName, dstTableName, whereStr, joinStr, cols, selectCols, debug):
  
  # create the sql string
  sqlString = ""
  sqlString += "merge into " + dstTableName + " t\n"
  sqlString += "using ( \n"
  sqlString += "select \n"
  for str in selectCols[:-1]:
    sqlString += "  " + str + ", \n"
  sqlString += "  " + selectCols[-1] + "\n"
  sqlString += "from " + srcTableName
  sqlString += whereStr + "\n"
  sqlString += ") s\n"
  sqlString += "on \n"
  sqlString += joinStr + "\n"
  sqlString += "when matched then update set \n"
  updateCols = [s + " = s." + s for s in cols]
  for str in updateCols[:-1]:
    sqlString += "  " + str + ", \n"
  sqlString += "  " + updateCols[-1]
  sqlString += """
  when not matched then 
  insert (
  """
  for str in cols[:-1]:
    sqlString += "  " + str + ", \n"
  sqlString += "  " + cols[-1]
  sqlString += """
  )
  values (
  """
  for str in cols[:-1]:
    sqlString += "  " + str + ", \n"
  sqlString += "  " + cols[-1] + "\n"
  sqlString += ")"

  if debug == True:
    print(sqlString)

  spark.sql(sqlString)
  return sqlString

Example Call:

# define the tables (FROM MY_PROJECT.DEMO TO DEMO)
srcTableName = "my_project.demo"
dstTableName = "demo"
# define the where clause (ONLY AC DATA)
whereStr = """
where 1=1
  and org = 'ac'
  and org_patient_id is not null
"""
# define the join (MATCH ON PATIENT_ID)
joinStr = "t.patient_id = s.patient_id"
# define the columns (JUST THESE COLUMNS)
cols = [
  'org',
  'data_lot',
  'raw_table',
  'org_patient_id',
  'patient_id',
  'sex',
  'sexual_orientation',
  'year_of_birth',
  'year_of_death',
  'race',
  'ethnicity',
  'city',
  'state',
  'zip'  
]
# create any needed aliases for query string
selectCols = cols
selectCols = [x if x != 'year_of_birth' else '(2020 - age) as year_of_birth' for x in selectCols]
# do the merge
merge(srcTableName, dstTableName, whereStr, joinStr, cols, selectCols)
print("Done.")
user_26705710
Related