From a11f5377a011cb4c46e1722da16a5fb4614e751f Mon Sep 17 00:00:00 2001
From: Devansh Mahant <86195162+devansh-m12@users.noreply.github.com>
Date: Fri, 27 Dec 2024 16:48:13 +0530
Subject: [PATCH] added datapoint count on dataset table (#288)

* eval time progression (#210)

* initial work to compare evals

* remove unnecessary div

* design

---------

Co-authored-by: Din <dinmukhamed.mailibay@gmail.com>

* added datapoint count on dataset table

* added types and empty check

---------

Co-authored-by: Dinmukhamed Mailibay <47117969+dinmukhamedm@users.noreply.github.com>
Co-authored-by: skull8888888 <skull8888888@gmail.com>
Co-authored-by: Din <dinmukhamed.mailibay@gmail.com>
---
 .../api/projects/[projectId]/datasets/route.ts   | 16 +++++++++++++---
 frontend/components/datasets/datasets.tsx        | 13 +++++++++----
 frontend/lib/dataset/types.ts                    |  4 ++++
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/frontend/app/api/projects/[projectId]/datasets/route.ts b/frontend/app/api/projects/[projectId]/datasets/route.ts
index 72502ffb..e627fc91 100644
--- a/frontend/app/api/projects/[projectId]/datasets/route.ts
+++ b/frontend/app/api/projects/[projectId]/datasets/route.ts
@@ -1,8 +1,8 @@
-import { and, desc, eq, inArray } from 'drizzle-orm';
+import { and, desc, eq, getTableColumns, inArray, sql } from 'drizzle-orm';
 import { NextRequest } from 'next/server';
 
 import { db } from '@/lib/db/drizzle';
-import { datasets } from '@/lib/db/migrations/schema';
+import { datasetDatapoints,datasets } from '@/lib/db/migrations/schema';
 import { paginatedGet } from '@/lib/db/utils';
 
 export async function POST(
@@ -42,12 +42,22 @@ export async function GET(
   const pageSize =
     parseInt(req.nextUrl.searchParams.get('pageSize') ?? '50') || 50;
   const filters = [eq(datasets.projectId, projectId)];
+  const { ...columns } = getTableColumns(datasets);
+
   const datasetsData = await paginatedGet({
     table: datasets,
     pageNumber,
     pageSize,
     filters,
-    orderBy: desc(datasets.createdAt)
+    orderBy: desc(datasets.createdAt),
+    columns: {
+      ...columns,
+      datapointsCount: sql<number>`COALESCE((
+        SELECT COUNT(*)
+        FROM ${datasetDatapoints} dp
+        WHERE dp.dataset_id = datasets.id
+      ), 0)::int`.as('datapointsCount')
+    }
   });
 
   return new Response(JSON.stringify(datasetsData), { status: 200 });
diff --git a/frontend/components/datasets/datasets.tsx b/frontend/components/datasets/datasets.tsx
index cd25a701..599950d7 100644
--- a/frontend/components/datasets/datasets.tsx
+++ b/frontend/components/datasets/datasets.tsx
@@ -17,7 +17,7 @@ import {
   DialogTrigger
 } from '@/components/ui/dialog';
 import { useProjectContext } from '@/contexts/project-context';
-import { Dataset } from '@/lib/dataset/types';
+import { DatasetInfo } from '@/lib/dataset/types';
 import { useToast } from '@/lib/hooks/use-toast';
 import { PaginatedResponse } from '@/lib/types';
 import { swrFetcher } from '@/lib/utils';
@@ -33,7 +33,7 @@ export default function Datasets() {
   const { projectId } = useProjectContext();
 
   const router = useRouter();
-  const { data, mutate } = useSWR<PaginatedResponse<Dataset>>(
+  const { data, mutate } = useSWR<PaginatedResponse<DatasetInfo>>(
     `/api/projects/${projectId}/datasets/`,
     swrFetcher
   );
@@ -72,7 +72,7 @@ export default function Datasets() {
     setIsDeleteDialogOpen(false);
   };
 
-  const columns: ColumnDef<Dataset>[] = [
+  const columns: ColumnDef<DatasetInfo>[] = [
     {
       cell: ({ row }) => <Mono>{row.original.id}</Mono>,
       size: 300,
@@ -83,6 +83,11 @@ export default function Datasets() {
       header: 'name',
       size: 300
     },
+    {
+      accessorKey: 'datapointsCount',
+      header: 'Datapoints Count',
+      size: 300
+    },
     {
       header: 'Created at',
       accessorKey: 'createdAt',
@@ -107,7 +112,7 @@ export default function Datasets() {
           onRowClick={(row) => {
             router.push(`/project/${projectId}/datasets/${row.original.id}`);
           }}
-          getRowId={(row: Dataset) => row.id}
+          getRowId={(row: DatasetInfo) => row.id}
           columns={columns}
           data={data?.items}
           selectionPanel={(selectedRowIds) => (
diff --git a/frontend/lib/dataset/types.ts b/frontend/lib/dataset/types.ts
index 20fe0582..e68e2bbd 100644
--- a/frontend/lib/dataset/types.ts
+++ b/frontend/lib/dataset/types.ts
@@ -5,6 +5,10 @@ export interface Dataset {
   indexedOn: string | null;
 }
 
+export interface DatasetInfo extends Dataset {
+  datapointsCount: number;
+}
+
 export interface Datapoint {
   id: string;
   createdAt: string;