Repository: oppia/oppia
File: core/jobs/batch_jobs/audit_threads_with_missing_suggestions_jobs.py (Python)

Code before:
            | 'Extract suggestion thread ids'
            >> beam.Map(lambda model: model.id)
        )

    def _get_invalid_threads(
        self,
        valid_suggestion_thread_ids: beam.PCollection[str],
    ) -> beam.PCollection[feedback_models.GeneralFeedbackThreadModel]:
        """Returns threads marked as having suggestions but without one."""
        return (
            self.pipeline
            | 'Get GeneralFeedbackThreadModels'
            >> ndb_io.GetModels(
                feedback_models.GeneralFeedbackThreadModel.get_all(
                    include_deleted=False
                )
            )
            | 'Filter threads with missing suggestions'
            >> beam.Filter(
                lambda thread, suggestion_thread_ids: (
                    thread.has_suggestion
                    and thread.id not in suggestion_thread_ids
                ),
                beam.pvalue.AsList(valid_suggestion_thread_ids),
            )
        )


class AuditThreadsWithMissingSuggestionsJob(
    BaseThreadsWithMissingSuggestionsJob
):
    """Audit job reporting invalid feedback threads."""

    DATASTORE_UPDATES_ALLOWED = False

    def run(self) -> beam.PCollection[job_run_result.JobRunResult]:
        suggestion_thread_ids = self._get_suggestion_thread_ids()
        invalid_threads = self._get_invalid_threads(suggestion_thread_ids)

        invalid_thread_logs = invalid_threads | 'Log invalid threads' >> beam.Map(
            lambda model: job_run_result.JobRunResult.as_stdout(
                (
                    'GeneralFeedbackThreadModel marked as has_suggestion=True '
                    'but no GeneralSuggestionModel exists: '
                    f'id={model.id}'
                )
            )
        )

        invalid_thread_count = (
            invalid_threads

Review comment (suggestion):
`AsList` works, but it scales very poorly as the number of threads grows. Since we don't have enough threads where this'd be an issue, **I won't block the PR on this comment and you can disregard it if you want**.

However, I still think it's valuable to show a solution that's able to scale to _billions_ of threads, since that's what Apache Beam is designed to do. TL;DR: instead of loading all of the ids into memory with `AsList`, we can use [`CoGroupByKey`](https://beam.apache.org/documentation/transforms/python/aggregation/cogroupbykey/) to have Apache Beam find the model <-> valid id relationship for us:

```python
def expand(self, valid_ids: beam.PCollection[str]) -> beam.PCollection[Model]:
    models: beam.PCollection[Model] = (
        self.pipeline
        | 'Get models' >> ndb_io.GetModels(...)
        | 'Filter models with suggestions' >> beam.Filter(lambda m: m.has_suggestion)
    )
    model_id_and_model_pairs: beam.PCollection[tuple[str, Model]] = (
        models | beam.Map(lambda model: (model.id, model))
    )
    valid_id_and_none_pairs: beam.PCollection[tuple[str, None]] = (
        valid_ids | beam.Map(lambda id: (id, None))
    )

    return (
        { 'models': model_id_and_model_pairs, 'ids': valid_id_and_none_pairs }
        | 'Group models and ids'
        >> beam.CoGroupByKey()
        | 'Keep models that DID NOT get grouped with a valid id'
        >> beam.FlatMap(
            # `group` is a tuple:
            # (id: str, { 'models': list[Model], 'ids': list[None] })
            lambda group: group[1]['models'] if not group[1]['ids'] else []
        )
    )
```

Note that this solution is analogous to how we'd solve this problem with SQL:

```sql
SELECT f.*
FROM feedback f
LEFT JOIN suggestions s ON f.id = s.id   # <- CoGroupByKey
WHERE f.has_suggestion AND s.id IS NULL  # <- FlatMap
```

And your original approach would look something like:

```sql
SELECT f.*
FROM feedback f
WHERE f.has_suggestion 
  AND f.id NOT IN (SELECT s.id FROM suggestions s)  # <- AsList
```

Code after:
        """

        suggestion_thread_ids = (
            self.pipeline
            | 'Get GeneralSuggestionModels'
            >> ndb_io.GetModels(
                suggestion_models.GeneralSuggestionModel.get_all(
                    include_deleted=False
                )
            )
            | 'Extract suggestion thread ids'
            >> beam.Map(lambda model: model.id)
        )

        feedback_threads = (
            self.pipeline
            | 'Get GeneralFeedbackThreadModels'
            >> ndb_io.GetModels(
                feedback_models.GeneralFeedbackThreadModel.get_all(
                    include_deleted=False
                )
            )
            | 'Keep threads with has_suggestion=True'
            >> beam.Filter(lambda thread: thread.has_suggestion)
        )

        feedback_id_to_model = (
            feedback_threads
            | 'Map feedback threads to (id, model)'
            >> beam.Map(lambda thread: (thread.id, thread))
        )

        suggestion_id_to_none = (
            suggestion_thread_ids
            | 'Map suggestion ids to None'
            >> beam.Map(lambda thread_id: (thread_id, None))
        )

        invalid_threads = (
            {
                'feedback': feedback_id_to_model,
                'suggestions': suggestion_id_to_none,
            }
            | 'CoGroup feedback threads with suggestion ids'
            >> beam.CoGroupByKey()
            | 'Select threads with missing suggestions'
            >> beam.FlatMap(
                lambda group: (
                    group[1]['feedback'] if not group[1]['suggestions'] else []
                )
            )