Skip to content

Commit

Permalink
Fix a spelling error and some docstrings
Browse files Browse the repository at this point in the history
  • Loading branch information
dagardner-nv committed Jun 3, 2024
1 parent a013cf7 commit 001bd7d
Showing 1 changed file with 17 additions and 3 deletions.
20 changes: 17 additions & 3 deletions morpheus/stages/preprocess/group_by_column_stage.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,21 +43,35 @@ def name(self) -> str:
return "group-by-column"

def accepted_types(self) -> tuple:
"""
Returns accepted input types for this stage.
"""
return (MessageMeta, )

def supports_cpp_node(self) -> bool:
"""
Indicates whether this stage supports C++ node.
"""
return False

def on_data(self, message: MessageMeta) -> list[MessageMeta]:
"""
Group the incoming message by a column in the DataFrame.
Parameters
----------
message : MessageMeta
Incoming message
"""
with message.mutable_dataframe() as df:
grouper = df.groupby(self._column_name)

ouptut_messages = []
output_messages = []
for group_name in sorted(grouper.groups.keys()):
group_df = grouper.get_group(group_name)
ouptut_messages.append(MessageMeta(group_df))
output_messages.append(MessageMeta(group_df))

return ouptut_messages
return output_messages

def _build_single(self, builder: mrc.Builder, input_node: mrc.SegmentObject) -> mrc.SegmentObject:
node = builder.make_node(self.unique_name, ops.map(self.on_data), ops.flatten())
Expand Down

0 comments on commit 001bd7d

Please sign in to comment.