[ESQL] Cleanup and bug fixes in NDJSON datasource (#143969)

* Cleanup constructors, add settings for the number of lines used for schema inference
* Move implicit schema inference to the top-level `NdJsonFormatReader`
* Fix schema inference (priorities were wrong)
* Fix unexpected value parsing leading to an imbalance in block sizes. A `null` is added.
* Fix nested arrays and structures in arrays parsing. All values are flattened.
* Fix null column handling.
* Use the "strict_date_optional_time" to parse date, to accept both date+time and date only.
* Fall back to double for big integers.

Author: swallez

docs/changelog/143969.yaml

@@ -0,0 +1,5 @@
+area: ES|QL
+issues: []
+pr: 143969
+summary: Cleanup and bug fixes in NDJSON datasource
+type: bug

x-pack/plugin/esql-datasource-ndjson/src/main/java/org/elasticsearch/xpack/esql/datasource/ndjson/NdJsonDataSourcePlugin.java

@@ -32,6 +32,6 @@ public Set<FormatSpec> formatSpecs() {
 
     @Override
     public Map<String, FormatReaderFactory> formatReaders(Settings settings) {
-        return Map.of("ndjson", (s, blockFactory) -> new NdJsonFormatReader(blockFactory));
+        return Map.of("ndjson", NdJsonFormatReader::new);
     }
 }

x-pack/plugin/esql-datasource-ndjson/src/main/java/org/elasticsearch/xpack/esql/datasource/ndjson/NdJsonFormatReader.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.datasource.ndjson;
 
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
@@ -27,28 +28,47 @@
  */
 public class NdJsonFormatReader implements SegmentableFormatReader {
 
+    public static final String SCHEMA_SAMPLE_SIZE_SETTING = "esql.datasource.ndjson.schema_sample_size";
+    public static final int DEFAULT_SCHEMA_SAMPLE_SIZE = 100;
+
     private final BlockFactory blockFactory;
+    private final Settings settings;
     private final List<Attribute> resolvedSchema;
 
-    public NdJsonFormatReader(BlockFactory blockFactory) {
-        this(blockFactory, null);
-    }
-
-    private NdJsonFormatReader(BlockFactory blockFactory, List<Attribute> resolvedSchema) {
+    public NdJsonFormatReader(Settings settings, BlockFactory blockFactory, List<Attribute> resolvedSchema) {
         this.blockFactory = blockFactory;
+        this.settings = settings == null ? Settings.EMPTY : settings;
         this.resolvedSchema = resolvedSchema;
     }
 
+    NdJsonFormatReader(Settings settings, BlockFactory blockFactory) {
+        this(settings, blockFactory, null);
+    }
+
     @Override
     public NdJsonFormatReader withSchema(List<Attribute> schema) {
-        return new NdJsonFormatReader(blockFactory, schema);
+        return new NdJsonFormatReader(settings, blockFactory, schema);
+    }
+
+    private List<Attribute> inferSchemaIfNeeded(List<Attribute> attributes, StorageObject object) throws IOException {
+        if (attributes != null && attributes.isEmpty() == false) {
+            return attributes;
+        }
+
+        try (var stream = object.newStream()) {
+            return NdJsonSchemaInferrer.inferSchema(stream, schemaSampleSize(settings));
+        }
+    }
+
+    private static int schemaSampleSize(Settings settings) {
+        return settings.getAsInt(SCHEMA_SAMPLE_SIZE_SETTING, DEFAULT_SCHEMA_SAMPLE_SIZE);
     }
 
     @Override
     public SourceMetadata metadata(StorageObject object) throws IOException {
         List<Attribute> schema;
         try (var stream = object.newStream()) {
-            schema = NdJsonSchemaInferrer.inferSchema(stream);
+            schema = NdJsonSchemaInferrer.inferSchema(stream, schemaSampleSize(settings));
         }
         return new SimpleSourceMetadata(schema, formatName(), object.path().toString());
     }
@@ -64,7 +84,7 @@ public CloseableIterator<Page> read(StorageObject object, FormatReadContext cont
             blockFactory,
             skipFirstLine,
             trimLastPartialLine,
-            resolvedSchema
+            inferSchemaIfNeeded(resolvedSchema, object)
         );
     }
 

x-pack/plugin/esql-datasource-ndjson/src/main/java/org/elasticsearch/xpack/esql/datasource/ndjson/NdJsonPageDecoder.java

@@ -10,6 +10,7 @@
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
+import com.fasterxml.jackson.core.exc.InputCoercionException;
 import com.fasterxml.jackson.core.io.JsonEOFException;
 
 import org.apache.lucene.util.BytesRef;
@@ -33,15 +34,14 @@
 import java.io.Closeable;
 import java.io.IOException;
 import java.io.InputStream;
-import java.time.Instant;
 import java.util.BitSet;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 public class NdJsonPageDecoder implements Closeable {
 
-    private static final Logger LOGGER = LogManager.getLogger(NdJsonPageDecoder.class);
+    private static final Logger logger = LogManager.getLogger(NdJsonPageDecoder.class);
 
     private InputStream input;
     private final BlockDecoder decoder;
@@ -64,7 +64,7 @@ public class NdJsonPageDecoder implements Closeable {
         this.input = input;
 
         var projectedAttributes = attributes;
-        if (projectedColumns.isEmpty() == false) {
+        if (projectedColumns != null && projectedColumns.isEmpty() == false) {
             // Keep projected columns in order, adding NULL for missing columns
             projectedAttributes = projectedColumns.stream()
                 .map(
@@ -99,9 +99,9 @@ Page decodePage() throws IOException {
                     }
                 } catch (JsonParseException e) {
                     if (e instanceof JsonEOFException) {
-                        LOGGER.debug("Truncated NDJSON at line {} (expected at split boundaries): {}", lineCount, e.getOriginalMessage());
+                        logger.debug("Truncated NDJSON at line {} (expected at split boundaries): {}", lineCount, e.getOriginalMessage());
                     } else {
-                        LOGGER.warn("Malformed NDJSON at line {}: {}", lineCount, e.getOriginalMessage());
+                        logger.debug("Malformed NDJSON at line {}: {}", lineCount, e.getOriginalMessage());
                     }
                     this.input = NdJsonUtils.moveToNextLine(parser, this.input);
                     parser = NdJsonUtils.JSON_FACTORY.createParser(this.input);
@@ -112,12 +112,12 @@ Page decodePage() throws IOException {
                 this.blockTracker.clear();
 
                 try {
-                    decoder.decodeObject(parser);
+                    decoder.decodeObject(parser, false);
                 } catch (JsonParseException e) {
                     if (e instanceof JsonEOFException) {
-                        LOGGER.debug("Truncated NDJSON at line {} (expected at split boundaries): {}", lineCount, e.getOriginalMessage());
+                        logger.debug("Truncated NDJSON at line {} (expected at split boundaries): {}", lineCount, e.getOriginalMessage());
                     } else {
-                        LOGGER.warn("Malformed NDJSON at line {}: {}", lineCount, e.getOriginalMessage());
+                        logger.debug("Malformed NDJSON at line {}: {}", lineCount, e.getOriginalMessage());
                     }
                     this.input = NdJsonUtils.moveToNextLine(parser, this.input);
                     parser = NdJsonUtils.JSON_FACTORY.createParser(this.input);
@@ -185,20 +185,22 @@ public void close() throws IOException {
     // A tree of decoders. Avoids path reconstruction when traversing nested objects.
     private class BlockDecoder {
         @Nullable
-        Attribute attribute;
-        Block.Builder blockBuilder;
+        DataType dataType;
+        String name;
         int blockIdx;
+        Block.Builder blockBuilder;
         Map<String, BlockDecoder> children;
 
         void setAttribute(Attribute attribute, int blockIdx) {
-            this.attribute = attribute;
+            this.dataType = attribute.dataType();
+            this.name = attribute.name();
             this.blockIdx = blockIdx;
         }
 
         // Builders setup independently as we need to create new ones for each page.
         void setupBuilders(Block.Builder[] blockBuilders) {
-            if (attribute != null) {
-                blockBuilder = switch (attribute.dataType()) {
+            if (dataType != null) {
+                blockBuilder = switch (dataType) {
                     // Keep in sync with NdJsonSchemaInferrer.inferValueSchema
                     case BOOLEAN -> blockFactory.newBooleanBlockBuilder(batchSize);
                     case NULL -> new ConstantNullBlock.Builder(blockFactory);
@@ -207,7 +209,7 @@ void setupBuilders(Block.Builder[] blockBuilders) {
                     case DOUBLE -> blockFactory.newDoubleBlockBuilder(batchSize);
                     case KEYWORD -> blockFactory.newBytesRefBlockBuilder(batchSize);
                     case DATETIME -> blockFactory.newLongBlockBuilder(batchSize); // milliseconds since epoch
-                    default -> throw new IllegalArgumentException("Unsupported data type: " + attribute.dataType());
+                    default -> throw new IllegalArgumentException("Unsupported data type: " + dataType);
                 };
                 blockBuilders[blockIdx] = blockBuilder;
             }
@@ -219,7 +221,7 @@ void setupBuilders(Block.Builder[] blockBuilders) {
             }
         }
 
-        private void decodeObject(JsonParser parser) throws IOException {
+        private void decodeObject(JsonParser parser, boolean inArray) throws IOException {
             JsonToken token = parser.currentToken();
             if (token != JsonToken.START_OBJECT) {
                 throw new NdJsonParseException(parser, "Expected JSON object");
@@ -234,69 +236,125 @@ private void decodeObject(JsonParser parser) throws IOException {
                     // Unknown field, skip it
                     parser.skipChildren();
                 } else {
-                    childDecoder.decodeValue(parser);
+                    childDecoder.decodeValue(parser, inArray);
                 }
             }
         }
 
-        private void decodeValue(JsonParser parser) throws IOException {
+        private void beginPositionEntry() {
+            // We may have DataType.NULL for unknown columns. And NullBlock.Builder throws on beginPositionEntry()
+            if (blockBuilder != null && dataType != DataType.NULL) {
+                blockBuilder.beginPositionEntry();
+            }
+            if (children != null) {
+                for (var child : children.values()) {
+                    child.beginPositionEntry();
+                }
+            }
+        }
+
+        private void endPositionEntry() {
+            if (blockBuilder != null && dataType != DataType.NULL) {
+                blockBuilder.endPositionEntry();
+            }
+            if (children != null) {
+                for (var child : children.values()) {
+                    child.endPositionEntry();
+                }
+            }
+        }
+
+        private void decodeValue(JsonParser parser, boolean inArray) throws IOException {
             JsonToken token = parser.currentToken();
-            blockTracker.set(blockIdx);
+
+            if (dataType == DataType.NULL) {
+                // Don't do anything. We must do a single appendNull() on null blocks, this will be done
+                // at the end of decodePage() when we check that all blocks have moved forward.
+                parser.skipChildren();
+                return;
+            }
+
             if (token == JsonToken.START_ARRAY) {
-                this.blockBuilder.beginPositionEntry();
+                // Start a multi-value entry on this decoder and all its children (nested arrays are flattened).
+                // Note: the `inArray` flag is needed because blockBuilder.beginPositionEntry() is not idempotent.
+                // Calling it twice implicitly calls endPositionEntry().
+                if (!inArray) {
+                    beginPositionEntry();
+                }
                 while (parser.nextToken() != JsonToken.END_ARRAY) {
-                    decodeValue(parser);
+                    decodeValue(parser, true);
                 }
-                this.blockBuilder.endPositionEntry();
+                if (!inArray) {
+                    endPositionEntry();
+                }
+                return;
+            }
+
+            if (token == JsonToken.START_OBJECT) {
+                decodeObject(parser, inArray);
                 return;
             }
 
-            if (token == JsonToken.VALUE_NULL) {
+            blockTracker.set(blockIdx);
+            if (token == JsonToken.VALUE_NULL && inArray == false) {
+                // Nulls in arrays aren't supported. Furthermore, appendNull will implicitly call endPositionEntry()
                 blockBuilder.appendNull();
                 return;
             }
 
-            switch (attribute.dataType()) {
+            switch (dataType) {
                 case BOOLEAN -> {
                     if (token == JsonToken.VALUE_TRUE) {
                         ((BooleanBlock.Builder) blockBuilder).appendBoolean(true);
                     } else if (token == JsonToken.VALUE_FALSE) {
                         ((BooleanBlock.Builder) blockBuilder).appendBoolean(false);
                     } else {
-                        unexpectedValue(parser);
+                        unexpectedValue(blockBuilder, parser, inArray);
                     }
                 }
                 case NULL -> {
                     // NULL handled above
-                    unexpectedValue(parser);
+                    unexpectedValue(blockBuilder, parser, inArray);
                 }
                 case INTEGER -> {
                     if (token == JsonToken.VALUE_NUMBER_INT || token == JsonToken.VALUE_NUMBER_FLOAT) {
-                        ((IntBlock.Builder) blockBuilder).appendInt(parser.getIntValue());
+                        try {
+                            ((IntBlock.Builder) blockBuilder).appendInt(parser.getIntValue());
+                        } catch (InputCoercionException e) {
+                            unexpectedValue(blockBuilder, parser, inArray);
+                        }
                     } else {
-                        unexpectedValue(parser);
+                        unexpectedValue(blockBuilder, parser, inArray);
                     }
                 }
                 case LONG -> {
                     if (token == JsonToken.VALUE_NUMBER_INT || token == JsonToken.VALUE_NUMBER_FLOAT) {
-                        ((LongBlock.Builder) blockBuilder).appendLong(parser.getLongValue());
+                        try {
+                            ((LongBlock.Builder) blockBuilder).appendLong(parser.getLongValue());
+                        } catch (InputCoercionException e) {
+                            unexpectedValue(blockBuilder, parser, inArray);
+                        }
                     } else {
-                        unexpectedValue(parser);
+                        unexpectedValue(blockBuilder, parser, inArray);
                     }
                 }
                 case DOUBLE -> {
                     if (token == JsonToken.VALUE_NUMBER_INT || token == JsonToken.VALUE_NUMBER_FLOAT) {
-                        ((DoubleBlock.Builder) blockBuilder).appendDouble(parser.getDoubleValue());
+                        try {
+                            ((DoubleBlock.Builder) blockBuilder).appendDouble(parser.getDoubleValue());
+                        } catch (InputCoercionException e) {
+                            unexpectedValue(blockBuilder, parser, inArray);
+                        }
                     } else {
-                        unexpectedValue(parser);
+                        unexpectedValue(blockBuilder, parser, inArray);
                     }
                 }
                 case DATETIME -> {
                     try {
-                        var millis = Instant.parse(parser.getValueAsString()).toEpochMilli();
+                        var millis = NdJsonSchemaInferrer.DATE_FORMATTER.parseMillis(parser.getValueAsString());
                         ((LongBlock.Builder) blockBuilder).appendLong(millis);
                     } catch (Exception e) {
-                        unexpectedValue(parser);
+                        unexpectedValue(blockBuilder, parser, inArray);
                     }
                 }
                 case KEYWORD -> {
@@ -305,20 +363,21 @@ private void decodeValue(JsonParser parser) throws IOException {
                     if (str != null) {
                         ((BytesRefBlock.Builder) blockBuilder).appendBytesRef(new BytesRef(str));
                     } else {
-                        unexpectedValue(parser);
+                        unexpectedValue(blockBuilder, parser, inArray);
                     }
                 }
-                default -> throw new IllegalArgumentException("Unsupported data type: " + attribute.dataType());
+                default -> throw new IllegalArgumentException("Unsupported data type: " + dataType);
             }
         }
 
-        private void unexpectedValue(JsonParser parser) throws IOException {
-            LOGGER.warn(
-                "Unexpected token type: {} for attribute: {} at {}",
-                parser.currentToken(),
-                attribute.name(),
-                parser.getTokenLocation()
-            );
+        private void unexpectedValue(Block.Builder builder, JsonParser parser, boolean inArray) throws IOException {
+            // Append a null and log the problem
+            if (inArray == false) {
+                // See previous comment about nulls and arrays
+                builder.appendNull();
+            }
+
+            logger.debug("Unexpected token type: {} for attribute: {} at {}", parser.currentToken(), name, parser.getTokenLocation());
             // Ignore any children to keep reading other values
             parser.skipChildren();
         }