@Test
  public void testArrayDoubleMap() throws Exception {
    ArrayType arrayType = new ArrayType(VARCHAR);
    MapType mapType = new MapType(arrayType, DOUBLE);
    InternalAggregationFunction aggFunc =
        metadata
            .getFunctionRegistry()
            .getAggregateFunctionImplementation(
                new Signature(
                    NAME,
                    AGGREGATE,
                    mapType.getTypeSignature().toString(),
                    arrayType.getTypeSignature().toString(),
                    StandardTypes.DOUBLE));

    assertAggregation(
        aggFunc,
        1.0,
        ImmutableMap.of(
            ImmutableList.of("a", "b"), 1.0,
            ImmutableList.of("c", "d"), 2.0,
            ImmutableList.of("e", "f"), 3.0),
        createStringArraysBlock(
            ImmutableList.of(
                ImmutableList.of("a", "b"),
                ImmutableList.of("c", "d"),
                ImmutableList.of("e", "f"))),
        createDoublesBlock(1.0, 2.0, 3.0));
  }
  @Test
  public void testDoubleRowMap() throws Exception {
    RowType innerRowType =
        new RowType(ImmutableList.of(BIGINT, DOUBLE), Optional.of(ImmutableList.of("f1", "f2")));
    MapType mapType = new MapType(DOUBLE, innerRowType);
    InternalAggregationFunction aggFunc =
        metadata
            .getFunctionRegistry()
            .getAggregateFunctionImplementation(
                new Signature(
                    NAME,
                    AGGREGATE,
                    mapType.getTypeSignature().toString(),
                    StandardTypes.DOUBLE,
                    innerRowType.getTypeSignature().toString()));

    BlockBuilder builder = innerRowType.createBlockBuilder(new BlockBuilderStatus(), 3);
    innerRowType.writeObject(builder, testRowBigintBigint(1L, 1.0));
    innerRowType.writeObject(builder, testRowBigintBigint(2L, 2.0));
    innerRowType.writeObject(builder, testRowBigintBigint(3L, 3.0));

    assertAggregation(
        aggFunc,
        1.0,
        ImmutableMap.of(
            1.0,
            ImmutableList.of(1L, 1.0),
            2.0,
            ImmutableList.of(2L, 2.0),
            3.0,
            ImmutableList.of(3L, 3.0)),
        createDoublesBlock(1.0, 2.0, 3.0),
        builder.build());
  }
  @Test
  public void testSimpleMaps() throws Exception {
    MapType mapType = new MapType(DOUBLE, VARCHAR);
    InternalAggregationFunction aggFunc =
        metadata
            .getFunctionRegistry()
            .getAggregateFunctionImplementation(
                new Signature(
                    NAME,
                    AGGREGATE,
                    mapType.getTypeSignature().toString(),
                    StandardTypes.DOUBLE,
                    StandardTypes.VARCHAR));
    assertAggregation(
        aggFunc,
        1.0,
        ImmutableMap.of(1.0, "a", 2.0, "b", 3.0, "c"),
        createDoublesBlock(1.0, 2.0, 3.0),
        createStringsBlock("a", "b", "c"));

    mapType = new MapType(DOUBLE, BIGINT);
    aggFunc =
        metadata
            .getFunctionRegistry()
            .getAggregateFunctionImplementation(
                new Signature(
                    NAME,
                    AGGREGATE,
                    mapType.getTypeSignature().toString(),
                    StandardTypes.DOUBLE,
                    StandardTypes.BIGINT));
    assertAggregation(
        aggFunc,
        1.0,
        ImmutableMap.of(1.0, 3L, 2.0, 2L, 3.0, 1L),
        createDoublesBlock(1.0, 2.0, 3.0),
        createLongsBlock(3L, 2L, 1L));

    mapType = new MapType(DOUBLE, BOOLEAN);
    aggFunc =
        metadata
            .getFunctionRegistry()
            .getAggregateFunctionImplementation(
                new Signature(
                    NAME,
                    AGGREGATE,
                    mapType.getTypeSignature().toString(),
                    StandardTypes.DOUBLE,
                    StandardTypes.BOOLEAN));
    assertAggregation(
        aggFunc,
        1.0,
        ImmutableMap.of(1.0, true, 2.0, false, 3.0, false),
        createDoublesBlock(1.0, 2.0, 3.0),
        createBooleansBlock(true, false, false));
  }
  @Test
  public void testDoubleMapMap() throws Exception {
    MapType innerMapType = new MapType(VARCHAR, VARCHAR);
    MapType mapType = new MapType(DOUBLE, innerMapType);
    InternalAggregationFunction aggFunc =
        metadata
            .getFunctionRegistry()
            .getAggregateFunctionImplementation(
                new Signature(
                    NAME,
                    AGGREGATE,
                    mapType.getTypeSignature().toString(),
                    StandardTypes.DOUBLE,
                    innerMapType.getTypeSignature().toString()));

    BlockBuilder builder = innerMapType.createBlockBuilder(new BlockBuilderStatus(), 3);
    innerMapType.writeObject(builder, mapBlockOf(VARCHAR, VARCHAR, ImmutableMap.of("a", "b")));
    innerMapType.writeObject(builder, mapBlockOf(VARCHAR, VARCHAR, ImmutableMap.of("c", "d")));
    innerMapType.writeObject(builder, mapBlockOf(VARCHAR, VARCHAR, ImmutableMap.of("e", "f")));

    assertAggregation(
        aggFunc,
        1.0,
        ImmutableMap.of(
            1.0,
            ImmutableMap.of("a", "b"),
            2.0,
            ImmutableMap.of("c", "d"),
            3.0,
            ImmutableMap.of("e", "f")),
        createDoublesBlock(1.0, 2.0, 3.0),
        builder.build());
  }
  @Test
  public void testDuplicateKeysValues() throws Exception {
    MapType mapType = new MapType(DOUBLE, VARCHAR);
    InternalAggregationFunction aggFunc =
        metadata
            .getFunctionRegistry()
            .getAggregateFunctionImplementation(
                new Signature(
                    NAME,
                    AGGREGATE,
                    mapType.getTypeSignature().toString(),
                    StandardTypes.DOUBLE,
                    StandardTypes.VARCHAR));
    assertAggregation(
        aggFunc,
        1.0,
        ImmutableMap.of(1.0, "a"),
        createDoublesBlock(1.0, 1.0, 1.0),
        createStringsBlock("a", "b", "c"));

    mapType = new MapType(DOUBLE, BIGINT);
    aggFunc =
        metadata
            .getFunctionRegistry()
            .getAggregateFunctionImplementation(
                new Signature(
                    NAME,
                    AGGREGATE,
                    mapType.getTypeSignature().toString(),
                    StandardTypes.DOUBLE,
                    StandardTypes.BIGINT));
    assertAggregation(
        aggFunc,
        1.0,
        ImmutableMap.of(1.0, 99L, 2.0, 99L, 3.0, 99L),
        createDoublesBlock(1.0, 2.0, 3.0),
        createLongsBlock(99L, 99L, 99L));
  }
Esempio n. 6
0
  @Test
  public void testWriter() throws Exception {
    List<Long> columnIds = ImmutableList.of(1L, 2L, 4L, 6L, 7L, 8L, 9L, 10L);
    ArrayType arrayType = new ArrayType(BIGINT);
    ArrayType arrayOfArrayType = new ArrayType(arrayType);
    MapType mapType = new MapType(createVarcharType(10), BOOLEAN);
    List<Type> columnTypes =
        ImmutableList.of(
            BIGINT,
            createVarcharType(10),
            VARBINARY,
            DOUBLE,
            BOOLEAN,
            arrayType,
            mapType,
            arrayOfArrayType);
    File file = new File(directory, System.nanoTime() + ".orc");

    byte[] bytes1 = octets(0x00, 0xFE, 0xFF);
    byte[] bytes3 = octets(0x01, 0x02, 0x19, 0x80);

    RowPagesBuilder rowPagesBuilder =
        RowPagesBuilder.rowPagesBuilder(columnTypes)
            .row(
                123L,
                "hello",
                wrappedBuffer(bytes1),
                123.456,
                true,
                arrayBlockOf(BIGINT, 1, 2),
                mapBlockOf(createVarcharType(5), BOOLEAN, "k1", true),
                arrayBlockOf(arrayType, arrayBlockOf(BIGINT, 5)))
            .row(
                null,
                "world",
                null,
                Double.POSITIVE_INFINITY,
                null,
                arrayBlockOf(BIGINT, 3, null),
                mapBlockOf(createVarcharType(5), BOOLEAN, "k2", null),
                arrayBlockOf(arrayType, null, arrayBlockOf(BIGINT, 6, 7)))
            .row(
                456L,
                "bye \u2603",
                wrappedBuffer(bytes3),
                Double.NaN,
                false,
                arrayBlockOf(BIGINT),
                mapBlockOf(createVarcharType(5), BOOLEAN, "k3", false),
                arrayBlockOf(arrayType, arrayBlockOf(BIGINT)));

    try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(new EmptyClassLoader());
        OrcFileWriter writer = new OrcFileWriter(columnIds, columnTypes, file)) {
      writer.appendPages(rowPagesBuilder.build());
    }

    try (OrcDataSource dataSource = fileOrcDataSource(file)) {
      OrcRecordReader reader = createReader(dataSource, columnIds, columnTypes);
      assertEquals(reader.getReaderRowCount(), 3);
      assertEquals(reader.getReaderPosition(), 0);
      assertEquals(reader.getFileRowCount(), reader.getReaderRowCount());
      assertEquals(reader.getFilePosition(), reader.getFilePosition());

      assertEquals(reader.nextBatch(), 3);
      assertEquals(reader.getReaderPosition(), 0);
      assertEquals(reader.getFilePosition(), reader.getFilePosition());

      Block column0 = reader.readBlock(BIGINT, 0);
      assertEquals(column0.isNull(0), false);
      assertEquals(column0.isNull(1), true);
      assertEquals(column0.isNull(2), false);
      assertEquals(BIGINT.getLong(column0, 0), 123L);
      assertEquals(BIGINT.getLong(column0, 2), 456L);

      Block column1 = reader.readBlock(createVarcharType(10), 1);
      assertEquals(createVarcharType(10).getSlice(column1, 0), utf8Slice("hello"));
      assertEquals(createVarcharType(10).getSlice(column1, 1), utf8Slice("world"));
      assertEquals(createVarcharType(10).getSlice(column1, 2), utf8Slice("bye \u2603"));

      Block column2 = reader.readBlock(VARBINARY, 2);
      assertEquals(VARBINARY.getSlice(column2, 0), wrappedBuffer(bytes1));
      assertEquals(column2.isNull(1), true);
      assertEquals(VARBINARY.getSlice(column2, 2), wrappedBuffer(bytes3));

      Block column3 = reader.readBlock(DOUBLE, 3);
      assertEquals(column3.isNull(0), false);
      assertEquals(column3.isNull(1), false);
      assertEquals(column3.isNull(2), false);
      assertEquals(DOUBLE.getDouble(column3, 0), 123.456);
      assertEquals(DOUBLE.getDouble(column3, 1), Double.POSITIVE_INFINITY);
      assertEquals(DOUBLE.getDouble(column3, 2), Double.NaN);

      Block column4 = reader.readBlock(BOOLEAN, 4);
      assertEquals(column4.isNull(0), false);
      assertEquals(column4.isNull(1), true);
      assertEquals(column4.isNull(2), false);
      assertEquals(BOOLEAN.getBoolean(column4, 0), true);
      assertEquals(BOOLEAN.getBoolean(column4, 2), false);

      Block column5 = reader.readBlock(arrayType, 5);
      assertEquals(column5.getPositionCount(), 3);

      assertTrue(
          arrayBlocksEqual(BIGINT, arrayType.getObject(column5, 0), arrayBlockOf(BIGINT, 1, 2)));
      assertTrue(
          arrayBlocksEqual(BIGINT, arrayType.getObject(column5, 1), arrayBlockOf(BIGINT, 3, null)));
      assertTrue(arrayBlocksEqual(BIGINT, arrayType.getObject(column5, 2), arrayBlockOf(BIGINT)));

      Block column6 = reader.readBlock(mapType, 6);
      assertEquals(column6.getPositionCount(), 3);

      assertTrue(
          mapBlocksEqual(
              createVarcharType(5),
              BOOLEAN,
              arrayType.getObject(column6, 0),
              mapBlockOf(createVarcharType(5), BOOLEAN, "k1", true)));
      assertTrue(
          mapBlocksEqual(
              createVarcharType(5),
              BOOLEAN,
              arrayType.getObject(column6, 1),
              mapBlockOf(createVarcharType(5), BOOLEAN, "k2", null)));
      assertTrue(
          mapBlocksEqual(
              createVarcharType(5),
              BOOLEAN,
              arrayType.getObject(column6, 2),
              mapBlockOf(createVarcharType(5), BOOLEAN, "k3", false)));

      Block column7 = reader.readBlock(arrayOfArrayType, 7);
      assertEquals(column7.getPositionCount(), 3);

      assertTrue(
          arrayBlocksEqual(
              arrayType,
              arrayOfArrayType.getObject(column7, 0),
              arrayBlockOf(arrayType, arrayBlockOf(BIGINT, 5))));
      assertTrue(
          arrayBlocksEqual(
              arrayType,
              arrayOfArrayType.getObject(column7, 1),
              arrayBlockOf(arrayType, null, arrayBlockOf(BIGINT, 6, 7))));
      assertTrue(
          arrayBlocksEqual(
              arrayType,
              arrayOfArrayType.getObject(column7, 2),
              arrayBlockOf(arrayType, arrayBlockOf(BIGINT))));

      assertEquals(reader.nextBatch(), -1);
      assertEquals(reader.getReaderPosition(), 3);
      assertEquals(reader.getFilePosition(), reader.getFilePosition());

      OrcFileMetadata orcFileMetadata =
          METADATA_CODEC.fromJson(reader.getUserMetadata().get(OrcFileMetadata.KEY).getBytes());
      assertEquals(
          orcFileMetadata,
          new OrcFileMetadata(
              ImmutableMap.<Long, TypeSignature>builder()
                  .put(1L, BIGINT.getTypeSignature())
                  .put(2L, createVarcharType(10).getTypeSignature())
                  .put(4L, VARBINARY.getTypeSignature())
                  .put(6L, DOUBLE.getTypeSignature())
                  .put(7L, BOOLEAN.getTypeSignature())
                  .put(8L, arrayType.getTypeSignature())
                  .put(9L, mapType.getTypeSignature())
                  .put(10L, arrayOfArrayType.getTypeSignature())
                  .build()));
    }

    File crcFile = new File(file.getParentFile(), "." + file.getName() + ".crc");
    assertFalse(crcFile.exists());
  }
    @Setup
    public void setup() {
      MetadataManager metadata = MetadataManager.createTestMetadataManager();
      ExpressionCompiler compiler = new ExpressionCompiler(metadata);

      List<String> keys;
      switch (mapSize) {
        case 1:
          keys = ImmutableList.of("do_not_use");
          break;
        case 13:
          keys =
              ImmutableList.of(
                  "is_inverted",
                  "device_model",
                  "country",
                  "carrier_id",
                  "network_type",
                  "os_version",
                  "device_brand",
                  "device_type",
                  "interface",
                  "device_os",
                  "app_version",
                  "device_type_class",
                  "browser");
          break;
        default:
          throw new UnsupportedOperationException();
      }
      verify(keys.size() == mapSize);

      MapType mapType;
      Block valueBlock;
      switch (name) {
        case "fix-width":
          mapType = new MapType(createUnboundedVarcharType(), DOUBLE);
          valueBlock = createFixWidthValueBlock(POSITIONS, mapSize);
          break;
        case "var-width":
          mapType = new MapType(createUnboundedVarcharType(), createUnboundedVarcharType());
          valueBlock = createVarWidthValueBlock(POSITIONS, mapSize);
          break;
        case "dictionary":
          mapType = new MapType(createUnboundedVarcharType(), createUnboundedVarcharType());
          valueBlock = createDictionaryValueBlock(POSITIONS, mapSize);
          break;
        default:
          throw new UnsupportedOperationException();
      }

      Block keyBlock = createKeyBlock(POSITIONS, keys);
      Block block = createMapBlock(POSITIONS, keyBlock, valueBlock);

      ImmutableList.Builder<RowExpression> projectionsBuilder = ImmutableList.builder();

      Signature signature =
          new Signature(
              "$operator$" + SUBSCRIPT.name(),
              FunctionKind.SCALAR,
              mapType.getValueType().getTypeSignature(),
              mapType.getTypeSignature(),
              mapType.getKeyType().getTypeSignature());
      for (int i = 0; i < mapSize; i++) {
        projectionsBuilder.add(
            new CallExpression(
                signature,
                mapType.getValueType(),
                ImmutableList.of(
                    new InputReferenceExpression(0, mapType),
                    new ConstantExpression(utf8Slice(keys.get(i)), createUnboundedVarcharType()))));
      }

      ImmutableList<RowExpression> projections = projectionsBuilder.build();
      pageProcessor =
          compiler
              .compilePageProcessor(new ConstantExpression(true, BooleanType.BOOLEAN), projections)
              .get();
      pageBuilder =
          new PageBuilder(
              projections.stream().map(RowExpression::getType).collect(Collectors.toList()));
      page = new Page(block);
    }