def _parse()

in tzrec/features/lookup_feature.py [0:0]

• 48 lines of code
• 10 McCabe index (conditional complexity)

    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:
            # input feature is already lookuped
            feat = input_data[self.name]
            if self.is_sparse:
                parsed_feat = _parse_fg_encoded_sparse_feature_impl(
                    self.name, feat, **self._fg_encoded_kwargs
                )
            else:
                parsed_feat = _parse_fg_encoded_dense_feature_impl(
                    self.name, feat, **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([])
                elif pa.types.is_map(x.type):
                    x = x.fill_null({})
                input_feats.append(x.tolist())
            if self.config.value_dim > 1:
                fgout, status = self._fg_op.process(dict(zip(self.inputs, input_feats)))
                assert status.ok(), status.message()
                if self.is_sparse:
                    values = np.asarray(fgout[self.name].values, np.int64)
                    lengths = np.asarray(fgout[self.name].lengths, np.int32)
                    parsed_feat = SparseData(
                        name=self.name, values=values, lengths=lengths
                    )
                else:
                    values = fgout[self.name].dense_values
                    parsed_feat = DenseData(name=self.name, values=values)
            else:
                if self.is_sparse:
                    values, lengths = self._fg_op.to_bucketized_jagged_tensor(
                        *input_feats
                    )
                    parsed_feat = SparseData(
                        name=self.name, values=values, lengths=lengths
                    )
                else:
                    values = self._fg_op.transform(*input_feats)
                    parsed_feat = DenseData(name=self.name, values=values)
        else:
            raise ValueError(
                f"fg_mode: {self.fg_mode} is not supported without fg handler."
            )
        return parsed_feat