Fix the cost model for a CSV export

Review the AI's codeEasy10 min

Your task

An AI's Django view exported every order and its line items to CSV. It ran one query per order to load the items (an N+1 query pattern) and built the whole file in memory, so in production it timed out and the worker ran out of memory. The rewrite streams orders in chunks and loads each chunk's items with one extra query. You model both versions' cost to show the difference.

An assistant wrote compare_export_cost(order_count, chunk_size). It is in the editor. It has a bug. Find it and fix it.

Input and output:

  • order_count: how many orders the export holds.
  • chunk_size: how many orders the streaming version loads at a time.
  • Return four numbers: [naive_queries, naive_peak, streamed_queries, streamed_peak].

Rules the tests check:

  • An order_count below 0 counts as 0. A chunk_size below 1 counts as 1.
  • naive_queries: 1 query for the orders, plus 1 per order for its items.
  • naive_peak: the number of orders, since all of them are held in memory at once.
  • streamed_queries: 2 per chunk, one for the chunk's orders and one for their items. The orders are split into chunks of chunk_size, and a last chunk with fewer orders still counts as a chunk. With no orders at all, it is 1: the export still runs one query.
  • streamed_peak: the most orders held at once, which is chunk_size, or the number of orders if that is smaller.

Examples

compare_export_cost(order_count: int, chunk_size: int) -> list[int]
Example 1
Inputcompare_export_cost(1000, 200)
Output[1001, 1000, 10, 200]
Example 2
Inputcompare_export_cost(0, 500)
Output[1, 0, 1, 0]
Example 3
Inputcompare_export_cost(1, 500)
Output[2, 1, 2, 1]