Skip to content

Commit

Permalink
support grouped collections
Browse files Browse the repository at this point in the history
  • Loading branch information
brimoor committed Jan 18, 2025
1 parent 7ed59b4 commit d11868c
Showing 1 changed file with 24 additions and 6 deletions.
30 changes: 24 additions & 6 deletions fiftyone/core/collections.py
Original file line number Diff line number Diff line change
Expand Up @@ -7728,17 +7728,19 @@ def count(self, field_or_expr=None, expr=None, safe=False):
"""

# Optimization: use estimated document count when possible
if (
isinstance(self, fod.Dataset)
and expr is None
if self._is_full_collection() and (
expr is None
and (
field_or_expr is None
or (field_or_expr == "frames" and self._has_frame_fields())
or (
etau.is_str(field_or_expr)
and field_or_expr == "frames"
and self._has_frame_fields()
)
)
):
frames = field_or_expr == "frames"
# pylint: disable=no-member
return self._estimated_count(frames=frames)
return self._dataset._estimated_count(frames=frames)

make = lambda field_or_expr: foa.Count(
field_or_expr, expr=expr, safe=safe
Expand Down Expand Up @@ -10616,6 +10618,22 @@ def _has_frame_fields(self):
def _handle_id_fields(self, field_name):
return _handle_id_fields(self, field_name)

def _is_full_collection(self):
if isinstance(self, fod.Dataset) and self.media_type != fom.GROUP:
return True

# pylint:disable=no-member
if (
isinstance(self, fov.DatasetView)
and self._dataset.media_type == fom.GROUP
and len(self._stages) == 1
and isinstance(self._stages[0], fos.SelectGroupSlices)
and self._pipeline() == []
):
return True

return False

def _is_label_field(self, field_name, label_type_or_types):
try:
label_type = self._get_label_field_type(field_name)
Expand Down

0 comments on commit d11868c

Please sign in to comment.