Support loading sequences in read_parquet (#250)

* support sequences in read parquet

* unit tests

* comments

* lint

Author: smcguire-cmu

src/nested_pandas/nestedframe/io.py

@@ -1,6 +1,8 @@
 # typing.Self and "|" union syntax don't exist in Python 3.9
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 import pandas as pd
 import pyarrow as pa
 import pyarrow.parquet as pq
@@ -82,9 +84,11 @@ def read_parquet(
         reject_nesting = [reject_nesting]
 
     # First load through pyarrow
-    # Check if `data` is a file-like object
-    if hasattr(data, "read"):
-        # If `data` is a file-like object, pass it directly to pyarrow
+    # Check if `data` is a file-like object or a sequence
+    if hasattr(data, "read") or (
+        isinstance(data, Sequence) and not isinstance(data, (str, bytes, bytearray))
+    ):
+        # If `data` is a file-like object or a sequence, pass it directly to pyarrow
         table = pq.read_table(data, columns=columns, **kwargs)
     else:
         # Otherwise, treat `data` as a file path and use UPath

tests/nested_pandas/nestedframe/test_io.py

@@ -27,6 +27,26 @@ def test_read_parquet():
     assert nf.lincc.nest.fields == ["band", "frameworks"]
 
 
+def test_read_parquet_list():
+    """Test reading a parquet file with no columns specified"""
+    # Load in the example files
+    single_file_nf = read_parquet("tests/test_data/nested.parquet")
+    nf = read_parquet(["tests/test_data/nested.parquet", "tests/test_data/nested.parquet"])
+
+    # Check the columns
+    assert nf.columns.tolist() == ["a", "flux", "nested", "lincc"]
+
+    # Make sure nested columns were recognized
+    assert nf.nested_columns == ["nested", "lincc"]
+
+    # Check the nested columns
+    assert nf.nested.nest.fields == ["t", "flux", "band"]
+    assert nf.lincc.nest.fields == ["band", "frameworks"]
+
+    # Check loading list works correctly
+    assert len(nf) == 2 * len(single_file_nf)
+
+
 def test_read_parquet_directory():
     """Test reading a parquet file with no columns specified"""
     # Load in the example file