5858import org .springframework .data .projection .EntityProjection ;
5959import org .springframework .data .projection .ProjectionFactory ;
6060import org .springframework .data .projection .SpelAwareProxyProjectionFactory ;
61+ import org .springframework .data .util .Lazy ;
6162import org .springframework .util .Assert ;
6263
6364import com .datastax .oss .driver .api .core .CqlIdentifier ;
@@ -353,10 +354,17 @@ public <T> List<T> select(Statement<?> statement, Class<T> entityClass) {
353354 Assert .notNull (statement , "Statement must not be null" );
354355 Assert .notNull (entityClass , "Entity type must not be null" );
355356
356- Function <Row , T > mapper = getMapper (EntityProjection .nonProjecting (entityClass ),
357- EntityQueryUtils .getTableName (statement ));
357+ return doSelect (statement , entityClass , getTableName (entityClass ), entityClass , QueryResultConverter .entity ());
358+ }
359+
360+ <T , R > List <R > doSelect (Statement <?> statement , Class <?> entityClass , CqlIdentifier tableName , Class <T > returnType ,
361+ QueryResultConverter <T , R > mappingFunction ) {
362+
363+ EntityProjection <T , ?> projection = entityOperations .introspectProjection (returnType , entityClass );
364+
365+ RowMapper <R > rowMapper = getRowMapper (projection , tableName , mappingFunction );
358366
359- return doQuery (statement , ( row , rowNum ) -> mapper . apply ( row ) );
367+ return doQuery (statement , rowMapper );
360368 }
361369
362370 @ Override
@@ -372,13 +380,14 @@ public <T> Slice<T> slice(Statement<?> statement, Class<T> entityClass) {
372380 Assert .notNull (statement , "Statement must not be null" );
373381 Assert .notNull (entityClass , "Entity type must not be null" );
374382
375- ResultSet resultSet = doQueryForResultSet (statement );
383+ return doSlice (statement ,
384+ getRowMapper (entityClass , EntityQueryUtils .getTableName (statement ), QueryResultConverter .entity ()));
385+ }
376386
377- Function <Row , T > mapper = getMapper (EntityProjection .nonProjecting (entityClass ),
378- EntityQueryUtils .getTableName (statement ));
387+ <T > Slice <T > doSlice (Statement <?> statement , RowMapper <T > mapper ) {
379388
380- return EntityQueryUtils . readSlice ( resultSet , ( row , rowNum ) -> mapper . apply ( row ), 0 ,
381- getEffectivePageSize (statement ));
389+ ResultSet resultSet = doQueryForResultSet ( statement );
390+ return EntityQueryUtils . readSlice ( resultSet , mapper , 0 , getEffectivePageSize (statement ));
382391 }
383392
384393 @ Override
@@ -387,9 +396,17 @@ public <T> Stream<T> stream(Statement<?> statement, Class<T> entityClass) throws
387396 Assert .notNull (statement , "Statement must not be null" );
388397 Assert .notNull (entityClass , "Entity type must not be null" );
389398
390- Function <Row , T > mapper = getMapper (EntityProjection .nonProjecting (entityClass ),
391- EntityQueryUtils .getTableName (statement ));
392- return doQueryForStream (statement , (row , rowNum ) -> mapper .apply (row ));
399+ return doStream (statement , entityClass , EntityQueryUtils .getTableName (statement ), entityClass ,
400+ QueryResultConverter .entity ());
401+ }
402+
403+ <T , R > Stream <R > doStream (Statement <?> statement , Class <?> entityClass , CqlIdentifier tableName , Class <T > returnType ,
404+ QueryResultConverter <T , R > mappingFunction ) {
405+
406+ EntityProjection <T , ?> projection = entityOperations .introspectProjection (returnType , entityClass );
407+
408+ RowMapper <R > rowMapper = getRowMapper (projection , tableName , mappingFunction );
409+ return doQueryForStream (statement , rowMapper );
393410 }
394411
395412 // -------------------------------------------------------------------------
@@ -402,10 +419,11 @@ public <T> List<T> select(Query query, Class<T> entityClass) throws DataAccessEx
402419 Assert .notNull (query , "Query must not be null" );
403420 Assert .notNull (entityClass , "Entity type must not be null" );
404421
405- return doSelect (query , entityClass , getTableName (entityClass ), entityClass );
422+ return doSelect (query , entityClass , getTableName (entityClass ), entityClass , QueryResultConverter . entity () );
406423 }
407424
408- <T > List <T > doSelect (Query query , Class <?> entityClass , CqlIdentifier tableName , Class <T > returnType ) {
425+ <T , R > List <R > doSelect (Query query , Class <?> entityClass , CqlIdentifier tableName , Class <T > returnType ,
426+ QueryResultConverter <? super T , ? extends R > mappingFunction ) {
409427
410428 CassandraPersistentEntity <?> entity = getRequiredPersistentEntity (entityClass );
411429 EntityProjection <T , ?> projection = entityOperations .introspectProjection (returnType , entityClass );
@@ -415,9 +433,9 @@ <T> List<T> doSelect(Query query, Class<?> entityClass, CqlIdentifier tableName,
415433 Query queryToUse = query .columns (columns );
416434
417435 StatementBuilder <Select > select = getStatementFactory ().select (queryToUse , entity , tableName );
418- Function < Row , T > mapper = getMapper (projection , tableName );
436+ RowMapper < R > rowMapper = getRowMapper (projection , tableName , mappingFunction );
419437
420- return doQuery (select .build (), ( row , rowNum ) -> mapper . apply ( row ) );
438+ return doQuery (select .build (), rowMapper );
421439 }
422440
423441 @ Override
@@ -434,9 +452,24 @@ public <T> Slice<T> slice(Query query, Class<T> entityClass) throws DataAccessEx
434452 Assert .notNull (query , "Query must not be null" );
435453 Assert .notNull (entityClass , "Entity type must not be null" );
436454
437- StatementBuilder <Select > select = getStatementFactory ().select (query , getRequiredPersistentEntity (entityClass ));
455+ return doSlice (query , entityClass , getRequiredPersistentEntity (entityClass ).getTableName (), entityClass ,
456+ QueryResultConverter .entity ());
457+ }
458+
459+ <T , R > Slice <R > doSlice (Query query , Class <?> entityClass , CqlIdentifier tableName , Class <T > returnType ,
460+ QueryResultConverter <? super T , ? extends R > mappingFunction ) {
461+
462+ CassandraPersistentEntity <?> entity = getRequiredPersistentEntity (entityClass );
463+ EntityProjection <T , ?> projection = entityOperations .introspectProjection (returnType , entityClass );
464+ Columns columns = getStatementFactory ().computeColumnsForProjection (projection , query .getColumns (), entity ,
465+ returnType );
466+
467+ Query queryToUse = query .columns (columns );
468+
469+ StatementBuilder <Select > select = getStatementFactory ().select (queryToUse , entity , tableName );
470+ RowMapper <R > rowMapper = getRowMapper (projection , tableName , mappingFunction );
438471
439- return slice (select .build (), entityClass );
472+ return doSlice (select .build (), rowMapper );
440473 }
441474
442475 @ Override
@@ -445,17 +478,19 @@ public <T> Stream<T> stream(Query query, Class<T> entityClass) throws DataAccess
445478 Assert .notNull (query , "Query must not be null" );
446479 Assert .notNull (entityClass , "Entity type must not be null" );
447480
448- return doStream (query , entityClass , getTableName (entityClass ), entityClass );
481+ return doStream (query , entityClass , getTableName (entityClass ), entityClass , QueryResultConverter . entity () );
449482 }
450483
451- <T > Stream <T > doStream (Query query , Class <?> entityClass , CqlIdentifier tableName , Class <T > returnType ) {
484+ <T , R > Stream <R > doStream (Query query , Class <?> entityClass , CqlIdentifier tableName , Class <T > returnType ,
485+ QueryResultConverter <? super T , ? extends R > mappingFunction ) {
452486
453487 StatementBuilder <Select > select = getStatementFactory ().select (query , getRequiredPersistentEntity (entityClass ),
454488 tableName );
455489 EntityProjection <T , ?> projection = entityOperations .introspectProjection (returnType , entityClass );
456490
457- Function <Row , T > mapper = getMapper (projection , tableName );
458- return doQueryForStream (select .build (), (row , rowNum ) -> mapper .apply (row ));
491+ RowMapper <R > rowMapper = getRowMapper (projection , tableName , mappingFunction );
492+
493+ return doQueryForStream (select .build (), rowMapper );
459494 }
460495
461496 @ Override
@@ -767,6 +802,16 @@ public <T> ExecutableSelect<T> query(Class<T> domainType) {
767802 return new ExecutableSelectOperationSupport (this ).query (domainType );
768803 }
769804
805+ @ Override
806+ public UntypedSelect query (String cql ) {
807+ return new ExecutableSelectOperationSupport (this ).query (cql );
808+ }
809+
810+ @ Override
811+ public UntypedSelect query (Statement <?> statement ) {
812+ return new ExecutableSelectOperationSupport (this ).query (statement );
813+ }
814+
770815 @ Override
771816 public <T > ExecutableInsert <T > insert (Class <T > domainType ) {
772817 return new ExecutableInsertOperationSupport (this ).insert (domainType );
@@ -909,6 +954,32 @@ public String getCql() {
909954 return getCqlOperations ().execute (new GetConfiguredPageSize ());
910955 }
911956
957+ @ SuppressWarnings ("unchecked" )
958+ <T , R > RowMapper <R > getRowMapper (EntityProjection <T , ?> projection , CqlIdentifier tableName ,
959+ QueryResultConverter <? super T , ? extends R > mappingFunction ) {
960+
961+ Function <Row , T > mapper = getMapper (projection , tableName );
962+
963+ return mappingFunction == QueryResultConverter .entity () ? (row , rowNum ) -> (R ) mapper .apply (row )
964+ : (row , rowNum ) -> {
965+ Lazy <T > reader = Lazy .of (() -> mapper .apply (row ));
966+ return mappingFunction .mapRow (row , reader ::get );
967+ };
968+ }
969+
970+ @ SuppressWarnings ("unchecked" )
971+ <T , R > RowMapper <R > getRowMapper (Class <T > domainClass , CqlIdentifier tableName ,
972+ QueryResultConverter <? super T , ? extends R > mappingFunction ) {
973+
974+ Function <Row , T > mapper = getMapper (EntityProjection .nonProjecting (domainClass ), tableName );
975+
976+ return mappingFunction == QueryResultConverter .entity () ? (row , rowNum ) -> (R ) mapper .apply (row )
977+ : (row , rowNum ) -> {
978+ Lazy <T > reader = Lazy .of (() -> mapper .apply (row ));
979+ return mappingFunction .mapRow (row , reader ::get );
980+ };
981+ }
982+
912983 @ SuppressWarnings ("unchecked" )
913984 private <T > Function <Row , T > getMapper (EntityProjection <T , ?> projection , CqlIdentifier tableName ) {
914985
0 commit comments