Support IGNORE_ERRORS when scanning from pyarrow/pandas (#4646)

royi-luo · andyfengHKU · web-flow · commit 6c40b466b710 · 2025-01-02T10:01:38.000+08:00
* Support IGNORE_ERRORS when scanning from pyarrow/pandas

* minor fix

---------

Co-authored-by: xiyang &lt;x74feng@uwaterloo.ca&gt;
diff --git a/src_cpp/include/pandas/pandas_scan.h b/src_cpp/include/pandas/pandas_scan.h
@@ -1,11 +1,9 @@
 #pragma once
 
-#include "function/scalar_function.h"
 #include "function/table/bind_data.h"
 #include "function/table/scan_functions.h"
 #include "function/table_functions.h"
 #include "pandas_bind.h"
-#include "pybind_include.h"
 
 namespace kuzu {
 
@@ -33,17 +31,24 @@ struct PandasScanFunction {
 struct PandasScanFunctionData : public function::TableFuncBindData {
     py::handle df;
     std::vector<std::unique_ptr<PandasColumnBindData>> columnBindData;
+    common::ReaderConfig config;
 
     PandasScanFunctionData(binder::expression_vector columns, py::handle df, uint64_t numRows,
-        std::vector<std::unique_ptr<PandasColumnBindData>> columnBindData)
+        std::vector<std::unique_ptr<PandasColumnBindData>> columnBindData,
+        common::ReaderConfig config)
         : TableFuncBindData{std::move(columns), 0 /* numWarningDataColumns */, numRows}, df{df},
-          columnBindData{std::move(columnBindData)} {}
+          columnBindData{std::move(columnBindData)}, config(std::move(config)) {}
 
     ~PandasScanFunctionData() override {
         py::gil_scoped_acquire acquire;
         columnBindData.clear();
     }
 
+    bool getIgnoreErrorsOption() const override {
+        return config.getOption(common::CopyConstants::IGNORE_ERRORS_OPTION_NAME,
+            common::CopyConstants::DEFAULT_IGNORE_ERRORS);
+    }
+
     std::vector<std::unique_ptr<PandasColumnBindData>> copyColumnBindData() const;
 
     std::unique_ptr<function::TableFuncBindData> copy() const override {
@@ -56,6 +61,7 @@ struct PandasScanFunctionData : public function::TableFuncBindData {
         for (const auto& i : other.columnBindData) {
             columnBindData.push_back(i->copy());
         }
+        config = other.config.copy();
     }
 };
 
diff --git a/src_cpp/include/pyarrow/pyarrow_scan.h b/src_cpp/include/pyarrow/pyarrow_scan.h
@@ -1,5 +1,7 @@
 #pragma once
 
+#include <utility>
+
 #include "common/arrow/arrow.h"
 #include "function/scalar_function.h"
 #include "function/table/bind_data.h"
@@ -12,6 +14,7 @@ namespace kuzu {
 struct PyArrowScanConfig {
     uint64_t skipNum;
     uint64_t limitNum;
+    bool ignoreErrors;
     explicit PyArrowScanConfig(const common::case_insensitive_map_t<common::Value>& options);
 };
 
@@ -35,19 +38,22 @@ struct PyArrowTableScanSharedState final : public function::BaseScanSharedStateW
 struct PyArrowTableScanFunctionData final : public function::TableFuncBindData {
     std::shared_ptr<ArrowSchemaWrapper> schema;
     std::vector<std::shared_ptr<ArrowArrayWrapper>> arrowArrayBatches;
+    bool ignoreErrors;
 
     PyArrowTableScanFunctionData(binder::expression_vector columns,
         std::shared_ptr<ArrowSchemaWrapper> schema,
-        std::vector<std::shared_ptr<ArrowArrayWrapper>> arrowArrayBatches, uint64_t numRows)
+        std::vector<std::shared_ptr<ArrowArrayWrapper>> arrowArrayBatches, uint64_t numRows,
+        bool ignoreErrors)
         : TableFuncBindData{std::move(columns), 0 /* numWarningDataColumns */, numRows},
-          schema{std::move(schema)}, arrowArrayBatches{std::move(arrowArrayBatches)} {}
+          schema{std::move(schema)}, arrowArrayBatches{std::move(arrowArrayBatches)},
+          ignoreErrors(ignoreErrors) {}
 
     ~PyArrowTableScanFunctionData() override {}
 
+    bool getIgnoreErrorsOption() const override { return ignoreErrors; }
+
 private:
-    PyArrowTableScanFunctionData(const PyArrowTableScanFunctionData& other)
-        : TableFuncBindData{other}, schema{other.schema},
-          arrowArrayBatches{other.arrowArrayBatches} {}
+    PyArrowTableScanFunctionData(const PyArrowTableScanFunctionData& other) = default;
 
 public:
     std::unique_ptr<function::TableFuncBindData> copy() const override {
diff --git a/src_cpp/pandas/pandas_scan.cpp b/src_cpp/pandas/pandas_scan.cpp
@@ -5,6 +5,7 @@
 #include "common/exception/runtime.h"
 #include "function/table/bind_input.h"
 #include "numpy/numpy_scan.h"
+#include "processor/execution_context.h"
 #include "py_connection.h"
 #include "pyarrow/pyarrow_scan.h"
 #include "pybind11/pytypes.h"
@@ -31,8 +32,9 @@ std::unique_ptr<TableFuncBindData> bindFunc(ClientContext* /*context*/,
     auto getFunc = df.attr("__getitem__");
     auto numRows = py::len(getFunc(columns[0]));
     auto returnColumns = input->binder->createVariables(names, returnTypes);
+    auto scanConfig = input->extraInput->constPtrCast<ExtraScanTableFuncBindInput>()->config.copy();
     return std::make_unique<PandasScanFunctionData>(std::move(returnColumns), df, numRows,
-        std::move(columnBindData));
+        std::move(columnBindData), std::move(scanConfig));
 }
 
 bool sharedStateNext(const TableFuncBindData* /*bindData*/, PandasScanLocalState* localState,
@@ -119,6 +121,10 @@ static double progressFunc(TableFuncSharedState* sharedState) {
     return static_cast<double>(pandasSharedState->numRowsRead) / pandasSharedState->numRows;
 }
 
+static void finalizeFunc(const processor::ExecutionContext* ctx, TableFuncSharedState*) {
+    ctx->clientContext->getWarningContextUnsafe().defaultPopulateAllWarnings(ctx->queryID);
+}
+
 function_set PandasScanFunction::getFunctionSet() {
     function_set functionSet;
     functionSet.push_back(getFunction().copy());
@@ -132,6 +138,7 @@ TableFunction PandasScanFunction::getFunction() {
     function.initSharedStateFunc = initSharedState;
     function.initLocalStateFunc = initLocalState;
     function.progressFunc = progressFunc;
+    function.finalizeFunc = finalizeFunc;
     return function;
 }
 
diff --git a/src_cpp/pyarrow/pyarrow_scan.cpp b/src_cpp/pyarrow/pyarrow_scan.cpp
@@ -5,6 +5,7 @@
 #include "common/arrow/arrow_converter.h"
 #include "function/cast/functions/numeric_limits.h"
 #include "function/table/bind_input.h"
+#include "processor/execution_context.h"
 #include "py_connection.h"
 #include "pyarrow/pyarrow_bind.h"
 #include "pybind11/pytypes.h"
@@ -18,6 +19,7 @@ namespace kuzu {
 PyArrowScanConfig::PyArrowScanConfig(const case_insensitive_map_t<Value>& options) {
     skipNum = 0;
     limitNum = NumericLimits<uint64_t>::maximum();
+    ignoreErrors = CopyConstants::DEFAULT_IGNORE_ERRORS;
     for (const auto& i : options) {
         if (i.first == "SKIP") {
             if (i.second.getDataType().getLogicalTypeID() != LogicalTypeID::INT64 ||
@@ -31,8 +33,14 @@ PyArrowScanConfig::PyArrowScanConfig(const case_insensitive_map_t<Value>& option
                 throw BinderException("LIMIT Option must be a positive integer literal.");
             }
             limitNum = i.second.val.int64Val;
+        } else if (i.first == CopyConstants::IGNORE_ERRORS_OPTION_NAME) {
+            if (i.second.getDataType().getLogicalTypeID() != LogicalTypeID::BOOL) {
+                throw BinderException("IGNORE_ERRORS Option must be a boolean.");
+            }
+            ignoreErrors = i.second.val.booleanVal;
         } else {
-            throw BinderException(stringFormat("{} Option not recognized by pyArrow scanner."));
+            throw BinderException(
+                stringFormat("{} Option not recognized by pyArrow scanner.", i.first));
         }
     }
 }
@@ -82,7 +90,7 @@ static std::unique_ptr<TableFuncBindData> bindFunc(ClientContext*,
 
     auto columns = input->binder->createVariables(names, returnTypes);
     return std::make_unique<PyArrowTableScanFunctionData>(std::move(columns), std::move(schema),
-        arrowArrayBatches, numRows);
+        arrowArrayBatches, numRows, config.ignoreErrors);
 }
 
 ArrowArrayWrapper* PyArrowTableScanSharedState::getNextChunk() {
@@ -143,13 +151,18 @@ function_set PyArrowTableScanFunction::getFunctionSet() {
     return functionSet;
 }
 
+static void finalizeFunc(const processor::ExecutionContext* ctx, TableFuncSharedState*) {
+    ctx->clientContext->getWarningContextUnsafe().defaultPopulateAllWarnings(ctx->queryID);
+}
+
 TableFunction PyArrowTableScanFunction::getFunction() {
     auto function = TableFunction(name, std::vector{LogicalTypeID::POINTER});
     function.tableFunc = tableFunc;
     function.bindFunc = bindFunc;
     function.initSharedStateFunc = initSharedState;
     function.initLocalStateFunc = initLocalState;
     function.progressFunc = progressFunc;
+    function.finalizeFunc = finalizeFunc;
     return function;
 }
 
diff --git a/test/test_scan_pandas.py b/test/test_scan_pandas.py
@@ -429,5 +429,29 @@ def test_scan_string_to_nested(tmp_path: Path) -> None:
     )
     conn.execute("COPY tab from df")
     result = conn.execute("match (t:tab) return t.*")
-    assert result.get_next() == [1, [1, 2, 3], {"'a'": 1, "'b'": 2}, {"a": 1, "b": 2}, [[], [1, 2, 3], [4, 5, 6]]]
+    assert result.get_next() == [
+        1,
+        [1, 2, 3],
+        {"'a'": 1, "'b'": 2},
+        {"a": 1, "b": 2},
+        [[], [1, 2, 3], [4, 5, 6]],
+    ]
     assert not result.has_next()
+
+
+def test_pandas_scan_ignore_errors(tmp_path: Path) -> None:
+    db = kuzu.Database(tmp_path)
+    conn = kuzu.Connection(db)
+    df = pd.DataFrame({"id": [1, 2, 3, 1]})
+    conn.execute("CREATE NODE TABLE person(id INT64, PRIMARY KEY(id))")
+    conn.execute("COPY person FROM df(IGNORE_ERRORS=true)")
+
+    people = conn.execute("MATCH (p:person) RETURN p.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()
diff --git a/test/test_scan_polars.py b/test/test_scan_polars.py
@@ -47,5 +47,25 @@ def test_polars_error(conn_db_readonly: ConnDB) -> None:
     with pytest.raises(RuntimeError, match="Binder exception: Variable df is not in scope."):
         conn.execute("LOAD FROM df RETURN *;")
     df = []
-    with pytest.raises(RuntimeError, match="Binder exception: Variable df found but no matches were scannable"):
+    with pytest.raises(
+        RuntimeError,
+        match="Binder exception: Variable df found but no matches were scannable",
+    ):
         conn.execute("LOAD FROM df RETURN *;")
+
+
+def test_polars_scan_ignore_errors(conn_db_readwrite: ConnDB) -> None:
+    conn, db = conn_db_readwrite
+    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()
diff --git a/test/test_scan_pyarrow.py b/test/test_scan_pyarrow.py
@@ -1,4 +1,5 @@
 import pyarrow as pa
+import pytest
 from type_aliases import ConnDB
 
 
@@ -64,3 +65,39 @@ def test_pyarrow_copy_from(conn_db_readwrite: ConnDB) -> None:
     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, db = 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, db = 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)")
