real-ladybug 0.0.1.dev1__cp311-cp311-win_amd64.whl

real_ladybug-source/tools/python_api/test/test_scan_polars.py

@@ -0,0 +1,165 @@
+from datetime import datetime
+
+import polars as pl
+import pytest
+from type_aliases import ConnDB
+
+
+def test_polars_basic(conn_db_empty: ConnDB) -> None:
+    conn, _ = conn_db_empty
+    df = pl.DataFrame([
+        pl.Series("col1", [1, 2], dtype=pl.Int64),
+        pl.Series("col2", ["a", "b"], dtype=pl.String),
+        pl.Series("col3", [1.0, 2.0], dtype=pl.Float64),
+        pl.Series("col4", [datetime(2020, 5, 5), datetime(2000, 1, 1)], dtype=pl.Datetime),
+    ])
+    result = conn.execute("LOAD FROM df RETURN *").get_as_pl()
+    equivalency = result == df
+    assert equivalency["col1"].all()
+    assert equivalency["col2"].all()
+    assert equivalency["col3"].all()
+    assert equivalency["col4"].all()
+    assert result["col1"].dtype == pl.Int64
+    assert result["col2"].dtype == pl.String
+    assert result["col3"].dtype == pl.Float64
+    assert result["col4"].dtype == pl.Datetime
+    # since polars just uses arrow as its backend, it's probably not necessary to extensively test this functionality
+    conn.execute(
+        "CREATE NODE TABLE polarstab(col1 INT64, col2 STRING, col3 DOUBLE, col4 TIMESTAMP, PRIMARY KEY(col1))"
+    )
+    conn.execute("COPY polarstab FROM df")
+    result = conn.execute(
+        "MATCH (t:polarstab) RETURN t.col1 AS col1, t.col2 AS col2, t.col3 AS col3, t.col4 AS col4 ORDER BY col1"
+    ).get_as_pl()
+    equivalency = result == df
+    assert equivalency["col1"].all()
+    assert equivalency["col2"].all()
+    assert equivalency["col3"].all()
+    assert equivalency["col4"].all()
+    assert result["col1"].dtype == pl.Int64
+    assert result["col2"].dtype == pl.String
+    assert result["col3"].dtype == pl.Float64
+    assert result["col4"].dtype == pl.Datetime
+
+
+def test_polars_basic_param(conn_db_empty: ConnDB) -> None:
+    conn, _ = conn_db_empty
+    df = pl.DataFrame([
+        pl.Series("col1", [1, 2], dtype=pl.Int64),
+        pl.Series("col2", ["a", "b"], dtype=pl.String),
+        pl.Series("col3", [1.0, 2.0], dtype=pl.Float64),
+        pl.Series("col4", [datetime(2020, 5, 5), datetime(2000, 1, 1)], dtype=pl.Datetime),
+    ])
+    result = conn.execute("LOAD FROM df RETURN *").get_as_pl()
+    equivalency = result == df
+    assert equivalency["col1"].all()
+    assert equivalency["col2"].all()
+    assert equivalency["col3"].all()
+    assert equivalency["col4"].all()
+    assert result["col1"].dtype == pl.Int64
+    assert result["col2"].dtype == pl.String
+    assert result["col3"].dtype == pl.Float64
+    assert result["col4"].dtype == pl.Datetime
+    # since polars just uses arrow as its backend, it's probably not necessary to extensively test this functionality
+    conn.execute(
+        "CREATE NODE TABLE polarstab(col1 INT64, col2 STRING, col3 DOUBLE, col4 TIMESTAMP, PRIMARY KEY(col1))"
+    )
+    conn.execute("COPY polarstab FROM $df", {"df": df})
+    result = conn.execute(
+        "MATCH (t:polarstab) RETURN t.col1 AS col1, t.col2 AS col2, t.col3 AS col3, t.col4 AS col4 ORDER BY col1"
+    ).get_as_pl()
+    equivalency = result == df
+    assert equivalency["col1"].all()
+    assert equivalency["col2"].all()
+    assert equivalency["col3"].all()
+    assert equivalency["col4"].all()
+    assert result["col1"].dtype == pl.Int64
+    assert result["col2"].dtype == pl.String
+    assert result["col3"].dtype == pl.Float64
+    assert result["col4"].dtype == pl.Datetime
+
+
+def test_polars_error(conn_db_readonly: ConnDB) -> None:
+    conn, _ = conn_db_readonly
+    with pytest.raises(RuntimeError, match=r"Binder exception: Variable df is not in scope."):
+        conn.execute("LOAD FROM df RETURN *;")
+    df = []
+    with pytest.raises(
+        RuntimeError,
+        match=r"Binder exception: Attempted to scan from unsupported python object. Can only scan from pandas/polars dataframes and pyarrow tables.",
+    ):
+        conn.execute("LOAD FROM df RETURN *;")
+
+
+def test_polars_scan_ignore_errors(conn_db_empty: ConnDB) -> None:
+    conn, _ = conn_db_empty
+    df = pl.DataFrame({"id": [1, 2, 3, 1]})
+    conn.execute("CREATE NODE TABLE ids(id INT64, PRIMARY KEY(id))")
+    conn.execute("COPY ids FROM df(IGNORE_ERRORS=true)")
+
+    people = conn.execute("MATCH (i:ids) RETURN i.id")
+    assert people.get_next() == [1]
+    assert people.get_next() == [2]
+    assert people.get_next() == [3]
+    assert not people.has_next()
+
+    warnings = conn.execute("CALL show_warnings() RETURN *")
+    assert warnings.get_next()[1].startswith("Found duplicated primary key value 1")
+    assert not warnings.has_next()
+
+
+def test_copy_from_polars_multi_pairs(conn_db_empty: ConnDB) -> None:
+    conn, _ = conn_db_empty
+    conn.execute("CREATE NODE TABLE prof(id INT64, PRIMARY KEY(id))")
+    conn.execute("CREATE (p:prof {id: 3});")
+    conn.execute("CREATE (p:prof {id: 4});")
+    conn.execute("CREATE NODE TABLE student(id INT64, PRIMARY KEY(id))")
+    conn.execute("CREATE (p:student {id: 2});")
+    conn.execute("CREATE REL TABLE teaches(from prof to prof, from prof to student, length int64)")
+    df = pl.DataFrame({"from": [3], "to": [4], "length": [252]})
+    conn.execute("COPY teaches from df (from = 'prof', to = 'prof');")
+    result = conn.execute("match (:prof)-[e:teaches]->(:prof) return e.*")
+    assert result.has_next()
+    assert result.get_next()[0] == 252
+    assert not result.has_next()
+
+
+def test_scan_from_empty_lst(conn_db_empty: ConnDB) -> None:
+    conn, _ = conn_db_empty
+    df = pl.DataFrame({"prop1": [3], "prop2": [[]]})
+    result = conn.execute("LOAD FROM df RETURN *")
+    assert result.has_next()
+    tp = result.get_next()
+    assert tp[0] == 3
+    assert tp[1] == []
+
+
+def test_scan_from_parameterized_df_docs_example_1(conn_db_empty: ConnDB):
+    conn, _ = conn_db_empty
+
+    conn.execute("CREATE NODE TABLE Person(name STRING, age INT64, PRIMARY KEY (name))")
+
+    df = pl.DataFrame({"name": ["Adam", "Karissa", "Zhang"], "age": [30, 40, 50]})
+
+    conn.execute("COPY Person FROM $dataframe", {"dataframe": df})
+
+
+def test_scan_from_parameterized_df_docs_example_2(conn_db_empty: ConnDB):
+    conn, _ = conn_db_empty
+
+    conn.execute("CREATE NODE TABLE Person(name STRING, age INT64, PRIMARY KEY (name))")
+
+    def get_df():
+        return pl.DataFrame({"name": ["Adam", "Karissa", "Zhang"], "age": [30, 40, 50]})
+
+    conn.execute("COPY Person FROM $dataframe", {"dataframe": get_df()})
+
+
+def test_scan_from_df_docs_example(conn_db_empty: ConnDB):
+    conn, _ = conn_db_empty
+
+    conn.execute("CREATE NODE TABLE Person(name STRING, age INT64, PRIMARY KEY (name))")
+
+    df = pl.DataFrame({"name": ["Adam", "Karissa", "Zhang"], "age": [30, 40, 50]})
+
+    conn.execute("COPY Person FROM df")

real_ladybug-source/tools/python_api/test/test_scan_pyarrow.py

@@ -0,0 +1,167 @@
+import pyarrow as pa
+import pytest
+from type_aliases import ConnDB
+
+
+def test_pyarrow_basic(conn_db_readonly: ConnDB) -> None:
+    conn, _ = conn_db_readonly
+    tab = pa.Table.from_arrays(
+        [
+            [1, 2, 3],
+            ["a", "b", "c"],
+            [True, False, None],
+        ],
+        names=["col1", "col2", "col3"],
+    )
+    result = conn.execute("LOAD FROM tab RETURN * ORDER BY col1")
+    assert result.get_next() == [1, "a", True]
+    assert result.get_next() == [2, "b", False]
+    assert result.get_next() == [3, "c", None]
+
+
+def test_pyarrow_copy_from_parameterized_df(conn_db_readwrite: ConnDB) -> None:
+    conn, _ = conn_db_readwrite
+
+    def get_tab_func():
+        return pa.Table.from_arrays(
+            [
+                pa.array([1, 2, 3], type=pa.int32()),
+                pa.array(["a", "b", "c"], type=pa.string()),
+                pa.array([True, False, None], type=pa.bool_()),
+            ],
+            names=["id", "A", "B"],
+        )
+
+    conn.execute("CREATE NODE TABLE pyarrowtab(id INT32, A STRING, B BOOL, PRIMARY KEY(id))")
+    conn.execute("COPY pyarrowtab FROM $tab", {"tab": get_tab_func()})
+    result = conn.execute("MATCH (t:pyarrowtab) RETURN t.id AS id, t.A AS A, t.B AS B ORDER BY t.id")
+    assert result.get_next() == [1, "a", True]
+    assert result.get_next() == [2, "b", False]
+    assert result.get_next() == [3, "c", None]
+
+    rels = pa.Table.from_arrays(
+        [pa.array([1, 2, 3], type=pa.int32()), pa.array([2, 3, 1], type=pa.int32())], names=["from", "to"]
+    )
+    conn.execute("CREATE REL TABLE pyarrowrel(FROM pyarrowtab TO pyarrowtab)")
+    conn.execute("COPY pyarrowrel FROM $tab", {"tab": rels})
+    result = conn.execute("MATCH (a:pyarrowtab)-[:pyarrowrel]->(b:pyarrowtab) RETURN a.id, b.id ORDER BY a.id")
+    assert result.get_next() == [1, 2]
+    assert result.get_next() == [2, 3]
+    assert result.get_next() == [3, 1]
+
+
+def test_pyarrow_copy_from_invalid_source(conn_db_readwrite: ConnDB) -> None:
+    conn, _ = conn_db_readwrite
+    conn.execute("CREATE NODE TABLE pyarrowtab(id INT32, A STRING, B BOOL, PRIMARY KEY(id))")
+    with pytest.raises(
+        RuntimeError,
+        match=r"Binder exception: Trying to scan from unsupported data type INT8\[\]. The only parameter types that can be scanned from are pandas/polars dataframes and pyarrow tables.",
+    ):
+        conn.execute("COPY pyarrowtab FROM $tab", {"tab": [1, 2, 3]})
+
+
+def test_pyarrow_copy_from(conn_db_readwrite: ConnDB) -> None:
+    conn, _ = conn_db_readwrite
+    tab = pa.Table.from_arrays(
+        [
+            pa.array([1, 2, 3], type=pa.int32()),
+            pa.array(["a", "b", "c"], type=pa.string()),
+            pa.array([True, False, None], type=pa.bool_()),
+        ],
+        names=["id", "A", "B"],
+    )
+    conn.execute("CREATE NODE TABLE pyarrowtab(id INT32, A STRING, B BOOL, PRIMARY KEY(id))")
+    conn.execute("COPY pyarrowtab FROM tab")
+    result = conn.execute("MATCH (t:pyarrowtab) RETURN t.id AS id, t.A AS A, t.B AS B ORDER BY t.id")
+    assert result.get_next() == [1, "a", True]
+    assert result.get_next() == [2, "b", False]
+    assert result.get_next() == [3, "c", None]
+
+    rels = pa.Table.from_arrays(
+        [pa.array([1, 2, 3], type=pa.int32()), pa.array([2, 3, 1], type=pa.int32())], names=["from", "to"]
+    )
+    conn.execute("CREATE REL TABLE pyarrowrel(FROM pyarrowtab TO pyarrowtab)")
+    conn.execute("COPY pyarrowrel FROM rels")
+    result = conn.execute("MATCH (a:pyarrowtab)-[:pyarrowrel]->(b:pyarrowtab) RETURN a.id, b.id ORDER BY a.id")
+    assert result.get_next() == [1, 2]
+    assert result.get_next() == [2, 3]
+    assert result.get_next() == [3, 1]
+
+    conn.execute("DROP TABLE pyarrowrel")
+
+    rels = pa.Table.from_arrays(
+        [
+            pa.array([1, 2, 3], type=pa.int32()),
+            pa.array([2, 3, 1], type=pa.int32()),
+            pa.array(["honk", "shoo", "mimimimimimimi"], type=pa.string()),
+        ],
+        names=["foo", "bar", "spongebobmeboy"],
+    )
+
+    conn.execute("CREATE REL TABLE pyarrowrel(FROM pyarrowtab TO pyarrowtab, a_distraction STRING)")
+    conn.execute("COPY pyarrowrel FROM rels")
+    result = conn.execute(
+        "MATCH (a:pyarrowtab)-[r:pyarrowrel]->(b:pyarrowtab) RETURN a.id, r.a_distraction, b.id ORDER BY a.id"
+    )
+    assert result.get_next() == [1, "honk", 2]
+    assert result.get_next() == [2, "shoo", 3]
+    assert result.get_next() == [3, "mimimimimimimi", 1]
+
+
+def test_pyarrow_scan_ignore_errors(conn_db_readwrite: ConnDB) -> None:
+    conn, _ = conn_db_readwrite
+    tab = pa.Table.from_arrays(
+        [
+            pa.array([1, 2, 3, 1], type=pa.int32()),
+        ],
+        names=["id"],
+    )
+    conn.execute("CREATE NODE TABLE ids(id INT64, PRIMARY KEY(id))")
+    conn.execute("COPY ids FROM tab(IGNORE_ERRORS=true)")
+
+    people = conn.execute("MATCH (i:ids) RETURN i.id")
+    assert people.get_next() == [1]
+    assert people.get_next() == [2]
+    assert people.get_next() == [3]
+    assert not people.has_next()
+
+    warnings = conn.execute("CALL show_warnings() RETURN *")
+    assert warnings.get_next()[1].startswith("Found duplicated primary key value 1")
+    assert not warnings.has_next()
+
+
+def test_pyarrow_scan_invalid_option(conn_db_readwrite: ConnDB) -> None:
+    conn, _ = conn_db_readwrite
+    tab = pa.Table.from_arrays(
+        [
+            pa.array([1, 2, 3], type=pa.int32()),
+        ],
+        names=["id"],
+    )
+    conn.execute("CREATE NODE TABLE ids(id INT64, PRIMARY KEY(id))")
+    error_message = "INVALID_OPTION Option not recognized by pyArrow scanner."
+    with pytest.raises(RuntimeError, match=error_message):
+        conn.execute("COPY ids FROM tab(INVALID_OPTION=1)")
+
+
+def test_copy_from_pyarrow_multi_pairs(conn_db_readwrite: ConnDB) -> None:
+    conn, _ = conn_db_readwrite
+    conn.execute("CREATE NODE TABLE prof(id INT64, PRIMARY KEY(id))")
+    conn.execute("CREATE (p:prof {id: 3});")
+    conn.execute("CREATE (p:prof {id: 4});")
+    conn.execute("CREATE NODE TABLE student(id INT64, PRIMARY KEY(id))")
+    conn.execute("CREATE (p:student {id: 2});")
+    conn.execute("CREATE REL TABLE teaches(from prof to prof, from prof to student, length int64)")
+    df = pa.Table.from_arrays(
+        [
+            pa.array([3], type=pa.int64()),
+            pa.array([4], type=pa.int64()),
+            pa.array([252], type=pa.int64()),
+        ],
+        names=["from", "to", "length"],
+    )
+    conn.execute("COPY teaches from df (from = 'prof', to = 'prof');")
+    result = conn.execute("match (:prof)-[e:teaches]->(:prof) return e.*")
+    assert result.has_next()
+    assert result.get_next()[0] == 252
+    assert not result.has_next()

real_ladybug-source/tools/python_api/test/test_timeout.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import pytest
+from type_aliases import ConnDB
+
+
+def test_timeout(conn_db_readonly: ConnDB) -> None:
+    conn, _ = conn_db_readonly
+    conn.set_query_timeout(1000)
+    with pytest.raises(RuntimeError, match=r"Interrupted."):
+        conn.execute("UNWIND RANGE(1,100000) AS x UNWIND RANGE(1, 100000) AS y RETURN COUNT(x + y);")