diff --git a/.github/workflows/python-integration.yml b/.github/workflows/python-integration.yml index 8f048852..4c51f586 100644 --- a/.github/workflows/python-integration.yml +++ b/.github/workflows/python-integration.yml @@ -31,6 +31,12 @@ jobs: run: | python -m pip install --upgrade pip pip install -r requirements/test.txt + - name: Install multicorn2 + run: | + git clone https://github.com/pgsql-io/multicorn2.git + cd multicorn2 + git checkout v2.5 + pip install . - name: Test with pytest env: SHILLELAGH_ADAPTER_KWARGS: ${{ secrets.SHILLELAGH_ADAPTER_KWARGS }} diff --git a/src/shillelagh/backends/multicorn/fdw.py b/src/shillelagh/backends/multicorn/fdw.py index 53dbbcc6..def86730 100644 --- a/src/shillelagh/backends/multicorn/fdw.py +++ b/src/shillelagh/backends/multicorn/fdw.py @@ -129,6 +129,26 @@ def update(self, oldvalues: Row, newvalues: Row) -> Row: def rowid_column(self): return "rowid" + def get_rel_size(self, quals: List[Qual], columns: List[str]) -> Tuple[int, int]: + """ + Estimate query cost. + """ + all_bounds = get_all_bounds(quals) + filtered_columns = [ + (column, operator[0]) + for column, operators in all_bounds.items() + for operator in operators + ] + + # the adapter returns an arbitrary cost that takes in consideration filtering and + # sorting; let's use that as an approximation for rows + rows = int(self.adapter.get_cost(filtered_columns, [])) + + # same assumption as the parent class + row_width = len(columns) * 100 + + return (rows, row_width) + @classmethod def import_schema( # pylint: disable=too-many-arguments cls, diff --git a/tests/backends/multicorn/fdw_test.py b/tests/backends/multicorn/fdw_test.py index aa8a6d6c..f4518664 100644 --- a/tests/backends/multicorn/fdw_test.py +++ b/tests/backends/multicorn/fdw_test.py @@ -194,3 +194,19 @@ def test_update(mocker: MockerFixture, registry: AdapterLoader) -> None: {"rowid": 1, "name": "Bob", "age": 23, "pets": 3}, {"rowid": 0, "name": "Alice", "age": 20, "pets": 1}, ] + + +def test_get_rel_Size(mocker: MockerFixture, registry: AdapterLoader) -> None: + """ + Test the ``get_rel_size`` method. + """ + mocker.patch("shillelagh.backends.multicorn.fdw.registry", registry) + + registry.add("dummy", FakeAdapter) + + wrapper = MulticornForeignDataWrapper( + {"adapter": "dummy", "args": "qQA="}, + {}, + ) + + assert wrapper.get_rel_size([Qual("age", ">", 21)], ["name", "age"]) == (666, 200)