agturley commented on code in PR #10981:
URL: https://github.com/apache/nifi/pull/10981#discussion_r2908817605
##########
nifi-extension-bundles/nifi-elasticsearch-bundle/nifi-elasticsearch-restapi-processors/src/main/java/org/apache/nifi/processors/elasticsearch/PutElasticsearchJson.java:
##########
@@ -153,174 +224,569 @@ public class PutElasticsearchJson extends
AbstractPutElasticsearch {
SCRIPT,
SCRIPTED_UPSERT,
DYNAMIC_TEMPLATES,
+ INPUT_FORMAT,
BATCH_SIZE,
+ MAX_BATCH_SIZE,
CHARSET,
MAX_JSON_FIELD_STRING_LENGTH,
CLIENT_SERVICE,
LOG_ERROR_RESPONSES,
OUTPUT_ERROR_RESPONSES,
+ OUTPUT_BULK_REQUEST,
NOT_FOUND_IS_SUCCESSFUL
);
+
static final Set<Relationship> BASE_RELATIONSHIPS =
Set.of(REL_ORIGINAL, REL_FAILURE, REL_RETRY, REL_SUCCESSFUL,
REL_ERRORS);
+ private static final int READER_BUFFER_SIZE = 65536;
+
+ private final AtomicBoolean bulkRequestOutputEnabled = new
AtomicBoolean(false);
+ private boolean outputBulkRequest;
+ private ObjectReader mapReader;
+
@Override
Set<Relationship> getBaseRelationships() {
return BASE_RELATIONSHIPS;
}
+ @Override
+ public Set<Relationship> getRelationships() {
+ final Set<Relationship> rels = new HashSet<>(super.getRelationships());
+ if (bulkRequestOutputEnabled.get()) {
+ rels.add(REL_BULK_REQUEST);
+ }
+ return rels;
+ }
+
@Override
public final List<PropertyDescriptor> getSupportedPropertyDescriptors() {
return DESCRIPTORS;
}
+ @Override
+ public void onPropertyModified(final PropertyDescriptor descriptor, final
String oldValue, final String newValue) {
+ super.onPropertyModified(descriptor, oldValue, newValue);
+ if (OUTPUT_BULK_REQUEST.equals(descriptor)) {
+ bulkRequestOutputEnabled.set(Boolean.parseBoolean(newValue));
+ }
+ }
+
@Override
public void migrateProperties(final PropertyConfiguration config) {
super.migrateProperties(config);
+ // Migrate legacy property names from PutElasticsearchJson
config.removeProperty("put-es-json-error-documents");
config.renameProperty("put-es-json-id-attr", ID_ATTRIBUTE.getName());
config.renameProperty("put-es-json-script", SCRIPT.getName());
config.renameProperty("put-es-json-scripted-upsert",
SCRIPTED_UPSERT.getName());
config.renameProperty("put-es-json-dynamic_templates",
DYNAMIC_TEMPLATES.getName());
config.renameProperty("put-es-json-charset", CHARSET.getName());
config.renameProperty("put-es-json-not_found-is-error",
AbstractPutElasticsearch.NOT_FOUND_IS_SUCCESSFUL.getName());
+
+ // Migrate "Batch Size" (from PutElasticsearchJson) to the new name
used in this processor
+ config.renameProperty(AbstractPutElasticsearch.BATCH_SIZE.getName(),
BATCH_SIZE.getName());
+
+ // If INPUT_FORMAT was not explicitly set, this flow was migrated from
PutElasticsearchJson — default to Single JSON
+ if (!config.hasProperty(INPUT_FORMAT.getName())) {
+ config.setProperty(INPUT_FORMAT.getName(), FORMAT_SINGLE_JSON);
+ }
}
@Override
public void migrateRelationships(final RelationshipConfiguration config) {
super.migrateRelationships(config);
+ // PutElasticsearchJson used "success" before it was renamed to
"original"
config.renameRelationship("success",
AbstractPutElasticsearch.REL_ORIGINAL.getName());
}
@Override
@OnScheduled
public void onScheduled(final ProcessContext context) {
super.onScheduled(context);
-
this.notFoundIsSuccessful =
context.getProperty(NOT_FOUND_IS_SUCCESSFUL).asBoolean();
+ this.outputBulkRequest =
context.getProperty(OUTPUT_BULK_REQUEST).asBoolean();
+ this.bulkRequestOutputEnabled.set(this.outputBulkRequest);
+ this.mapReader = mapper.readerFor(Map.class);
}
@Override
public void onTrigger(final ProcessContext context, final ProcessSession
session) {
- final int batchSize =
context.getProperty(BATCH_SIZE).evaluateAttributeExpressions().asInteger();
-
- final List<FlowFile> flowFiles = session.get(batchSize);
- if (flowFiles.isEmpty()) {
- return;
- }
-
+ final long maxBatchBytes =
context.getProperty(MAX_BATCH_SIZE).evaluateAttributeExpressions().asDataSize(DataUnit.B).longValue();
final String idAttribute =
context.getProperty(ID_ATTRIBUTE).getValue();
+ final String inputFormat =
context.getProperty(INPUT_FORMAT).getValue();
+ final int batchSize = FORMAT_SINGLE_JSON.equals(inputFormat)
+ ?
context.getProperty(BATCH_SIZE).evaluateAttributeExpressions().asInteger()
+ : Integer.MAX_VALUE;
+ int flowFilesProcessed = 0;
+
+ // Tracks all FlowFiles that were successfully parsed and submitted
(even partially)
+ final Set<FlowFile> allProcessedFlowFiles = new LinkedHashSet<>();
+ // Tracks FlowFiles that had at least one Elasticsearch document error
+ final Set<FlowFile> errorFlowFiles = new LinkedHashSet<>();
+ // Deferred bulk-error attributes: applied after each FlowFile's
InputStream is closed
+ final Map<FlowFile, List<Map<String, Object>>> pendingBulkErrors = new
HashMap<>();
+
+ // Current chunk accumulation — operationFlowFiles is parallel to
operations (same index)
+ final List<FlowFile> operationFlowFiles = new ArrayList<>();
+ final List<IndexOperationRequest> operations = new ArrayList<>();
+ long totalBytesAccumulated = 0;
+ long chunkBytes = 0;
+
+ FlowFile flowFile;
+ while ((flowFile = session.get()) != null) {
+ final String indexOp =
context.getProperty(INDEX_OP).evaluateAttributeExpressions(flowFile).getValue();
+ final String index =
context.getProperty(INDEX).evaluateAttributeExpressions(flowFile).getValue();
+ final String type =
context.getProperty(TYPE).evaluateAttributeExpressions(flowFile).getValue();
+ final String charset =
context.getProperty(CHARSET).evaluateAttributeExpressions(flowFile).getValue();
+ final String flowFileIdAttribute =
StringUtils.isNotBlank(idAttribute) ? flowFile.getAttribute(idAttribute) : null;
+
+ final Map<String, Object> scriptMap = getMapFromAttribute(SCRIPT,
context, flowFile);
+ final boolean scriptedUpsert =
context.getProperty(SCRIPTED_UPSERT).evaluateAttributeExpressions(flowFile).asBoolean();
+ final Map<String, Object> dynamicTemplatesMap =
getMapFromAttribute(DYNAMIC_TEMPLATES, context, flowFile);
+ final Map<String, String> dynamicProperties =
getRequestParametersFromDynamicProperties(context, flowFile);
+ final Map<String, String> bulkHeaderFields =
getBulkHeaderParameters(dynamicProperties);
+
+ boolean parseError = false;
+ try {
+ final IndexOperationRequest.Operation o =
IndexOperationRequest.Operation.forValue(indexOp);
+ if (FORMAT_NDJSON.equals(inputFormat)) {
+ // NDJSON: each non-blank line is one JSON document.
+ // Index/Create operations pass raw UTF-8 bytes directly
to avoid Map allocation.
+ // Update/Delete/Upsert parse into a Map so the bulk
serializer can wrap the payload.
+ try (final BufferedReader reader = new BufferedReader(
+ new InputStreamReader(session.read(flowFile),
charset), READER_BUFFER_SIZE)) {
+ String line;
+ while ((line = reader.readLine()) != null) {
+ final String trimmedLine = line.trim();
+ if (trimmedLine.isEmpty()) {
+ continue;
+ }
+ final IndexOperationRequest opRequest;
+ final long docBytes;
+ if (o == IndexOperationRequest.Operation.Index ||
o == IndexOperationRequest.Operation.Create) {
+ final String id = extractId(trimmedLine,
idAttribute, flowFileIdAttribute);
+ final byte[] rawJsonBytes =
trimmedLine.getBytes(StandardCharsets.UTF_8);
+ opRequest = new IndexOperationRequest(index,
type, id, rawJsonBytes, o, scriptMap, scriptedUpsert, dynamicTemplatesMap,
bulkHeaderFields);
+ docBytes = rawJsonBytes.length;
+ } else {
+ @SuppressWarnings("unchecked")
+ final Map<String, Object> contentMap =
mapReader.readValue(trimmedLine);
+ final String id = resolveId(contentMap,
idAttribute, flowFileIdAttribute);
+ opRequest = new IndexOperationRequest(index,
type, id, contentMap, o, scriptMap, scriptedUpsert, dynamicTemplatesMap,
bulkHeaderFields);
+ docBytes = trimmedLine.length();
+ }
+ operations.add(opRequest);
+ operationFlowFiles.add(flowFile);
+ chunkBytes += docBytes;
+ totalBytesAccumulated += docBytes;
+ if (chunkBytes >= maxBatchBytes) {
+ flushChunk(operations, operationFlowFiles,
errorFlowFiles, flowFile, pendingBulkErrors, context, session);
Review Comment:
I'm not a huge fan of the potential to have duplicates, this seems to
work.... When a bulk request comes back from Elasticsearch, we look at which
individual documents failed. Instead of buffering the raw bytes of every
document as we go (which would hold the entire file in memory twice), we just
record the index of each failed document — a small integer per error,
regardless of document size.
Once all chunks for a FlowFile are processed, we route it:
**No errors:** clone the original FlowFile straight to REL_SUCCESSFUL. Zero
extra I/O.
**Single JSON input (one doc per FlowFile):** if it errored, clone it to
REL_ERRORS. If not, clone to REL_SUCCESSFUL. Again, no re-read needed.
**NDJSON or JSON Array input with at least one error:** we do a single
re-read of the original FlowFile and split it in one pass into two streams: one
for the failed records (REL_ERRORS) and one for the successful records
(REL_SUCCESSFUL). Both outputs are written as clean NDJSON with no trailing
newline.
The re-read only happens for FlowFiles that actually had partial failures in
NDJSON/JSON Array format, so the common happy path (all docs succeed) is just a
cheap clone with no extra I/O.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
