Remove select IDs optimization

src/datachain/query/dispatch.py

@@ -20,7 +20,6 @@
     process_udf_outputs,
 )
 from datachain.query.queue import (
-    EMPTY_SIGNAL,
     get_from_queue,
     msgpack_pack,
     msgpack_unpack,
@@ -408,16 +407,7 @@ def get_inputs(self):
                 ids = [row[0] for row in batch.rows]
                 rows = warehouse.dataset_rows_select(self.query.where(col_id.in_(ids)))
                 yield RowsOutputBatch(list(rows))
-            return
-
-        ids = []
-        while (batch := get_from_queue(self.task_queue, timeout=1)) != STOP_SIGNAL:
-            if batch != EMPTY_SIGNAL:
-                ids.append(batch[0])
-            if len(ids) >= 1000 or (ids and batch == EMPTY_SIGNAL):
-                rows = warehouse.dataset_rows_select(self.query.where(col_id.in_(ids)))
+        else:
+            while (row := get_from_queue(self.task_queue)) != STOP_SIGNAL:
+                rows = warehouse.dataset_rows_select(self.query.where(col_id == row[0]))
                 yield from rows
-                ids.clear()
-        if ids:
-            rows = warehouse.dataset_rows_select(self.query.where(col_id.in_(ids)))
-            yield from rows

src/datachain/query/queue.py

@@ -1,5 +1,4 @@
 import datetime
-import time
 from queue import Empty, Full, Queue
 from struct import pack, unpack
 from time import sleep
@@ -11,7 +10,6 @@
 
 DEFAULT_BATCH_SIZE = 10000
 STOP_SIGNAL = "STOP"
-EMPTY_SIGNAL = "EMPTY"
 OK_STATUS = "OK"
 FINISHED_STATUS = "FINISHED"
 FAILED_STATUS = "FAILED"
@@ -27,20 +25,17 @@
 # https://github.com/python/cpython/issues/108645
 
 
-def get_from_queue(queue: Queue, timeout: float = 0) -> Any:
+def get_from_queue(queue: Queue) -> Any:
     """
     Gets an item from a queue.
     This is required to handle signals, such as KeyboardInterrupt exceptions
     while waiting for items to be available, although only on certain installations.
     (See the above comment for more context.)
     """
-    started_at = time.time()
     while True:
         try:
             return queue.get_nowait()
         except Empty:
-            if 0 < timeout <= time.time() - started_at:
-                return EMPTY_SIGNAL
             sleep(0.01)