3434import org .apache .pinot .common .datatable .DataTable ;
3535import org .apache .pinot .common .proto .Server ;
3636import org .apache .pinot .common .utils .DataSchema ;
37- import org .apache .pinot .common .utils .grpc .GrpcRequestBuilder ;
37+ import org .apache .pinot .common .utils .grpc .ServerGrpcRequestBuilder ;
3838import org .apache .pinot .core .common .datatable .DataTableBuilder ;
3939import org .apache .pinot .core .common .datatable .DataTableBuilderV4 ;
4040import org .apache .pinot .spi .data .DimensionFieldSpec ;
4141import org .apache .pinot .spi .data .FieldSpec ;
42+ import org .apache .pinot .spi .utils .ByteArray ;
4243import org .apache .pinot .spi .utils .CommonConstants ;
4344import org .testng .annotations .Test ;
4445
@@ -172,9 +173,6 @@ protected static DataTable createDataTableWithAllTypes()
172173 case STRING :
173174 dataTableBuilder .setColumn (colId , generateRandomStringWithLength (RANDOM .nextInt (20 )));
174175 break ;
175- case OBJECT :
176- dataTableBuilder .setColumn (colId , (Object ) RANDOM .nextDouble ());
177- break ;
178176 case BOOLEAN_ARRAY :
179177 int length = RANDOM .nextInt (20 );
180178 int [] booleanArray = new int [length ];
@@ -233,7 +231,7 @@ protected static DataTable createDataTableWithAllTypes()
233231 case BYTES :
234232 try {
235233 dataTableBuilder .setColumn (colId ,
236- Hex .decodeHex ("0DE0B6B3A7640000" .toCharArray ())); // Hex of BigDecimal.ONE
234+ new ByteArray ( Hex .decodeHex ("0DE0B6B3A7640000" .toCharArray () ))); // Hex of BigDecimal.ONE
237235 }
238236 catch (DecoderException e ) {
239237 throw new RuntimeException (e );
@@ -390,7 +388,6 @@ public void testPinotProxyGrpcRequest()
390388 Assert .assertEquals (grpcRequest .getSql (), "SELECT * FROM myTable" );
391389 Assert .assertEquals (grpcRequest .getSegmentsCount (), 1 );
392390 Assert .assertEquals (grpcRequest .getSegments (0 ), "segment1" );
393- Assert .assertEquals (grpcRequest .getMetadataCount (), 9 );
394391 Assert .assertEquals (grpcRequest .getMetadataOrThrow ("k1" ), "v1" );
395392 Assert .assertEquals (grpcRequest .getMetadataOrThrow ("k2" ), "v2" );
396393 Assert .assertEquals (grpcRequest .getMetadataOrThrow ("FORWARD_HOST" ), "localhost" );
@@ -401,6 +398,7 @@ public void testPinotProxyGrpcRequest()
401398 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .ENABLE_STREAMING ), "true" );
402399 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .PAYLOAD_TYPE ), "sql" );
403400
401+ Assert .assertTrue (grpcRequest .getMetadataCount () >= 9 , "Expected at least 9 metadata entries, but got " + grpcRequest .getMetadataCount ());
404402 grpcRequest = new PinotProxyGrpcRequestBuilder ()
405403 .setSegments (ImmutableList .of ("segment1" ))
406404 .setEnableStreaming (true )
@@ -412,20 +410,23 @@ public void testPinotProxyGrpcRequest()
412410 Assert .assertEquals (grpcRequest .getSql (), "SELECT * FROM myTable" );
413411 Assert .assertEquals (grpcRequest .getSegmentsCount (), 1 );
414412 Assert .assertEquals (grpcRequest .getSegments (0 ), "segment1" );
415- Assert .assertEquals (grpcRequest .getMetadataCount (), 7 );
416413 Assert .assertEquals (grpcRequest .getMetadataOrThrow ("k1" ), "v1" );
417414 Assert .assertEquals (grpcRequest .getMetadataOrThrow ("k2" ), "v2" );
418415 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .REQUEST_ID ), "121" );
419416 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .BROKER_ID ), "presto-coordinator-grpc" );
420417 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .ENABLE_TRACE ), "false" );
421418 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .ENABLE_STREAMING ), "true" );
422419 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .PAYLOAD_TYPE ), "sql" );
420+
421+ // Verifying minimum metadata count (7 fields: k1, k2, requestId, brokerId, enableTrace, enableStreaming, payloadType)
422+ Assert .assertTrue (grpcRequest .getMetadataCount () >= 7 ,
423+ "Expected at least 7 metadata entries, but got " + grpcRequest .getMetadataCount ());
423424 }
424425
425426 @ Test
426427 public void testPinotGrpcRequest ()
427428 {
428- final Server .ServerRequest grpcRequest = new GrpcRequestBuilder ()
429+ final Server .ServerRequest grpcRequest = new ServerGrpcRequestBuilder ()
429430 .setSegments (ImmutableList .of ("segment1" ))
430431 .setEnableStreaming (true )
431432 .setRequestId (121 )
@@ -435,7 +436,7 @@ public void testPinotGrpcRequest()
435436 Assert .assertEquals (grpcRequest .getSql (), "SELECT * FROM myTable" );
436437 Assert .assertEquals (grpcRequest .getSegmentsCount (), 1 );
437438 Assert .assertEquals (grpcRequest .getSegments (0 ), "segment1" );
438- Assert .assertEquals (grpcRequest .getMetadataCount (), 5 );
439+ Assert .assertEquals (grpcRequest .getMetadataCount (), 6 );
439440 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .REQUEST_ID ), "121" );
440441 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .BROKER_ID ), "presto-coordinator-grpc" );
441442 Assert .assertEquals (grpcRequest .getMetadataOrThrow (CommonConstants .Query .Request .MetadataKeys .ENABLE_TRACE ), "false" );
@@ -455,7 +456,7 @@ private static final class TestingPinotStreamingQueryClient
455456 }
456457
457458 @ Override
458- public Iterator <Server .ServerResponse > submit (String host , int port , GrpcRequestBuilder requestBuilder )
459+ public Iterator <Server .ServerResponse > submit (String host , int port , ServerGrpcRequestBuilder requestBuilder )
459460 {
460461 return new Iterator <Server .ServerResponse >()
461462 {
0 commit comments