Browse Source

support download json file of sample data

Signed-off-by: ryjiang <jiangruiyi@gmail.com>
ryjiang 1 year ago
parent
commit
00013bfda1

+ 3 - 3
client/src/http/Collection.ts

@@ -28,7 +28,7 @@ export class CollectionHttp extends BaseModel implements CollectionView {
   private id!: string;
   private loadedPercentage!: string;
   private createdTime!: string;
-  private csv!: string;
+  private sampleFile!: string;
   private schema!: {
     fields: Field[];
     autoID: boolean;
@@ -253,7 +253,7 @@ export class CollectionHttp extends BaseModel implements CollectionView {
     return this.schema;
   }
 
-  get _csv() {
-    return this.csv;
+  get _sampleFile() {
+    return this.sampleFile;
   }
 }

+ 3 - 2
client/src/i18n/en/insert.ts

@@ -31,8 +31,9 @@ const insertTrans = {
 
   importSampleData: 'Import sample data into {{collection}}',
   sampleDataSize: 'Choose sample data size',
-  importSampleDataDesc: `Import random data based on the collection's schema.`,
-  downloadSampleDataCSV: `Download Sample Data CSV`,
+  importSampleDataDesc: `This function imports randomly generated data matching the collection schema. Useful for testing and development. Click the download button to get the data.`,
+  downloadSampleDataCSV: `Download Sample CSV Data`,
+  downloadSampleDataJSON: `Download Sample JSON Data`,
 };
 
 export default insertTrans;

+ 76 - 41
client/src/pages/dialogs/ImportSampleDialog.tsx

@@ -1,10 +1,9 @@
-import { makeStyles, Theme, Typography } from '@material-ui/core';
+import { makeStyles, Theme, Typography, Chip } from '@material-ui/core';
 import { FC, useState, useContext } from 'react';
 import { useTranslation } from 'react-i18next';
 import { saveAs } from 'file-saver';
 import DialogTemplate from '@/components/customDialog/DialogTemplate';
 import CustomSelector from '@/components/customSelector/CustomSelector';
-import CustomIconButton from '@/components/customButton/CustomIconButton';
 import { rootContext } from '@/context';
 import { InsertStatusEnum } from './insert/Types';
 import { CollectionHttp, MilvusHttp } from '@/http';
@@ -18,9 +17,9 @@ const getStyles = makeStyles((theme: Theme) => {
       fontSize: '16px',
     },
     downloadBtn: {
-      margin: theme.spacing(1.5, 1),
+      maxWidth: '240px',
+      margin: theme.spacing(1.5, 1.5, 0, 0),
     },
-
     selectors: {
       '& .selectorWrapper': {
         display: 'flex',
@@ -36,26 +35,56 @@ const getStyles = makeStyles((theme: Theme) => {
 
         '& .description': {
           color: theme.palette.attuGrey.dark,
-          marginBottom: theme.spacing(1),
-          fontSize: 12,
+          marginBottom: theme.spacing(2),
+          fontSize: 13,
+          lineHeight: 1.5,
+          width: '35vw',
         },
       },
-
       '& .actions': {
+        display: 'flex',
+        flexDirection: 'column',
+      },
+      '& .download-actions': {
         display: 'flex',
         flexDirection: 'row',
       },
-
       '& .selector': {
-        minWidth: theme.spacing(48),
+        maxWidth: '35vw',
       },
     },
   };
 });
 
+const sizeOptions = [
+  {
+    label: '100',
+    value: '100',
+  },
+  {
+    label: '1000',
+    value: '1000',
+  },
+  {
+    label: '5000',
+    value: '5000',
+  },
+  {
+    label: '10k',
+    value: '10000',
+  },
+];
+
 const ImportSampleDialog: FC<{ collection: string }> = props => {
   const classes = getStyles();
-  const [size, setSize] = useState<string>('100');
+  const { collection } = props;
+  const [size, setSize] = useState<string>(sizeOptions[0].value);
+  const [csvFileName, setCsvFileName] = useState<string>(
+    `${collection}.sample.${size}.csv`
+  );
+  const [jsonFileName, setJsonFileName] = useState<string>(
+    `${collection}.sample.${size}.json`
+  );
   const [insertStatus, setInsertStatus] = useState<InsertStatusEnum>(
     InsertStatusEnum.init
   );
@@ -65,34 +94,17 @@ const ImportSampleDialog: FC<{ collection: string }> = props => {
   const { handleCloseDialog, openSnackBar } = useContext(rootContext);
   // selected collection name
 
-  const sizeOptions = [
-    {
-      label: '100',
-      value: '100',
-    },
-    {
-      label: '1k',
-      value: '1000',
-    },
-    {
-      label: '5k',
-      value: '5000',
-    },
-    {
-      label: '10k',
-      value: '10000',
-    },
-  ];
-
   const handleImportSample = async (
     collectionName: string,
     size: string,
-    download: boolean = false
+    download: boolean = false,
+    format: 'csv' | 'json' = 'csv'
   ): Promise<{ result: string | boolean; msg: string }> => {
     const param: LoadSampleParam = {
       collection_name: collectionName,
       size: size,
       download,
+      format: format,
     };
     try {
       const res = (await CollectionHttp.importSample(
@@ -100,9 +112,12 @@ const ImportSampleDialog: FC<{ collection: string }> = props => {
         param
       )) as CollectionHttp;
       if (download) {
-        const blob = new Blob([res._csv], { type: 'text/csv;charset=utf-8;' });
-        saveAs(blob, `${collectionName}.sample.${size}.csv`);
-        return { result: res._csv, msg: '' };
+        const fileName = format === 'csv' ? csvFileName : jsonFileName;
+        const type =
+          format === 'csv' ? 'text/csv;charset=utf-8;' : 'application/json';
+        const blob = new Blob([res._sampleFile], { type });
+        saveAs(blob, fileName);
+        return { result: res._sampleFile, msg: '' };
       }
       await MilvusHttp.flush(collectionName);
       return { result: true, msg: '' };
@@ -117,7 +132,11 @@ const ImportSampleDialog: FC<{ collection: string }> = props => {
   };
 
   const onDownloadCSVClicked = async () => {
-    return await handleImportSample(props.collection, size, true);
+    return await handleImportSample(collection, size, true, 'csv');
+  };
+
+  const onDownloadJSONClicked = async () => {
+    return await handleImportSample(collection, size, true, 'json');
   };
 
   const importData = async () => {
@@ -127,7 +146,7 @@ const ImportSampleDialog: FC<{ collection: string }> = props => {
     }
     // start loading
     setInsertStatus(InsertStatusEnum.loading);
-    const { result, msg } = await handleImportSample(props.collection, size);
+    const { result, msg } = await handleImportSample(collection, size);
 
     if (!result) {
       openSnackBar(msg, 'error');
@@ -142,7 +161,7 @@ const ImportSampleDialog: FC<{ collection: string }> = props => {
   return (
     <DialogTemplate
       title={insertTrans('importSampleData', {
-        collection: props.collection,
+        collection,
       })}
       handleClose={handleCloseDialog}
       confirmLabel={
@@ -180,15 +199,31 @@ const ImportSampleDialog: FC<{ collection: string }> = props => {
               onChange={(e: { target: { value: unknown } }) => {
                 const size = e.target.value;
                 setSize(size as string);
+                setCsvFileName(`${collection}.sample.${size}.csv`);
+                setJsonFileName(`${collection}.sample.${size}.json`);
               }}
             />
-            <CustomIconButton
+          </div>
+
+          <div className="download-actions">
+            <Chip
               className={classes.downloadBtn}
-              tooltip={insertTrans('downloadSampleDataCSV')}
+              icon={<DownloadIcon />}
+              label={csvFileName}
+              title={csvFileName}
+              variant="outlined"
+              size="small"
               onClick={onDownloadCSVClicked}
-            >
-              <DownloadIcon />
-            </CustomIconButton>
+            />
+            <Chip
+              className={classes.downloadBtn}
+              icon={<DownloadIcon />}
+              label={jsonFileName}
+              title={jsonFileName}
+              variant="outlined"
+              size="small"
+              onClick={onDownloadJSONClicked}
+            />
           </div>
         </div>
       </form>

+ 1 - 0
client/src/pages/dialogs/Types.ts

@@ -38,4 +38,5 @@ export interface LoadSampleParam {
   // e.g. [{vector: [1,2,3], age: 10}]
   size: string;
   download?: boolean;
+  format?: 'csv' | 'json';
 }

+ 11 - 3
server/src/collections/collections.service.ts

@@ -308,7 +308,12 @@ export class CollectionsService {
   /**
    * Load sample data into collection
    */
-  async importSample({ collection_name, size, download }: ImportSampleDto) {
+  async importSample({
+    collection_name,
+    size,
+    download,
+    format,
+  }: ImportSampleDto) {
     const collectionInfo = await this.describeCollection({ collection_name });
     const fields_data = genRows(
       collectionInfo.schema.fields,
@@ -318,9 +323,12 @@ export class CollectionsService {
 
     if (download) {
       const parser = new Parser({});
-      const csv = parser.parse(fields_data);
+      const sampleFile =
+        format === 'csv'
+          ? parser.parse(fields_data)
+          : JSON.stringify(fields_data);
       // If download is true, return the generated data directly
-      return { csv };
+      return { sampleFile };
     } else {
       // Otherwise, insert the data into the collection
       return await this.insert({ collection_name, fields_data });

+ 7 - 0
server/src/collections/dto.ts

@@ -55,10 +55,17 @@ export class InsertDataDto {
 export class ImportSampleDto {
   @IsOptional()
   readonly collection_name: string;
+
   @IsString()
   readonly size: string;
+
   @IsBoolean()
+  @IsOptional()
   readonly download?: boolean;
+
+  @IsString()
+  @IsOptional()
+  readonly format?: string;
 }
 
 export class GetReplicasDto {