NodeTask

Defines a node classification or regression 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 with ICLEstimator. Set this instead of train_table — the two cannot both be set, and setting context_table is what lets you skip train_table/validation_table entirely. Not used by GNNEstimator.	Yes
label_column	str	The name of the column in the train and validation and (optionally) test tables containing the labels. The label column is the column holding the values that will be used to train the model. (Users can choose not to provide labels for the test data.)	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 nodes in the task for which the model will make predictions. 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
task_type	TaskType	The type of the node task, it can be one of TaskType.BINARY_CLASSIFICATION, TaskType.MULTICLASS_CLASSIFICATION, TaskType.MULTILABEL_CLASSIFICATION or TaskType.REGRESSION	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.

context_table and train_table are mutually exclusive

train_table and context_table cannot both be set — the constructor raises a ValueError if they are. Use train_table to train a GNNEstimator; use context_table for zero-shot inference with ICLEstimator instead. validation_table can be paired with either one: with train_table for GNN training, or with context_table alone (no train_table) for ICLEstimator.score().

Examples

Binary Classification 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 (churn_task) is a node-level binary classification task (binary_classification): given a customer and a date, we want to predict whether they are likely to churn in the next month. The time column is required so that the model does not see future transactions when making predictions.

In this example, the node of interest is the customers table, since we are predicting churn for customers. Therefore, the target_entity_column links to the customers table.

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

from relationalai_predictive import NodeTask, TaskType, ForeignKey

binary_clf_task = NodeTask(
    connector=connector,
    name="churn_task",
    train_table="DATABASE.SCHEMA.TRAIN",
    validation_table="DATABASE.SCHEMA.VALIDATION",
    test_table="DATABASE.SCHEMA.TEST",
    target_entity_column=ForeignKey(column_name="id", link_to="TableWithCKey.Id"),
    time_column="timestamp",
    label_column="label",
    task_type=TaskType.BINARY_CLASSIFICATION
)

Regression Task and Evaluation Metric Setup

from relationalai_predictive import NodeTask, TaskType, ForeignKey
from relationalai_predictive import EvaluationMetric

regr_task = NodeTask(
    connector=connector,
    name="my_node_task",
    train_table="DATABASE.REGRESSION_SCHEMA.TRAIN",
    validation_table="DATABASE.REGRESSION_SCHEMA.VALIDATION",
    test_table="DATABASE.REGRESSION_SCHEMA.TEST",
    target_entity_column=ForeignKey(column_name="target_id", link_to="TableWithCKey.Id"),
    label_column="value",
    task_type=TaskType.REGRESSION,
    evaluation_metric=EvaluationMetric(name="r2")
)

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 NodeTask, TaskType, ForeignKey

inference_task = NodeTask(
    connector=connector,
    name="my_inference_task",
    test_table="DATABASE.SCHEMA.TEST",
    target_entity_column=ForeignKey(
        column_name="id",
        link_to="TableWithCKey.Id"
    ),
    task_type=TaskType.BINARY_CLASSIFICATION
)

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().

Zero-Shot ICL Task

To use a NodeTask with ICLEstimator instead of training a GNN, set context_table instead of train_table:

from relationalai_predictive import NodeTask, TaskType, ForeignKey

icl_task = NodeTask(
    connector=connector,
    name="churn_task_icl",
    context_table="DATABASE.SCHEMA.CONTEXT",
    test_table="DATABASE.SCHEMA.TEST",
    target_entity_column=ForeignKey(column_name="id", link_to="TableWithCKey.Id"),
    label_column="label",
    task_type=TaskType.BINARY_CLASSIFICATION,
)

train_table and context_table cannot both be set on the same task. See ICLEstimator for the full zero-shot inference workflow.

Methods

NodeTask 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

binary_clf_task.show_task()

Attributes

.label_column

Retrieves label_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.

.context_table

Retrieves the ICLEstimator context table path. Cannot be set after initialization. It is read-only.