fix: change statistics to use pyarrow result type

rustyconover · rustyconover · commit 44a024ab803a · 2025-07-13T10:54:49.000-04:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,7 +8,6 @@ authors = [
 dependencies = [
     "pyarrow>=20.0.0",
     "query-farm-flight-server",
-#    "query-farm-flight-server @ git+https://github.com/Query-farm/server_base.git@master",
     "duckdb>=1.3.1",
     "query-farm-duckdb-json-serialization>=0.1.1",
 ]
diff --git a/requirements-dev.lock b/requirements-dev.lock
@@ -12,20 +12,20 @@
 -e file:.
 annotated-types==0.7.0
     # via pydantic
-boto3==1.39.0
+boto3==1.39.4
     # via query-farm-flight-server
-botocore==1.39.0
+botocore==1.39.4
     # via boto3
     # via s3transfer
 cache3==0.4.3
     # via query-farm-flight-server
-certifi==2025.6.15
+certifi==2025.7.9
     # via sentry-sdk
 click==8.2.1
     # via query-farm-flight-server
-coverage==7.9.1
+coverage==7.9.2
     # via pytest-cov
-duckdb==1.3.1
+duckdb==1.3.2
     # via query-farm-airport-test-server
 execnet==2.1.1
     # via pytest-xdist
@@ -46,7 +46,7 @@ mypy==1.16.1
     # via pytest-mypy
 mypy-boto3-dynamodb==1.39.0
     # via query-farm-flight-server
-mypy-boto3-s3==1.39.0
+mypy-boto3-s3==1.39.2
     # via query-farm-flight-server
 mypy-extensions==1.1.0
     # via mypy
@@ -79,14 +79,14 @@ pytest-cov==6.2.1
 pytest-env==1.1.5
 pytest-mypy==1.0.1
 pytest-parallel==0.1.1
-pytest-xdist==3.7.0
+pytest-xdist==3.8.0
 python-dateutil==2.9.0.post0
     # via botocore
 python-levenshtein==0.27.1
     # via query-farm-flight-server
-query-farm-duckdb-json-serialization==0.1.1
+query-farm-duckdb-json-serialization==0.1.2
     # via query-farm-airport-test-server
-query-farm-flight-server==0.1.2
+query-farm-flight-server==0.1.4
     # via query-farm-airport-test-server
 rapidfuzz==3.13.0
     # via levenshtein
@@ -101,7 +101,7 @@ structlog==25.4.0
     # via query-farm-flight-server
 tblib==3.1.0
     # via pytest-parallel
-typing-extensions==4.14.0
+typing-extensions==4.14.1
     # via mypy
     # via pydantic
     # via pydantic-core
diff --git a/requirements.lock b/requirements.lock
@@ -12,18 +12,18 @@
 -e file:.
 annotated-types==0.7.0
     # via pydantic
-boto3==1.39.0
+boto3==1.39.4
     # via query-farm-flight-server
-botocore==1.39.0
+botocore==1.39.4
     # via boto3
     # via s3transfer
 cache3==0.4.3
     # via query-farm-flight-server
-certifi==2025.6.15
+certifi==2025.7.9
     # via sentry-sdk
 click==8.2.1
     # via query-farm-flight-server
-duckdb==1.3.1
+duckdb==1.3.2
     # via query-farm-airport-test-server
 fuzzywuzzy==0.18.0
     # via query-farm-flight-server
@@ -36,7 +36,7 @@ msgpack==1.1.1
     # via query-farm-flight-server
 mypy-boto3-dynamodb==1.39.0
     # via query-farm-flight-server
-mypy-boto3-s3==1.39.0
+mypy-boto3-s3==1.39.2
     # via query-farm-flight-server
 prettytable==3.16.0
     # via query-farm-flight-server
@@ -52,9 +52,9 @@ python-dateutil==2.9.0.post0
     # via botocore
 python-levenshtein==0.27.1
     # via query-farm-flight-server
-query-farm-duckdb-json-serialization==0.1.1
+query-farm-duckdb-json-serialization==0.1.2
     # via query-farm-airport-test-server
-query-farm-flight-server==0.1.2
+query-farm-flight-server==0.1.4
     # via query-farm-airport-test-server
 rapidfuzz==3.13.0
     # via levenshtein
@@ -66,7 +66,7 @@ six==1.17.0
     # via python-dateutil
 structlog==25.4.0
     # via query-farm-flight-server
-typing-extensions==4.14.0
+typing-extensions==4.14.1
     # via pydantic
     # via pydantic-core
     # via typing-inspection
diff --git a/src/query_farm_airport_test_server/server.py b/src/query_farm_airport_test_server/server.py
@@ -1,8 +1,11 @@
+import datetime
 import hashlib
 import json
 import re
+import uuid
 from collections.abc import Callable, Generator, Iterator
 from dataclasses import dataclass, field
+from decimal import Decimal
 from typing import Any, Literal, TypeVar, overload
 
 import click
@@ -889,6 +892,49 @@ def in_out_handler(
 
                 return pa.RecordBatch.from_arrays([["last"], ["row"]], schema=output_schema)
 
+            static_data_schema = SchemaCollection(
+                scalar_functions_by_name=CaseInsensitiveDict(),
+                table_functions_by_name=CaseInsensitiveDict(),
+                tables_by_name=CaseInsensitiveDict(
+                    {
+                        "employees": TableInfo(
+                            table_versions=[
+                                pa.Table.from_arrays(
+                                    [
+                                        ["Emily", "Amy"],
+                                        [30, 32],
+                                        [datetime.datetime(2023, 10, 1), datetime.datetime(2024, 10, 2)],
+                                        ["{}", "[1,2,3]"],
+                                        [uuid.uuid4().bytes, uuid.uuid4().bytes],
+                                        [datetime.date(2023, 10, 1), datetime.date(2024, 10, 2)],
+                                        [True, False],
+                                        ["Ann", None],
+                                        [1234.123, 5678.123],
+                                        [Decimal("12345.678790"), Decimal("67890.123456")],
+                                    ],
+                                    schema=pa.schema(
+                                        [
+                                            pa.field("name", pa.string()),
+                                            pa.field("age", pa.int32()),
+                                            pa.field("start_date", pa.timestamp("ms")),
+                                            pa.field("json_data", pa.json_(pa.string())),
+                                            pa.field("id", pa.uuid()),
+                                            pa.field("birthdate", pa.date32()),
+                                            pa.field("is_active", pa.bool_()),
+                                            pa.field("nickname", pa.string()),
+                                            pa.field("salary", pa.float64()),
+                                            pa.field("balance", pa.decimal128(12, 6)),
+                                        ],
+                                        metadata={"can_produce_statistics": "1"},
+                                    ),
+                                )
+                            ],
+                            row_id_counter=2,
+                        )
+                    }
+                ),
+            )
+
             util_schema = SchemaCollection(
                 scalar_functions_by_name=CaseInsensitiveDict(
                     {
@@ -1065,6 +1111,7 @@ def in_out_handler(
             )
 
             library.databases_by_name[database_name].schemas_by_name["utils"] = util_schema
+            library.databases_by_name[database_name].schemas_by_name["static_data"] = static_data_schema
 
             return iter([])
         elif action.type == "drop_database":
@@ -1611,6 +1658,58 @@ def action_change_column_type(
             schema_name=parameters.schema_name,
         )[0]
 
+    def action_column_statistics(
+        self,
+        *,
+        context: base_server.CallContext[auth.Account, auth.AccountToken],
+        parameters: parameter_types.ColumnStatistics,
+    ) -> pa.Table:
+        assert context.caller is not None
+
+        descriptor_parts = descriptor_unpack_(parameters.flight_descriptor)
+        library = self.contents[context.caller.token.token]
+        database = library.by_name(descriptor_parts.catalog_name)
+        schema = database.by_name(descriptor_parts.schema_name)
+
+        assert descriptor_parts.type == "table"
+        table = schema.by_name("table", descriptor_parts.name)
+
+        contents = table.version().column(parameters.column_name)
+        # Since the table is a Pyarrow table we need to produce some values.
+        not_null_count = pc.count(contents, "only_valid").as_py()
+        null_count = pc.count(contents, "only_null").as_py()
+        distinct_count = len(set(contents.to_pylist()))
+        sorted_contents = sorted(filter(lambda x: x is not None, contents.to_pylist()))
+        min_value = sorted_contents[0]
+        max_value = sorted_contents[-1]
+
+        if contents.type == pa.uuid():
+            # For UUIDs, we need to convert them to strings for the output.
+            min_value = min_value.bytes
+            max_value = max_value.bytes
+
+        result_table = pa.Table.from_pylist(
+            [
+                {
+                    "has_not_null": not_null_count > 0,
+                    "has_null": null_count > 0,
+                    "distinct_count": distinct_count,
+                    "min": min_value,
+                    "max": max_value,
+                }
+            ],
+            schema=pa.schema(
+                [
+                    pa.field("has_not_null", pa.bool_()),
+                    pa.field("has_null", pa.bool_()),
+                    pa.field("distinct_count", pa.uint64()),
+                    pa.field("min", contents.type),
+                    pa.field("max", contents.type),
+                ]
+            ),
+        )
+        return result_table
+
     def impl_do_get(
         self,
         *,

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,6 @@ authors = [`
`8`	`8`	`dependencies = [`
`9`	`9`	`"pyarrow>=20.0.0",`
`10`	`10`	`"query-farm-flight-server",`
`11`		`-# "query-farm-flight-server @ git+https://github.com/Query-farm/server_base.git@master",`
`12`	`11`	`"duckdb>=1.3.1",`
`13`	`12`	`"query-farm-duckdb-json-serialization>=0.1.1",`
`14`	`13`	`]`