Add ib.collect support for raw records (#36516)

Author: ian-Liaozy

sdks/python/apache_beam/runners/interactive/interactive_beam.py

@@ -879,7 +879,8 @@ def collect(
     runner=None,
     options=None,
     force_compute=False,
-    force_tuple=False):
+    force_tuple=False,
+    raw_records=False):
   """Materializes the elements from a PCollection into a Dataframe.
 
   This reads each element from file and reads only the amount that it needs
@@ -901,6 +902,8 @@ def collect(
         cached PCollections
     force_tuple: (optional) if True, return a 1-tuple or results rather than
         the bare results if only one PCollection is computed
+    raw_records: (optional) if True, return a list of collected records
+        without converting to a DataFrame. Default False.
 
   For example::
 
@@ -910,6 +913,9 @@ def collect(
 
     # Run the pipeline and bring the PCollection into memory as a Dataframe.
     in_memory_square = head(square, n=5)
+    
+    # Run the pipeline and get the raw list of elements.
+    raw_squares = collect(square, n=5, raw_records=True)
   """
   if len(pcolls) == 0:
     return ()
@@ -986,15 +992,19 @@ def as_pcollection(pcoll_or_df):
   if n == float('inf'):
     n = None
 
-  # Collecting DataFrames may have a length > n, so slice again to be sure. Note
-  # that array[:None] returns everything.
-  empty = pd.DataFrame()
-  result_tuple = tuple(
-      elements_to_df(
-          computed[pcoll],
-          include_window_info=include_window_info,
-          element_type=pcolls_to_element_types[pcoll])[:n] if pcoll in
-      computed else empty for pcoll in pcolls)
+  if raw_records:
+    result_tuple = tuple([el.value for el in computed.get(pcoll, [])][:n]
+                         for pcoll in pcolls)
+  else:
+    # Collecting DataFrames may have a length > n, so slice again to be sure.
+    # Note that array[:None] returns everything.
+    empty = pd.DataFrame()
+    result_tuple = tuple(
+        elements_to_df(
+            computed.get(pcoll, []),
+            include_window_info=include_window_info,
+            element_type=pcolls_to_element_types[pcoll])[:n] if pcoll in
+        computed else empty for pcoll in pcolls)
 
   if len(result_tuple) == 1 and not force_tuple:
     return result_tuple[0]

sdks/python/apache_beam/runners/interactive/interactive_beam_test.py

@@ -293,6 +293,91 @@ def is_triggered(self):
     self.assertTrue(ib.recordings.record(p))
     ib.recordings.stop(p)
 
+  def test_collect_raw_records_true(self):
+    p = beam.Pipeline(ir.InteractiveRunner())
+    data = list(range(5))
+    pcoll = p | 'Create' >> beam.Create(data)
+    ib.watch(locals())
+    ie.current_env().track_user_pipelines()
+
+    result = ib.collect(pcoll, raw_records=True)
+    self.assertIsInstance(result, list)
+    self.assertEqual(result, data)
+
+    result_n = ib.collect(pcoll, n=3, raw_records=True)
+    self.assertIsInstance(result_n, list)
+    self.assertEqual(result_n, data[:3])
+
+  def test_collect_raw_records_false(self):
+    p = beam.Pipeline(ir.InteractiveRunner())
+    data = list(range(5))
+    pcoll = p | 'Create' >> beam.Create(data)
+    ib.watch(locals())
+    ie.current_env().track_user_pipelines()
+
+    result = ib.collect(pcoll)
+    self.assertNotIsInstance(result, list)
+    self.assertTrue(
+        hasattr(result, 'columns'), "Result should have 'columns' attribute")
+    self.assertTrue(
+        hasattr(result, 'values'), "Result should have 'values' attribute")
+
+    result_n = ib.collect(pcoll, n=3)
+    self.assertNotIsInstance(result_n, list)
+    self.assertTrue(
+        hasattr(result_n, 'columns'),
+        "Result (n=3) should have 'columns' attribute")
+    self.assertTrue(
+        hasattr(result_n, 'values'),
+        "Result (n=3) should have 'values' attribute")
+
+  def test_collect_raw_records_true_multiple_pcolls(self):
+    p = beam.Pipeline(ir.InteractiveRunner())
+    data1 = list(range(3))
+    data2 = [x * x for x in range(3)]
+    pcoll1 = p | 'Create1' >> beam.Create(data1)
+    pcoll2 = p | 'Create2' >> beam.Create(data2)
+    ib.watch(locals())
+    ie.current_env().track_user_pipelines()
+
+    result = ib.collect(pcoll1, pcoll2, raw_records=True)
+    self.assertIsInstance(result, tuple)
+    self.assertEqual(len(result), 2)
+    self.assertIsInstance(result[0], list)
+    self.assertEqual(result[0], data1)
+    self.assertIsInstance(result[1], list)
+    self.assertEqual(result[1], data2)
+
+  def test_collect_raw_records_false_multiple_pcolls(self):
+    p = beam.Pipeline(ir.InteractiveRunner())
+    data1 = list(range(3))
+    data2 = [x * x for x in range(3)]
+    pcoll1 = p | 'Create1' >> beam.Create(data1)
+    pcoll2 = p | 'Create2' >> beam.Create(data2)
+    ib.watch(locals())
+    ie.current_env().track_user_pipelines()
+
+    result = ib.collect(pcoll1, pcoll2)
+    self.assertIsInstance(result, tuple)
+    self.assertEqual(len(result), 2)
+    self.assertNotIsInstance(result[0], list)
+    self.assertTrue(hasattr(result[0], 'columns'))
+    self.assertNotIsInstance(result[1], list)
+    self.assertTrue(hasattr(result[1], 'columns'))
+
+  def test_collect_raw_records_true_force_tuple(self):
+    p = beam.Pipeline(ir.InteractiveRunner())
+    data = list(range(5))
+    pcoll = p | 'Create' >> beam.Create(data)
+    ib.watch(locals())
+    ie.current_env().track_user_pipelines()
+
+    result = ib.collect(pcoll, raw_records=True, force_tuple=True)
+    self.assertIsInstance(result, tuple)
+    self.assertEqual(len(result), 1)
+    self.assertIsInstance(result[0], list)
+    self.assertEqual(result[0], data)
+
 
 @unittest.skipIf(
     not ie.current_env().is_interactive_ready,