amarnathkarthik commented on a change in pull request #6382: URL: https://github.com/apache/incubator-pinot/pull/6382#discussion_r548765615
########## File path: pinot-integration-tests/src/test/java/org/apache/pinot/compat/tests/SegmentOp.java ########## @@ -82,14 +107,175 @@ public void setTableConfigFileName(String tableConfigFileName) { _tableConfigFileName = tableConfigFileName; } + public void setSchemaFileName(String schemaFileName) { + _schemaFileName = schemaFileName; + } + + public String getSchemaFileName() { + return _schemaFileName; + } + + public void setRecordReaderConfigFileName(String recordReaderConfigFileName) { + _recordReaderConfigFileName = recordReaderConfigFileName; + } + + public String getRecordReaderConfigFileName() { + return _recordReaderConfigFileName; + } + + public void setSegmentName(String segmentName) { + _segmentName = segmentName; + } + + public String getSegmentName() { + return _segmentName; + } + @Override boolean runOp() { - switch(_op) { + switch (_op) { case UPLOAD: - System.out.println("Generating segment " + _segmentName + " from " + _inputDataFileName + " and uploading to " + - _tableConfigFileName); + return createAndUploadSegments(); case DELETE: + return deleteSegment(); } return true; } + + /** + * Create Segment file, compress to TarGz, and upload the files to controller. + * @return true if all successful, false in case of failure. + */ + private boolean createAndUploadSegments() { + File localTempDir = new File(FileUtils.getTempDirectory(), "pinot-compat-test-" + UUID.randomUUID()); + File localOutputTempDir = new File(localTempDir, "output"); + try { + FileUtils.forceMkdir(localOutputTempDir); + File segmentTarFile = generateSegment(localOutputTempDir); + uploadSegment(segmentTarFile); + + Pair<Long, Long> onlineSegmentCount = getOnlineSegmentCount(getTableExternalView()); + if (onlineSegmentCount.getFirst() <= 0 && onlineSegmentCount.getSecond() <= 0) { + LOGGER.error("Uploaded segment {} not found or not in {} state.", _segmentName, STATE_ONLINE); + return false; + } + LOGGER.info("Successfully verified segment {} and its current status is {}.", _segmentName, STATE_ONLINE); + + return true; + } catch (Exception e) { + LOGGER.error("Failed to create and upload segment for input data file {}.", _inputDataFileName, e); + return false; + } finally { + FileUtils.deleteQuietly(localTempDir); + } + } + + /** + * Generate the Segment(s) and then compress to TarGz file. Supports generation of segment files for one input data + * file. + * @param outputDir to generate the Segment file(s). + * @return File object of the TarGz compressed segment file. + * @throws Exception while generating segment files and/or compressing to TarGz. + */ + private File generateSegment(File outputDir) + throws Exception { + TableConfig tableConfig = JsonUtils.fileToObject(new File(_tableConfigFileName), TableConfig.class); + _tableName = tableConfig.getTableName(); + + Schema schema = JsonUtils.fileToObject(new File(_schemaFileName), Schema.class); + RecordReaderConfig recordReaderConfig = + RecordReaderFactory.getRecordReaderConfig(DEFAULT_FILE_FORMAT, _recordReaderConfigFileName); + + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, schema); + segmentGeneratorConfig.setInputFilePath(_inputDataFileName); + segmentGeneratorConfig.setFormat(DEFAULT_FILE_FORMAT); + segmentGeneratorConfig.setOutDir(outputDir.getAbsolutePath()); + segmentGeneratorConfig.setReaderConfig(recordReaderConfig); + segmentGeneratorConfig.setTableName(_tableName); + segmentGeneratorConfig.setSegmentName(_segmentName); + + SegmentIndexCreationDriver driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig); + driver.build(); + String segmentName = driver.getSegmentName(); Review comment: Both should be the same since the segment name is passed when generating the segment but above one should give the actual segment name used when generating. In the case of segment naming logic change to use either sequence id or min/max value, this implementation would always return the correct segment name. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org