Native capability to write to Unity catalog

[anilmenon14]

This is more of an idea rather than a feature request. So, thought of initiating it here .

Since Daft has good native capability to read from Databricks Unity catalog using the credential vending in UC OSS ( 0.1.1 version , as of Daft 0.3.13) and has native integration to DeltaLake , implemented through df.write_deltalake , does it make sense for your team to have write capability into UC as well?

I have been playing around with using the existing capabilities in Daft (version 0.3.11) and could use the existing APIs to write out to a new/existing Unity catalog table as shown in this sample snippet.

import unitycatalog # 0.1.1 version , which is distributed along with Daft 0.3.x versions
uc_writer_client = unitycatalog.Unitycatalog(
            base_url=DATABRICKS_HOST.rstrip("/") + "/api/2.1/unity-catalog/",
            default_headers={"Authorization": f"Bearer {PAT_TOKEN_AZURE}"},
        )
table_name = "main_catalog.taxi_samples.taxi_data_adf_uc_v1_silver_agg_ext" 
external_table_storage_location = "abfss://my-container@storage_acc.dfs.core.windows.net/some_folder/some_sub_folder/table_sub_folder"

# The below function will take the Daft Dataframe and along with the UC OSS client, UC table name and the storage path to write to, can write to a table in 'overwrite' / 'append' mode
daft_uc_utilities.daft_write_to_unity(my_daft_df ,uc_writer_client, table_name, external_table_storage_location, mode="overwrite")   

@jaychia , I know you had some past conversations on delta-rs vs delta-kernel-rs that might play into this discussion , hence tagging you for your thoughts on this.
I also know @kevinzwang has been contributing towards the new UC OSS 0.2 being available as a Python client and tagging you as well, in case you have some context of whether there is a dependency on using the new Python client.

[kevinzwang]

Hi @anilmenon14 , thanks for bringing this up! I think what we could do is also allow users to pass in a daft.unity_catalog.UnityCatalogTable into DataFrame.write_deltalake, similar to how they can currently do so for daft.read_deltalake. I'm curious if there is any unity-specific functionality that would be enabled via a method specifically for unity?

[jaychia]

I just started work on a df.write_table(...) API that could be interesting: #3323

We should think very hard about the API here. Spark's DataFrame writer v2 API could be a good reference point: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrameWriterV2.html

It's a pretty complex writer that handles a ton of corner-cases. I think we can have maybe a simpler API but we should definitely think about it 😛

[uditrana]

I think it could be worth it to support the solution that @kevinzwang suggested (where you allow the passing of a UnityCatalogTable to write_deltalake in the same way you support read. We have been looking for a solution to efficiently read/write from Databricks UC on our external GPU machines, and just discovered Daft as a potential great solution.

Then the full featured write syntax to match v2 could come later?

[anilmenon14]

Hi @kevinzwang , that was exactly my line of thinking too and how I have written the sample utility function that performs the writes (I have pasted a snippet of the code I have used in my utility function in my Daft project to write to UC).

@jaychia has a very valid point about needing to be able to handle a variety of cases. For e.g., here we are making an assumption that the table is of Delta format when the user provides an arbitrary table name while calling any function to write to UC. It is highly likely the table is a Delta table; however, that cannot be guaranteed, and the engine should be able to handle different table formats, especially on overwrites on existing tables (e.g., CSV, JSON, ORC, etc.)
I suppose a way to handle this is to initially support UC writes (and overwrites) for only a subset of the table formats that Daft natively supports channelled through the df.write_table API and expand that out over time.

# Retrieve existing table or create new table
# 'uc_writer_client' is an instance of 'unitycatalog.Unitycatalog'
try:
    table_info = uc_writer_client.tables.retrieve(table_name) # Find table, if existing one.
except unitycatalog.NotFoundError as e:
    print("{} : {} \n Not an existing table. Attempting to create new table....".format(e.body['error_code'],e.body['message']))
    table_info = uc_writer_client.tables.create(**params) # Create a new table since it was not found
    print("Table {} has been successfully created!".format(table_name))

credentials_generated = uc_writer_client.temporary_table_credentials.create(operation="READ_WRITE", table_id=table_info.table_id)

# TO-DO : Conditional logic for building IOConfig based on choosing Azure or AWS
io_config = IOConfig(azure=AzureConfig(sas_token=credentials_generated.azure_user_delegation_sas['sas_token']))

df.write_deltalake(table=table_info.storage_location,io_config=io_config,mode=write_mode)

[kevinzwang]

Adding it to write_deltalake seems pretty doable. Created an issue for it!
#3336

[anilmenon14]

Nice @kevinzwang 👏.
If there is something I can help to collaborate on the feature, happy to help.

[kevinzwang]

Actually, if you're interested, feel free to take the issue on! I believe it is doable solely on the Python side

[anilmenon14]

Cool. I'll give it a go. I'm a bit strapped for time for a couple of days but will surely get to this and raise a PR.