From e8e3812d287eeae5b4275883b0927928932280bd Mon Sep 17 00:00:00 2001 From: liyaxing Date: Fri, 8 Sep 2023 15:52:05 +0800 Subject: [PATCH] Add chunk store and chunk list helper --- src/knowledge/data/__init__.py | 4 +- src/knowledge/data/chunk.py | 2 +- src/knowledge/data/chunk_store.py | 26 +++++++++ src/knowledge/data/gen.py | 94 +++++++++++++++++++++++++++++++ src/knowledge/data/tracker.py | 40 ++++++------- test/test_chunk.py | 26 ++++++++- 6 files changed, 164 insertions(+), 28 deletions(-) create mode 100644 src/knowledge/data/chunk_store.py create mode 100644 src/knowledge/data/gen.py diff --git a/src/knowledge/data/__init__.py b/src/knowledge/data/__init__.py index 81c672a..f8e69aa 100644 --- a/src/knowledge/data/__init__.py +++ b/src/knowledge/data/__init__.py @@ -1,2 +1,4 @@ from .chunk import ChunkID, PositionType, PositionFileRange -from .tracker import ChunkTracker \ No newline at end of file +from .tracker import ChunkTracker +from .chunk_store import ChunkStore +from .gen import ChunkListGenerator \ No newline at end of file diff --git a/src/knowledge/data/chunk.py b/src/knowledge/data/chunk.py index ff355bd..5267dc3 100644 --- a/src/knowledge/data/chunk.py +++ b/src/knowledge/data/chunk.py @@ -8,7 +8,7 @@ class PositionType(Enum): Device = 2 File = 3 FileRange = 4 - ChunkManager = 5 + ChunkStore = 5 class PositionFileRange: diff --git a/src/knowledge/data/chunk_store.py b/src/knowledge/data/chunk_store.py new file mode 100644 index 0000000..d06eb38 --- /dev/null +++ b/src/knowledge/data/chunk_store.py @@ -0,0 +1,26 @@ +import os +import logging +from ..object import FileBlobStorage +from .chunk import ChunkID + + +class ChunkStore: + def __init__(self, root_dir: str): + logging.info(f"will init chunk store, root_dir={root_dir}") + + if not os.path.exists(root_dir): + os.makedirs(root_dir) + + self.root = root_dir + self.blob = FileBlobStorage(root_dir) + + def put_chunk(self, chunk_id: ChunkID, contents: bytes): + self.blob.put(chunk_id, contents) + + def get_chunk(self, chunk_id: ChunkID) -> bytes: + return self.blob.get(chunk_id) + + def delete_chunk(self, chunk_id: ChunkID): + self.blob.delete(chunk_id) + + \ No newline at end of file diff --git a/src/knowledge/data/gen.py b/src/knowledge/data/gen.py new file mode 100644 index 0000000..a6d9425 --- /dev/null +++ b/src/knowledge/data/gen.py @@ -0,0 +1,94 @@ +import os +import hashlib +import re +from typing import Tuple, List +from .chunk_store import ChunkStore +from .chunk import ChunkID, PositionFileRange, PositionType +from ..object import HashValue +from .tracker import ChunkTracker + + +class ChunkListGenerator: + def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker): + self.chunk_store = chunk_store + self.chunk_tracker = chunk_tracker + + def create_chunk_list_from_file( + self, file_path: str, chunk_size: int, restore: bool + ) -> Tuple[List[ChunkID], HashValue]: + assert ( + chunk_size % (1024 * 1024) == 0 + ), "chunk size should be an integral multiple of 1MB" + chunk_list = [] + hash_obj = hashlib.sha256() + + with open(file_path, "rb") as file: + while True: + chunk = file.read(chunk_size) + if not chunk: + break + chunk_id = ChunkID.hash_data(chunk) + chunk_list.append(chunk_id) + + hash_obj.update(chunk) + + if restore: + self.chunk_tracker.add_position( + chunk_id, file_path, PositionType.ChunkStore + ) + self.chunk_store.put_chunk(chunk_id, chunk) + else: + file_range = PositionFileRange( + file_path, file.tell() - chunk_size, chunk_size + ) + self.chunk_tracker.add_position( + chunk_id, file_range, PositionType.FileRange + ) + + file_hash = HashValue(hash_obj.digest()) + print(f"calc file hash: {file_path}, {file_hash}") + return chunk_list, file_hash + + def create_chunk_list_from_text( + self, text: str, chunk_max_words: int, separator_chars: str = ".," + ) -> Tuple[List[ChunkID], HashValue]: + text_list = self._split_text_list(text, chunk_max_words, separator_chars) + chunk_list = [] + hash_obj = hashlib.sha256() + + for text in text_list: + chunk_bytes = text.encode("utf-8") + hash_obj.update(chunk_bytes) + + chunk_id = ChunkID.hash_data(chunk_bytes) + chunk_list.append(chunk_id) + self.chunk_tracker.add_position( + chunk_id, "", PositionType.ChunkStore + ) + self.chunk_store.put_chunk(chunk_id, chunk_bytes) + + hash = HashValue(hash_obj.digest()) + return chunk_list, hash + + @staticmethod + def _split_text_list( + text: str, chunk_max_words: int, separator_chars: str = ".," + ) -> List[str]: + sentences = re.split(f"[{separator_chars}]", text) + chunk_list = [] + chunk = [] + word_count = 0 + for sentence in sentences: + words = sentence.split() + for word in words: + if word_count < chunk_max_words: + chunk.append(word) + word_count += 1 + else: + chunk_list.append(" ".join(chunk)) + chunk = [word] + word_count = 1 + if chunk: + chunk_list.append(" ".join(chunk)) + return chunk_list + diff --git a/src/knowledge/data/tracker.py b/src/knowledge/data/tracker.py index efe4ede..2905d11 100644 --- a/src/knowledge/data/tracker.py +++ b/src/knowledge/data/tracker.py @@ -2,31 +2,17 @@ import sqlite3 import time import logging import os -from .chunk import ChunkID, PositionType, PositionFileRange +from .chunk import ChunkID, PositionType, PositionFileRange class ChunkTracker: - _instance = None - - def __new__(cls): - if cls._instance is None: - cls._instance = super().__new__(cls) - directory = os.path.join( - os.path.dirname(__file__), "../../../rootfs/data/chunk/" - ) - - if not os.path.exists(directory): - os.mkdir(directory) - file = os.path.join(directory, "chunk_tracker.db") - logging.info(f"will init chunk tracker, db={file}") - cls._instance.__singleton_init__(file) - - return cls._instance + def __init__(self, root_dir: str): + if not os.path.exists(root_dir): + os.makedirs(root_dir) + file = os.path.join(root_dir, "chunk_tracker.db") + logging.info(f"will init chunk tracker, db={file}") - def __singleton_init__(self, db_path: str): - logging.info(f"will init chunk tracker, db={db_path}") - - self.conn = sqlite3.connect(db_path) + self.conn = sqlite3.connect(file) self.cursor = self.conn.cursor() self.cursor.execute( """ @@ -43,7 +29,9 @@ class ChunkTracker: ) self.conn.commit() - def add_position(self, chunk_id: ChunkID, position: str, position_type: PositionType): + def add_position( + self, chunk_id: ChunkID, position: str, position_type: PositionType + ): logging.debug(f"add chunk position: {chunk_id}, {position}, {position_type}") insert_time = update_time = int(time.time()) @@ -52,7 +40,13 @@ class ChunkTracker: INSERT OR REPLACE INTO chunks (id, pos, pos_type, insert_time, update_time) VALUES (?, ?, ?, ?, ?) """, - (str(chunk_id), position, position_type.value, insert_time, update_time,), + ( + str(chunk_id), + position, + position_type.value, + insert_time, + update_time, + ), ) self.conn.commit() diff --git a/test/test_chunk.py b/test/test_chunk.py index 40a7aeb..b2f2e7e 100644 --- a/test/test_chunk.py +++ b/test/test_chunk.py @@ -7,16 +7,27 @@ print(dir_path) sys.path.append("{}/../src/".format(dir_path)) print(sys.path) -from knowledge import ChunkTracker, ChunkID, HashValue, PositionType + +import logging +import sys + +root = logging.getLogger() +root.setLevel(logging.DEBUG) +handler = logging.StreamHandler(sys.stdout) +handler.setLevel(logging.DEBUG) +formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') +handler.setFormatter(formatter) +root.addHandler(handler) +from knowledge import ChunkTracker, ChunkID, HashValue, PositionType, KnowledgeStore, ChunkListGenerator import asyncio import unittest class TestChunk(unittest.TestCase): def test_chunk_tracker(self): - tracker = ChunkTracker() + tracker = KnowledgeStore().get_chunk_tracker() hash = HashValue.hash_data("1234567890".encode("utf-8")); cid = ChunkID.new_chunk_id(hash) @@ -29,8 +40,17 @@ class TestChunk(unittest.TestCase): tracker.remove_position(cid) ret = tracker.get_position(cid) self.assertEqual(ret, None) - + + def test_chunk(self): + gen = ChunkListGenerator(KnowledgeStore().get_chunk_store(), KnowledgeStore().get_chunk_tracker()) + gen.create_chunk_list_from_file("H:/test", 1024 * 1024, True) + + # Read the file + text_file = "H:/test.txt" + with open(text_file, 'r', encoding = "utf-8") as file: + text = file.read() + gen.create_chunk_list_from_text(text, 1024) if __name__ == "__main__": unittest.main()