Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

alias metadata ops from core #1843

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 10 additions & 57 deletions nvtabular/ops/add_metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,60 +13,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
from merlin.core.dispatch import DataFrameType
from merlin.schema.tags import Tags
from nvtabular.ops.operator import ColumnSelector, Operator


class AddMetadata(Operator):
"""
This operator will add user defined tags and properties
to a Schema.
"""

def __init__(self, tags=None, properties=None):
super().__init__()
self.tags = tags or []
self.properties = properties or {}

def transform(self, col_selector: ColumnSelector, df: DataFrameType) -> DataFrameType:
return df

@property
def output_tags(self):
return self.tags

@property
def output_properties(self):
return self.properties


class AddTags(AddMetadata):
def __init__(self, tags=None):
super().__init__(tags=tags)


class AddProperties(AddMetadata):
def __init__(self, properties=None):
super().__init__(properties=properties)


# Wrappers for common features
class TagAsUserID(AddTags):
def __init__(self, tags=None):
super().__init__(tags=[Tags.ID, Tags.USER])


class TagAsItemID(AddTags):
def __init__(self, tags=None):
super().__init__(tags=[Tags.ID, Tags.ITEM])


class TagAsUserFeatures(AddTags):
def __init__(self, tags=None):
super().__init__(tags=[Tags.USER])


class TagAsItemFeatures(AddTags):
def __init__(self, tags=None):
super().__init__(tags=[Tags.ITEM])
# keep this file as alias, users may be apply import from here
from merlin.dag.ops.add_metadata import ( # noqa pylint: disable=unused-import
AddMetadata,
AddProperties,
AddTags,
TagAsItemFeatures,
TagAsItemID,
TagAsUserFeatures,
TagAsUserID,
)