Repository: mabel-dev/opteryx

Issue:
🪲 SQL Projection Push `COUNT(*)`

**Sample Code/Statement** _If you can, please submit the SQL statement or Python code snippet, or a representative example using the sample datasets._

~~~sql
SELECT COUNT(*)
~~~

pushes `SELECT *` to the SQL source

**Additional context** _Add any other context about the problem here, for example what you have done to try to diagnose or workaround the problem._



Discussion:

Resolving patch:
diff --git a/opteryx/__version__.py b/opteryx/__version__.py
--- a/opteryx/__version__.py
+++ b/opteryx/__version__.py
@@ -1,4 +1,4 @@
-__build__ = 346
+__build__ = 348
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/opteryx/connectors/sql_connector.py b/opteryx/connectors/sql_connector.py
--- a/opteryx/connectors/sql_connector.py
+++ b/opteryx/connectors/sql_connector.py
@@ -139,11 +139,9 @@ def read_dataset(  # type:ignore
             result_schema.columns = [  # type:ignore
                 col for col in self.schema.columns if col.name in column_names  # type:ignore
             ]
-        elif self.schema.columns:  # type:ignore
-            query_builder.add("SELECT", "*")
         else:
             query_builder.add("SELECT", "1")
-            self.schema.columns.append(ConstantColumn(name="1", value=1))  # type:ignore
+            self.schema.columns = [ConstantColumn(name="1", value=1)]  # type:ignore
 
         # Update SQL if we've pushed predicates
         parameters: dict = {}
diff --git a/opteryx/utils/file_decoders.py b/opteryx/utils/file_decoders.py
--- a/opteryx/utils/file_decoders.py
+++ b/opteryx/utils/file_decoders.py
@@ -132,12 +132,6 @@ def parquet_decoder(buffer, projection: List = None, selection=None, just_schema
         if just_schema:
             return convert_arrow_schema_to_orso_schema(parquet_file.schema_arrow)
 
-        # Special handling for projection of ["count_*"]
-        if projection == ["count_*"]:
-            return pyarrow.Table.from_pydict(
-                {"_": numpy.full(parquet_file.metadata.num_rows, True, dtype=numpy.bool_)}
-            )
-
         # Projection processing
         schema_columns_set = set(parquet_file.schema_arrow.names)
         projection_names = {name for proj_col in projection for name in proj_col.all_names}
@@ -146,6 +140,12 @@ def parquet_decoder(buffer, projection: List = None, selection=None, just_schema
         # If no columns are selected, set to None
         if not selected_columns:
             selected_columns = None
+    # Special handling for projection of [] (usually means COUNT(*))
+    if projection == []:
+        parquet_file = parquet.ParquetFile(stream)
+        return pyarrow.Table.from_pydict(
+            {"_": numpy.full(parquet_file.metadata.num_rows, True, dtype=numpy.bool_)}
+        )
     else:
         selected_columns = None
 


Test patch:
diff --git a/tests/storage/test_blob_local.py b/tests/storage/test_blob_local.py
new file mode 100644
--- /dev/null
+++ b/tests/storage/test_blob_local.py
@@ -0,0 +1,128 @@
+import os
+import sys
+import pytest
+
+sys.path.insert(1, os.path.join(sys.path[0], "../.."))
+
+import opteryx
+from opteryx.utils.formatter import format_sql
+
+# fmt: off
+STATEMENTS = [
+    ("SELECT * FROM testdata.planets", 9, 20, None),
+    ("SELECT * FROM testdata.satellites", 177, 8, None),
+    ("SELECT COUNT(*) FROM testdata.planets;", 1, 1, None),
+    ("SELECT COUNT(*) FROM testdata.satellites;", 1, 1, None),
+    ("SELECT COUNT(*) FROM (SELECT * FROM testdata.planets) AS p", 1, 1, None),
+    ("SELECT COUNT(*) FROM (SELECT COUNT(*) FROM testdata.planets) AS p", 1, 1, None),
+    ("SELECT COUNT(*) FROM (SELECT * FROM testdata.planets WHERE id > 4) AS p", 1, 1, None),
+    ("SELECT COUNT(*) FROM (SELECT * FROM testdata.planets) AS p WHERE id > 4", 1, 1, None),
+    ("SELECT name FROM testdata.planets;", 9, 1, None),
+    ("SELECT name FROM testdata.satellites;", 177, 1, None),
+    ("SELECT * FROM testdata.planets INNER JOIN $satellites ON testdata.planets.id = $satellites.planetId;", 177, 28, None),
+    ("SELECT * FROM testdata.planets, $satellites WHERE testdata.planets.id = $satellites.planetId;", 177, 28, None),
+    ("SELECT * FROM testdata.planets CROSS JOIN $satellites WHERE testdata.planets.id = $satellites.planetId;", 177, 28, None),
+    ("SELECT * FROM testdata.planets INNER JOIN testdata.satellites ON testdata.planets.id = testdata.satellites.planetId;", 177, 28, None),
+    ("SELECT name FROM testdata.planets WHERE name LIKE 'Earth';", 1, 1, None),
+    ("SELECT * FROM testdata.planets WHERE id > gravity", 2, 20, None),
+    ("SELECT * FROM testdata.planets WHERE surfacePressure IS NULL", 4, 20, None),
+    ("SELECT * FROM testdata.planets WHERE surfacePressure IS NOT NULL", 5, 20, None),
+    ("SELECT * FROM testdata.planets, testdata.satellites WHERE testdata.planets.id = 5 AND testdata.satellites.planetId = 5;", 67, 28, None),
+    ("SELECT * FROM testdata.planets, testdata.satellites WHERE testdata.planets.id - testdata.satellites.planetId = 0;", 177, 28, None),
+    ("SELECT * FROM testdata.planets, testdata.satellites WHERE testdata.planets.id - testdata.satellites.planetId != 0;", 1416, 28, None),
+    ("SELECT * FROM testdata.planets WHERE testdata.planets.id - testdata.planets.numberOfMoons < 0;", 4, 20, None),
+]
+# fmt: on
+
+
+@pytest.mark.parametrize("statement, rows, columns, exception", STATEMENTS)
+def test_sql_battery(statement, rows, columns, exception):
+    """
+    Test an battery of statements
+    """
+    try:
+        # query to arrow is the fastest way to query
+        result = opteryx.query_to_arrow(statement)
+        actual_rows, actual_columns = result.shape
+        assert (
+            rows == actual_rows
+        ), f"\n\033[38;5;203mQuery returned {actual_rows} rows but {rows} were expected.\033[0m\n{statement}"
+        assert (
+            columns == actual_columns
+        ), f"\n\033[38;5;203mQuery returned {actual_columns} cols but {columns} were expected.\033[0m\n{statement}"
+        assert (
+            exception is None
+        ), f"Exception {exception} not raised but expected\n{format_sql(statement)}"
+    except AssertionError as err:
+        raise Exception(err) from err
+    except Exception as err:
+        if type(err) != exception:
+            raise Exception(
+                f"{format_sql(statement)}\nQuery failed with error {type(err)} but error {exception} was expected"
+            ) from err
+
+
+if __name__ == "__main__":  # pragma: no cover
+    """
+    Running in the IDE we do some formatting - it's not functional but helps
+    when reading the outputs.
+    """
+
+    import shutil
+    import time
+
+    from tests.tools import trunc_printable
+
+    start_suite = time.monotonic_ns()
+
+    width = shutil.get_terminal_size((80, 20))[0] - 15
+
+    passed = 0
+    failed = 0
+
+    nl = "\n"
+
+    failures = []
+
+    print(f"RUNNING BATTERY OF {len(STATEMENTS)} testdata TESTS")
+    for index, (statement, rows, cols, err) in enumerate(STATEMENTS):
+
+        printable = statement
+        if hasattr(printable, "decode"):
+            printable = printable.decode()
+        print(
+            f"\033[38;2;255;184;108m{(index + 1):04}\033[0m"
+            f" {trunc_printable(format_sql(printable), width - 1)}",
+            end="",
+            flush=True,
+        )
+        try:
+            start = time.monotonic_ns()
+            test_sql_battery(statement, rows, cols, err)
+            print(
+                f"\033[38;2;26;185;67m{str(int((time.monotonic_ns() - start)/1e6)).rjust(4)}ms\033[0m ✅",
+                end="",
+            )
+            passed += 1
+            if failed > 0:
+                print(" \033[0;31m*\033[0m")
+            else:
+                print()
+        except Exception as err:
+            print(f"\033[0;31m{str(int((time.monotonic_ns() - start)/1e6)).rjust(4)}ms ❌ *\033[0m")
+            print(">", err)
+            failed += 1
+            failures.append((statement, err))
+
+    print("--- ✅ \033[0;32mdone\033[0m")
+
+    if failed > 0:
+        print("\n\033[38;2;139;233;253m\033[3mFAILURES\033[0m")
+        for statement, err in failures:
+            print(err)
+
+    print(
+        f"\n\033[38;2;139;233;253m\033[3mCOMPLETE\033[0m ({((time.monotonic_ns() - start_suite) / 1e9):.2f} seconds)\n"
+        f"  \033[38;2;26;185;67m{passed} passed ({(passed * 100) // (passed + failed)}%)\033[0m\n"
+        f"  \033[38;2;255;121;198m{failed} failed\033[0m"
+    )
diff --git a/tests/storage/test_sql_sqlite.py b/tests/storage/test_sql_sqlite.py
--- a/tests/storage/test_sql_sqlite.py
+++ b/tests/storage/test_sql_sqlite.py
@@ -21,7 +21,7 @@
 from opteryx.connectors import SqlConnector
 from opteryx.utils.formatter import format_sql
 
-#fmt: off
+# fmt: off
 STATEMENTS = [
     ("SELECT * FROM sqlite.planets", 9, 20, None),
     ("SELECT * FROM sqlite.satellites", 177, 8, None),
@@ -29,6 +29,10 @@
     ("SELECT COUNT(*) FROM sqlite.planets;", 1, 1, None),
     ("SELECT COUNT(*) FROM sqlite.satellites;", 1, 1, None),
     ("SELECT COUNT(*) FROM sqlite_tweets.tweets", 1, 1, None),
+    ("SELECT COUNT(*) FROM (SELECT * FROM sqlite.planets) AS p", 1, 1, None),
+    ("SELECT COUNT(*) FROM (SELECT COUNT(*) FROM sqlite.planets) AS p", 1, 1, None),
+    ("SELECT COUNT(*) FROM (SELECT * FROM sqlite.planets WHERE id > 4) AS p", 1, 1, None),
+    ("SELECT COUNT(*) FROM (SELECT * FROM sqlite.planets) AS p WHERE id > 4", 1, 1, None),
     ("SELECT name FROM sqlite.planets;", 9, 1, None),
     ("SELECT name FROM sqlite.satellites;", 177, 1, None),
     ("SELECT user_name FROM sqlite_tweets.tweets;", 100000, 1, None),
@@ -46,7 +50,8 @@
     ("SELECT * FROM sqlite.planets, sqlite.satellites WHERE sqlite.planets.id - sqlite.satellites.planetId != 0;", 1416, 28, None),
     ("SELECT * FROM sqlite.planets WHERE sqlite.planets.id - sqlite.planets.numberOfMoons < 0;", 4, 20, None),
 ]
-#fmt: on
+# fmt: on
+
 
 @pytest.mark.parametrize("statement, rows, columns, exception", STATEMENTS)
 def test_sql_battery(statement, rows, columns, exception):