fix: it wouldn't be a demo without collatz

Author: rustyconover

pyproject.toml

@@ -9,7 +9,7 @@ dependencies = [
     "pyarrow>=20.0.0",
     "query-farm-flight-server",
     "duckdb>=1.3.1",
-    "query-farm-duckdb-json-serialization>=0.1.1",
+    "query-farm-duckdb-json-serialization>=0.1.1"
 ]
 readme = "README.md"
 requires-python = ">= 3.12"

src/query_farm_airport_test_server/database_impl.py

@@ -780,6 +780,36 @@ def yellow_taxi_endpoint_generator(ticket_data: Any) -> list[flight.FlightEndpoi
     ),
 )
 
+
+def collatz_step_count(n: int) -> int:
+    steps = 0
+    while n != 1:
+        if n % 2 == 0:
+            n //= 2
+        else:
+            n = 3 * n + 1
+        steps += 1
+    return steps
+
+
+def collatz(inputs: pa.Array) -> pa.Array:
+    results = [collatz_step_count(n) for n in inputs.to_pylist()]
+    return pa.array(results, type=pa.int64())
+
+
+def collatz_steps(n: int) -> list[int]:
+    steps = 0
+    results = []
+    while n != 1:
+        if n % 2 == 0:
+            n //= 2
+        else:
+            n = 3 * n + 1
+        results.append(n)
+        steps += 1
+    return results
+
+
 util_schema = SchemaCollection(
     scalar_functions_by_name=CaseInsensitiveDict(
         {
@@ -798,6 +828,18 @@ def yellow_taxi_endpoint_generator(ticket_data: Any) -> list[flight.FlightEndpoi
                 output_schema=pa.schema([pa.field("result", pa.int64())]),
                 handler=add_handler,
             ),
+            "collatz": ScalarFunction(
+                input_schema=pa.schema([pa.field("n", pa.int64())]),
+                output_schema=pa.schema([pa.field("result", pa.int64())]),
+                handler=lambda table: collatz(table.column(0)),
+            ),
+            "collatz_sequence": ScalarFunction(
+                input_schema=pa.schema([pa.field("n", pa.int64())]),
+                output_schema=pa.schema([pa.field("result", pa.list_(pa.int64()))]),
+                handler=lambda table: pa.array(
+                    [collatz_steps(n) for n in table.column(0).to_pylist()], type=pa.list_(pa.int64())
+                ),
+            ),
         }
     ),
     table_functions_by_name=CaseInsensitiveDict(