Read column descriptions when using lightweight transforms

When using PySpark, we can both, read and write the column descriptions from/to Foundry datasets. When using lightweight transforms, in Duck DB, there is a way to write column descriptions, but how can we read them?

Our organization mostly transforms small-medium sized datasets. The probability for the datasets to become large one day is very low. So we are in a perfect position to use lightweight transforms. But it is critical for us to have column descriptions. So it seems, we cannot employ lightweight…

How to read the column descriptions from Foundry datasets when using lightweight transforms?

We recently also ran into this limitation. I put my +1 to Palantir supporting this in the LightweightInput.

Here is code that you can use to achieve it in the meantime:

from transforms.api import LightweightContext

# import polars as pl
from transforms.api import Input, Output, transform, LightweightInput, LightweightOutput
from foundry.transforms import (
    LightweightFoundryConnector,
)
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from urllib.parse import quote_plus


def _get_session(retries=3, backoff_factor=0.3):
    session = requests.Session()
    retry = Retry(
        total=retries,
        backoff_factor=backoff_factor,
    )
    adapter = HTTPAdapter(max_retries=retry)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    return session


SESSION = _get_session()


def _request(
    method: str,
    auth_header: str,
    url: str,
    params: dict,
    expected_status_codes: list[int] = [200, 204],
    json: dict = None,
) -> requests.Response:
    response = SESSION.request(
        method=method,
        url=url,
        headers={"Authorization": auth_header},
        json=json,
        params=params,
        timeout=30,
    )
    response.raise_for_status()
    if response.status_code not in expected_status_codes:
        raise ValueError(
            f"Url: {url}\n"
            + f"Fields: {params}\n"
            + f"Body: {json}\n"
            + f"Response Code: {response.status_code}\n"
            + f"Expected Status Code: {expected_status_codes}\n"
            + f"Response Body: {response.text}"
        )
    return response


def _get_dataset_column_description(
    auth_header: str,
    foundry_url: str,
    dataset_rid: str,
    transaction_rid: str,
    branch: str,
) -> dict:
    response = _request(
        method="GET",
        auth_header=auth_header,
        url=f"{foundry_url}/metadata/datasets/"
        f"{dataset_rid}/branches/{quote_plus(branch)}/view/namespace/foundryColumnMetadata",
        params={"endTransactionRid": transaction_rid},
        expected_status_codes=[200, 204],
    )
    if response.status_code == 204:
        return {}
    maybe_descriptions = response.json()
    return {
        column_name: description_pack["description"]
        for column_name, description_pack in maybe_descriptions.items()
        if "description" in description_pack
    }


@transform.using(
    output_dataset1=Output(
        "/.../column-test"
    ),
    output_dataset2=Output(
        "/.../column-test2"
    ),
    lineitem=Input("ri.foundry.main.dataset.d137c770-137c-441c-9029-420db0078c33"),
    no_desc=Input("ri.foundry.main.dataset.b73a4565-f7db-4ea1-bbc8-2ff4444de878"),
)
def compute(
    ctx: LightweightContext,
    lineitem: LightweightInput,
    no_desc: LightweightInput,
    output_dataset1: LightweightOutput,
    output_dataset2: LightweightOutput,
) -> None:
    service_url = LightweightFoundryConnector()._service_provider.services_config[
        "foundry-metadata-api"
    ][0]
    cd = _get_dataset_column_description(
        ctx.auth_header,
        foundry_url=service_url,
        dataset_rid=lineitem.rid,
        transaction_rid=lineitem.end_transaction_rid,
        branch=lineitem.branch,
    )
    print(cd)
    output_dataset1.write_table(
        lineitem.polars(lazy=True).limit(1), column_descriptions=cd
    )
    cd2 = _get_dataset_column_description(
        ctx.auth_header,
        foundry_url=service_url,
        dataset_rid=no_desc.rid,
        transaction_rid=no_desc.end_transaction_rid,
        branch=no_desc.branch,
    )
    print(cd2)
    output_dataset2.write_table(
        no_desc.polars(lazy=True).limit(1), column_descriptions=cd2
    )

Returning a couple of years later…
Is the call to foundryColumnMetadata undocumented API endpoint still the only and best option to get column descriptions of lightweight inputs?

:waving_hand: We are looking into providing a first-class solution for this. Thank you all for raising! We’ll update this thread when available.