User Guide: SDKData & Graph

Connectors

View as Markdown

After you have initialized the SDK, the first step is to define a connector to your source tables. The SDK supports connectors for the following sources:

  • Amazon S3 (S3Connector)
  • Google Cloud Storage (GCSConnector)
  • Azure Data Lake Storage (ADLSConnector)
  • AWS Glue (GlueConnector)
  • Snowflake (SnowflakeConnector)
  • BigQuery (BigQueryConnector)
  • Databricks (DatabricksConnector) The examples here use S3, but equivalent steps apply to other supported data warehouses. You can also connect multiple tables across multiple connectors: for example, using S3 and Snowflake together.

If you are using the Kumo Snowpark Container Services edition, the SDK supports only SnowflakeConnector.

Creating a Connector

To create a connector to a dataset on S3, specify the root directory of your data:

connector = kumo.S3Connector(root_dir="s3://kumo-public-datasets/customerltv_mini/")

You can access tables with Python indexing semantics or with the table() method:

# Access the 'customer' table by indexing into the connector:
customer_src = connector['customer']
# Access the 'transaction' table by explicitly calling the `.table`
# method on the connector:
transaction_src = connector.table('transaction')
# Create a connector without a root directory, and obtain a table by
# passing the full table path:
stock_src = kumo.S3Connector().table('s3://kumo-public-datasets/customerltv_mini/stock')

Inspecting Source Tables

The tables customer_src, transaction_src, and stock_src are SourceTable objects, which support basic operations to verify the types and raw data you have connected to Kumo. For example, you can view a sample of the source data (as a pandas.DataFrame) or inspect the source columns and their data types:

print(customer_src.head())
>>
    CustomerID
428    16909.0
312    14002.0
306    17101.0
141    13385.0
273    14390.0
print(len(transaction_src.columns))
>> 8

For tables with semantically meaningful text columns, Kumo supports a language model integration that allows for modeling to utilize powerful large language model embeddings, for example from OpenAI’s GPT. Refer to add_llm() for more details.

Data Transformations

In addition to viewing raw source table data, you can perform data transformations using your own data platform alongside the Kumo RDL SDK. For example, with pyspark:

import kumoai as kumo
from pyspark.sql.functions import col
root_dir = "s3://kumo-public-datasets/customerltv_mini/"
output_dir = ...  # An output directory that you can write to
# Perform transformation with Spark
spark.read.parquet(f"{root_dir}/transaction") \
    .withColumn("TotalPrice", col("Quantity") * col("UnitPrice")) \
    .write.format("parquet").option("header","true").mode("Overwrite") \
    .save(f"{output_dir}/transaction_altered/")
# Access the altered table from the same connector:
assert kumo.S3Connector(root_dir=output_dir).has_table("transaction_altered")

Uploading Local Tables

For local files, you can use upload_table() to upload Parquet or CSV files directly to Kumo. Kumo supports files larger than 1 GB by default through automatic partitioning. After uploading, access tables through FileUploadConnector.

from kumoai.connector import upload_table
# Upload local file (supports >1GB automatically)
upload_table(name="my_table", path="/path/to/local/file.parquet")
# Access uploaded table
connector = kumo.FileUploadConnector(file_type="parquet")
my_table_src = connector["my_table"]

Key parameters: name (table name), path (local file path), auto_partition (defaults to True for files larger than 1 GB), and partition_size_mb (default 250 MB).