Resolved conflicts and merged changes from master

Author: inv-jishnu

.github/workflows/release-snapshot.yaml

@@ -30,7 +30,7 @@ jobs:
           VERSION=$(./gradlew :core:properties -q | grep "version:" | awk '{print $2}')
           echo "version=${VERSION}" >> $GITHUB_OUTPUT
 
-      - name: Upload SNAPSHOT versions for scalardb, scalardb-schema-loader, and scalardb-integration-test to Maven Snapshot Repository
+      - name: Upload SNAPSHOT versions for scalardb, scalardb-schema-loader, scalardb-data-loader-core, and scalardb-integration-test to Maven Snapshot Repository
         if: contains(steps.version.outputs.version, '-SNAPSHOT')
         run: |
           echo "${{secrets.SIGNING_SECRET_KEY_RING}}" | base64 -d > ~/.gradle/secring.gpg

.github/workflows/upload-artifacts.yaml

@@ -57,7 +57,7 @@ jobs:
         run: |
           docker push ghcr.io/scalar-labs/scalardb-schema-loader:${{ steps.version.outputs.version }}
 
-      - name: Upload scalardb, scalardb-schema-loader, and scalardb-integration-test to Maven Central Repository
+      - name: Upload scalardb, scalardb-schema-loader, scalardb-data-loader-core, and scalardb-integration-test to Maven Central Repository
         run: |
           echo "${{secrets.SIGNING_SECRET_KEY_RING}}" | base64 -d > ~/.gradle/secring.gpg
           ./gradlew publish \

core/src/main/java/com/scalar/db/common/error/CoreError.java

@@ -923,6 +923,14 @@ public enum CoreError implements ScalarDbError {
       "Some scanners were not closed. All scanners must be closed before preparing the transaction.",
       "",
       ""),
+  DATA_LOADER_INVALID_DATA_CHUNK_SIZE(
+      Category.USER_ERROR, "0207", "Data chunk size must be greater than 0", "", ""),
+  DATA_LOADER_INVALID_TRANSACTION_SIZE(
+      Category.USER_ERROR, "0208", "Transaction size must be greater than 0", "", ""),
+  DATA_LOADER_INVALID_MAX_THREADS(
+      Category.USER_ERROR, "0209", "Number of max threads must be greater than 0", "", ""),
+  DATA_LOADER_INVALID_DATA_CHUNK_QUEUE_SIZE(
+      Category.USER_ERROR, "0210", "Data chunk queue size must be greater than 0", "", ""),
 
   //
   // Errors for the concurrency error category
@@ -1117,7 +1125,7 @@ public enum CoreError implements ScalarDbError {
   JDBC_ERROR_OCCURRED_IN_SELECTION(
       Category.INTERNAL_ERROR, "0027", "An error occurred in the selection. Details: %s", "", ""),
   JDBC_FETCHING_NEXT_RESULT_FAILED(
-      Category.INTERNAL_ERROR, "0028", "Fetching the next result failed", "", ""),
+      Category.INTERNAL_ERROR, "0028", "Fetching the next result failed. Details: %s", "", ""),
   JDBC_TRANSACTION_ROLLING_BACK_TRANSACTION_FAILED(
       Category.INTERNAL_ERROR, "0029", "Rolling back the transaction failed. Details: %s", "", ""),
   JDBC_TRANSACTION_COMMITTING_TRANSACTION_FAILED(
@@ -1196,6 +1204,8 @@ public enum CoreError implements ScalarDbError {
       Category.INTERNAL_ERROR, "0053", "Failed to read JSON Lines file. Details: %s.", "", ""),
   JDBC_TRANSACTION_GETTING_SCANNER_FAILED(
       Category.INTERNAL_ERROR, "0054", "Getting the scanner failed. Details: %s", "", ""),
+  JDBC_CLOSING_SCANNER_FAILED(
+      Category.INTERNAL_ERROR, "0055", "Closing the scanner failed. Details: %s", "", ""),
 
   //
   // Errors for the unknown transaction status error category

core/src/main/java/com/scalar/db/config/DatabaseConfig.java

@@ -38,6 +38,7 @@ public class DatabaseConfig {
   private boolean crossPartitionScanFilteringEnabled;
   private boolean crossPartitionScanOrderingEnabled;
   private String systemNamespaceName;
+  private int scanFetchSize;
 
   public static final String PREFIX = "scalar.db.";
   public static final String CONTACT_POINTS = PREFIX + "contact_points";
@@ -56,9 +57,11 @@ public class DatabaseConfig {
   public static final String CROSS_PARTITION_SCAN_FILTERING = SCAN_PREFIX + "filtering.enabled";
   public static final String CROSS_PARTITION_SCAN_ORDERING = SCAN_PREFIX + "ordering.enabled";
   public static final String SYSTEM_NAMESPACE_NAME = PREFIX + "system_namespace_name";
+  public static final String SCAN_FETCH_SIZE = PREFIX + "scan_fetch_size";
 
   public static final int DEFAULT_METADATA_CACHE_EXPIRATION_TIME_SECS = 60;
   public static final String DEFAULT_SYSTEM_NAMESPACE_NAME = "scalardb";
+  public static final int DEFAULT_SCAN_FETCH_SIZE = 10;
 
   public DatabaseConfig(File propertiesFile) throws IOException {
     try (FileInputStream stream = new FileInputStream(propertiesFile)) {
@@ -118,6 +121,8 @@ protected void load() {
     }
 
     systemNamespaceName = getSystemNamespaceName(getProperties());
+
+    scanFetchSize = getInt(getProperties(), SCAN_FETCH_SIZE, DEFAULT_SCAN_FETCH_SIZE);
   }
 
   public List<String> getContactPoints() {
@@ -172,6 +177,10 @@ public String getSystemNamespaceName() {
     return systemNamespaceName;
   }
 
+  public int getScanFetchSize() {
+    return scanFetchSize;
+  }
+
   public static String getTransactionManager(Properties properties) {
     return getString(properties, TRANSACTION_MANAGER, "consensus-commit");
   }

core/src/main/java/com/scalar/db/storage/cassandra/Cassandra.java

@@ -56,7 +56,7 @@ public Cassandra(DatabaseConfig config) {
 
     handlers =
         StatementHandlerManager.builder()
-            .select(new SelectStatementHandler(session))
+            .select(new SelectStatementHandler(session, config.getScanFetchSize()))
             .insert(new InsertStatementHandler(session))
             .update(new UpdateStatementHandler(session))
             .delete(new DeleteStatementHandler(session))

core/src/main/java/com/scalar/db/storage/cassandra/SelectStatementHandler.java

@@ -41,13 +41,18 @@
  */
 @ThreadSafe
 public class SelectStatementHandler extends StatementHandler {
+
+  private final int fetchSize;
+
   /**
    * Constructs {@code SelectStatementHandler} with the specified {@code Session}
    *
    * @param session session to be used with this statement
+   * @param fetchSize the number of rows to be fetched at once
    */
-  public SelectStatementHandler(Session session) {
+  public SelectStatementHandler(Session session, int fetchSize) {
     super(session);
+    this.fetchSize = fetchSize;
   }
 
   @Override
@@ -94,6 +99,7 @@ protected BoundStatement bind(PreparedStatement prepared, Operation operation) {
   @Override
   @Nonnull
   protected ResultSet execute(BoundStatement bound, Operation operation) {
+    bound.setFetchSize(fetchSize);
     return session.execute(bound);
   }
 

core/src/main/java/com/scalar/db/storage/cosmos/Cosmos.java

@@ -62,7 +62,8 @@ public Cosmos(DatabaseConfig databaseConfig) {
             new CosmosAdmin(client, config), databaseConfig.getMetadataCacheExpirationTimeSecs());
     operationChecker = new CosmosOperationChecker(databaseConfig, metadataManager);
 
-    selectStatementHandler = new SelectStatementHandler(client, metadataManager);
+    selectStatementHandler =
+        new SelectStatementHandler(client, metadataManager, databaseConfig.getScanFetchSize());
     putStatementHandler = new PutStatementHandler(client, metadataManager);
     deleteStatementHandler = new DeleteStatementHandler(client, metadataManager);
     batchHandler = new BatchHandler(client, metadataManager);

core/src/main/java/com/scalar/db/storage/cosmos/SelectStatementHandler.java

@@ -44,8 +44,12 @@
 @ThreadSafe
 public class SelectStatementHandler extends StatementHandler {
 
-  public SelectStatementHandler(CosmosClient client, TableMetadataManager metadataManager) {
+  private final int fetchSize;
+
+  public SelectStatementHandler(
+      CosmosClient client, TableMetadataManager metadataManager, int fetchSize) {
     super(client, metadataManager);
+    this.fetchSize = fetchSize;
   }
 
   /**
@@ -85,9 +89,10 @@ private Scanner executeRead(Get get, TableMetadata tableMetadata) throws CosmosE
       return executeReadWithIndex(get, tableMetadata);
     }
 
+    PartitionKey partitionKey = cosmosOperation.getCosmosPartitionKey();
+
     if (get.getProjections().isEmpty()) {
       String id = cosmosOperation.getId();
-      PartitionKey partitionKey = cosmosOperation.getCosmosPartitionKey();
       Record record = getContainer(get).readItem(id, partitionKey, Record.class).getItem();
       return new SingleRecordScanner(
           record, new ResultInterpreter(get.getProjections(), tableMetadata));
@@ -100,8 +105,10 @@ private Scanner executeRead(Get get, TableMetadata tableMetadata) throws CosmosE
                     .eq(cosmosOperation.getConcatenatedPartitionKey()),
                 DSL.field("r.id").eq(cosmosOperation.getId()))
             .getSQL(ParamType.INLINED);
+    CosmosQueryRequestOptions options =
+        new CosmosQueryRequestOptions().setPartitionKey(partitionKey);
 
-    return executeQuery(get, tableMetadata, query);
+    return executeQuery(get, tableMetadata, query, options);
   }
 
   private Scanner executeReadWithIndex(Selection selection, TableMetadata tableMetadata)
@@ -327,8 +334,8 @@ private Scanner executeQuery(
       CosmosQueryRequestOptions queryOptions) {
     Iterator<FeedResponse<Record>> pagesIterator =
         getContainer(selection)
-            .queryItems(query, queryOptions, Record.class)
-            .iterableByPage()
+            .queryItems(query, queryOptions.setMaxBufferedItemCount(fetchSize), Record.class)
+            .iterableByPage(fetchSize)
             .iterator();
 
     return new ScannerImpl(

core/src/main/java/com/scalar/db/storage/dynamo/Dynamo.java

@@ -77,7 +77,11 @@ public Dynamo(DatabaseConfig databaseConfig) {
     operationChecker = new DynamoOperationChecker(databaseConfig, metadataManager);
 
     selectStatementHandler =
-        new SelectStatementHandler(client, metadataManager, config.getNamespacePrefix());
+        new SelectStatementHandler(
+            client,
+            metadataManager,
+            config.getNamespacePrefix(),
+            databaseConfig.getScanFetchSize());
     putStatementHandler =
         new PutStatementHandler(client, metadataManager, config.getNamespacePrefix());
     deleteStatementHandler =

core/src/main/java/com/scalar/db/storage/dynamo/QueryScanner.java

@@ -22,19 +22,22 @@ public class QueryScanner extends AbstractScanner {
   private final ResultInterpreter resultInterpreter;
 
   private Iterator<Map<String, AttributeValue>> itemsIterator;
+  private final int fetchSize;
   @Nullable private Integer remainingLimit;
   @Nullable private Map<String, AttributeValue> lastEvaluatedKey;
 
   @SuppressFBWarnings("EI_EXPOSE_REP2")
-  public QueryScanner(PaginatedRequest request, int limit, ResultInterpreter resultInterpreter) {
+  public QueryScanner(
+      PaginatedRequest request, int fetchSize, int limit, ResultInterpreter resultInterpreter) {
     this.request = request;
+    this.fetchSize = fetchSize;
 
     if (limit > 0) {
       remainingLimit = limit;
-      handleResponse(request.execute(limit));
+      handleResponse(request.execute(Math.min(fetchSize, limit)));
     } else {
       remainingLimit = null;
-      handleResponse(request.execute());
+      handleResponse(request.execute(fetchSize));
     }
 
     this.resultInterpreter = resultInterpreter;
@@ -54,15 +57,13 @@ private boolean hasNext() {
     if (itemsIterator.hasNext()) {
       return true;
     }
-    if (lastEvaluatedKey != null) {
-      if (remainingLimit != null) {
-        handleResponse(request.execute(lastEvaluatedKey, remainingLimit));
-      } else {
-        handleResponse(request.execute(lastEvaluatedKey));
-      }
-      return itemsIterator.hasNext();
+    if (lastEvaluatedKey == null) {
+      return false;
     }
-    return false;
+
+    int nextFetchSize = remainingLimit != null ? Math.min(fetchSize, remainingLimit) : fetchSize;
+    handleResponse(request.execute(lastEvaluatedKey, nextFetchSize));
+    return itemsIterator.hasNext();
   }
 
   private void handleResponse(PaginatedRequestResponse response) {
@@ -71,25 +72,21 @@ private void handleResponse(PaginatedRequestResponse response) {
       remainingLimit -= items.size();
     }
     itemsIterator = items.iterator();
-    if ((remainingLimit == null || remainingLimit > 0) && response.hasLastEvaluatedKey()) {
-      lastEvaluatedKey = response.lastEvaluatedKey();
-    } else {
-      lastEvaluatedKey = null;
-    }
+
+    boolean shouldContinue =
+        (remainingLimit == null || remainingLimit > 0) && response.hasLastEvaluatedKey();
+    lastEvaluatedKey = shouldContinue ? response.lastEvaluatedKey() : null;
   }
 
   @Override
   @Nonnull
   public List<Result> all() {
-    List<Result> ret = new ArrayList<>();
-    while (true) {
-      Optional<Result> one = one();
-      if (!one.isPresent()) {
-        break;
-      }
-      ret.add(one.get());
+    List<Result> results = new ArrayList<>();
+    Optional<Result> next;
+    while ((next = one()).isPresent()) {
+      results.add(next.get());
     }
-    return ret;
+    return results;
   }
 
   @Override