feat: update parquet file handling and documentation for GraphRAG v2.x.x (#40)

Author: noworneverev

README.md

@@ -85,12 +85,22 @@ The logic for creating relationships for text units, documents, communities, and
 
 To load `.parquet` files automatically when the application starts, place your Parquet files in the `public/artifacts` directory. These files will be loaded into the application for visualization and data table display. The files can be organized as follows:
 
-- `public/artifacts/create_final_entities.parquet`
-- `public/artifacts/create_final_relationships.parquet`
-- `public/artifacts/create_final_documents.parquet`
-- `public/artifacts/create_final_text_units.parquet`
-- `public/artifacts/create_final_communities.parquet`
-- `public/artifacts/create_final_community_reports.parquet`
-- `public/artifacts/create_final_covariates.parquet`
+- GraphRAG v2.x.x
+  - `public/artifacts/entities.parquet`
+  - `public/artifacts/relationships.parquet`
+  - `public/artifacts/documents.parquet`
+  - `public/artifacts/text_units.parquet`
+  - `public/artifacts/communities.parquet`
+  - `public/artifacts/community_reports.parquet`
+  - `public/artifacts/covariates.parquet`
+
+- GraphRAG v1.x.x
+  - `public/artifacts/create_final_entities.parquet`
+  - `public/artifacts/create_final_relationships.parquet`
+  - `public/artifacts/create_final_documents.parquet`
+  - `public/artifacts/create_final_text_units.parquet`
+  - `public/artifacts/create_final_communities.parquet`
+  - `public/artifacts/create_final_community_reports.parquet`
+  - `public/artifacts/create_final_covariates.parquet`
 
 If the files are placed in the `public/artifacts` folder, the app will automatically load and display them on startup.

src/app/components/DataTableContainer.tsx

@@ -114,47 +114,47 @@ const DataTableContainer: React.FC<DataTableContainerProps> = ({
         {selectedTable === "entities" && (
           <>
             <Typography variant="h4" gutterBottom>
-              Entities (create_final_entities.parquet)
+              Entities (entities.parquet)
             </Typography>
             <DataTable columns={entityColumns} data={entities} />
           </>
         )}
         {selectedTable === "relationships" && (
           <>
             <Typography variant="h4" gutterBottom>
-              Relationships (create_final_relationships.parquet)
+              Relationships (relationships.parquet)
             </Typography>
             <DataTable columns={relationshipColumns} data={relationships} />
           </>
         )}
         {selectedTable === "documents" && (
           <>
             <Typography variant="h4" gutterBottom>
-              Documents (create_final_documents.parquet)
+              Documents (documents.parquet)
             </Typography>
             <DataTable columns={documentColumns} data={documents} />
           </>
         )}
         {selectedTable === "textunits" && (
           <>
             <Typography variant="h4" gutterBottom>
-              TextUnits (create_final_text_units.parquet)
+              TextUnits (text_units.parquet)
             </Typography>
             <DataTable columns={textUnitColumns} data={textunits} />
           </>
         )}
         {selectedTable === "communities" && (
           <>
             <Typography variant="h4" gutterBottom>
-              Communities (create_final_communities.parquet)
+              Communities (communities.parquet)
             </Typography>
             <DataTable columns={communityColumns} data={communities} />
           </>
         )}
         {selectedTable === "communityReports" && (
           <>
             <Typography variant="h4" gutterBottom>
-              Community Reports (create_final_community_reports.parquet)
+              Community Reports (community_reports.parquet)
             </Typography>
             <DataTable
               columns={communityReportColumns}
@@ -165,7 +165,7 @@ const DataTableContainer: React.FC<DataTableContainerProps> = ({
         {selectedTable === "covariates" && (
           <>
             <Typography variant="h4" gutterBottom>
-              Covariates (create_final_covariates.parquet)
+              Covariates (covariates.parquet)
             </Typography>
             <DataTable columns={covariateColumns} data={covariates} />
           </>

src/app/hooks/useFileHandler.ts

@@ -9,27 +9,32 @@ import { CommunityReport } from "../models/community-report";
 import { Covariate } from "../models/covariate";
 import { readParquetFile } from "../utils/parquet-utils";
 
-// Paths to default files in the public folder
-const defaultFiles = [
-  process.env.PUBLIC_URL + "/artifacts/create_final_entities.parquet",
-  process.env.PUBLIC_URL + "/artifacts/create_final_relationships.parquet",
-  process.env.PUBLIC_URL + "/artifacts/create_final_documents.parquet",
-  process.env.PUBLIC_URL + "/artifacts/create_final_text_units.parquet",
-  process.env.PUBLIC_URL + "/artifacts/create_final_communities.parquet",
-  process.env.PUBLIC_URL + "/artifacts/create_final_community_reports.parquet",
-  process.env.PUBLIC_URL + "/artifacts/create_final_covariates.parquet",
+const baseFileNames = [
+  "entities.parquet",
+  "relationships.parquet",
+  "documents.parquet",
+  "text_units.parquet",
+  "communities.parquet",
+  "community_reports.parquet",
+  "covariates.parquet",
 ];
 
-const fileSchemas: { [key: string]: string } = {
-  "create_final_entities.parquet": "entity",
-  "create_final_relationships.parquet": "relationship",
-  "create_final_text_units.parquet": "text_unit",
-  "create_final_communities.parquet": "community",
-  "create_final_community_reports.parquet": "community_report",
-  "create_final_documents.parquet": "document",
-  "create_final_covariates.parquet": "covariate",
+const baseMapping: { [key: string]: string } = {
+  "entities.parquet": "entity",
+  "relationships.parquet": "relationship",
+  "documents.parquet": "document",
+  "text_units.parquet": "text_unit",
+  "communities.parquet": "community",
+  "community_reports.parquet": "community_report",
+  "covariates.parquet": "covariate",
 };
 
+const fileSchemas: { [key: string]: string } = {};
+Object.entries(baseMapping).forEach(([key, value]) => {  
+  fileSchemas[key] = value;  
+  fileSchemas[`create_final_${key}`] = value;
+});
+
 const useFileHandler = () => {
   const navigate = useNavigate();
   const [entities, setEntities] = useState<Entity[]>([]);
@@ -57,8 +62,8 @@ const useFileHandler = () => {
 
     for (const file of files) {
       const fileName =
-        typeof file === "string" ? file.split("/").pop()! : file.name;
-      const schema = fileSchemas[fileName];
+        typeof file === "string" ? file.split("/").pop()! : file.name;      
+      const schema = fileSchemas[fileName] || fileSchemas[`create_final_${fileName}`];
 
       let data;
       if (typeof file === "string") {
@@ -144,12 +149,17 @@ const useFileHandler = () => {
   const loadDefaultFiles = async () => {
     const filesToLoad = [];
 
-    for (const file of defaultFiles) {
-      const fileExists = await checkFileExists(file);
-      if (fileExists) {
-        filesToLoad.push(file); // Add to load queue if the file exists
+    for (const baseName of baseFileNames) {
+      const prefixedPath = process.env.PUBLIC_URL + `/artifacts/create_final_${baseName}`;
+      const unprefixedPath = process.env.PUBLIC_URL + `/artifacts/${baseName}`;
+  
+      if (await checkFileExists(prefixedPath)) {
+        filesToLoad.push(prefixedPath);
+      } else if (await checkFileExists(unprefixedPath)) {
+        filesToLoad.push(unprefixedPath);
       }
     }
+    
     if (filesToLoad.length > 0) {
       await loadFiles(filesToLoad);
       navigate("/graph", { replace: true });