Repository: microsoft/graphrag
File: graphrag/query/cli.py (Python)

Code before:
        )
    )
    reporter.success(f"Local Search Response:\n{response}")
    # NOTE: we return the response and context data here purely as a complete demonstration of the API.
    # External users should use the API directly to get the response and context data.
    return response, context_data


def _resolve_parquet_files(
    data_path: Path, config: GraphRagConfig
) -> dict[str, pd.DataFrame]:
    """Read parquet files to a dataframe dict."""
    dataframe_dict = {}
    match config.storage.type:
        case StorageType.blob:
            if config.storage.container_name is None:
                msg = "Container name required for querying from blob storage"
                raise ValueError(msg)
            if (
                config.storage.connection_string is None
                and config.storage.storage_account_blob_url is None
            ):
                msg = "Connection string or storage account blob url required for querying blob storage"
                raise ValueError(msg)

            if config.storage.storage_account_blob_url is not None:
                storage_account_blob_url = str(config.storage.storage_account_blob_url)
                storage_account_name = storage_account_blob_url.split("//")[1].split(
                    "."
                )[0]

                storage_options = {
                    "account_name": storage_account_name,
                    "credential": DefaultAzureCredentialAsync(),
                }

                blob_service_client = BlobServiceClient(
                    account_url=storage_account_blob_url,
                    credential=DefaultAzureCredentialSync(),
                )

            else:
                connection_string = str(config.storage.connection_string)
                storage_account_name = connection_string.split("AccountName=")[1].split(
                    ";"
                )[0]
                storage_options = {
                    "account_name": storage_account_name,
                    "connection_string": connection_string,
                }


Review comment (refactor):
Can we extract this function to an util? Can we also add an optional parameter to pass a list of the files one is interested on? 
To provide an example, I'm definitely needing this for incremental indexing, so I can load a few cherry picked outputs :)

Code after:
            print()  # noqa: T201
            return full_response, context_data

        return asyncio.run(run_streaming_search())
    # not streaming
    response, context_data = asyncio.run(
        api.local_search(
            config=config,
            nodes=final_nodes,
            entities=final_entities,
            community_reports=final_community_reports,
            text_units=final_text_units,
            relationships=final_relationships,
            covariates=final_covariates,
            community_level=community_level,
            response_type=response_type,
            query=query,
        )
    )
    reporter.success(f"Local Search Response:\n{response}")
    # NOTE: we return the response and context data here purely as a complete demonstration of the API.
    # External users should use the API directly to get the response and context data.
    return response, context_data


def _resolve_parquet_files(
    root_dir: str,
    config: GraphRagConfig,
    parquet_list: list[str],
    optional_list: list[str],
) -> dict[str, pd.DataFrame]:
    """Read parquet files to a dataframe dict."""
    dataframe_dict = {}
    pipeline_config = create_pipeline_config(config)
    storage_obj = _create_storage(root_dir=root_dir, config=pipeline_config.storage)
    for parquet_file in parquet_list:
        df_key = parquet_file.split(".")[0]
        df_value = asyncio.run(
            _load_table_from_storage(name=parquet_file, storage=storage_obj)
        )
        dataframe_dict[df_key] = df_value

    # for optional parquet files, set the dict entry to None instead of erroring out if it does not exist
    for optional_file in optional_list:
        file_exists = asyncio.run(storage_obj.has(optional_file))
        df_key = optional_file.split(".")[0]
        if file_exists:
            df_value = asyncio.run(
                _load_table_from_storage(name=optional_file, storage=storage_obj)
            )
            dataframe_dict[df_key] = df_value