gigl.common.utils.feature_quantization.lossless_compression#

Build-time lossless compression of feature columns into the packed feature sidecar.

Feature quantization already stores selected columns of a node or edge type as packed uint8 rows next to the raw float matrix, with a FeatureQuantizationMetadata describing them. Every later stage is generic over that pair:

  • partitioning

  • GLT storage

  • sampling

  • graph-store metadata fetching

  • loader decode

This module adds columns to the same sidecar: every column that is integral on all ranks (one-hot columns, categorical ids) is stored as the narrowest integer encoding that holds its values (FeatureCodec). The encoding is inferred from the data when the dataset is built instead of declared in the preprocessor, so existing preprocessed data needs no rerun. Columns that preprocessing already quantized are not in the raw matrix and are left alone.

compress_lossless_columns is the pure tensor operation; compress_loaded_features applies it to the tensors of all node or edge types and makes the codec identical on every rank by merging per-rank column statistics over RPC.

Compression is lossless and bit-exact: a column holding a non-integral value, a value no integer encoding holds, or a -0.0 stays float32.

Attributes#

FEATURE_COMPRESSION_ENV

logger

Classes#

FeatureCompression

How node and edge features are compressed when a dataset is built.

Functions#

compress_loaded_features(features, quantized_features, ...)

Applies compress_lossless_columns to every node or edge type of the loaded tensors of this rank.

compress_lossless_columns(features, ...)

Moves the integral columns of one node or edge type from the raw float matrix into the packed sidecar.

resolve_feature_compression()

Returns the mode named by the GIGL_FEATURE_COMPRESSION environment variable.

Module Contents#

class gigl.common.utils.feature_quantization.lossless_compression.FeatureCompression[source]#

Bases: str, enum.Enum

How node and edge features are compressed when a dataset is built.

Initialize self. See help(type(self)) for accurate signature.

AUTO = 'auto'[source]#
OFF = 'off'[source]#
gigl.common.utils.feature_quantization.lossless_compression.compress_loaded_features(features, quantized_features, metadata, is_homogeneous, default_key)[source]#

Applies compress_lossless_columns to every node or edge type of the loaded tensors of this rank.

Every rank must call this (it joins one RPC all_gather of the column statistics), so that all ranks end up with identical metadata even if a column only fails to fit on one rank’s partition or a rank has no rows.

Every type with a raw feature matrix on any rank is a candidate. Columns that preprocessing already quantized are not in the raw matrix and stay as they are.

Logs on each rank the seconds spent scanning the statistics, merging them across ranks, encoding, and in total, and the number of torch threads those steps ran with.

Peak host memory is the input tensors of all compressed types, plus their remaining raw columns, plus the packed sidecars, until the caller releases the inputs.

Parameters:

  • features (Optional[Union[torch.Tensor, Mapping[_Key, torch.Tensor]]]) – Raw float matrix, or one per type for heterogeneous graphs. A type missing here has no rows on this rank.

  • quantized_features (Optional[Union[torch.Tensor, Mapping[_Key, torch.Tensor]]]) – Existing packed sidecar(s) from preprocessing.

  • metadata (Optional[Union[gigl.types.graph.FeatureQuantizationMetadata, Mapping[_Key, gigl.types.graph.FeatureQuantizationMetadata]]]) – Existing quantization metadata, a single value for homogeneous graphs.

  • is_homogeneous (bool) – Whether the graph is homogeneous, which decides the form of metadata.

  • default_key (_Key) – Key the homogeneous graph’s single type uses internally.

Returns:

features, quantized_features and metadata after compression, in the same form as the inputs (single values for homogeneous graphs, dicts for heterogeneous graphs, None if empty).

Raises:

ValueError – If a homogeneous input is a mapping, or a heterogeneous input is a single value.

Return type:

tuple[Optional[Union[torch.Tensor, dict[_Key, torch.Tensor]]], Optional[Union[torch.Tensor, dict[_Key, torch.Tensor]]], Optional[Union[gigl.types.graph.FeatureQuantizationMetadata, dict[_Key, gigl.types.graph.FeatureQuantizationMetadata]]]]

gigl.common.utils.feature_quantization.lossless_compression.compress_lossless_columns(features, quantized_features, metadata, stats, name)[source]#

Moves the integral columns of one node or edge type from the raw float matrix into the packed sidecar.

The encoding of each raw column comes from stats (FeatureCodec.from_stats), the statistics of the data merged over all partitions. Nothing is changed if no column is moved or moving them would not make rows smaller.

Peak host memory is the input raw matrix, plus the remaining raw columns, plus the packed sidecar; the encode works in row chunks, so it needs no other full-size temporary. The input matrix can be released once the caller drops it.

Parameters:

  • features (Optional[jaxtyping.Float[torch.Tensor, rows raw_dim]]) – Raw float matrix [N, raw_dim] holding every column that is not already quantized; None if this partition has no rows.

  • quantized_features (Optional[jaxtyping.UInt8[torch.Tensor, rows packed_dim]]) – Existing packed sidecar [N, packed_dim] from preprocessing, if any.

  • metadata (Optional[gigl.types.graph.FeatureQuantizationMetadata]) – Quantization metadata describing quantized_features, if any.

  • stats (Optional[gigl.common.utils.feature_quantization.lossless.ColumnStats]) – ColumnStats of the raw matrix, merged over partitions; None if no partition has a raw matrix.

  • name (str) – Name of the node or edge type, for log messages.

Returns:

The remaining raw matrix (None if empty), the new sidecar and the metadata describing it.

Raises:

ValueError – If the metadata already has a lossless region, or stats does not match the raw column count.

Return type:

tuple[Optional[jaxtyping.Float[torch.Tensor, rows _remaining_dim]], Optional[jaxtyping.UInt8[torch.Tensor, rows _new_packed_dim]], Optional[gigl.types.graph.FeatureQuantizationMetadata]]

gigl.common.utils.feature_quantization.lossless_compression.resolve_feature_compression()[source]#

Returns the mode named by the GIGL_FEATURE_COMPRESSION environment variable.

An unset or empty variable means AUTO. The value is case-insensitive.

Raises:

ValueError – If the variable holds a value that is not a FeatureCompression mode.

Return type:

FeatureCompression

gigl.common.utils.feature_quantization.lossless_compression.FEATURE_COMPRESSION_ENV: Final[str] = 'GIGL_FEATURE_COMPRESSION'[source]#
gigl.common.utils.feature_quantization.lossless_compression.logger[source]#