• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""
16This module is to  write data into mindrecord.
17"""
18import mindspore._c_mindrecord as ms
19from mindspore import log as logger
20from .common.exceptions import MRMIndexGeneratorError, MRMGenerateIndexError
21
22__all__ = ['ShardIndexGenerator']
23
24
25class ShardIndexGenerator:
26    """
27    Wrapper class which is represent ShardIndexGenerator class in c++ module.
28
29    The class would generate db files for accelerating reading.
30
31    Args:
32        path (str): Absolute path of MindRecord File.
33        append (bool): If True, open existed MindRecord Files for appending, or create new MindRecord Files.
34
35    Raises:
36        MRMIndexGeneratorError: If failed to create index generator.
37    """
38    def __init__(self, path, append=False):
39        self._generator = ms.ShardIndexGenerator(path, append)
40        if not self._generator:
41            logger.critical("Failed to create index generator.")
42            raise MRMIndexGeneratorError
43
44    def build(self):
45        """
46        Build index generator.
47
48        Returns:
49            MSRStatus, SUCCESS or FAILED.
50
51        Raises:
52            MRMGenerateIndexError: If failed to build index generator.
53        """
54        ret = self._generator.build()
55        if ret != ms.MSRStatus.SUCCESS:
56            logger.critical("Failed to build index generator.")
57            raise MRMGenerateIndexError
58        return ret
59
60    def write_to_db(self):
61        """
62        Create index field in table for reading data.
63
64        Returns:
65            MSRStatus, SUCCESS or FAILED.
66
67        Raises:
68            MRMGenerateIndexError: If failed to write to database.
69        """
70        ret = self._generator.write_to_db()
71        if ret != ms.MSRStatus.SUCCESS:
72            logger.critical("Failed to write to database.")
73            raise MRMGenerateIndexError
74        return ret
75