Use bulk ingest in AdbcHook.insert_rows when fast_executemany is set (#76363)

* Use bulk ingest in AdbcHook.insert_rows when fast_executemany is set

AdbcHook.insert_rows built an Arrow RecordBatch per commit_every chunk
and bound it with cursor.executemany. Several ADBC drivers run a bound
stream once per row, so the previous fast_executemany path was a no-op
that set cursor.fast_executemany inside suppress(AttributeError).

Route fast_executemany (now the default) to cursor.adbc_ingest in
append mode. Keep executemany when fast_executemany is False, and
always bind when replace=True because bulk ingest cannot carry the
dialect's conflict clause. Chunking, commits, target_fields, and
schema.table splitting are unchanged.

Co-authored-by: Aldo <aoelvp94@gmail.com>

* Fall back to binding when the ADBC driver has no bulk ingest

Co-authored-by: Aldo <aoelvp94@gmail.com>

* Use Logger spec and patch decorators in ADBC insert_rows tests

Match AGENTS.md: spec every mock, and prefer @mock.patch.object over
a with-block for the replace=True insert_rows tests.

Co-authored-by: Aldo <aoelvp94@gmail.com>

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

providers/apache/arrow/docs/connections/adbc.rst

@@ -186,7 +186,9 @@ Use :class:`~airflow.providers.apache.arrow.hooks.adbc.AdbcHook` in a task to qu
         # Run a query and get results as a list of tuples
         records = hook.get_records("SELECT id, name FROM users WHERE active = 1")
 
-        # Bulk-insert rows in chunks of 5 000 using Arrow-native transfer
+        # Insert rows in chunks of 5 000. By default each chunk is loaded with the
+        # driver's bulk ingest (e.g. COPY on PostgreSQL); pass fast_executemany=False
+        # to bind the rows to an INSERT statement instead.
         rows = [(1, "Alice"), (2, "Bob")]
         hook.insert_rows(
             table="users",

providers/apache/arrow/src/airflow/providers/apache/arrow/hooks/adbc.py

@@ -17,17 +17,16 @@
 # under the License.
 from __future__ import annotations
 
-import contextlib
 import pathlib
 import re
 import sys
 from collections.abc import Callable, Iterable, Mapping
 from contextlib import closing
-from functools import cached_property
+from functools import cached_property, partial
 from importlib import resources as importlib_resources
 from typing import TYPE_CHECKING, Any
 
-from adbc_driver_manager.dbapi import Connection, connect
+from adbc_driver_manager.dbapi import Connection, NotSupportedError, connect
 from more_itertools import chunked
 from pyarrow import RecordBatch, Schema, array, schema
 
@@ -78,7 +77,8 @@ class AdbcHook(DbApiHook):
         - Handles dialect-specific connection URIs and driver entrypoints.
         - Integrates with Airflow's connection system (conn_id, extras, etc.).
         - Provides custom placeholder replacement for parameterized SQL queries.
-        - Supports both native Arrow binding and DBAPI ``executemany`` for inserts.
+        - Loads rows with the driver's bulk ingest (``adbc_ingest``, e.g. ``COPY`` on PostgreSQL)
+          by default, or binds them with DBAPI ``executemany`` (``fast_executemany=False``).
         - Exposes configuration via connection extras: driver, entrypoint, db_kwargs, conn_kwargs, dialect.
 
     Connection Extras:
@@ -238,34 +238,58 @@ class AdbcHook(DbApiHook):
             schema=schema,
         )
 
-    @classmethod
-    def _execute_executemany(cls, cursor, statement: str, record_batch: RecordBatch) -> None:
-        """
-        Execute a statement using cursor.executemany with a RecordBatch.
-
-        ADBC's executemany() accepts Arrow data directly and routes it through
-        bind_stream internally, making this the Arrow-native fast path.
-        """
-        cursor.executemany(statement, record_batch)
-
     def _resolve_execute_batch(
-        self, cursor, executemany: bool, fast_executemany: bool
-    ) -> Callable[[Cursor, str, RecordBatch], None]:
+        self,
+        cursor: Cursor,
+        *,
+        table_name: str,
+        schema_name: str | None,
+        sql: str,
+        replace: bool,
+        fast_executemany: bool,
+    ) -> Callable[[RecordBatch], object]:
         """
-        Return the batch-execute callable for the given cursor.
+        Return the callable that loads one RecordBatch through ``cursor``.
 
-        ADBC's executemany() natively accepts a RecordBatch, so it is always
-        the fast path.  fast_executemany is passed through to drivers that
-        expose that knob (silently ignored otherwise).
+        With ``fast_executemany`` the batch goes through the driver's bulk ingest
+        (:meth:`adbc_driver_manager.dbapi.Cursor.adbc_ingest`, ``COPY`` on PostgreSQL).
+        Otherwise it is bound to the generated ``INSERT`` statement with
+        ``cursor.executemany``, which most drivers execute once per row.
+        ``replace=True`` always binds, because bulk ingest cannot carry the
+        dialect's conflict clause.
+
+        ``fast_executemany`` is a hint: if the driver does not implement bulk ingest
+        (the first ``adbc_ingest`` call raises ``NotSupportedError``), that batch and
+        all following ones are bound with ``cursor.executemany`` instead.
         """
-        if fast_executemany:
-            with contextlib.suppress(AttributeError):
-                cursor.fast_executemany = True
+        bind = partial(cursor.executemany, sql)
+        if replace or not fast_executemany:
+            if replace and fast_executemany:
+                self.log.debug("Bulk ingest cannot replace rows; binding them to the statement instead")
+            return bind
+
+        self.log.info("Using bulk ingest (adbc_ingest) for conn_id '%s'", self.get_conn_id())
+        ingest = partial(cursor.adbc_ingest, table_name, mode="append", db_schema_name=schema_name)
+        loader: Callable[[RecordBatch], object] | None = None
+
+        def execute_batch(batch: RecordBatch) -> object:
+            nonlocal loader
+            if loader is not None:
+                return loader(batch)
+            try:
+                result = ingest(batch)
+            except NotSupportedError:
                 self.log.info(
-                    "Fast_executemany is enabled for conn_id '%s'!",
+                    "Driver for conn_id '%s' does not support bulk ingest; "
+                    "binding rows with executemany instead",
                     self.get_conn_id(),
                 )
-        return self._execute_executemany
+                loader = bind
+                return bind(batch)
+            loader = ingest
+            return result
+
+        return execute_batch
 
     def insert_rows(
         self,
@@ -276,27 +300,31 @@ class AdbcHook(DbApiHook):
         replace=False,
         *,
         executemany=False,
-        fast_executemany=False,
+        fast_executemany=True,
         autocommit=False,
         **kwargs,
     ):
         """
         Insert a collection of tuples into a table.
 
-        Rows are inserted in chunks, each chunk (of size ``commit_every``) is
-        done in a new transaction.
+        Rows are converted to Arrow ``RecordBatch`` objects and inserted in chunks;
+        each chunk (of size ``commit_every``) is done in a new transaction.
 
-        :param table: Name of the target table
+        :param table: Name of the target table, optionally prefixed with its schema
+            (``schema.table``).
         :param rows: The rows to insert into the table
         :param target_fields: The names of the columns to fill in the table
         :param commit_every: The maximum number of rows to insert in one
             transaction. Set to 0 to insert all rows in one transaction.
-        :param replace: Whether to replace instead of insert
-        :param executemany: If True, all rows are inserted at once in
-            chunks defined by the commit_every parameter. This only works if all rows
-            have same number of column names, but leads to better performance.
-        :param fast_executemany: If True, the `fast_executemany` parameter will be set on the
-            cursor used by `executemany` which leads to better performance, if supported by driver.
+        :param replace: Whether to replace instead of insert. Replacing always binds
+            the rows to the generated statement, because bulk ingest cannot express it.
+        :param executemany: Kept for signature compatibility with
+            :meth:`~airflow.providers.common.sql.hooks.sql.DbApiHook.insert_rows`;
+            rows are always sent as one Arrow batch per chunk.
+        :param fast_executemany: If True (default), load each batch with the driver's
+            bulk ingest (``cursor.adbc_ingest``, ``COPY`` on PostgreSQL). If the driver
+            does not implement bulk ingest, falls back to binding. If False, bind
+            each batch to the ``INSERT`` statement with ``cursor.executemany``.
         :param autocommit: What to set the connection's autocommit setting to
             before executing the query.
         """
@@ -329,12 +357,17 @@ class AdbcHook(DbApiHook):
 
             with closing(conn.cursor()) as cur:
                 execute_batch = self._resolve_execute_batch(
-                    cur, executemany=executemany, fast_executemany=fast_executemany
+                    cur,
+                    table_name=table_name,
+                    schema_name=schema_name,
+                    sql=sql,
+                    replace=replace,
+                    fast_executemany=fast_executemany,
                 )
 
                 for chunked_rows in chunked(rows, commit_every or None):
                     batch = self._to_record_batch(rows=chunked_rows, schema=table_schema)
-                    execute_batch(cur, sql, batch)
+                    execute_batch(batch)
 
                     if not autocommit:
                         conn.commit()

providers/apache/arrow/tests/unit/apache/arrow/hooks/test_adbc.py

@@ -22,9 +22,9 @@ import logging
 from unittest import mock
 
 import pytest
-from adbc_driver_manager import dbapi
+from adbc_driver_manager import AdbcStatusCode, dbapi
 from adbc_driver_manager.dbapi import Cursor
-from pyarrow import field, schema, string
+from pyarrow import RecordBatch, field, schema, string
 
 from airflow.models import Connection
 from airflow.providers.apache.arrow.hooks.adbc import AdbcHook
@@ -97,22 +97,121 @@ class TestAdbcHook:
 
         assert result == [(10, 30), (20, 40)]
 
-    def test_insert_rows_uses_executemany_with_record_batch(self):
-        """insert_rows must call executemany with a RecordBatch — the Arrow-native fast path."""
+    def test_insert_rows_uses_adbc_ingest_by_default(self):
+        """By default insert_rows must hand each RecordBatch to the driver's bulk ingest."""
         rows = [("a",), ("b",)]
         self.hook.insert_rows("table", rows)
 
-        assert self.cur.executemany.called
-        batch = self.cur.executemany.call_args[0][1]
+        self.cur.executemany.assert_not_called()
+        self.cur.adbc_ingest.assert_called_once()
+        args, kwargs = self.cur.adbc_ingest.call_args
+        assert args[0] == "table"
+        assert isinstance(args[1], RecordBatch)
+        assert args[1].num_rows == 2
+        assert kwargs == {"mode": "append", "db_schema_name": None}
+        assert self.conn.commit.call_count == 1
+
+    def test_insert_rows_fast_executemany_false_uses_executemany(self):
+        """fast_executemany=False keeps the previous bind path (one INSERT bound to the batch)."""
+        rows = [("a",), ("b",)]
+        self.hook.insert_rows("table", rows, fast_executemany=False)
+
+        self.cur.adbc_ingest.assert_not_called()
+        self.cur.executemany.assert_called_once()
+        sql, batch = self.cur.executemany.call_args[0]
+        assert sql.startswith("INSERT INTO table")
         assert batch.num_rows == 2
-        assert self.conn.commit.call_count >= 1
+
+    @mock.patch.object(AdbcHook, "_generate_insert_sql", return_value="REPLACE INTO table (col) VALUES ($1)")
+    def test_insert_rows_replace_binds_even_with_fast_executemany(self, generate_insert_sql):
+        """replace=True cannot go through bulk ingest, so it must bind to the generated statement."""
+        rows = [("a",), ("b",)]
+        self.hook.insert_rows("table", rows, replace=True, fast_executemany=True)
+
+        assert generate_insert_sql.call_args[0][3] is True
+        self.cur.adbc_ingest.assert_not_called()
+        self.cur.executemany.assert_called_once()
+        assert self.cur.executemany.call_args[0][0] == "REPLACE INTO table (col) VALUES ($1)"
+
+    @mock.patch.object(AdbcHook, "_generate_insert_sql", return_value="REPLACE INTO table (col) VALUES ($1)")
+    def test_insert_rows_replace_with_fast_executemany_logs_debug_not_info(self, _generate_insert_sql):
+        """The caller did not ask for ingest explicitly, so binding for replace=True is only a debug note."""
+        self.hook._log = mock.MagicMock(spec=logging.Logger)
+        self.hook.insert_rows("table", [("a",)], replace=True)
+
+        self.hook._log.debug.assert_any_call(
+            "Bulk ingest cannot replace rows; binding them to the statement instead"
+        )
+        info_messages = [c.args[0] for c in self.hook._log.info.call_args_list]
+        assert not any("cannot replace" in m for m in info_messages)
+
+    def test_insert_rows_falls_back_to_executemany_when_ingest_not_supported(self):
+        """A driver without bulk ingest must fall back to binding, once, for every remaining batch."""
+        self.cur.adbc_ingest.side_effect = dbapi.NotSupportedError("NOT_IMPLEMENTED: ingest")
+        self.hook._log = mock.MagicMock(spec=logging.Logger)
+        rows = [("a",), ("b",), ("c",)]
+
+        self.hook.insert_rows("table", rows, commit_every=2)
+
+        self.cur.adbc_ingest.assert_called_once()
+        assert self.cur.executemany.call_count == 2
+        assert [c.args[0].startswith("INSERT INTO table") for c in self.cur.executemany.call_args_list] == [
+            True,
+            True,
+        ]
+        assert [c.args[1].num_rows for c in self.cur.executemany.call_args_list] == [2, 1]
+        assert self.conn.commit.call_count == 2
+        fallback_logs = [
+            c for c in self.hook._log.info.call_args_list if "does not support bulk ingest" in c.args[0]
+        ]
+        assert len(fallback_logs) == 1
+
+    def test_insert_rows_ingest_error_after_first_batch_is_not_swallowed(self):
+        """Only an unsupported first ingest falls back; later NotSupportedErrors propagate."""
+        self.cur.adbc_ingest.side_effect = [None, dbapi.NotSupportedError("boom")]
+
+        with pytest.raises(dbapi.NotSupportedError, match="boom"):
+            self.hook.insert_rows("table", [("a",), ("b",)], commit_every=1)
+
+        self.cur.executemany.assert_not_called()
+
+    def test_insert_rows_ingest_other_errors_propagate(self):
+        """Errors other than NotSupportedError from ingest are not hidden by the fallback."""
+        self.cur.adbc_ingest.side_effect = dbapi.ProgrammingError(
+            "bad table", status_code=AdbcStatusCode.INVALID_ARGUMENT
+        )
+
+        with pytest.raises(dbapi.ProgrammingError, match="bad table"):
+            self.hook.insert_rows("table", [("a",)])
+
+        self.cur.executemany.assert_not_called()
+
+    def test_insert_rows_schema_qualified_table_uses_db_schema_name(self):
+        """``schema.table`` must be split into the ingest target table and db_schema_name."""
+        rows = [("a",)]
+        self.hook.insert_rows("my_schema.my_table", rows)
+
+        self.conn.adbc_get_table_schema.assert_called_once_with(
+            table_name="my_table", db_schema_filter="my_schema"
+        )
+        args, kwargs = self.cur.adbc_ingest.call_args
+        assert args[0] == "my_table"
+        assert kwargs == {"mode": "append", "db_schema_name": "my_schema"}
+
+    def test_insert_rows_commit_every_chunks_ingest(self):
+        """Chunking and per-chunk commits must be the same on the ingest path."""
+        rows = [("a",), ("b",), ("c",)]
+        self.hook.insert_rows("table", rows, commit_every=2)
+
+        assert [c.args[1].num_rows for c in self.cur.adbc_ingest.call_args_list] == [2, 1]
+        assert self.conn.commit.call_count == 2
 
     def test_insert_rows_autocommit_skips_explicit_commit(self):
         """With autocommit=True the chunk loop must not call conn.commit() — drivers raise on it."""
         rows = [("a",), ("b",)]
         self.hook.insert_rows("table", rows, autocommit=True)
 
-        assert self.cur.executemany.called
+        assert self.cur.adbc_ingest.called
         assert self.conn.commit.call_count == 0
 
     def test_set_autocommit_applies_adbc_option(self):
@@ -132,36 +231,13 @@ class TestAdbcHook:
         rows = [("a",), ("b",), ("c",)]
         self.hook.insert_rows("table", rows, commit_every=0)
 
-        assert self.cur.executemany.called
-        batch = self.cur.executemany.call_args[0][1]
+        assert self.cur.adbc_ingest.call_count == 1
+        batch = self.cur.adbc_ingest.call_args[0][1]
         assert batch.num_rows == 3
         assert self.conn.commit.call_count == 1
 
-    def test_insert_rows_fast_executemany_not_supported(self):
-        # Cursor that doesn't support setting fast_executemany
-        class NoFastExecCursor(mock.MagicMock):
-            def __setattr__(self, name, value):
-                if name == "fast_executemany":
-                    raise AttributeError("fast_executemany not supported")
-                super().__setattr__(name, value)
-
-        cur = NoFastExecCursor(spec=Cursor)
-        conn = mock.MagicMock()
-        conn.cursor.return_value = cur
-        conn.adbc_get_table_schema.return_value = schema([field("col", string())])
-        conn.extra_dejson = {}
-        hook = self.make_hook_for_conn(conn)
-
-        table = "table"
-        rows = [("x",), ("y",)]
-
-        hook.insert_rows(table, rows, executemany=True, fast_executemany=True)
-
-        assert cur.executemany.called
-        assert conn.commit.call_count >= 1
-
-    def test_insert_rows_fast_executemany_supported(self):
-        # Cursor that supports setting fast_executemany
+    def test_insert_rows_does_not_set_fast_executemany_on_cursor(self):
+        """ADBC cursors have no fast_executemany attribute; the flag selects the path instead."""
         cur = mock.MagicMock(spec=Cursor)
         conn = mock.MagicMock()
         conn.cursor.return_value = cur
@@ -169,14 +245,11 @@ class TestAdbcHook:
         conn.extra_dejson = {}
         hook = self.make_hook_for_conn(conn)
 
-        table = "table"
-        rows = [("x",), ("y",)]
+        hook.insert_rows("table", [("x",), ("y",)], executemany=True, fast_executemany=True)
 
-        hook.insert_rows(table, rows, executemany=True, fast_executemany=True)
-
-        assert cur.fast_executemany
-        assert cur.executemany.called
-        assert conn.commit.call_count >= 1
+        assert not hasattr(cur, "fast_executemany")
+        assert cur.adbc_ingest.called
+        assert not cur.executemany.called
 
     def test_insert_rows_target_fields_order_preserved(self):
         """Schema must follow target_fields order so SQL and RecordBatch agree."""
@@ -196,8 +269,8 @@ class TestAdbcHook:
         rows = [("Smith", "Alice"), ("Jones", "Bob")]
         hook.insert_rows("t", rows, target_fields=["last", "first"])
 
-        assert cur.executemany.called
-        batch = cur.executemany.call_args[0][1]  # second positional arg is the RecordBatch
+        assert cur.adbc_ingest.called
+        batch = cur.adbc_ingest.call_args[0][1]  # second positional arg is the RecordBatch
         # Schema order must match target_fields, not the table declaration.
         assert batch.schema.names == ["last", "first"]
         # Values must land in the declared column order.
@@ -225,6 +298,36 @@ class TestAdbcHook:
         with AdbcHook()._create_autocommit_connection() as adbc_conn:
             assert isinstance(adbc_conn, dbapi.Connection)
 
+    @pytest.mark.skipif(
+        importlib.util.find_spec("adbc_driver_sqlite") is None,
+        reason="adbc_driver_sqlite not installed",
+    )
+    @pytest.mark.parametrize("fast_executemany", [True, False])
+    def test_insert_rows_sqlite_round_trip(self, create_connection_without_db, tmp_path, fast_executemany):
+        """Both load paths must write the same rows, including a target_fields subset."""
+        create_connection_without_db(
+            Connection(
+                conn_id="adbc_default",
+                conn_type="adbc",
+                extra=json.dumps({"driver": "adbc_driver_sqlite"}),
+            )
+        )
+        hook = AdbcHook()
+        # A file database, so the rows outlive the connection that inserted them.
+        with mock.patch.object(AdbcHook, "uri", str(tmp_path / "db.sqlite")):
+            # SQLite reports column types from the stored values, so an empty table only
+            # yields a usable schema for integer columns.
+            hook.run("CREATE TABLE t (id INTEGER, qty INTEGER, note TEXT DEFAULT 'dflt')", autocommit=True)
+
+            hook.insert_rows(
+                "t", [(10, 1), (20, 2)], target_fields=["qty", "id"], fast_executemany=fast_executemany
+            )
+
+            assert hook.get_records("SELECT id, qty, note FROM t ORDER BY id") == [
+                (1, 10, "dflt"),
+                (2, 20, "dflt"),
+            ]
+
     @pytest.mark.skipif(
         importlib.util.find_spec("adbc_driver_sqlite") is None,
         reason="adbc_driver_sqlite not installed",