Merge pull request #84 from palantir/ssanjay/aggCount

Add count to aggregations

examples/basic/cli/src/index.ts

@@ -181,6 +181,32 @@ async function runTests() {
 
     console.log(intersectResultbbox.data.map(data => data.usState));
 
+    const testAggregateCountNoGroup = await client.objects.BoundariesUsState
+      .aggregateOrThrow({
+        select: { $count: true, latitude: ["min", "max", "avg"] },
+      });
+
+    // Should be 51 because it includes DC
+    console.log(
+      testAggregateCountNoGroup.$count,
+      testAggregateCountNoGroup.latitude.avg,
+      testAggregateCountNoGroup.latitude.max,
+      testAggregateCountNoGroup.latitude.min,
+    );
+    const testAggregateCountWithGroups = await client.objects.BoundariesUsState
+      .aggregateOrThrow({
+        select: { $count: true, latitude: ["min", "max", "avg"] },
+        groupBy: { usState: "exact" },
+      });
+
+    console.log(
+      testAggregateCountWithGroups[0].$group.usState,
+      testAggregateCountWithGroups[0].$count,
+      testAggregateCountWithGroups[0].latitude.avg,
+      testAggregateCountWithGroups[0].latitude.max,
+      testAggregateCountWithGroups[0].latitude.min,
+    );
+
     if (runOld) await typeChecks(client);
   } catch (e) {
     console.error("Caught an error we did not expect", typeof e);

examples/basic/sdk/ontology.json

@@ -156,6 +156,16 @@
               "type": "string"
             }
           },
+          "latitude": {
+            "dataType": {
+              "type": "double"
+            }
+          },
+          "longitude": {
+            "dataType": {
+              "type": "double"
+            }
+          },
           "geometry10M": {
             "dataType": {
               "type": "geoshape"

examples/basic/sdk/src/generatedNoCheck/ontology/objects/BoundariesUsState.ts

@@ -7,6 +7,8 @@ export interface BoundariesUsState extends ObjectTypeDefinition<'BoundariesUsSta
   links: {};
   properties: {
     usState: PropertyDef<'string', 'non-nullable', 'single'>;
+    latitude: PropertyDef<'double', 'nullable', 'single'>;
+    longitude: PropertyDef<'double', 'nullable', 'single'>;
     geometry10M: PropertyDef<'geoshape', 'nullable', 'single'>;
   };
 }
@@ -24,6 +26,16 @@ export const BoundariesUsState: BoundariesUsState = {
       type: 'string',
       nullable: false,
     },
+    latitude: {
+      multiplicity: false,
+      type: 'double',
+      nullable: true,
+    },
+    longitude: {
+      multiplicity: false,
+      type: 'double',
+      nullable: true,
+    },
     geometry10M: {
       multiplicity: false,
       description: 'geoshape',

packages/client/changelog/@unreleased/pr-84.v2.yml

@@ -0,0 +1,5 @@
+type: feature
+feature:
+  description: Add count to aggregations
+  links:
+  - https://github.com/palantir/osdk-ts/pull/84

packages/client/src/internal/conversions/legacyToModernSingleAggregationResult.ts

@@ -30,6 +30,9 @@ export function legacyToModernSingleAggregationResult<
   return entry.metrics.reduce(
     (accumulator, curValue) => {
       const parts = curValue.name.split(".");
+      if (parts[0] === "count") {
+        return accumulator;
+      }
       invariant(
         parts.length === 2,
         "assumed we were getting a `${key}.${type}`",

packages/client/src/internal/conversions/modernToLegacyAggregationClause.ts

@@ -21,7 +21,10 @@ export function modernToLegacyAggregationClause<
   AC extends AggregationClause<any>,
 >(select: AC) {
   return Object.entries(select).flatMap<Aggregation>(([k, v]) => {
-    if (Array.isArray(v)) {
+    if (k === "$count") {
+      if (v) return { type: "count", name: "count" };
+      return [];
+    } else if (Array.isArray(v)) {
       return v.map((v2) => {
         return {
           type: v2,

packages/client/src/object/aggregateOrThrow.test.ts

@@ -141,6 +141,7 @@ describe("aggregateOrThrow", () => {
           text: "approximateDistinct",
           priority: "avg",
           id: ["max", "avg"],
+          $count: true,
         },
       },
     );
@@ -149,6 +150,7 @@ describe("aggregateOrThrow", () => {
     expectType<number | undefined>(notGrouped.priority.avg);
     expectType<number | undefined>(notGrouped.id.max);
     expectType<number | undefined>(notGrouped.id.avg);
+    expectType<number>(notGrouped.$count);
     expectType<
       TypeOf<
         {
@@ -168,6 +170,7 @@ describe("aggregateOrThrow", () => {
       {
         select: {
           id: "approximateDistinct",
+          $count: true,
         },
         groupBy: {
           text: "exact",
@@ -177,6 +180,7 @@ describe("aggregateOrThrow", () => {
     expectType<Array<any>>(grouped);
     expectType<string | undefined>(grouped[0].$group.text);
     expectType<number>(grouped[0].id.approximateDistinct);
+    expectType<number>(grouped[0].$count);
   });
 
   it("works with where: todo", async () => {
@@ -197,17 +201,17 @@ describe("aggregateOrThrow", () => {
               id: {
                 type: "double";
               };
+              locationCity: {
+                type: "string";
+              };
             };
           };
         };
         actions: {};
         queries: {};
-      }["objects"]["Todo"],
-      {
-        locationCity: "approximateDistinct";
-        text: "approximateDistinct";
-      }
-    >;
+      }["objects"]["Todo"]
+    > // "locationCity" | "text"
+    ;
 
     const f: AggregateOpts<
       {
@@ -226,22 +230,22 @@ describe("aggregateOrThrow", () => {
               id: {
                 type: "double";
               };
+              locationCity: {
+                type: "string";
+              };
             };
           };
         };
         actions: {};
         queries: {};
-      }["objects"]["Todo"],
-      {
-        locationCity: "approximateDistinct";
-        text: "approximateDistinct";
-      }
+      }["objects"]["Todo"]
     > = {
       select: {
         locationCity: "approximateDistinct",
+        text: "approximateDistinct",
       },
-    } as any;
+    };
 
-    expectType<"approximateDistinct">(f.select.locationCity);
+    // expectType<"approximateDistinct">(f.select.locationCity);
   });
 });

packages/client/src/object/aggregateOrThrow.ts

@@ -16,7 +16,11 @@
 
 import type { ObjectOrInterfaceDefinition } from "@osdk/api";
 import { aggregateObjectSetV2 } from "@osdk/gateway/requests";
-import type { AggregateObjectsRequestV2, ObjectSet } from "@osdk/gateway/types";
+import type {
+  AggregateObjectsRequestV2,
+  AggregateObjectsResponseV2,
+  ObjectSet,
+} from "@osdk/gateway/types";
 import { createOpenApiRequest } from "@osdk/shared.net";
 import type { ClientContext } from "@osdk/shared.net";
 import invariant from "tiny-invariant";
@@ -31,10 +35,11 @@ import type {
   AggregationResultsWithGroups,
   AggregationsResults,
 } from "../query/index.js";
+import type { ArrayElement } from "../util/ArrayElement.js";
 
 export async function aggregateOrThrow<
   Q extends ObjectOrInterfaceDefinition,
-  const AO extends AggregateOpts<Q, any>,
+  AO extends AggregateOpts<Q>,
 >(
   clientCtx: ClientContext<any>,
   objectType: Q,
@@ -77,18 +82,32 @@ export async function aggregateOrThrow<
       "no group by clause should mean only one data result",
     );
 
-    return legacyToModernSingleAggregationResult<AO["select"]>(
-      result.data[0],
-    ) as any;
+    return {
+      ...aggregationToCountResult(result.data[0]),
+      ...legacyToModernSingleAggregationResult<AO["select"]>(
+        result.data[0],
+      ),
+    } as any;
   }
 
   const ret: AggregationResultsWithGroups<Q, AO["select"], any> = result.data
     .map((entry) => {
       return {
         $group: entry.group as any,
+        ...aggregationToCountResult(entry),
         ...legacyToModernSingleAggregationResult(entry),
       };
     }) as any; // fixme
 
   return ret as any; // FIXME
 }
+
+function aggregationToCountResult(
+  entry: ArrayElement<AggregateObjectsResponseV2["data"]>,
+): { $count: number } | undefined {
+  for (const aggregateResult of entry.metrics) {
+    if (aggregateResult.name === "count") {
+      return { $count: aggregateResult.value };
+    }
+  }
+}

packages/client/src/objectSet/ObjectSet.ts

@@ -73,7 +73,7 @@ export interface ObjectSet<Q extends ObjectOrInterfaceDefinition> {
   //   Osdk<K, O, PropertyKeysFrom<O, K>>
   // >;
 
-  aggregateOrThrow: <const AO extends AggregateOpts<Q, any>>(
+  aggregateOrThrow: <AO extends AggregateOpts<Q>>(
     req: AO,
   ) => Promise<AggregationsResults<Q, AO>>;
 

packages/client/src/objectSet/createObjectSet.ts

@@ -30,7 +30,7 @@ import { fetchSingle } from "../object/fetchSingle.js";
 import { aggregateOrThrow, fetchPageOrThrow } from "../object/index.js";
 import type { Osdk } from "../OsdkObjectFrom.js";
 import type { AggregateOpts } from "../query/aggregations/AggregateOpts.js";
-import type { AggregationClause, AggregationsResults } from "../query/index.js";
+import type { AggregationsResults } from "../query/index.js";
 import type { LinkedType, LinkNames } from "./LinkUtils.js";
 import type { BaseObjectSet, ObjectSet } from "./ObjectSet.js";
 import { ObjectSetListenerWebsocket } from "./ObjectSetListenerWebsocket.js";
@@ -64,9 +64,9 @@ export function createObjectSet<Q extends ObjectOrInterfaceDefinition>(
     //   throw "TODO";
     // },
     aggregateOrThrow: async <
-      AC extends AggregationClause<Q>,
+      // AC extends AggregationClause<Q>,
       // GBC extends GroupByClause<O, K>,
-      AO extends AggregateOpts<Q, AC>,
+      AO extends AggregateOpts<Q>,
     >(
       req: AO,
     ): Promise<AggregationsResults<Q, AO>> => {

packages/client/src/query/aggregations/AggregatableKeys.ts

@@ -17,7 +17,7 @@
 import type { ObjectOrInterfaceDefinition } from "@osdk/api";
 
 export type AggregatableKeys<
-  Q extends ObjectOrInterfaceDefinition<any, any>,
+  Q extends ObjectOrInterfaceDefinition,
 > = keyof {
   [P in keyof Q["properties"]]: any;
 };

packages/client/src/query/aggregations/AggregateOpts.ts

@@ -16,16 +16,17 @@
 
 import type { ObjectOrInterfaceDefinition } from "@osdk/api";
 import type {
+  AggregatableKeys,
   AggregationClause,
   GroupByClause,
   WhereClause,
 } from "../../query/index.js";
 
 export type AggregateOpts<
-  Q extends ObjectOrInterfaceDefinition<any, any>,
-  AC extends AggregationClause<Q>,
+  Q extends ObjectOrInterfaceDefinition,
+  KK extends AggregatableKeys<Q> = AggregatableKeys<Q>,
 > = {
-  select: AC;
+  select: AggregationClause<Q, KK>;
   where?: WhereClause<Q>;
   groupBy?: GroupByClause<Q>;
 };

packages/client/src/query/aggregations/AggregationResultsWithGroups.ts

@@ -16,18 +16,25 @@
 
 import type { ObjectOrInterfaceDefinition } from "@osdk/api";
 import type { OsdkObjectPropertyType } from "../../Definitions.js";
-import type { AggregationResultsWithoutGroups } from "./AggregationResultsWithoutGroups.js";
+import type {
+  AggregationCountResult,
+  AggregationResultsWithoutGroups,
+} from "./AggregationResultsWithoutGroups.js";
 import type { AggregationClause } from "./AggregationsClause.js";
 import type { GroupByClause } from "./GroupByClause.js";
 
 export type AggregationResultsWithGroups<
   Q extends ObjectOrInterfaceDefinition<any, any>,
   A extends AggregationClause<Q>,
   G extends GroupByClause<Q> | undefined,
-> = ({
-  $group: {
-    [P in keyof G & keyof Q["properties"]]: OsdkObjectPropertyType<
-      Q["properties"][P]
-    >;
-  };
-} & AggregationResultsWithoutGroups<Q, A>)[];
+> = (
+  & {
+    $group: {
+      [P in keyof G & keyof Q["properties"]]: OsdkObjectPropertyType<
+        Q["properties"][P]
+      >;
+    };
+  }
+  & AggregationCountResult<Q, A>
+  & AggregationResultsWithoutGroups<Q, A>
+)[];

packages/client/src/query/aggregations/AggregationResultsWithoutGroups.ts

@@ -24,7 +24,7 @@ type SubselectKeys<AC extends AggregationClause<any>, P extends keyof AC> =
 
 export type AggregationResultsWithoutGroups<
   Q extends ObjectOrInterfaceDefinition<any, any>,
-  AC extends AggregationClause<any>,
+  AC extends AggregationClause<Q>,
 > = {
   [P in keyof Q["properties"] as SubselectKeys<AC, P>]: AC[P] extends
     readonly string[] | string ? {
@@ -33,3 +33,8 @@ export type AggregationResultsWithoutGroups<
     }
     : never;
 };
+
+export type AggregationCountResult<
+  Q extends ObjectOrInterfaceDefinition<any, any>,
+  A extends AggregationClause<Q>,
+> = "$count" extends keyof A ? { $count: number } : {};

packages/client/src/query/aggregations/AggregationsClause.ts

@@ -26,13 +26,18 @@ type NumericAggregateOption =
   | "avg"
   | "approximateDistinct";
 
+type totalCountOption = { $count?: true };
+
 export type AggregationClause<
-  Q extends ObjectOrInterfaceDefinition<any, any>,
-> = {
-  [P in AggregatableKeys<Q>]?: Q["properties"][P]["type"] extends "string"
-    ? StringAggregateOption | StringAggregateOption[]
-    : Q["properties"][P]["type"] extends
-      "double" | "integer" | "float" | "decimal" | "byte" | "long" | "short"
-      ? NumericAggregateOption | NumericAggregateOption[]
-    : never;
-};
+  Q extends ObjectOrInterfaceDefinition,
+  K extends AggregatableKeys<Q> = AggregatableKeys<Q>,
+> =
+  & {
+    [P in K]?: Q["properties"][P]["type"] extends "string"
+      ? StringAggregateOption | StringAggregateOption[]
+      : Q["properties"][P]["type"] extends
+        "double" | "integer" | "float" | "decimal" | "byte" | "long" | "short"
+        ? NumericAggregateOption | NumericAggregateOption[]
+      : never;
+  }
+  & totalCountOption;