Пример #1
0
    def commit(self, overwrite=False):
        partitions = list(set(p for p in self._blocks.values() if p is not None))
        if not partitions:
            partitions = ['']
        channel = CupidTaskServiceRpcChannel(self._session)
        stub = task_service_pb.CupidTaskService_Stub(channel)

        part_specs = [pt.replace("'", '') for pt in partitions]
        req = task_service_pb.CommitTableRequest(
            outputTableHandle=self._handle,
            projectName=self._project_name,
            tableName=self._table_name,
            isOverWrite=overwrite,
            lookupName=self._session.lookup_name,
            partSpecs=part_specs,
        )

        controller = CupidRpcController()
        resp = None
        for _ in range(options.retry_times):
            resp = stub.CommitTable(controller, req, None)
            if controller.Failed():
                time.sleep(0.1)
                controller = CupidRpcController()
            else:
                break
        if controller.Failed():
            raise CupidError(controller.ErrorText())

        logger.info(
            "[CupidTask] commitTable call, CurrentInstanceId: %s, "
            "request: %s, response: %s", self._session.lookup_name, req, resp,
        )
Пример #2
0
def create_upload_session(session, table):
    controller = CupidRpcController()
    channel = CupidTaskServiceRpcChannel(session)
    stub = task_service_pb.CupidTaskService_Stub(channel)

    req = task_service_pb.WriteTableRequest(lookupName=session.lookup_name, tableName=table.name,
                                            projectName=table.project.name)
    resp = stub.WriteTable(controller, req, None)
    if controller.Failed():
        raise CupidError(controller.ErrorText())
    logger.info(
        "[CupidTask] writeTable call, CurrentInstanceId: %s, "
        "request: %s, response: %s", session.lookup_name, req, resp,
    )
    return CupidTableUploadSession(
        session=session, table_name=table.name, project_name=table.project.name, handle=resp.outputTableHandle)
Пример #3
0
def query_table_meta(session, table):
    controller = CupidRpcController()
    channel = CupidTaskServiceRpcChannel(session)
    stub = task_service_pb.CupidTaskService_Stub(channel)

    table_info = task_service_pb.TableInfo(projectName=table.project.name, tableName=table.name)
    req = task_service_pb.GetTableMetaRequest(lookupName=session.lookup_name, tableInfo=table_info,
                                              needContent=True, uploadFile='')
    resp = stub.GetTableMeta(controller, req, None)
    if controller.Failed():
        raise CupidError(controller.ErrorText())
    logger.info(
        "[CupidTask] getTableMeta call, CurrentInstanceId: %s, "
        "request: %s, response: %s", session.lookup_name, req, resp,
    )
    return json.loads(resp.getTableMetaContent)
Пример #4
0
def create_download_session(session,
                            table_or_parts,
                            split_size=None,
                            split_count=None,
                            columns=None,
                            with_split_meta=False):
    channel = CupidTaskServiceRpcChannel(session)
    stub = task_service_pb.CupidTaskService_Stub(channel)

    if not isinstance(table_or_parts, (list, tuple, set, GeneratorType)):
        table_or_parts = [table_or_parts]

    if split_size is None and split_count is None:
        split_count = 1
    split_count = split_count or 0
    split_size = (split_size or 1024**2) // 1024**2

    table_pbs = []
    for t in table_or_parts:
        if isinstance(t, Table):
            if not columns:
                columns = t.schema.names
            table_kw = dict(
                projectName=t.project.name,
                tableName=t.name,
                columns=','.join(columns),
            )
        elif isinstance(t, TablePartition):
            if not columns:
                columns = t.table.schema.names
            table_kw = dict(
                projectName=t.table.project.name,
                tableName=t.table.name,
                columns=','.join(columns),
                partSpec=str(t.partition_spec).replace("'", '').strip(),
            )
        else:
            raise NotImplementedError
        table_pbs.append(task_service_pb.TableInputInfo(**table_kw))

    request = task_service_pb.SplitTablesRequest(
        lookupName=session.lookup_name,
        splitSize=split_size,
        splitCount=split_count,
        tableInputInfos=table_pbs,
        allowNoColumns=True,
        requireSplitMeta=with_split_meta,
    )

    controller = CupidRpcController()
    resp = stub.SplitTables(controller, request, None)
    if controller.Failed():
        raise CupidError(controller.ErrorText())
    logger.info("[CupidTask] splitTables call, CurrentInstanceId: %s, "
                "request: %s, response: %s" % (
                    session.lookup_name,
                    str(request),
                    str(resp),
                ))
    handle = resp.inputTableHandle

    channel = SandboxRpcChannel()
    stub = subprocess_pb.CupidSubProcessService_Stub(channel)

    if not with_split_meta:
        split_meta = itertools.repeat(None)
    else:
        req = subprocess_pb.GetSplitsMetaRequest(inputTableHandle=handle, )
        controller = CupidRpcController()
        resp = stub.GetSplitsMeta(controller, req, None)
        logger.info("[CupidTask] getSplitsMeta call, CurrentInstanceId: %s, "
                    "request: %s, response: %s" % (
                        session.lookup_name,
                        str(request),
                        str(resp),
                    ))
        if controller.Failed():
            split_meta = itertools.repeat(None)
            logger.warning('Failed to get results of getSplitsMeta, '
                           'may running on an old service')
        else:
            split_meta = resp.inputSplitsMeta

    req = subprocess_pb.GetSplitsRequest(inputTableHandle=handle)
    controller = CupidRpcController()
    resp = stub.GetSplits(controller, req, None)
    if controller.Failed():
        raise CupidError(controller.ErrorText())

    input_splits = []
    for info, meta in zip(resp.inputSplits, split_meta):
        input_splits.append(
            TableSplit(split_proto=info,
                       meta_proto=meta,
                       handle=handle,
                       columns=columns))
    logger.info("[SubProcess] getSplits call, CurrentInstanceId: %s, "
                "request: %s, response: %s" % (
                    session.lookup_name,
                    str(req),
                    str(resp),
                ))
    return CupidTableDownloadSession(session=session,
                                     handle=handle,
                                     splits=input_splits)