feat(api): count nulls with topk (#8531)

Default to count star instead of count(expr) in topk API.

BREAKING CHANGE: expr.topk(...) now includes null counts. The row count of the topk call will not differ, but the number of nulls counted will no longer be zero. To drop the null row use the dropna method.

ci/schema/clickhouse.sql

@@ -94,3 +94,6 @@ INSERT INTO ibis_testing.win VALUES
     ('a', 2, 0),
     ('a', 3, 1),
     ('a', 4, 1);
+
+CREATE OR REPLACE TABLE ibis_testing.topk (x Nullable(Int64)) ENGINE = Memory;
+INSERT INTO ibis_testing.topk VALUES (1), (1), (NULL);

ci/schema/duckdb.sql

@@ -51,3 +51,7 @@ CREATE OR REPLACE TABLE map (idx BIGINT, kv MAP(STRING, BIGINT));
 INSERT INTO map VALUES
     (1, MAP(['a', 'b', 'c'], [1, 2, 3])),
     (2, MAP(['d', 'e', 'f'], [4, 5, 6]));
+
+
+CREATE OR REPLACE TABLE topk (x BIGINT);
+INSERT INTO topk VALUES (1), (1), (NULL);

ci/schema/exasol.sql

@@ -87,3 +87,6 @@ INSERT INTO "win" VALUES
     ('a', 2, 0),
     ('a', 3, 1),
     ('a', 4, 1);
+
+CREATE OR REPLACE TABLE EXASOL."topk" ("x" BIGINT);
+INSERT INTO "topk" VALUES (1), (1), (NULL);

ci/schema/mssql.sql

@@ -131,3 +131,8 @@ INSERT INTO win VALUES
     ('a', 2, 0),
     ('a', 3, 1),
     ('a', 4, 1);
+
+DROP TABLE IF EXISTS topk;
+
+CREATE TABLE topk (x BIGINT);
+INSERT INTO topk VALUES (1), (1), (NULL);

ci/schema/mysql.sql

@@ -119,3 +119,8 @@ INSERT INTO win VALUES
     ('a', 2, 0),
     ('a', 3, 1),
     ('a', 4, 1);
+
+DROP TABLE IF EXISTS topk CASCADE;
+
+CREATE TABLE topk (x BIGINT);
+INSERT INTO topk VALUES (1), (1), (NULL);

ci/schema/oracle.sql

@@ -116,4 +116,9 @@ INSERT INTO "win" VALUES
     ('a', 3, 1),
     ('a', 4, 1);
 
+DROP TABLE IF EXISTS "topk";
+
+CREATE TABLE "topk" ("x" NUMBER(18));
+INSERT INTO "topk" VALUES (1), (1), (NULL);
+
 COMMIT;

ci/schema/postgres.sql

@@ -289,3 +289,8 @@ CREATE TABLE map (idx BIGINT, kv HSTORE);
 INSERT INTO map VALUES
     (1, 'a=>1,b=>2,c=>3'),
     (2, 'd=>4,e=>5,c=>6');
+
+DROP TABLE IF EXISTS topk;
+
+CREATE TABLE topk (x BIGINT);
+INSERT INTO topk VALUES (1), (1), (NULL);

ci/schema/risingwave.sql

@@ -175,3 +175,8 @@ INSERT INTO "win" VALUES
     ('a', 2, 0),
     ('a', 3, 1),
     ('a', 4, 1);
+
+DROP TABLE IF EXISTS "topk";
+
+CREATE TABLE "topk" ("x" BIGINT);
+INSERT INTO "topk" VALUES (1), (1), (NULL);

ci/schema/snowflake.sql

@@ -140,3 +140,6 @@ INSERT INTO "win" VALUES
     ('a', 2, 0),
     ('a', 3, 1),
     ('a', 4, 1);
+
+CREATE OR REPLACE TABLE "topk" ("x" BIGINT);
+INSERT INTO "topk" VALUES (1), (1), (NULL);

ci/schema/sqlite.sql

@@ -119,3 +119,7 @@ INSERT INTO win VALUES
     ('a', 2, 0),
     ('a', 3, 1),
     ('a', 4, 1);
+
+DROP TABLE IF EXISTS topk;
+CREATE TABLE "topk" ("x" BIGINT);
+INSERT INTO "topk" VALUES (1), (1), (NULL);

ci/schema/trino.sql

@@ -178,3 +178,7 @@ INSERT INTO win VALUES
     ('a', 2, 0),
     ('a', 3, 1),
     ('a', 4, 1);
+
+DROP TABLE IF EXISTS topk;
+CREATE TABLE topk (x BIGINT);
+INSERT INTO topk VALUES (1), (1), (NULL);

ibis/backends/bigquery/tests/conftest.py

@@ -16,7 +16,13 @@
 from ibis.backends.bigquery.datatypes import BigQuerySchema
 from ibis.backends.conftest import TEST_TABLES
 from ibis.backends.tests.base import BackendTest
-from ibis.backends.tests.data import json_types, non_null_array_types, struct_types, win
+from ibis.backends.tests.data import (
+    json_types,
+    non_null_array_types,
+    struct_types,
+    topk,
+    win,
+)
 
 DATASET_ID = "ibis_gbq_testing"
 DATASET_ID_TOKYO = "ibis_gbq_testing_tokyo"
@@ -215,6 +221,19 @@ def _load_data(self, **_: Any) -> None:
                 )
             )
 
+            futures.append(
+                e.submit(
+                    make_job,
+                    client.load_table_from_dataframe,
+                    topk.to_pandas(),
+                    bq.TableReference(testing_dataset, "topk"),
+                    job_config=bq.LoadJobConfig(
+                        write_disposition=write_disposition,
+                        schema=BigQuerySchema.from_ibis(ibis.schema(dict(x="int64"))),
+                    ),
+                )
+            )
+
             futures.append(
                 e.submit(
                     make_job,

ibis/backends/dask/tests/conftest.py

@@ -12,7 +12,7 @@
 import ibis.expr.datatypes as dt
 from ibis.backends.conftest import TEST_TABLES
 from ibis.backends.pandas.tests.conftest import TestConf as PandasTest
-from ibis.backends.tests.data import array_types, json_types, win
+from ibis.backends.tests.data import array_types, json_types, topk, win
 
 dd = pytest.importorskip("dask.dataframe")
 
@@ -66,6 +66,11 @@ def _load_data(self, **_: Any) -> None:
             dd.from_pandas(json_types, npartitions=NPARTITIONS),
             overwrite=True,
         )
+        con.create_table(
+            "topk",
+            dd.from_pandas(topk.to_pandas(), npartitions=NPARTITIONS),
+            overwrite=True,
+        )
 
     @classmethod
     def assert_series_equal(

ibis/backends/datafusion/tests/conftest.py

@@ -7,7 +7,7 @@
 import ibis
 from ibis.backends.conftest import TEST_TABLES
 from ibis.backends.tests.base import BackendTest
-from ibis.backends.tests.data import array_types, win
+from ibis.backends.tests.data import array_types, topk, win
 
 
 class TestConf(BackendTest):
@@ -27,6 +27,7 @@ def _load_data(self, **_: Any) -> None:
             con.register(path, table_name=table_name)
         con.register(array_types, table_name="array_types")
         con.register(win, table_name="win")
+        con.register(topk, table_name="topk")
 
     @staticmethod
     def connect(*, tmpdir, worker_id, **kw):

ibis/backends/flink/tests/conftest.py

@@ -8,7 +8,7 @@
 import ibis
 from ibis.backends.conftest import TEST_TABLES
 from ibis.backends.tests.base import BackendTest
-from ibis.backends.tests.data import array_types, json_types, struct_types, win
+from ibis.backends.tests.data import array_types, json_types, struct_types, topk, win
 
 
 class TestConf(BackendTest):
@@ -75,6 +75,7 @@ def _load_data(self, **_: Any) -> None:
             schema=ibis.schema({"idx": "int64", "kv": "map<string, int64>"}),
             temp=True,
         )
+        con.create_table("topk", topk, temp=True)
 
 
 class TestConfForStreaming(TestConf):

ibis/backends/impala/tests/conftest.py

@@ -113,6 +113,10 @@ def _load_data(self, **_: Any) -> None:
                 ('a', 4, 1)
             """
         )
+        con.create_table(
+            "topk", schema=ibis.schema({"x": "int64"}), database="ibis_testing"
+        )
+        con.raw_sql("INSERT INTO ibis_testing.topk VALUES (1), (1), (NULL)")
         assert con.list_tables(database="ibis_testing")
 
     def postload(self, **kw):

ibis/backends/pandas/tests/conftest.py

@@ -12,7 +12,7 @@
 from ibis.backends.conftest import TEST_TABLES
 from ibis.backends.pandas import Backend
 from ibis.backends.tests.base import BackendTest
-from ibis.backends.tests.data import array_types, json_types, struct_types, win
+from ibis.backends.tests.data import array_types, json_types, struct_types, topk, win
 
 
 class TestConf(BackendTest):
@@ -35,6 +35,7 @@ def _load_data(self, **_: Any) -> None:
         con.create_table("struct", struct_types, overwrite=True)
         con.create_table("win", win, overwrite=True)
         con.create_table("json_t", json_types, overwrite=True)
+        con.create_table("topk", topk.to_pandas(), overwrite=True)
 
     @staticmethod
     def connect(*, tmpdir, worker_id, **kw):

ibis/backends/polars/tests/conftest.py

@@ -3,12 +3,13 @@
 from typing import Any
 
 import numpy as np
+import polars as pl
 import pytest
 
 import ibis
 from ibis.backends.conftest import TEST_TABLES
 from ibis.backends.tests.base import BackendTest
-from ibis.backends.tests.data import array_types, struct_types, win
+from ibis.backends.tests.data import array_types, struct_types, topk, win
 
 
 class TestConf(BackendTest):
@@ -27,6 +28,9 @@ def _load_data(self, **_: Any) -> None:
         con.register(struct_types, table_name="struct")
         con.register(win, table_name="win")
 
+        # TODO: remove when pyarrow inputs are supported
+        con._add_table("topk", pl.from_arrow(topk).lazy())
+
     @staticmethod
     def connect(*, tmpdir, worker_id, **kw):
         return ibis.polars.connect(**kw)

ibis/backends/pyspark/tests/conftest.py

@@ -12,7 +12,7 @@
 from ibis import util
 from ibis.backends.conftest import TEST_TABLES
 from ibis.backends.tests.base import BackendTest
-from ibis.backends.tests.data import json_types, win
+from ibis.backends.tests.data import json_types, topk, win
 
 
 def set_pyspark_database(con, database):
@@ -129,6 +129,7 @@ def _load_data(self, **_: Any) -> None:
 
         s.createDataFrame(json_types).createOrReplaceTempView("json_t")
         s.createDataFrame(win).createOrReplaceTempView("win")
+        s.createDataFrame(topk.to_pandas()).createOrReplaceTempView("topk")
 
     @staticmethod
     def connect(*, tmpdir, worker_id, **kw):

ibis/backends/tests/data.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 
 array_types = pd.DataFrame(
     [
@@ -124,3 +125,5 @@
         "y": [3, 2, 0, 1, 1],
     }
 )
+
+topk = pa.Table.from_pydict({"x": [1, 1, None]})

ibis/backends/tests/sql/snapshots/test_select_sql/test_chain_limit_doesnt_collapse/result.sql

@@ -3,17 +3,17 @@ SELECT
 FROM (
   SELECT
     "t1"."city",
-    "t1"."Count(city)"
+    "t1"."CountStar(tbl)"
   FROM (
     SELECT
       "t0"."city",
-      COUNT("t0"."city") AS "Count(city)"
+      COUNT(*) AS "CountStar(tbl)"
     FROM "tbl" AS "t0"
     GROUP BY
       1
   ) AS "t1"
   ORDER BY
-    "t1"."Count(city)" DESC
+    "t1"."CountStar(tbl)" DESC
   LIMIT 10
 ) AS "t3"
 LIMIT 5
@@ -23,17 +23,17 @@ OFFSET (
   FROM (
     SELECT
       "t1"."city",
-      "t1"."Count(city)"
+      "t1"."CountStar(tbl)"
     FROM (
       SELECT
         "t0"."city",
-        COUNT("t0"."city") AS "Count(city)"
+        COUNT(*) AS "CountStar(tbl)"
       FROM "tbl" AS "t0"
       GROUP BY
         1
     ) AS "t1"
     ORDER BY
-      "t1"."Count(city)" DESC
+      "t1"."CountStar(tbl)" DESC
     LIMIT 10
   ) AS "t3"
 )

ibis/backends/tests/sql/snapshots/test_select_sql/test_topk_operation/e2.sql

@@ -8,17 +8,17 @@ FROM "tbl" AS "t1"
 SEMI JOIN (
   SELECT
     "t2"."city",
-    "t2"."Count(city)"
+    "t2"."CountStar(tbl)"
   FROM (
     SELECT
       "t0"."city",
-      COUNT("t0"."city") AS "Count(city)"
+      COUNT(*) AS "CountStar(tbl)"
     FROM "tbl" AS "t0"
     GROUP BY
       1
   ) AS "t2"
   ORDER BY
-    "t2"."Count(city)" DESC
+    "t2"."CountStar(tbl)" DESC
   LIMIT 10
 ) AS "t5"
   ON "t1"."city" = "t5"."city"

ibis/backends/tests/test_generic.py

@@ -1884,7 +1884,7 @@ def test_isnull_equality(con, backend, monkeypatch):
     ["druid"],
     raises=PyDruidProgrammingError,
     reason=(
-        "Query could not be planned. A possible reason is [SQL query requires ordering a "
+        "Query could not be planned. A possible reason is SQL query requires ordering a "
         "table by non-time column [[id]], which is not supported."
     ),
 )
@@ -1898,3 +1898,24 @@ def test_subsequent_overlapping_order_by(con, backend, alltypes, df):
     result = con.execute(ts2)
     expected = df.sort_values("id", ascending=False).reset_index(drop=True)
     backend.assert_frame_equal(result, expected)
+
+
+@pytest.mark.broken(
+    ["clickhouse"],
+    raises=AssertionError,
+    reason="https://github.com/ClickHouse/ClickHouse/issues/61313",
+)
+@pytest.mark.notimpl(
+    ["pandas", "dask"], raises=IndexError, reason="NaN isn't treated as NULL"
+)
+@pytest.mark.notimpl(
+    ["druid"],
+    raises=AttributeError,
+    reason="inserting three rows into druid is difficult",
+)
+def test_topk_counts_null(con):
+    t = con.tables.topk
+    tk = t.x.topk(10)
+    tkf = tk.filter(_.x.isnull())[1]
+    result = con.to_pyarrow(tkf)
+    assert result[0].as_py() == 1

ibis/expr/types/generic.py

@@ -1823,10 +1823,11 @@ def topk(
             raise com.IbisTypeError("TopK must depend on exactly one table.")
 
         table = table.to_expr()
+
         if by is None:
-            metric = self.count()
-        else:
-            (metric,) = bind(table, by)
+            by = lambda t: t.count()
+
+        (metric,) = bind(table, by)
 
         return table.aggregate(metric, by=[self]).order_by(metric.desc()).limit(k)