Fix arrow stream construction (duckdb#294)

Author: evertlammerts

src/duckdb_py/arrow/arrow_array_stream.cpp

@@ -27,15 +27,15 @@ void VerifyArrowDatasetLoaded() {
 	}
 }
 
-py::object PythonTableArrowArrayStreamFactory::ProduceScanner(DBConfig &config, py::object &arrow_scanner,
-                                                              py::handle &arrow_obj_handle,
+py::object PythonTableArrowArrayStreamFactory::ProduceScanner(py::object &arrow_scanner, py::handle &arrow_obj_handle,
                                                               ArrowStreamParameters &parameters,
                                                               const ClientProperties &client_properties) {
 	D_ASSERT(!py::isinstance<py::capsule>(arrow_obj_handle));
 	ArrowSchemaWrapper schema;
 	PythonTableArrowArrayStreamFactory::GetSchemaInternal(arrow_obj_handle, schema);
 	ArrowTableSchema arrow_table;
-	ArrowTableFunction::PopulateArrowTableSchema(config, arrow_table, schema.arrow_schema);
+	ArrowTableFunction::PopulateArrowTableSchema(*client_properties.client_context.get_mutable(), arrow_table,
+	                                             schema.arrow_schema);
 
 	auto filters = parameters.filters;
 	auto &column_list = parameters.projected_columns.columns;
@@ -86,26 +86,23 @@ unique_ptr<ArrowArrayStreamWrapper> PythonTableArrowArrayStreamFactory::Produce(
 		auto arrow_dataset = import_cache.pyarrow.dataset().attr("dataset");
 		auto dataset = arrow_dataset(arrow_obj_handle);
 		py::object arrow_scanner = dataset.attr("__class__").attr("scanner");
-		scanner = ProduceScanner(factory->config, arrow_scanner, dataset, parameters, factory->client_properties);
+		scanner = ProduceScanner(arrow_scanner, dataset, parameters, factory->client_properties);
 		break;
 	}
 	case PyArrowObjectType::RecordBatchReader: {
-		scanner = ProduceScanner(factory->config, arrow_batch_scanner, arrow_obj_handle, parameters,
-		                         factory->client_properties);
+		scanner = ProduceScanner(arrow_batch_scanner, arrow_obj_handle, parameters, factory->client_properties);
 		break;
 	}
 	case PyArrowObjectType::Scanner: {
 		// If it's a scanner we have to turn it to a record batch reader, and then a scanner again since we can't stack
 		// scanners on arrow Otherwise pushed-down projections and filters will disappear like tears in the rain
 		auto record_batches = arrow_obj_handle.attr("to_reader")();
-		scanner = ProduceScanner(factory->config, arrow_batch_scanner, record_batches, parameters,
-		                         factory->client_properties);
+		scanner = ProduceScanner(arrow_batch_scanner, record_batches, parameters, factory->client_properties);
 		break;
 	}
 	case PyArrowObjectType::Dataset: {
 		py::object arrow_scanner = arrow_obj_handle.attr("__class__").attr("scanner");
-		scanner =
-		    ProduceScanner(factory->config, arrow_scanner, arrow_obj_handle, parameters, factory->client_properties);
+		scanner = ProduceScanner(arrow_scanner, arrow_obj_handle, parameters, factory->client_properties);
 		break;
 	}
 	default: {

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

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

src/duckdb_py/python_replacement_scan.cpp

@@ -18,7 +18,7 @@ namespace duckdb {
 
 static void CreateArrowScan(const string &name, py::object entry, TableFunctionRef &table_function,
                             vector<unique_ptr<ParsedExpression>> &children, ClientProperties &client_properties,
-                            PyArrowObjectType type, DBConfig &config, DatabaseInstance &db) {
+                            PyArrowObjectType type, DatabaseInstance &db) {
 	shared_ptr<ExternalDependency> external_dependency = make_shared_ptr<ExternalDependency>();
 	if (type == PyArrowObjectType::MessageReader) {
 		if (!db.ExtensionIsLoaded("nanoarrow")) {
@@ -56,7 +56,7 @@ static void CreateArrowScan(const string &name, py::object entry, TableFunctionR
 			type = PyArrowObjectType::PyCapsule;
 		}
 
-		auto stream_factory = make_uniq<PythonTableArrowArrayStreamFactory>(entry.ptr(), client_properties, config);
+		auto stream_factory = make_uniq<PythonTableArrowArrayStreamFactory>(entry.ptr(), client_properties);
 		auto stream_factory_produce = PythonTableArrowArrayStreamFactory::Produce;
 		auto stream_factory_get_schema = PythonTableArrowArrayStreamFactory::GetSchema;
 
@@ -113,7 +113,7 @@ unique_ptr<TableRef> PythonReplacementScan::TryReplacementObject(const py::objec
 		if (PandasDataFrame::IsPyArrowBacked(entry)) {
 			auto table = PandasDataFrame::ToArrowTable(entry);
 			CreateArrowScan(name, table, *table_function, children, client_properties, PyArrowObjectType::Table,
-			                DBConfig::GetConfig(context), *context.db);
+			                *context.db);
 		} else {
 			string name = "df_" + StringUtil::GenerateRandomName();
 			auto new_df = PandasScanFunction::PandasReplaceCopiedNames(entry);
@@ -143,17 +143,16 @@ unique_ptr<TableRef> PythonReplacementScan::TryReplacementObject(const py::objec
 	} else if (PolarsDataFrame::IsDataFrame(entry)) {
 		auto arrow_dataset = entry.attr("to_arrow")();
 		CreateArrowScan(name, arrow_dataset, *table_function, children, client_properties, PyArrowObjectType::Table,
-		                DBConfig::GetConfig(context), *context.db);
+		                *context.db);
 	} else if (PolarsDataFrame::IsLazyFrame(entry)) {
 		auto materialized = entry.attr("collect")();
 		auto arrow_dataset = materialized.attr("to_arrow")();
 		CreateArrowScan(name, arrow_dataset, *table_function, children, client_properties, PyArrowObjectType::Table,
-		                DBConfig::GetConfig(context), *context.db);
+		                *context.db);
 	} else if (DuckDBPyConnection::GetArrowType(entry) != PyArrowObjectType::Invalid &&
 	           !(DuckDBPyConnection::GetArrowType(entry) == PyArrowObjectType::MessageReader && !relation)) {
 		arrow_type = DuckDBPyConnection::GetArrowType(entry);
-		CreateArrowScan(name, entry, *table_function, children, client_properties, arrow_type,
-		                DBConfig::GetConfig(context), *context.db);
+		CreateArrowScan(name, entry, *table_function, children, client_properties, arrow_type, *context.db);
 	} else if (DuckDBPyConnection::IsAcceptedNumpyObject(entry) != NumpyObjectType::INVALID) {
 		numpytype = DuckDBPyConnection::IsAcceptedNumpyObject(entry);
 		string np_name = "np_" + StringUtil::GenerateRandomName();

src/duckdb_py/python_udf.cpp

@@ -74,8 +74,7 @@ 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(), DBConfig::GetConfig(context));
+	auto stream_factory = make_uniq<PythonTableArrowArrayStreamFactory>(ptr, context.GetClientProperties());
 	auto stream_factory_produce = PythonTableArrowArrayStreamFactory::Produce;
 	auto stream_factory_get_schema = PythonTableArrowArrayStreamFactory::GetSchema;