add arrow stream test

Author: wangxiaoying

connectorx-python/connectorx/tests/test_arrow.py

@@ -44,6 +44,73 @@ def test_arrow(postgres_url: str) -> None:
     df.sort_values(by="test_int", inplace=True, ignore_index=True)
     assert_frame_equal(df, expected, check_names=True)
 
+def test_arrow_stream(postgres_url: str) -> None:
+    import pyarrow as pa
+    query = "SELECT * FROM test_table"
+    reader = read_sql(
+        postgres_url,
+        query,
+        return_type="arrow_record_batches",
+        record_batch_size=2,
+    )
+    batches = []
+    for batch in reader:
+        batches.append(batch)
+    table = pa.Table.from_batches(batches)
+    df = table.to_pandas()
+    df.sort_values(by="test_int", inplace=True, ignore_index=True)
+
+    expected = pd.DataFrame(
+        index=range(6),
+        data={
+            "test_int": pd.Series([0, 1, 2, 3, 4, 1314], dtype="int64"),
+            "test_nullint": pd.Series([5, 3, None, 7, 9, 2], dtype="float64"),
+            "test_str": pd.Series(
+                ["a", "str1", "str2", "b", "c", None], dtype="object"
+            ),
+            "test_float": pd.Series([3.1, None, 2.2, 3, 7.8, -10], dtype="float64"),
+            "test_bool": pd.Series(
+                [None, True, False, False, None, True], dtype="object"
+            ),
+        },
+    )
+    assert_frame_equal(df, expected, check_names=True)
+
+def test_arrow_stream_with_partition(postgres_url: str) -> None:
+    import pyarrow as pa
+    query = "SELECT * FROM test_table"
+    reader = read_sql(
+        postgres_url,
+        query,
+        partition_on="test_int",
+        partition_range=(0, 2000),
+        partition_num=3,
+        return_type="arrow_record_batches",
+        record_batch_size=2,
+    )
+    batches = []
+    for batch in reader:
+        batches.append(batch)
+    table = pa.Table.from_batches(batches)
+    df = table.to_pandas()
+    df.sort_values(by="test_int", inplace=True, ignore_index=True)
+
+    expected = pd.DataFrame(
+        index=range(6),
+        data={
+            "test_int": pd.Series([0, 1, 2, 3, 4, 1314], dtype="int64"),
+            "test_nullint": pd.Series([5, 3, None, 7, 9, 2], dtype="float64"),
+            "test_str": pd.Series(
+                ["a", "str1", "str2", "b", "c", None], dtype="object"
+            ),
+            "test_float": pd.Series([3.1, None, 2.2, 3, 7.8, -10], dtype="float64"),
+            "test_bool": pd.Series(
+                [None, True, False, False, None, True], dtype="object"
+            ),
+        },
+    )
+    assert_frame_equal(df, expected, check_names=True)
+
 def decimal_s10(val):
     return Decimal(val).quantize(Decimal("0.0000000001"))