| /** |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| */ |
| package org.apache.pinot.integration.tests; |
| |
| import com.fasterxml.jackson.databind.JsonNode; |
| import com.google.common.collect.ImmutableList; |
| import com.google.common.collect.ImmutableMap; |
| import java.io.File; |
| import java.io.IOException; |
| import java.nio.ByteBuffer; |
| import java.util.ArrayList; |
| import java.util.HashMap; |
| import java.util.List; |
| import java.util.Map; |
| import org.apache.avro.Schema.Field; |
| import org.apache.avro.Schema.Type; |
| import org.apache.avro.file.DataFileWriter; |
| import org.apache.avro.generic.GenericData; |
| import org.apache.avro.generic.GenericDatumWriter; |
| import org.apache.commons.io.FileUtils; |
| import org.apache.commons.lang3.tuple.ImmutablePair; |
| import org.apache.commons.lang3.tuple.Pair; |
| import org.apache.datasketches.theta.UpdateSketch; |
| import org.apache.datasketches.theta.UpdateSketchBuilder; |
| import org.apache.pinot.spi.config.table.TableConfig; |
| import org.apache.pinot.spi.config.table.TableType; |
| import org.apache.pinot.spi.data.FieldSpec; |
| import org.apache.pinot.spi.data.Schema; |
| import org.apache.pinot.spi.utils.builder.TableConfigBuilder; |
| import org.apache.pinot.util.TestUtils; |
| import org.testng.annotations.AfterClass; |
| import org.testng.annotations.BeforeClass; |
| import org.testng.annotations.Test; |
| |
| import static org.testng.Assert.assertEquals; |
| |
| |
| public class ThetaSketchIntegrationTest extends BaseClusterIntegrationTest { |
| private static final String DIM_NAME = "dimName"; |
| private static final String DIM_VALUE = "dimValue"; |
| private static final String SHARD_ID = "shardId"; |
| private static final String THETA_SKETCH = "thetaSketchCol"; |
| |
| @BeforeClass |
| public void setup() |
| throws Exception { |
| TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir); |
| |
| // Start the Pinot cluster |
| startZk(); |
| startController(); |
| startBroker(); |
| startServer(); |
| |
| // create & upload schema AND table config |
| Schema schema = new Schema.SchemaBuilder().setSchemaName(DEFAULT_SCHEMA_NAME) |
| .addSingleValueDimension(DIM_NAME, FieldSpec.DataType.STRING) |
| .addSingleValueDimension(DIM_VALUE, FieldSpec.DataType.STRING) |
| .addSingleValueDimension(SHARD_ID, FieldSpec.DataType.INT) |
| .addSingleValueDimension(THETA_SKETCH, FieldSpec.DataType.BYTES).build(); |
| addSchema(schema); |
| TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(DEFAULT_TABLE_NAME).build(); |
| addTableConfig(tableConfig); |
| |
| // create & upload segments |
| File avroFile = createAvroFile(); |
| ClusterIntegrationTestUtils.buildSegmentFromAvro(avroFile, tableConfig, schema, 0, _segmentDir, _tarDir); |
| uploadSegments(DEFAULT_TABLE_NAME, _tarDir); |
| |
| waitForAllDocsLoaded(60_000); |
| } |
| |
| @Override |
| protected long getCountStarResult() { |
| /* |
| Uploaded table content: |
| |
| row# dimName dimValue shardId thetaSketchCol |
| ---- ======= ======== ======= ============== |
| 1 Course Math 1 ... |
| 2 Course History 1 ... |
| 3 Course Biology 1 ... |
| 4 Gender Female 1 ... |
| 5 Gender Male 1 ... |
| 6 Course Math 2 ... |
| 7 Course History 2 ... |
| 8 Course Biology 2 ... |
| 9 Gender Female 2 ... |
| 10 Gender Male 2 ... |
| */ |
| return 10; |
| } |
| |
| @Test |
| public void testThetaSketchQuery() |
| throws Exception { |
| /* |
| Original data: |
| |
| Gender Course Shard#1 Shard#2 |
| -------- ------- ------- ------- |
| Female Math 50 110 |
| Female History 60 120 |
| Female Biology 70 130 |
| Male Math 80 140 |
| Male History 90 150 |
| Male Biology 100 160 |
| */ |
| |
| // gender = female |
| { |
| String query = "select distinctCountThetaSketch(thetaSketchCol) from " + DEFAULT_TABLE_NAME |
| + " where dimName = 'gender' and dimValue = 'Female'"; |
| int expected = 50 + 60 + 70 + 110 + 120 + 130; |
| runAndAssert(query, expected); |
| |
| query = "select distinctCountThetaSketch(thetaSketchCol, '', 'dimName = ''gender'' and dimValue = ''Female''', " |
| + "'$1') from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| |
| query = "select distinctCountThetaSketch(thetaSketchCol, '', " |
| + "'dimName = ''gender''', 'dimValue = ''Female''', 'SET_INTERSECT($1, $2)') from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| } |
| |
| // gender = male |
| { |
| String query = "select distinctCountThetaSketch(thetaSketchCol) from " + DEFAULT_TABLE_NAME |
| + " where dimName = 'gender' and dimValue = 'Male'"; |
| int expected = 80 + 90 + 100 + 140 + 150 + 160; |
| runAndAssert(query, expected); |
| |
| query = |
| "select distinctCountThetaSketch(thetaSketchCol, '', 'dimName = ''gender'' and dimValue = ''Male''', '$1') " |
| + "from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| |
| query = "select distinctCountThetaSketch(thetaSketchCol, '', " |
| + "'dimName = ''gender''', 'dimValue = ''Male''', 'SET_INTERSECT($1, $2)') from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| } |
| |
| // course = math |
| { |
| String query = "select distinctCountThetaSketch(thetaSketchCol) from " + DEFAULT_TABLE_NAME |
| + " where dimName = 'course' AND dimValue = 'Math'"; |
| int expected = 50 + 80 + 110 + 140; |
| runAndAssert(query, expected); |
| |
| query = |
| "select distinctCountThetaSketch(thetaSketchCol, '', 'dimName = ''course'' and dimValue = ''Math''', '$1') " |
| + "from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| |
| query = "select distinctCountThetaSketch(thetaSketchCol, '', " |
| + "'dimName = ''course''', 'dimValue = ''Math''', 'SET_INTERSECT($1, $2)') from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| } |
| |
| // gender = female INTERSECT course = math |
| { |
| String query = "select distinctCountThetaSketch(thetaSketchCol, '', " |
| + "'dimName = ''gender'' and dimValue = ''Female''', 'dimName = ''course'' and dimValue = ''Math''', " |
| + "'SET_INTERSECT($1, $2)') from " + DEFAULT_TABLE_NAME; |
| int expected = 50 + 110; |
| runAndAssert(query, expected); |
| |
| query = "select distinctCountThetaSketch(thetaSketchCol, '', " |
| + "'dimName = ''gender''', 'dimValue = ''Female''', 'dimName = ''course''', 'dimValue = ''Math''', " |
| + "'SET_INTERSECT($1, $2, $3, $4)') from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| |
| query = "select distinctCountThetaSketch(thetaSketchCol, '', " |
| + "'dimName = ''gender''', 'dimValue = ''Female''', 'dimName = ''course''', 'dimValue = ''Math''', " |
| + "'SET_INTERSECT(SET_INTERSECT($1, $2), SET_INTERSECT($3, $4))') from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| } |
| |
| // gender = male UNION course = biology |
| { |
| String query = "select distinctCountThetaSketch(thetaSketchCol, '', " |
| + "'dimName = ''gender'' and dimValue = ''Male''', 'dimName = ''course'' and dimValue = ''Biology''', " |
| + "'SET_UNION($1, $2)') from " + DEFAULT_TABLE_NAME; |
| int expected = 70 + 80 + 90 + 100 + 130 + 140 + 150 + 160; |
| runAndAssert(query, expected); |
| |
| query = "select distinctCountThetaSketch(thetaSketchCol, '', " |
| + "'dimName = ''gender''', 'dimValue = ''Male''', 'dimName = ''course''', 'dimValue = ''Biology''', " |
| + "'SET_UNION(SET_INTERSECT($1, $2), SET_INTERSECT($3, $4))') from " + DEFAULT_TABLE_NAME; |
| runAndAssert(query, expected); |
| } |
| |
| // group by gender |
| { |
| String query = "select dimValue, distinctCountThetaSketch(thetaSketchCol) from " + DEFAULT_TABLE_NAME |
| + " where dimName = 'gender' group by dimValue"; |
| ImmutableMap<String, Integer> expected = |
| ImmutableMap.of("Female", 50 + 60 + 70 + 110 + 120 + 130, "Male", 80 + 90 + 100 + 140 + 150 + 160); |
| runAndAssert(query, expected); |
| } |
| } |
| |
| private void runAndAssert(String query, int expected) |
| throws Exception { |
| JsonNode jsonNode = postQuery(query, _brokerBaseApiUrl); |
| int actual = Integer.parseInt(jsonNode.get("resultTable").get("rows").get(0).get(0).asText()); |
| assertEquals(actual, expected); |
| } |
| |
| private void runAndAssert(String query, Map<String, Integer> expectedGroupToValueMap) |
| throws Exception { |
| Map<String, Integer> actualGroupToValueMap = new HashMap<>(); |
| JsonNode jsonNode = postQuery(query, _brokerBaseApiUrl); |
| jsonNode.get("resultTable").get("rows").forEach(node -> { |
| String group = node.get(0).textValue(); |
| int value = node.get(1).intValue(); |
| actualGroupToValueMap.put(group, value); |
| }); |
| assertEquals(actualGroupToValueMap, expectedGroupToValueMap); |
| } |
| |
| private File createAvroFile() |
| throws IOException { |
| |
| // create avro schema |
| org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false); |
| avroSchema.setFields(ImmutableList.of(new Field(DIM_NAME, org.apache.avro.Schema.create(Type.STRING), null, null), |
| new Field(DIM_VALUE, org.apache.avro.Schema.create(Type.STRING), null, null), |
| new Field(SHARD_ID, org.apache.avro.Schema.create(Type.INT), null, null), |
| new Field(THETA_SKETCH, org.apache.avro.Schema.create(Type.BYTES), null, null))); |
| |
| // create avro file |
| File avroFile = new File(_tempDir, "data.avro"); |
| try (DataFileWriter<GenericData.Record> fileWriter = new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) { |
| fileWriter.create(avroSchema, avroFile); |
| |
| int studentId = 0; |
| int cardinality = 50; |
| for (int shardId = 0; shardId < 2; shardId++) { |
| |
| // populate student-course data (studentId, gender, course) for this shard id |
| String[] allGenders = {"Female", "Male"}; |
| String[] allCountries = {"Math", "History", "Biology"}; |
| Map<Pair<String, String>, List<Integer>> genderCourseToStudentIds = new HashMap<>(); |
| for (String gender : allGenders) { |
| for (String course : allCountries) { |
| List<Integer> studentIds = |
| genderCourseToStudentIds.computeIfAbsent(ImmutablePair.of(gender, course), key -> new ArrayList<>()); |
| for (int i = 0; i < cardinality; i++) { |
| studentIds.add(studentId++); |
| } |
| cardinality += 10; |
| } |
| } |
| |
| // [gender dimension] calculate theta sketches & add them to avro file |
| for (String gender : allGenders) { |
| |
| // calculate theta sketch |
| UpdateSketch sketch = new UpdateSketchBuilder().build(); |
| genderCourseToStudentIds.forEach((genderCourse, studentIds) -> { |
| if (gender.equals(genderCourse.getLeft())) { |
| studentIds.forEach(sketch::update); |
| } |
| }); |
| |
| // create avro record |
| GenericData.Record record = new GenericData.Record(avroSchema); |
| record.put(DIM_NAME, "gender"); |
| record.put(DIM_VALUE, gender); |
| record.put(SHARD_ID, shardId); |
| record.put(THETA_SKETCH, ByteBuffer.wrap(sketch.compact().toByteArray())); |
| |
| // add avro record to file |
| fileWriter.append(record); |
| } |
| |
| // [course dimension] calculate theta sketches & add them to avro file |
| for (String course : allCountries) { |
| |
| // calculate theta sketch |
| UpdateSketch sketch = new UpdateSketchBuilder().build(); |
| genderCourseToStudentIds.forEach((genderCourse, studentIds) -> { |
| if (course.equals(genderCourse.getRight())) { |
| studentIds.forEach(sketch::update); |
| } |
| }); |
| |
| // create avro record |
| GenericData.Record record = new GenericData.Record(avroSchema); |
| record.put(DIM_NAME, "course"); |
| record.put(DIM_VALUE, course); |
| record.put(SHARD_ID, shardId); |
| record.put(THETA_SKETCH, ByteBuffer.wrap(sketch.compact().toByteArray())); |
| |
| // add avro record to file |
| fileWriter.append(record); |
| } |
| } |
| } |
| |
| return avroFile; |
| } |
| |
| @AfterClass |
| public void tearDown() |
| throws IOException { |
| dropOfflineTable(DEFAULT_TABLE_NAME); |
| |
| stopServer(); |
| stopBroker(); |
| stopController(); |
| stopZk(); |
| |
| FileUtils.deleteDirectory(_tempDir); |
| } |
| } |
