add in a full text index creation on bulk load (#39)

Author: maguec

README.md

@@ -55,6 +55,7 @@ python3 redisgraph_bulk_loader/bulk_insert.py GRAPHNAME [OPTIONS]
 |  -b   | --max-buffer-size INT      |               (Debug argument) Max batch size (MBs) of each Redis query (default 4096)               |
 |  -c   | --max-token-size INT       |              (Debug argument) Max size (MBs) of each token sent to Redis (default 500)               |
 |  -i   | --index Label:Property     |              After bulk import, create an Index on provided Label:Property pair (optional)           |
+|  -f   | --full-text-index Label:Property     |              After bulk import, create an full text index on provided Label:Property pair (optional)           |
 
 
 The only required arguments are the name to give the newly-created graph (which can appear anywhere) and at least one node CSV file.

redisgraph_bulk_loader/bulk_insert.py

@@ -63,7 +63,8 @@ def process_entities(entities):
 @click.option('--max-buffer-size', '-b', default=2048, help='max buffer size in megabytes (default 2048)')
 @click.option('--max-token-size', '-t', default=500, help='max size of each token in megabytes (default 500, max 512)')
 @click.option('--index', '-i', multiple=True, help='Label:Propery on which to create an index')
-def bulk_insert(graph, host, port, password, nodes, nodes_with_label, relations, relations_with_type, separator, enforce_schema, skip_invalid_nodes, skip_invalid_edges, quote, max_token_count, max_buffer_size, max_token_size, index):
+@click.option('--full-text-index', '-f', multiple=True, help='Label:Propery on which to create an full text search index')
+def bulk_insert(graph, host, port, password, nodes, nodes_with_label, relations, relations_with_type, separator, enforce_schema, skip_invalid_nodes, skip_invalid_edges, quote, max_token_count, max_buffer_size, max_token_size, index, full_text_index):
     if sys.version_info[0] < 3:
         raise Exception("Python 3 is required for the RedisGraph bulk loader.")
 
@@ -116,6 +117,7 @@ def bulk_insert(graph, host, port, password, nodes, nodes_with_label, relations,
     end_time = timer()
     query_buf.report_completion(end_time - start_time)
 
+    # Add in Graph Indices after graph creation
     for i in index:
         l, p = i.split(":")
         print("Creating Index on Label: %s, Property: %s" %(l, p))
@@ -127,5 +129,18 @@ def bulk_insert(graph, host, port, password, nodes, nodes_with_label, relations,
             print("Unable to create Index on Label: %s, Property %s" %(l, p))
             print(e)
 
+    # Add in Full Text Search Indices after graph creation
+    for i in full_text_index:
+        l, p = i.split(":")
+        print("Creating Full Text Search Index on Label: %s, Property: %s" %(l, p))
+        try:
+            index_create = client.execute_command("GRAPH.QUERY", graph, "CALL db.idx.fulltext.createNodeIndex('%s', '%s')" %(l, p))
+            print(index_create[-1][0].decode("utf-8"))
+        except redis.exceptions.ResponseError as e:
+            print("Unable to create Full Text Search Index on Label: %s, Property %s" %(l, p))
+            print(e)
+        except:
+            print("Unknown Error: Unable to create Full Text Search Index on Label: %s, Property %s" %(l, p))
+
 if __name__ == '__main__':
     bulk_insert()

test/test_bulk_loader.py

@@ -41,6 +41,8 @@ def tearDownClass(cls):
         """Delete temporary files"""
         os.remove('/tmp/nodes.tmp')
         os.remove('/tmp/relations.tmp')
+        os.remove('/tmp/nodes_index.tmp')
+        os.remove('/tmp/nodes_full_text_index.tmp')
         cls.redis_con.flushall()
 
     def validate_exception(self, res, expected_msg):
@@ -685,6 +687,35 @@ def test17_ensure_index_is_created(self):
         res = r.execute_command("GRAPH.EXPLAIN", graphname, 'MATCH (p:Person) WHERE p.age > 16 RETURN p')
         self.assertIn('        Index Scan | (p:Person)', res)
 
+    def test18_ensure_full_text_index_is_created(self):
+        graphname = "index_full_text_test"
+        with open('/tmp/nodes_full_text_index.tmp', mode='w') as csv_file:
+            out = csv.writer(csv_file, delimiter='|')
+            out.writerow(['name:STRING'])
+            out.writerow(['Emperor Tamarin'])
+            out.writerow(['Golden Lion Tamarin'])
+            out.writerow(['Cotton-top Tamarin'])
+            out.writerow(['Olive Baboon'])
+        csv_file.close()
+
+        runner = CliRunner()
+        res = runner.invoke(bulk_insert, ['--nodes-with-label', 'Monkeys', '/tmp/nodes_full_text_index.tmp',
+                                          '--full-text-index', 'Monkeys:name',
+                                          '--enforce-schema',
+                                          graphname], catch_exceptions=False)
+
+        self.assertEqual(res.exit_code, 0)
+        self.assertIn('4 nodes created', res.output)
+        self.assertIn('Indices created: 1', res.output)
+
+        graph = Graph(graphname, self.redis_con)
+        query_result = graph.query("CALL db.idx.fulltext.queryNodes('Monkeys', 'tamarin') YIELD node RETURN node.name")
+        expected_result = [ ['Emperor Tamarin'],['Golden Lion Tamarin'], ['Cotton-top Tamarin'] ]
+
+        # We should find only the tamarins
+        self.assertEqual(query_result.result_set, expected_result)
+
+
 
 if __name__ == '__main__':
     unittest.main()