LinkTask

Defines a link prediction GNN learning task.

Parameters

Name	Type	Description	Optional
connector	SnowflakeConnector	The connector object used for sending requests to the GNN engine.	No
name	str	The name of the task, can be anything describing the task at hand. The name must comply with Snowflake object identifier rules.	No
train_table	str	Fully qualified Snowflake table name (Database.Schema.Table) for the training split. Must be provided together with validation_table.	Yes
validation_table	str	Fully qualified Snowflake table name for the validation split. Must be provided together with train_table.	Yes
test_table	str	Fully qualified Snowflake table name for the test split. For inference-only tasks (no training), this is the only table you provide.	Yes
context_table	str	Fully qualified Snowflake table name used as in-context examples for zero-shot inference. Set this instead of train_table — the two cannot both be set. Accepted by the constructor for consistency with NodeTask, but ICLEstimator does not currently support LINK_PREDICTION/REPEATED_LINK_PREDICTION task types, so setting it on a LinkTask has no estimator to use it with yet.	Yes
source_entity_column	ForeignKey	A foreign key that specifies the name of the source entity column in the training, validation, and test tables, and references the corresponding source entity RelationalTable and its column. The column identified by this foreign key represents the source node in the task. The node IDs contained in the foreign key’s column_name must match, or be a subset of, the values in the column specified by the foreign key’s link_to attribute.	No
target_entity_column	ForeignKey	A foreign key that specifies the name of the target entity column in the training, validation, and test tables, and references the corresponding target entity RelationalTable and its column. The column identified by this foreign key represents the target node in the task. The node IDs contained in the foreign key’s column_name must match, or be a subset of, the values in the column specified by the foreign key’s link_to attribute. Required when train_table/validation_table are provided; optional for inference-only tasks.	Yes
task_type	TaskType	The type of the link task, it can be one of TaskType.LINK_PREDICTION or TaskType.REPEATED_LINK_PREDICTION	No
time_column	str	If the dataset includes a time-based dimension, you can specify a timestamp column to incorporate temporal dependencies. Only one time column is supported. For details, see the Time Columns section.	Yes
evaluation_metric	EvaluationMetric	The name of the evaluation metric that we want to optimize for	Yes
current_time	bool	If set to False the current time of the task table will be reduced by one time unit. Useful when the time column at the task table does not need to see the values from the database tables at the same timestamp	Yes
column_dtypes	Dict{str, Union[str, ColumnDType]}	A mapping of task table column name to data type. It is used to override inferred dtypes or assign types to columns whose dtype could not be inferred.	Yes

Note

Provide at least one of train_table, test_table, or context_table so the task can infer its schema. train_table and validation_table must be provided together; supplying one without the other raises a ValueError. For training, provide train_table, validation_table, and (optionally) test_table. For inference only, provide just test_table. train_table and context_table cannot both be set — the constructor raises a ValueError if they are.

Examples

Link Prediction Task With Time

As shown in the figure, this dataset contains three tables:

• customers with candidate key customer_id
• articles (products) with candidate key article_id
• transactions with two foreign keys: customer_id linking to the customers table, and article_id linking to the articles table, as well as a time column t_dat.

Each row in the transactions table shows that a specific customer (customer_id) buying a specific product (article_id) on a specific date (t_dat).

Our task (purchase_task) is a recommendation task (link_prediction): given a customer and a date, we want to recommend articles the customer is likely to purchase. The time column is required so that the model does not see future transactions of a customer.

In this example, the source_entity_column links to the customers table, since we are making predictions about customers, and the target_entity_column links to the articles table, since we are predicting which articles the customers are likely to purchase next.

In this case, the task can be defined as follows:

from relationalai_predictive import LinkTask, TaskType, ForeignKey

link_task = LinkTask(
    connector=connector,
    name="recommendation_task",
    train_table="DATABASE.SCHEMA.TRAIN",
    validation_table="DATABASE.SCHEMA.VALIDATION",
    test_table="DATABASE.SCHEMA.TEST",
    source_entity_column=ForeignKey(column_name='customer_id', link_to='customers.customer_id'),
    target_entity_column=ForeignKey(column_name='article_id', link_to='articles.article_id'),
    time_column="timestamp",
    task_type=TaskType.LINK_PREDICTION
)

Repeated Link Prediction Task Setting Evaluation Metric

from relationalai_predictive import LinkTask, TaskType, ForeignKey
from relationalai_predictive import EvaluationMetric

rep_link_task = LinkTask(
    connector=connector,
    name="my_link_task",
    train_table="DATABASE.SCHEMA.TRAIN",
    validation_table="DATABASE.SCHEMA.VALIDATION",
    test_table="DATABASE.SCHEMA.TEST",
    source_entity_column=ForeignKey(column_name='source_ids', link_to='TableWithCKey1.Id1'),
    target_entity_column=ForeignKey(column_name='target_ids', link_to='TableWithCKey2.Id2'),
    task_type=TaskType.REPEATED_LINK_PREDICTION,
    evaluation_metric=EvaluationMetric(name="link_prediction_map", eval_at_k=12)
)

Inference-Only Task

If you only need to run inference (no training), you can create a task with just a test_table:

from relationalai_predictive import LinkTask, TaskType, ForeignKey

link_task = LinkTask(
    connector=connector,
    name="inference_recommendation_task",
    test_table="DATABASE.SCHEMA.TEST",
    source_entity_column=ForeignKey(column_name='source_ids', link_to='TableWithCKey1.Id1'),
    time_column="timestamp",
    task_type=TaskType.LINK_PREDICTION
)

Inference-only tasks can be used with estimator.predict() (by passing the task inside a new RelationalDataset) but cannot be passed to estimator.fit() or estimator.fit_predict().

Methods

LinkTask inherits from the RelationalTable, so it has the same methods. It additionally provides a show_task() method:

.show_task()

Prints the task metadata schema and task details.

Example

rep_link_task.show_task()

Attributes

.source_entity_column

Retrieves source_entity_column. Cannot be set after initialization. It is read-only.

.target_entity_column

Retrieves target_entity_column. Cannot be set after initialization. It is read-only.

.current_time

Retrieves current_time. Cannot be set after initialization. It is read-only.

.train_table

Retrieves the training split table path. Cannot be set after initialization. It is read-only.

.validation_table

Retrieves the validation split table path. Cannot be set after initialization. It is read-only.

.test_table

Retrieves the test split table path. Cannot be set after initialization. It is read-only.