# Copyright (c) 2024, Alibaba Group; # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os from typing import Any, Dict, List, Optional, Tuple import numpy as np import pyarrow as pa import pyfg from tzrec.datasets.utils import ( ParsedData, SequenceDenseData, SequenceSparseData, ) from tzrec.features.custom_feature import CustomFeature from tzrec.features.feature import MAX_HASH_BUCKET_SIZE from tzrec.features.id_feature import FgMode, IdFeature from tzrec.features.raw_feature import RawFeature from tzrec.protos import feature_pb2 from tzrec.protos.feature_pb2 import FeatureConfig from tzrec.utils.logging_util import logger def _parse_fg_encoded_sequence_sparse_feature_impl( name: str, feat: pa.Array, sequence_delim: str = ";", multival_sep: str = chr(3), default_value: Optional[List[int]] = None, ) -> SequenceSparseData: """Parse fg encoded sequence sparse feature. Args: name (str): feature name. feat (pa.Array): input feature data. sequence_delim (str): sequence delimiter. multival_sep (str): string separator for multi-val data. default_value (int): default value. Returns: an instance of SequenceSparseFeature. """ if pa.types.is_string(feat.type): # dtype = string is_empty = pa.compute.equal(feat, pa.scalar("")) feat = pa.compute.if_else(is_empty, pa.nulls(len(feat)), feat) if default_value is not None: feat = feat.fill_null(multival_sep.join(map(str, default_value))) list_seq_feat = pa.compute.split_pattern(feat, sequence_delim) list_feat = pa.compute.split_pattern(list_seq_feat.values, multival_sep) seq_offsets = list_seq_feat.offsets.to_numpy() seq_lengths = seq_offsets[1:] - seq_offsets[:-1] feat_values = list_feat.values.cast(pa.int64()).to_numpy() feat_offsets = list_feat.offsets.to_numpy() feat_lengths = feat_offsets[1:] - feat_offsets[:-1] elif pa.types.is_list(feat.type): if pa.types.is_list(feat.type.value_type): # dtype = list> or others can cast to list> feat = feat.cast(pa.list_(pa.list_(pa.int64())), safe=False) if default_value is not None: is_empty = pa.compute.equal(pa.compute.list_value_length(feat), 0) feat = pa.compute.if_else(is_empty, pa.nulls(len(feat)), feat) feat = feat.fill_null([default_value]) seq_offsets = feat.offsets.to_numpy() seq_lengths = seq_offsets[1:] - seq_offsets[:-1] feat_values = feat.values.values.to_numpy() feat_offsets = feat.values.offsets.to_numpy() feat_lengths = feat_offsets[1:] - feat_offsets[:-1] else: # dtype = list or others can cast to list feat = feat.cast(pa.list_(pa.int64()), safe=False) if default_value is not None: is_empty = pa.compute.equal(pa.compute.list_value_length(feat), 0) feat = pa.compute.if_else(is_empty, pa.nulls(len(feat)), feat) feat = feat.fill_null(default_value) seq_offsets = feat.offsets.to_numpy() seq_lengths = seq_offsets[1:] - seq_offsets[:-1] feat_values = feat.values.to_numpy() feat_lengths = np.ones_like(feat_values, dtype=np.int32) else: raise ValueError( f"{name} only support str|list|list> dtype input," f" but get {feat.dtype}." ) return SequenceSparseData(name, feat_values, feat_lengths, seq_lengths) def _parse_fg_encoded_sequence_dense_feature_impl( name: str, feat: pa.Array, sequence_delim: str = ";", multival_sep: str = chr(3), value_dim: int = 1, default_value: Optional[List[float]] = None, ) -> SequenceDenseData: """Parse fg encoded sequence dense feature. Args: name (str): feature name. feat (pa.Array): input feature data. sequence_delim (str): sequence delimiter. multival_sep (str): string separator for multi-val data. value_dtype (pa.DataType): value dtype. value_dim (int): value dimension. default_value (list): default value. Returns: an instance of SequenceSparseFeature. """ if pa.types.is_string(feat.type): is_empty = pa.compute.equal(feat, pa.scalar("")) feat = pa.compute.if_else(is_empty, pa.nulls(len(feat)), feat) if default_value is not None: feat = feat.fill_null(multival_sep.join(map(str, default_value))) list_seq_feat = pa.compute.split_pattern(feat, sequence_delim) list_feat = pa.compute.split_pattern(list_seq_feat.values, multival_sep) seq_offsets = list_seq_feat.offsets.to_numpy() seq_lengths = seq_offsets[1:] - seq_offsets[:-1] feat_values = ( list_feat.values.cast(pa.float32(), safe=False) .to_numpy() .reshape(-1, value_dim) ) elif pa.types.is_list(feat.type): if pa.types.is_list(feat.type.value_type): # dtype = list> or others can cast to list> feat = feat.cast(pa.list_(pa.list_(pa.float32())), safe=False) if default_value is not None: is_empty = pa.compute.equal(pa.compute.list_value_length(feat), 0) feat = pa.compute.if_else(is_empty, pa.nulls(len(feat)), feat) feat = feat.fill_null([default_value]) seq_offsets = feat.offsets.to_numpy() seq_lengths = seq_offsets[1:] - seq_offsets[:-1] feat_values = feat.values.values.to_numpy().reshape(-1, value_dim) else: # dtype = list or others can cast to list feat = feat.cast(pa.list_(pa.float32()), safe=False) if default_value is not None: is_empty = pa.compute.equal(pa.compute.list_value_length(feat), 0) feat = pa.compute.if_else(is_empty, pa.nulls(len(feat)), feat) feat = feat.fill_null(default_value) seq_offsets = feat.offsets.to_numpy() seq_lengths = seq_offsets[1:] - seq_offsets[:-1] feat_values = feat.values.to_numpy().reshape(-1, value_dim) else: raise ValueError( f"{name} only support str|list|list> dtype input," f" but get {feat.dtype}." ) return SequenceDenseData(name, feat_values, seq_lengths) class SequenceIdFeature(IdFeature): """SequenceIdFeature class. Args: feature_config (FeatureConfig): a instance of feature config. sequence_name (str): sequence group name. sequence_delim (str): separator for sequence feature. sequence_length (int): max sequence length. sequence_pk (str): sequence primary key name for serving. fg_mode (FgMode): input data fg mode. fg_encoded_multival_sep (str, optional): multival_sep when fg_mode=FG_NONE """ def __init__( self, feature_config: FeatureConfig, sequence_name: Optional[str] = None, sequence_delim: Optional[str] = None, sequence_length: Optional[int] = None, sequence_pk: Optional[str] = None, fg_mode: FgMode = FgMode.FG_NONE, fg_encoded_multival_sep: Optional[str] = None, ) -> None: fc_type = feature_config.WhichOneof("feature") config = getattr(feature_config, fc_type) self._is_grouped_seq = False self.sequence_name = None self.sequence_delim = None self.sequence_length = None self.sequence_pk = None if isinstance(config, feature_pb2.IdFeature): self._is_grouped_seq = True self.sequence_name = sequence_name self.sequence_delim = sequence_delim self.sequence_length = sequence_length if not sequence_pk: self.sequence_pk = f"user:{sequence_name}" else: self.sequence_pk = sequence_pk else: self.sequence_delim = config.sequence_delim self.sequence_length = config.sequence_length super().__init__(feature_config, fg_mode, fg_encoded_multival_sep) @property def name(self) -> str: """Feature name.""" if self._is_grouped_seq: return f"{self.sequence_name}__{self.config.feature_name}" else: return self.config.feature_name @property def value_dim(self) -> int: """Fg value dimension of the feature.""" if self.config.HasField("value_dim"): return self.config.value_dim else: return 1 @property def is_sequence(self) -> bool: """Feature is sequence or not.""" return True @property def is_grouped_sequence(self) -> bool: """Feature is grouped sequence or not.""" return self._is_grouped_seq def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]: """Input field names with side.""" if self.config.HasField("expression"): if self._is_grouped_seq: side, name = self.config.expression.split(":") return [(side, f"{self.sequence_name}__{name}")] else: return [tuple(self.config.expression.split(":"))] else: return None def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: """Parse input data for the feature impl. Args: input_data (dict): raw input feature data. Return: parsed feature data. """ if self.fg_mode == FgMode.FG_NONE: feat = input_data[self.name] parsed_feat = _parse_fg_encoded_sequence_sparse_feature_impl( self.name, feat, sequence_delim=self.sequence_delim, **self._fg_encoded_kwargs, ) elif self.fg_mode == FgMode.FG_NORMAL: input_feat = input_data[self.inputs[0]] if pa.types.is_list(input_feat.type): input_feat = input_feat.fill_null([]) input_feat = input_feat.tolist() values, key_lengths, seq_lengths = self._fg_op.to_bucketized_jagged_tensor( input_feat ) parsed_feat = SequenceSparseData( name=self.name, values=values, key_lengths=key_lengths, seq_lengths=seq_lengths, ) else: raise ValueError( f"fg_mode: {self.fg_mode} is not supported without fg handler." ) return parsed_feat def init_fg(self) -> None: """Init fg op.""" cfgs = self.fg_json() is_rank_zero = os.environ.get("RANK", "0") == "0" if self._is_grouped_seq: # pyre-ignore [16] self._fg_op = pyfg.FeatureFactory.create( cfgs[0], self.sequence_name, self.sequence_delim, self.sequence_length, is_rank_zero, ) else: # pyre-ignore [16] self._fg_op = pyfg.FeatureFactory.create( cfgs[0], is_rank_zero, ) def fg_json(self) -> List[Dict[str, Any]]: """Get fg json config.""" if self.config.default_value == "": logger.warning( "SequenceIdFeature not support empty default value now. reset to zero." ) self.config.default_value = "0" fg_cfg = { "feature_type": "id_feature" if self._is_grouped_seq else "sequence_id_feature", "feature_name": self.config.feature_name, "default_value": self.config.default_value, "expression": self.config.expression, "value_type": "string", "need_prefix": False, } if not self._is_grouped_seq: fg_cfg["sequence_delim"] = self.config.sequence_delim fg_cfg["sequence_length"] = self.config.sequence_length if self.config.separator != "\x1d": fg_cfg["separator"] = self.config.separator if self.config.HasField("zch"): fg_cfg["hash_bucket_size"] = MAX_HASH_BUCKET_SIZE elif self.config.HasField("hash_bucket_size"): fg_cfg["hash_bucket_size"] = self.config.hash_bucket_size elif self.config.HasField("num_buckets"): fg_cfg["num_buckets"] = self.config.num_buckets elif len(self.vocab_list) > 0: fg_cfg["vocab_list"] = self.vocab_list fg_cfg["default_bucketize_value"] = self.default_bucketize_value elif len(self.config.vocab_dict) > 0: fg_cfg["vocab_dict"] = self.vocab_dict fg_cfg["default_bucketize_value"] = self.default_bucketize_value elif len(self.vocab_file) > 0: fg_cfg["vocab_file"] = self.vocab_file fg_cfg["default_bucketize_value"] = self.default_bucketize_value if self.config.HasField("value_dim"): fg_cfg["value_dim"] = self.config.value_dim else: fg_cfg["value_dim"] = 1 return [fg_cfg] class SequenceRawFeature(RawFeature): """SequenceIdFeature class. Args: feature_config (FeatureConfig): a instance of feature config. sequence_name (str): sequence group name. sequence_delim (str): separator for sequence feature. sequence_length (int): max sequence length. sequence_pk (str): sequence primary key name for serving. fg_mode (FgMode): input data fg mode. fg_encoded_multival_sep (str, optional): multival_sep when fg_mode=FG_NONE """ def __init__( self, feature_config: FeatureConfig, sequence_name: Optional[str] = None, sequence_delim: Optional[str] = None, sequence_length: Optional[int] = None, sequence_pk: Optional[str] = None, fg_mode: FgMode = FgMode.FG_NONE, fg_encoded_multival_sep: Optional[str] = None, ) -> None: fc_type = feature_config.WhichOneof("feature") config = getattr(feature_config, fc_type) self._is_grouped_seq = False self.sequence_name = None self.sequence_delim = None self.sequence_length = None self.sequence_pk = None if isinstance(config, feature_pb2.RawFeature): self._is_grouped_seq = True self.sequence_name = sequence_name self.sequence_delim = sequence_delim self.sequence_length = sequence_length if not sequence_pk: self.sequence_pk = f"user:{sequence_name}" else: self.sequence_pk = sequence_pk else: self.sequence_delim = config.sequence_delim self.sequence_length = config.sequence_length super().__init__(feature_config, fg_mode, fg_encoded_multival_sep) @property def name(self) -> str: """Feature name.""" if self._is_grouped_seq: return f"{self.sequence_name}__{self.config.feature_name}" else: return self.config.feature_name @property def is_sequence(self) -> bool: """Feature is sequence or not.""" return True @property def is_grouped_sequence(self) -> bool: """Feature is grouped sequence or not.""" return self._is_grouped_seq @property def _dense_emb_type(self) -> Optional[str]: # TODO: support dense embedding for sequence raw feature. return None def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]: """Input field names with side.""" if self.config.HasField("expression"): if self._is_grouped_seq: side, name = self.config.expression.split(":") return [(side, f"{self.sequence_name}__{name}")] else: return [tuple(self.config.expression.split(":"))] else: return None def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: """Parse input data for the feature impl. Args: input_data (dict): raw input feature data. Return: parsed feature data. """ if self.fg_mode == FgMode.FG_NONE: feat = input_data[self.name] if self.is_sparse: parsed_feat = _parse_fg_encoded_sequence_sparse_feature_impl( self.name, feat, sequence_delim=self.sequence_delim, **self._fg_encoded_kwargs, ) else: parsed_feat = _parse_fg_encoded_sequence_dense_feature_impl( self.name, feat, sequence_delim=self.sequence_delim, value_dim=self.config.value_dim, **self._fg_encoded_kwargs, ) elif self.fg_mode == FgMode.FG_NORMAL: input_feat = input_data[self.inputs[0]] if pa.types.is_list(input_feat.type): input_feat = input_feat.fill_null([]) input_feat = input_feat.tolist() if self._fg_op.is_sparse: values, lengths = self._fg_op.to_bucketized_jagged_tensor(input_feat) parsed_feat = SequenceSparseData( name=self.name, values=values, key_lengths=np.array( [self._fg_op.value_dimension()] * sum(lengths) ), seq_lengths=lengths, ) else: values, lengths = self._fg_op.to_jagged_tensor(input_feat) parsed_feat = SequenceDenseData( name=self.name, values=values, seq_lengths=lengths, ) else: raise ValueError( f"fg_mode: {self.fg_mode} is not supported without fg handler." ) return parsed_feat def init_fg(self) -> None: """Init fg op.""" cfgs = self.fg_json() is_rank_zero = os.environ.get("RANK", "0") == "0" if self._is_grouped_seq: # pyre-ignore [16] self._fg_op = pyfg.FeatureFactory.create( cfgs[0], self.sequence_name, self.sequence_delim, self.sequence_length, is_rank_zero, ) else: # pyre-ignore [16] self._fg_op = pyfg.FeatureFactory.create( cfgs[0], is_rank_zero, ) def fg_json(self) -> List[Dict[str, Any]]: """Get fg json config.""" fg_cfg = { "feature_type": "raw_feature" if self._is_grouped_seq else "sequence_raw_feature", "feature_name": self.config.feature_name, "default_value": self.config.default_value, "expression": self.config.expression, "value_type": "float", } if not self._is_grouped_seq: fg_cfg["sequence_delim"] = self.config.sequence_delim fg_cfg["sequence_length"] = self.config.sequence_length if self.config.value_dim > 1: if self.config.separator != "\x1d": fg_cfg["separator"] = self.config.separator fg_cfg["value_dim"] = self.config.value_dim if self.config.normalizer != "": fg_cfg["normalizer"] = self.config.normalizer if len(self.config.boundaries) > 0: fg_cfg["boundaries"] = list(self.config.boundaries) return [fg_cfg] class SequenceCustomFeature(CustomFeature): """SequenceCustomFeature class. Args: feature_config (FeatureConfig): a instance of feature config. sequence_name (str): sequence group name. sequence_delim (str): separator for sequence feature. sequence_length (int): max sequence length. sequence_pk (str): sequence primary key name for serving. fg_mode (FgMode): input data fg mode. fg_encoded_multival_sep (str, optional): multival_sep when fg_mode=FG_NONE """ def __init__( self, feature_config: FeatureConfig, sequence_name: Optional[str] = None, sequence_delim: Optional[str] = None, sequence_length: Optional[int] = None, sequence_pk: Optional[str] = None, fg_mode: FgMode = FgMode.FG_NONE, fg_encoded_multival_sep: Optional[str] = None, ) -> None: fc_type = feature_config.WhichOneof("feature") config = getattr(feature_config, fc_type) self._is_grouped_seq = False self.sequence_name = None self.sequence_delim = None self.sequence_length = None self.sequence_pk = None if isinstance(config, feature_pb2.CustomFeature): self._is_grouped_seq = True self.sequence_name = sequence_name self.sequence_delim = sequence_delim self.sequence_length = sequence_length if not sequence_pk: self.sequence_pk = f"user:{sequence_name}" else: self.sequence_pk = sequence_pk else: self.sequence_delim = config.sequence_delim self.sequence_length = config.sequence_length super().__init__(feature_config, fg_mode, fg_encoded_multival_sep) @property def name(self) -> str: """Feature name.""" if self._is_grouped_seq: return f"{self.sequence_name}__{self.config.feature_name}" else: return self.config.feature_name @property def is_sequence(self) -> bool: """Feature is sequence or not.""" return True @property def is_grouped_sequence(self) -> bool: """Feature is grouped sequence or not.""" return self._is_grouped_seq def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]: """Input field names with side.""" if len(self.config.expression) > 0: side_inputs = [] if self._is_grouped_seq: for expression in self.config.expression: side, name = expression.split(":") side_inputs.append((side, f"{self.sequence_name}__{name}")) else: for expression in self.config.expression: side_inputs.append(tuple(expression.split(":"))) return side_inputs else: return None @property def _dense_emb_type(self) -> Optional[str]: return None def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: """Parse input data for the feature impl. Args: input_data (dict): raw input feature data. Return: parsed feature data. """ if self.fg_mode == FgMode.FG_NONE: feat = input_data[self.name] if self.is_sparse: parsed_feat = _parse_fg_encoded_sequence_sparse_feature_impl( self.name, feat, sequence_delim=self.sequence_delim, **self._fg_encoded_kwargs, ) else: parsed_feat = _parse_fg_encoded_sequence_dense_feature_impl( self.name, feat, sequence_delim=self.sequence_delim, value_dim=self.config.value_dim, **self._fg_encoded_kwargs, ) elif self.fg_mode == FgMode.FG_NORMAL: input_feats = [] for name in self.inputs: x = input_data[name] if pa.types.is_list(x.type): x = x.fill_null([]) x = x.tolist() input_feats.append(x) if self._fg_op.is_sparse: values, key_lengths, seq_lengths = ( self._fg_op.to_bucketized_jagged_tensor(input_feats) ) parsed_feat = SequenceSparseData( name=self.name, values=values, key_lengths=key_lengths, seq_lengths=seq_lengths, ) else: values, lengths = self._fg_op.to_jagged_tensor(input_feats) parsed_feat = SequenceDenseData( name=self.name, values=values, seq_lengths=lengths, ) else: raise ValueError( f"fg_mode: {self.fg_mode} is not supported without fg handler." ) return parsed_feat def init_fg(self) -> None: """Init fg op.""" cfgs = self.fg_json() is_rank_zero = os.environ.get("RANK", "0") == "0" if self._is_grouped_seq: # pyre-ignore [16] self._fg_op = pyfg.FeatureFactory.create( cfgs[0], self.sequence_name, self.sequence_delim, self.sequence_length, is_rank_zero, ) else: # pyre-ignore [16] self._fg_op = pyfg.FeatureFactory.create( cfgs[0], is_rank_zero, ) def fg_json(self) -> List[Dict[str, Any]]: """Get fg json config.""" if self.config.default_value == "": logger.warning( "SequenceCustomFeature not support empty default value now. " "reset to zero." ) self.config.default_value = "0" fg_cfg = super().fg_json()[0] fg_cfg["feature_name"] = self.config.feature_name fg_cfg["is_sequence"] = True if not self._is_grouped_seq: fg_cfg["sequence_delim"] = self.config.sequence_delim fg_cfg["sequence_length"] = self.config.sequence_length return [fg_cfg]