Cache arrow object type to avoid repeated detection (#321)

This caches the arrow type to avoid detecting it every time a stream is
created.

Author: evertlammerts

src/duckdb_py/arrow/arrow_array_stream.cpp

@@ -64,9 +64,9 @@ unique_ptr<ArrowArrayStreamWrapper> PythonTableArrowArrayStreamFactory::Produce(
 	auto factory = static_cast<PythonTableArrowArrayStreamFactory *>(reinterpret_cast<void *>(factory_ptr)); // NOLINT
 	D_ASSERT(factory->arrow_object);
 	py::handle arrow_obj_handle(factory->arrow_object);
-	auto arrow_object_type = DuckDBPyConnection::GetArrowType(arrow_obj_handle);
+	auto arrow_object_type = factory->cached_arrow_type;
 
-	if (arrow_object_type == PyArrowObjectType::PyCapsuleInterface) {
+	if (arrow_object_type == PyArrowObjectType::PyCapsuleInterface || arrow_object_type == PyArrowObjectType::Table) {
 		py::object capsule_obj = arrow_obj_handle.attr("__arrow_c_stream__")();
 		auto capsule = py::reinterpret_borrow<py::capsule>(capsule_obj);
 		auto stream = capsule.get_pointer<struct ArrowArrayStream>();
@@ -181,8 +181,8 @@ void PythonTableArrowArrayStreamFactory::GetSchema(uintptr_t factory_ptr, ArrowS
 	D_ASSERT(factory->arrow_object);
 	py::handle arrow_obj_handle(factory->arrow_object);
 
-	auto type = DuckDBPyConnection::GetArrowType(arrow_obj_handle);
-	if (type == PyArrowObjectType::PyCapsuleInterface) {
+	auto type = factory->cached_arrow_type;
+	if (type == PyArrowObjectType::PyCapsuleInterface || type == PyArrowObjectType::Table) {
 		// Get __arrow_c_schema__ if it exists
 		if (py::hasattr(arrow_obj_handle, "__arrow_c_schema__")) {
 			auto schema_capsule = arrow_obj_handle.attr("__arrow_c_schema__")();

src/duckdb_py/include/duckdb_python/arrow/arrow_array_stream.hpp

@@ -59,8 +59,9 @@ PyArrowObjectType GetArrowType(const py::handle &obj);
 
 class PythonTableArrowArrayStreamFactory {
 public:
-	explicit PythonTableArrowArrayStreamFactory(PyObject *arrow_table, const ClientProperties &client_properties_p)
-	    : arrow_object(arrow_table), client_properties(client_properties_p) {};
+	explicit PythonTableArrowArrayStreamFactory(PyObject *arrow_table, const ClientProperties &client_properties_p,
+	                                            PyArrowObjectType arrow_type_p)
+	    : arrow_object(arrow_table), client_properties(client_properties_p), cached_arrow_type(arrow_type_p) {};
 
 	//! Produces an Arrow Scanner, should be only called once when initializing Scan States
 	static unique_ptr<ArrowArrayStreamWrapper> Produce(uintptr_t factory, ArrowStreamParameters &parameters);
@@ -73,6 +74,7 @@ class PythonTableArrowArrayStreamFactory {
 	PyObject *arrow_object;
 
 	const ClientProperties client_properties;
+	const PyArrowObjectType cached_arrow_type;
 
 private:
 	static py::object ProduceScanner(py::object &arrow_scanner, py::handle &arrow_obj_handle,

src/duckdb_py/python_replacement_scan.cpp

@@ -51,7 +51,7 @@ static void CreateArrowScan(const string &name, py::object entry, TableFunctionR
 		auto dependency_item = PythonDependencyItem::Create(stream_messages);
 		external_dependency->AddDependency("replacement_cache", std::move(dependency_item));
 	} else {
-		auto stream_factory = make_uniq<PythonTableArrowArrayStreamFactory>(entry.ptr(), client_properties);
+		auto stream_factory = make_uniq<PythonTableArrowArrayStreamFactory>(entry.ptr(), client_properties, type);
 		auto stream_factory_produce = PythonTableArrowArrayStreamFactory::Produce;
 		auto stream_factory_get_schema = PythonTableArrowArrayStreamFactory::GetSchema;
 

src/duckdb_py/python_udf.cpp

@@ -74,7 +74,8 @@ static void ConvertArrowTableToVector(const py::object &table, Vector &out, Clie
 	D_ASSERT(py::gil_check());
 	py::gil_scoped_release gil;
 
-	auto stream_factory = make_uniq<PythonTableArrowArrayStreamFactory>(ptr, context.GetClientProperties());
+	auto stream_factory =
+	    make_uniq<PythonTableArrowArrayStreamFactory>(ptr, context.GetClientProperties(), PyArrowObjectType::Table);
 	auto stream_factory_produce = PythonTableArrowArrayStreamFactory::Produce;
 	auto stream_factory_get_schema = PythonTableArrowArrayStreamFactory::GetSchema;