def write_gremlin_df()

in awswrangler/neptune/_utils.py [0:0]

• 18 lines of code
• 7 McCabe index (conditional complexity)

def write_gremlin_df(client: "NeptuneClient", df: pd.DataFrame, mode: WriteDFType, batch_size: int) -> bool:
    """Write the provided DataFrame using Gremlin.

    Parameters
    ----------
    client
        The Neptune client to write the DataFrame
    df
        The DataFrame to write
    mode
        The type of DataFrame to write
    batch_size
        The size of the batch to write

    Returns
    -------
        True if the write operation succeeded
    """
    g = Graph().traversal()
    # Loop through items in the DF
    for index, row in df.iterrows():
        # build up a query
        if mode == WriteDFType.EDGE:
            g = _build_gremlin_edges(g, row.to_dict())
        elif mode == WriteDFType.VERTEX:
            g = _build_gremlin_vertices(g, row.to_dict())
        else:
            g = _build_gremlin_update(g, row.to_dict())
        # run the query
        if index > 0 and index % batch_size == 0:
            res = _run_gremlin_insert(client, g)
            if res:
                g = Graph().traversal()
            else:
                _logger.debug(res)
                raise exceptions.QueryFailed(
                    """Failed to insert part or all of the data in the DataFrame, please check the log output."""
                )

    return _run_gremlin_insert(client, g)