Neo4jCsvPublisher Speed Optimization (Parallelism)

[chonyy]

Hi Team,
I’m wondering if there’s a plan to apply multiprocessing on the publishers. We have a large amount of metadata in our production, which ended up running 3 million queries on neo4j . It takes about 90 minutes to finish.

To investigate the bottleneck, I looked into the code and logged the time elapsed for each step in a single iteration in the _publish_node function. This is the result

• Neo4j query: 0.1ms
• Create statement: 1ms
• Others: super fast, doesn’t matter
  

Surprisingly, the bottleneck is not the db query, it’s the statement creation. The process is basically

1. loop each row in csv
2. parse the row into a dictionary
3. loop through each key value pair in the dictionary to get the props
4. fill the statement Jinjia template with the props
5. execute the query with the statement

I’m thinking that instead of read a row => create a node in graph db one by one, maybe we could use multiprocessing to speed up the process. I believe there will be no dependency issue as long as we publish all the nodes before publishing relations, which is already handled in the current codebase.
I’m planning on implementing multiprocessing for this, is there any potential problem? Like dependency, graph db load, etc..

Expected Behavior or Use Case

Speed up the performance of the publisher. Currently, a 90 min sync is not acceptable for our use case 😢

Service or Ingestion ETL

Ingestion ETL, publisher

Possible Implementation

Thanks to @dkunitsk 's idea, I think there are three possible implementations

1. Multiprocessing on call side
2. Multiprocessing on Neo4j publisher
3. Neo4j UNWIND (Batch processing)

class HiveParallelIndexer:
    # Shim for adding all node labels to the NEO4J_DEADLOCK_NODE_LABELS config
    # which enables retries for those node labels. This is important for parallel writing
    # since we see intermittent Neo4j deadlock errors relatively often.
    class ContainsAllList(list):
        def __contains__(self, item):
            return True

    def __init__(self, publish_tag: str, parallelism: int):
        self.publish_tag = publish_tag
        self.parallelism = parallelism

    def __call__(self, worker_index: int):
        # Sharding:
        #   - take the md5 hash of the schema.table_name
        #   - convert the first 3 characters of the hash to decimal (3 chosen arbitrarily)
        #   - mod by total number of processes
        where_clause_suffix = """
            WHERE MOD(CONV(LEFT(MD5(CONCAT(d.NAME, '.', t.TBL_NAME)), 3), 16, 10), {total_parallelism}) = {worker_index}
            AND t.TBL_TYPE IN ('EXTERNAL_TABLE', 'MANAGED_TABLE', 'VIRTUAL_VIEW')
            AND (t.VIEW_EXPANDED_TEXT != '/* Presto View */' OR t.VIEW_EXPANDED_TEXT is NULL)
        """.format(total_parallelism=self.parallelism,
            worker_index=worker_index)

        # configs relevant for multiprocessing
        job_config = ConfigFactory.from_dict({
            'extractor.hive_table_metadata.{}'.format(HiveTableMetadataExtractor.WHERE_CLAUSE_SUFFIX_KEY):
                where_clause_suffix,
            # keeping this relatively low, in our experience, reduces neo4j deadlocks
            'publisher.neo4j.{}'.format(neo4j_csv_publisher.NEO4J_TRANSACTION_SIZE):
                100,
            'publisher.neo4j.{}'.format(neo4j_csv_publisher.NEO4J_DEADLOCK_NODE_LABELS):
                HiveParallelIndexer.ContainsAllList(),
        })
        job = DefaultJob(conf=job_config,
                         task=DefaultTask(
                             extractor=HiveTableMetadataExtractor(),
                             loader=FsNeo4jCSVLoader()),
                         publisher=Neo4jCsvPublisher())
        job.launch()


parallelism = 16
indexer = HiveParallelIndexer(
    publish_tag='2021-12-03'
    parallelism=parallelism)

with multiprocessing.Pool(processes=parallelism) as pool:
    def callback(_):
        # fast fail in case of exception in any process
        print('terminating due to exception')
        pool.terminate()
    res = pool.map_async(indexer, [i for i in range(parallelism)], error_callback=callback)
    res.get()

Screenshots of Slack Discussion

[chonyy]

The where_clause_suffix provided in the sample code above is for MySQL. The syntax is not compatible for PostgreSQL. For Postgres it should be replaced to

# PostgreSQL where clause
where_clause_suffix = """
    WHERE MOD(('x' || lpad(LEFT(MD5(CONCAT(d."NAME", '.', t."TBL_NAME")), 3), 8, '0'))::bit(32)::int, {total_parallelism}) = {worker_index}
    AND t."TBL_TYPE" IN ('EXTERNAL_TABLE', 'MANAGED_TABLE', 'VIRTUAL_VIEW')
    AND (t."VIEW_EXPANDED_TEXT" != '/* Presto View */' OR t."VIEW_EXPANDED_TEXT" is NULL)
""".format(total_parallelism=self.parallelism, worker_index=worker_index)

Reference

[chonyy]

When working with Airflow pipelines, it would be better to split it into multiple tasks in a dag instead of running multiprocess in a single task. The sample code is provided below

# Extract metastore table/column metadata from metastore and publish it to neo4j
def run_multi_hive_metastore_job(parallelism, worker_index):
    # Sharding:
    #   - take the md5 hash of the schema.table_name
    #   - convert the first 3 characters of the hash to decimal (3 chosen arbitrarily)
    #   - mod by total number of processes
    where_clause_suffix = """
        WHERE MOD(('x' || lpad(LEFT(MD5(CONCAT(d."NAME", '.', t."TBL_NAME")), 3), 8, '0'))::bit(32)::int, {total_parallelism}) = {worker_index}
        AND t."TBL_TYPE" IN ('EXTERNAL_TABLE', 'MANAGED_TABLE', 'VIRTUAL_VIEW')
        AND (t."VIEW_EXPANDED_TEXT" != '/* Presto View */' OR t."VIEW_EXPANDED_TEXT" is NULL)
    """.format(total_parallelism=parallelism, worker_index=worker_index)

    tmp_folder = '/var/tmp/amundsen/table_metadata' + str(worker_index)
    node_files_folder = f'{tmp_folder}/nodes/'
    relationship_files_folder = f'{tmp_folder}/relationships/'

    job_config = ConfigFactory.from_dict({
        'extractor.hive_table_metadata.{}'.format(HiveTableMetadataExtractor.WHERE_CLAUSE_SUFFIX_KEY): where_clause_suffix,
        # keeping this relatively low, in our experience, reduces neo4j deadlocks
        'publisher.neo4j.{}'.format(neo4j_csv_publisher.NEO4J_TRANSACTION_SIZE): 300,
        f'extractor.hive_table_metadata.extractor.sqlalchemy.{SQLAlchemyExtractor.CONN_STRING}': metastore_postgresql_connection_string(),
        f'loader.filesystem_csv_neo4j.{FsNeo4jCSVLoader.NODE_DIR_PATH}': node_files_folder,
        f'loader.filesystem_csv_neo4j.{FsNeo4jCSVLoader.RELATION_DIR_PATH}': relationship_files_folder,
        f'publisher.neo4j.{neo4j_csv_publisher.NODE_FILES_DIR}': node_files_folder,
        f'publisher.neo4j.{neo4j_csv_publisher.RELATION_FILES_DIR}': relationship_files_folder,
        f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_END_POINT_KEY}': neo4j_endpoint,
        f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_USER}': neo4j_user,
        f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_PASSWORD}': neo4j_password,
        f'publisher.neo4j.{neo4j_csv_publisher.JOB_PUBLISH_TAG}': timestamp,  # set the timestamp on each node in neo4j for later staleness removal
    })
    job = DefaultJob(conf=job_config,
                        task=DefaultTask(
                            extractor=HiveTableMetadataExtractor(),
                            loader=FsNeo4jCSVLoader()),
                        publisher=Neo4jCsvPublisher())
    job.launch()

parallelism = 16

for idx in range(parallelism):
    hive_metastore_task = PythonOperator(
        task_id='hive_metastore_task' + str(idx),
        python_callable=run_multi_hive_metastore_job,
        dag=dag,
        op_kwargs={'parallelism': parallelism, 'worker_index': idx})
    kick_off_dummy >> hive_metastore_task >> trino_view_task >> staleness_removal_task >> publish_elastic_search_task

[chonyy]

Based on my experience, the ingestion speed is about 2.5 times faster by applying multiprocessing on the caller side. Since the performance improvement will be bounded by the neo4j instance, it wouldn't help much if we also implement multiprocessing on the publisher side. This means that the only other possible optimization that might help is neo4j UNWIND.

However, as @dkunitsk pointed out, applying UNWIND in the publisher could be a tricky job. The query plan has to be well designed in order to gain the max performance improvement. Whether we should totally rewrite the publisher is still a question. I think I could provide some thought after my experiments.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs.

[mgorsk1]

Thanks @chonyy and @dkunitsk for bringing this up, I've stumbled upon this as well recently and have following thoughs:

1. I think looking into parallelization on caller (executor) side is not scalable as it'd require rework on every extractor. The benefit, however, would be that it's automatically available for any publisher.
2. I am totally in favor of improving neo4j publisher as it would support every extractor. I took initial spin with using current code with multiprocessing, however it still results in some locks. The improvement was 2.5 times faster with 15k tables (500k columns). The code is here 40791ad - if you have any hints on how to avoid locks it could be short-term solution. Bear in mind this implementation is very raw as I didn't want to invest more time unless it's a direction we want to take.
3. long-term unwind seems like a way to go, we already use bulk loads in ES publisher and should introduce it for neo4j too.

[zacr]

@mgorsk1, @chonyy, @dkunitsk I created a version of the publisher that uses UNWIND and APOC library. It's 5x faster for me on my test dataset (1500 tables, 650k columns). Still rough around the edges and I did it based on the CsvTable*Extractors so not sure if I baked in any assumptions that won't hold with other extractors. It will need some testing and I likely introduced some issues - but workable updated cypher queries should be in place and planning on running it for a few weeks to shake it out. There is additional room for optimization possibly; not doing anything in parallel to avoid locks though apoc.iterate supports parallelism and retries. If anybody wants to play with it and try to get it over the line happy to toss you the code.

[chonyy]

Hi @zacr , thanks for looking into this!

I'm not super familiar with the capabilities of UNWIND, but the 5x speedup sounds promising 🚀
One thing that I would like to clarify is currently there's no any kind of parallelism in your version right? Does this mean it's possible to have at least 10x speedup if we combined both multiprocessing and UNWIND?

Also, I'm interested in your implementation, could you point me to the commit or the branch? Thanks in advance 🙏

[zacr]

Wasn't sure if I am allowed to make a branch in the main repo or if I am supposed to fork first. Hopefully you can see this, I did it as a new file so could run both side by side: https://github.com/SaltIO/amundsen/blob/neo4j_csv_publisher_apoc/databuilder/databuilder/publisher/neo4j_csv_publisher_apoc.py

I didn't do any query optimization or tuning, the goal was to just try to get functional cypher queries using apoc.periodic.iterate which has support for parallelization. I suspect there is some meat on the bone by investigating and tuning these new cypher queries. The most recent version looks more like 3.5x faster instead of 5x, but I didn't do any serious timing. If you turn on Parallel:true and retries: 1 it goes 5x but locks will cause you to drop maybe 0.3% of the inserts. These cypher queries may be on the path to 7x+ with tuning unless I did them wrong.

After you investigate if you think this APOC based solution has promise and you or someone else wants to run with it that is great by me. In the meantime I am just going to try to let it run in a few projects with different extractors over the next few weeks to surface issues.

[afuzzyriff]

Hello - I stumbled across this issue as I was facing a similar problem of poor performance with large Neo4j writes.

@zacr Thank you for sharing your Neo4j CSV Publisher with APOC code. I reviewed the class and started to test this in my team's Amundsen POC implementation for Neo4j CSV publishing. The custom class provided seems to be working nicely.

I have tested a batch of data which resulted in writing 4,848,906 nodes and 9,715,678 relationships in 5704.64 seconds (~95 minutes) using an 8k batch size. I don't have good non-APOC query timing results as that scenario was running rather slowly, but I do know this is a vast timing improvement.

[zacr]

Great to hear! Our testing has been going well - we have been using this class for all publishes the last month and haven't seen any issues. I was hoping to find time to try a Neo4J version upgrade and do a little testing on that. I think the INDEX creation code in that publisher would need to be tweaked for a newer version of Neo4J. Was also thinking about looking at the parallelism one more time to see if there is any speed up there. Lots of good info in this article to try out still: https://neo4j.com/blog/nodes-2019-best-practices-to-make-large-updates-in-neo4j/

[VoLKyyyOG]

Regarding this, would be a simple merge into the main repo or can I download it and add it to mine. I presume the change would be as simple as replacing neo4j_csv_publisher with neo4j_csv_publisher_apoc?

There were also some issues with getting APOC installed alongside neo4j so if anyone has issues, let me know as I have been able to resolve it.