diff --git a/docs/src/performance.md b/docs/src/performance.md index 47f0aa482..0cc13ef2e 100644 --- a/docs/src/performance.md +++ b/docs/src/performance.md @@ -157,6 +157,43 @@ For large scans, increasing this to match your CPU core count enables more concu export LANCE_IO_THREADS=128 ``` +### Custom Read Metrics + +Lance Spark reports per-task custom metrics on the Spark UI Scan node, viewable in the SQL tab's +physical plan. These are useful for diagnosing read-path performance — pruning effectiveness, JNI +overhead, and where time is being spent. + +Naming conventions: + +- Counters use the `num*` prefix. +- Durations use the `*TimeNs` suffix and are displayed in the UI as formatted strings + (e.g. `1.2 s`, `350 ms`, `47 us`) rather than raw nanoseconds. + +| Metric | Type | Description | +|---|---|---| +| `numFragmentsScanned` | counter | Lance fragments actually opened by this task. Compare against the table fragment count to verify pruning is working. | +| `numBatchesLoaded` | counter | Arrow batches returned from the JNI scanner. | +| `numRowsScanned` | counter | Rows read from storage before filter evaluation. Pair with Spark's built-in `numOutputRows` to compute filter selectivity (`numOutputRows / numRowsScanned`). | +| `datasetOpenTimeNs` | duration | Time spent in `Dataset.open(...)` — manifest load, namespace lookup, credential fetch. High values indicate catalog or metadata cache misses. | +| `scannerCreateTimeNs` | duration | Time spent in `fragment.newScan(...)` — scan planning, predicate compilation, index lookup setup. | +| `batchLoadTimeNs` | duration | Wall-clock time inside `loadNextBatch` — JNI crossing, IO, and Arrow IPC deserialization. Divide by `numBatchesLoaded` to get per-batch cost. | + +How to read these together: + +- **Pushdown verification**: a low `numRowsScanned / numOutputRows` ratio means filters are being + pushed down effectively. A 1:1 ratio means every row is being scanned and filtered in Spark — + check whether your predicate is supported for pushdown. +- **Per-batch JNI cost**: `batchLoadTimeNs / numBatchesLoaded` gives the average cost of one batch + crossing JNI. If this is high relative to total query time, consider increasing + `LanceSparkReadOptions.batchSize` to amortize JNI overhead over more rows. +- **Catalog overhead**: `datasetOpenTimeNs` accumulates per fragment opened. If many fragments are + opened per task, this can dominate; metadata cache size and namespace caching matter most here. + +!!!note + Additional metrics covering bytes read, fragments pruned, index lookups, and IO/decode + time breakdown require new APIs in `lance-jni` and will be added once that surface is + available upstream. + ## Caching Lance Spark uses a multi-level caching strategy to minimize redundant I/O and improve query performance. diff --git a/lance-spark-3.4_2.12/src/test/java/org/lance/spark/read/metric/LanceCustomMetricsTest.java b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/read/metric/LanceCustomMetricsTest.java new file mode 100644 index 000000000..bb2c5b2e9 --- /dev/null +++ b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/read/metric/LanceCustomMetricsTest.java @@ -0,0 +1,16 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +public class LanceCustomMetricsTest extends BaseLanceCustomMetricsTest {} diff --git a/lance-spark-3.4_2.12/src/test/java/org/lance/spark/read/metric/LanceReadMetricsIntegrationTest.java b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/read/metric/LanceReadMetricsIntegrationTest.java new file mode 100644 index 000000000..9b35ac528 --- /dev/null +++ b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/read/metric/LanceReadMetricsIntegrationTest.java @@ -0,0 +1,16 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +public class LanceReadMetricsIntegrationTest extends BaseLanceReadMetricsIntegrationTest {} diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/read/metric/LanceCustomMetricsTest.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/read/metric/LanceCustomMetricsTest.java new file mode 100644 index 000000000..bb2c5b2e9 --- /dev/null +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/read/metric/LanceCustomMetricsTest.java @@ -0,0 +1,16 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +public class LanceCustomMetricsTest extends BaseLanceCustomMetricsTest {} diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/read/metric/LanceReadMetricsIntegrationTest.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/read/metric/LanceReadMetricsIntegrationTest.java new file mode 100644 index 000000000..9b35ac528 --- /dev/null +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/read/metric/LanceReadMetricsIntegrationTest.java @@ -0,0 +1,16 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +public class LanceReadMetricsIntegrationTest extends BaseLanceReadMetricsIntegrationTest {} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentColumnarBatchScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentColumnarBatchScanner.java index 2880e0211..ebc7b81e6 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentColumnarBatchScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentColumnarBatchScanner.java @@ -41,6 +41,7 @@ public class LanceFragmentColumnarBatchScanner implements AutoCloseable { private final LanceFragmentScanner fragmentScanner; private final ArrowReader arrowReader; private ColumnarBatch currentColumnarBatch; + private long lastBatchLoadTimeNs; public LanceFragmentColumnarBatchScanner( LanceFragmentScanner fragmentScanner, ArrowReader arrowReader) { @@ -55,8 +56,13 @@ public static LanceFragmentColumnarBatchScanner create( } public boolean loadNextBatch() throws IOException { - if (arrowReader.loadNextBatch()) { + long start = System.nanoTime(); + boolean hasNext = arrowReader.loadNextBatch(); + lastBatchLoadTimeNs = System.nanoTime() - start; + + if (hasNext) { VectorSchemaRoot root = arrowReader.getVectorSchemaRoot(); + List fieldVectors = root.getFieldVectors().stream() .map(LanceArrowColumnVector::new) @@ -86,6 +92,18 @@ public ColumnarBatch getCurrentBatch() { return currentColumnarBatch; } + public long getLastBatchLoadTimeNs() { + return lastBatchLoadTimeNs; + } + + public long getDatasetOpenTimeNs() { + return fragmentScanner.getDatasetOpenTimeNs(); + } + + public long getScannerCreateTimeNs() { + return fragmentScanner.getScannerCreateTimeNs(); + } + @Override public void close() throws IOException { try { diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java index 90199a9d7..6dd2b06fb 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java @@ -38,22 +38,29 @@ public class LanceFragmentScanner implements AutoCloseable { private final int fragmentId; private final boolean withFragemtId; private final LanceInputPartition inputPartition; + private final long datasetOpenTimeNs; + private final long scannerCreateTimeNs; private LanceFragmentScanner( Dataset dataset, LanceScanner scanner, int fragmentId, boolean withFragmentId, - LanceInputPartition inputPartition) { + LanceInputPartition inputPartition, + long datasetOpenTimeNs, + long scannerCreateTimeNs) { this.dataset = dataset; this.scanner = scanner; this.fragmentId = fragmentId; this.withFragemtId = withFragmentId; this.inputPartition = inputPartition; + this.datasetOpenTimeNs = datasetOpenTimeNs; + this.scannerCreateTimeNs = scannerCreateTimeNs; } public static LanceFragmentScanner create(int fragmentId, LanceInputPartition inputPartition) { Dataset dataset = null; + LanceScanner lanceScanner = null; try { LanceSparkReadOptions readOptions = inputPartition.getReadOptions(); if (inputPartition.getNamespaceImpl() != null) { @@ -61,10 +68,12 @@ public static LanceFragmentScanner create(int fragmentId, LanceInputPartition in LanceRuntime.getOrCreateNamespace( inputPartition.getNamespaceImpl(), inputPartition.getNamespaceProperties())); } + long dsOpenStart = System.nanoTime(); dataset = Utils.openDatasetBuilder(readOptions) .initialStorageOptions(inputPartition.getInitialStorageOptions()) .build(); + long dsOpenTimeNs = System.nanoTime() - dsOpenStart; Fragment fragment = dataset.getFragment(fragmentId); if (fragment == null) { throw new IllegalStateException( @@ -109,13 +118,25 @@ public static LanceFragmentScanner create(int fragmentId, LanceInputPartition in } boolean withFragmentId = inputPartition.getSchema().getFieldIndex(LanceConstant.FRAGMENT_ID).nonEmpty(); + long scanCreateStart = System.nanoTime(); + lanceScanner = fragment.newScan(scanOptions.build()); + long scanCreateTimeNs = System.nanoTime() - scanCreateStart; return new LanceFragmentScanner( dataset, - fragment.newScan(scanOptions.build()), + lanceScanner, fragmentId, withFragmentId, - inputPartition); + inputPartition, + dsOpenTimeNs, + scanCreateTimeNs); } catch (Throwable throwable) { + if (lanceScanner != null) { + try { + lanceScanner.close(); + } catch (Throwable closeError) { + throwable.addSuppressed(closeError); + } + } if (dataset != null) { try { dataset.close(); @@ -181,6 +202,14 @@ public LanceInputPartition getInputPartition() { return inputPartition; } + public long getDatasetOpenTimeNs() { + return datasetOpenTimeNs; + } + + public long getScannerCreateTimeNs() { + return scannerCreateTimeNs; + } + /** * Builds the projection column list for the scanner. Regular data columns come first, followed by * special metadata columns in the order matching {@link diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceColumnarPartitionReader.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceColumnarPartitionReader.java index df0addfd6..ef37756fc 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceColumnarPartitionReader.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceColumnarPartitionReader.java @@ -14,7 +14,9 @@ package org.lance.spark.read; import org.lance.spark.internal.LanceFragmentColumnarBatchScanner; +import org.lance.spark.read.metric.LanceReadMetricsTracker; +import org.apache.spark.sql.connector.metric.CustomTaskMetric; import org.apache.spark.sql.connector.read.PartitionReader; import org.apache.spark.sql.vectorized.ColumnarBatch; @@ -25,6 +27,7 @@ public class LanceColumnarPartitionReader implements PartitionReader fragmentIds = inputPartition.getLanceSplit().getFragments(); if (fragmentIds.isEmpty()) { return 0; } + metricsTracker.addNumFragmentsScanned(fragmentIds.size()); ScanOptions.Builder scanOptionsBuilder = new ScanOptions.Builder(); if (inputPartition.getWhereCondition().isPresent()) { @@ -80,10 +87,23 @@ private long computeCount() { scanOptionsBuilder.withRowId(true); scanOptionsBuilder.columns(Lists.newArrayList()); scanOptionsBuilder.fragmentIds(fragmentIds); + + long scanCreateStart = System.nanoTime(); try (LanceScanner scanner = dataset.newScan(scanOptionsBuilder.build())) { + metricsTracker.addScannerCreateTimeNs(System.nanoTime() - scanCreateStart); try (ArrowReader reader = scanner.scanBatches()) { - while (reader.loadNextBatch()) { - totalCount += reader.getVectorSchemaRoot().getRowCount(); + while (true) { + long batchStart = System.nanoTime(); + boolean hasNext = reader.loadNextBatch(); + long batchTimeNs = System.nanoTime() - batchStart; + if (!hasNext) { + break; + } + metricsTracker.addBatchLoadTimeNs(batchTimeNs); + long rowCount = reader.getVectorSchemaRoot().getRowCount(); + totalCount += rowCount; + metricsTracker.addNumBatchesLoaded(1); + metricsTracker.addNumRowsScanned(rowCount); } } } catch (Exception e) { @@ -118,13 +138,20 @@ private ColumnarBatch createCountResultBatch(long count, StructType resultSchema @Override public ColumnarBatch get() { - long rowCount = computeCount(); - StructType countSchema = - new StructType().add("count", org.apache.spark.sql.types.DataTypes.LongType); - currentBatch = createCountResultBatch(rowCount, countSchema); + if (currentBatch == null) { + long rowCount = computeCount(); + StructType countSchema = + new StructType().add("count", org.apache.spark.sql.types.DataTypes.LongType); + currentBatch = createCountResultBatch(rowCount, countSchema); + } return currentBatch; } + @Override + public CustomTaskMetric[] currentMetricsValues() { + return metricsTracker.currentMetricsValues(); + } + @Override public void close() throws IOException { if (currentBatch != null) { diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceRowPartitionReader.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceRowPartitionReader.java index b9ff09e3c..93ad43190 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceRowPartitionReader.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceRowPartitionReader.java @@ -14,6 +14,7 @@ package org.lance.spark.read; import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.metric.CustomTaskMetric; import org.apache.spark.sql.connector.read.PartitionReader; import org.apache.spark.sql.vectorized.ColumnarBatch; @@ -62,4 +63,9 @@ public InternalRow get() { public void close() throws IOException { reader.close(); } + + @Override + public CustomTaskMetric[] currentMetricsValues() { + return reader.currentMetricsValues(); + } } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScan.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScan.java index 9cfb41f9b..6b6d8d362 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScan.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScan.java @@ -16,6 +16,7 @@ import org.lance.index.scalar.ZoneStats; import org.lance.ipc.ColumnOrdering; import org.lance.spark.LanceSparkReadOptions; +import org.lance.spark.read.metric.LanceCustomMetrics; import org.lance.spark.utils.Optional; import org.apache.arrow.util.Preconditions; @@ -25,6 +26,7 @@ import org.apache.spark.sql.connector.expressions.aggregate.AggregateFunc; import org.apache.spark.sql.connector.expressions.aggregate.Aggregation; import org.apache.spark.sql.connector.expressions.aggregate.CountStar; +import org.apache.spark.sql.connector.metric.CustomMetric; import org.apache.spark.sql.connector.read.Batch; import org.apache.spark.sql.connector.read.InputPartition; import org.apache.spark.sql.connector.read.PartitionReader; @@ -410,6 +412,11 @@ public Statistics estimateStatistics() { return statistics; } + @Override + public CustomMetric[] supportedCustomMetrics() { + return LanceCustomMetrics.allMetrics(); + } + /** * Required for Spark's ReusedExchange: {@code BatchScanExec.equals()} compares {@code batch} * objects, which delegate to this method since LanceScan implements Batch. diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/CustomNsTimeMetric.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/CustomNsTimeMetric.java new file mode 100644 index 000000000..5f03807dd --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/CustomNsTimeMetric.java @@ -0,0 +1,53 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +import org.apache.spark.sql.connector.metric.CustomSumMetric; + +/** + * Base class for custom metrics whose task values are nanoseconds. Aggregation behaves like {@link + * CustomSumMetric} (sum of task values), but {@link #aggregateTaskMetrics(long[])} returns a + * human-readable duration string ("1.2 s", "350 ms", "47 us") instead of a raw long, so the Spark + * UI displays a usable value. Spark's {@code SQLMetrics.stringValue} formatter is {@code + * private[sql]}, so we reimplement it here. + */ +public abstract class CustomNsTimeMetric extends CustomSumMetric { + @Override + public String aggregateTaskMetrics(long[] taskMetrics) { + long total = 0; + for (long v : taskMetrics) { + total += v; + } + return formatDurationNs(total); + } + + /** Visible for testing. Formats nanoseconds as "X ns" / "X us" / "X ms" / "X.Y s". */ + static String formatDurationNs(long ns) { + if (ns < 0) { + ns = 0; + } + if (ns < 1_000L) { + return ns + " ns"; + } + if (ns < 1_000_000L) { + return (ns / 1_000L) + " us"; + } + if (ns < 1_000_000_000L) { + return (ns / 1_000_000L) + " ms"; + } + long secWhole = ns / 1_000_000_000L; + long secTenths = (ns % 1_000_000_000L) / 100_000_000L; + return secWhole + "." + secTenths + " s"; + } +} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/LanceCustomMetrics.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/LanceCustomMetrics.java new file mode 100644 index 000000000..fe85a6377 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/LanceCustomMetrics.java @@ -0,0 +1,133 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +import org.apache.spark.sql.connector.metric.CustomMetric; +import org.apache.spark.sql.connector.metric.CustomSumMetric; + +/** + * Custom metrics for the Lance read path, displayed on the Spark UI Scan node. + * + *

Naming conventions: + * + *

+ * + *

Future metrics gated on upstream {@code lance-jni} surface (not implementable in pure Java + * today): {@code numFragmentsPruned}, {@code bytesRead}, {@code numIndexLookups}, {@code + * ioWaitTimeNs}, {@code decodeTimeNs}. + */ +public final class LanceCustomMetrics { + public static final String NUM_FRAGMENTS_SCANNED = "numFragmentsScanned"; + public static final String NUM_BATCHES_LOADED = "numBatchesLoaded"; + public static final String NUM_ROWS_SCANNED = "numRowsScanned"; + + public static final String DATASET_OPEN_TIME_NS = "datasetOpenTimeNs"; + public static final String SCANNER_CREATE_TIME_NS = "scannerCreateTimeNs"; + public static final String BATCH_LOAD_TIME_NS = "batchLoadTimeNs"; + + private LanceCustomMetrics() {} + + // Each inner class MUST have a public no-arg constructor (Spark instantiates via reflection). + + public static class NumFragmentsScannedMetric extends CustomSumMetric { + @Override + public String name() { + return NUM_FRAGMENTS_SCANNED; + } + + @Override + public String description() { + return "number of Lance fragments scanned"; + } + } + + public static class NumBatchesLoadedMetric extends CustomSumMetric { + @Override + public String name() { + return NUM_BATCHES_LOADED; + } + + @Override + public String description() { + return "number of Arrow batches loaded"; + } + } + + public static class NumRowsScannedMetric extends CustomSumMetric { + @Override + public String name() { + return NUM_ROWS_SCANNED; + } + + @Override + public String description() { + return "number of rows read from storage (before filter evaluation)"; + } + } + + public static class DatasetOpenTimeNsMetric extends CustomNsTimeMetric { + @Override + public String name() { + return DATASET_OPEN_TIME_NS; + } + + @Override + public String description() { + return "time to open Lance dataset"; + } + } + + public static class ScannerCreateTimeNsMetric extends CustomNsTimeMetric { + @Override + public String name() { + return SCANNER_CREATE_TIME_NS; + } + + @Override + public String description() { + return "time to create fragment scanner"; + } + } + + public static class BatchLoadTimeNsMetric extends CustomNsTimeMetric { + @Override + public String name() { + return BATCH_LOAD_TIME_NS; + } + + @Override + public String description() { + return "time to load Arrow batch (JNI + IPC deserialization)"; + } + } + + private static final CustomMetric[] ALL_METRICS = { + new NumFragmentsScannedMetric(), + new NumBatchesLoadedMetric(), + new NumRowsScannedMetric(), + new DatasetOpenTimeNsMetric(), + new ScannerCreateTimeNsMetric(), + new BatchLoadTimeNsMetric(), + }; + + /** Returns all supported custom metrics, used by LanceScan.supportedCustomMetrics(). */ + public static CustomMetric[] allMetrics() { + return ALL_METRICS.clone(); + } +} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/LanceReadMetricsTracker.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/LanceReadMetricsTracker.java new file mode 100644 index 000000000..4e1b9df86 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/metric/LanceReadMetricsTracker.java @@ -0,0 +1,159 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +import org.apache.spark.sql.connector.metric.CustomTaskMetric; + +/** + * Accumulates read-path metrics on the executor side. Thread-confined (one instance per + * PartitionReader, single-threaded access). Returns snapshot values via {@link + * #currentMetricsValues()} — Spark calls this once per {@code next()} invocation on the + * PartitionReader. + * + *

The {@link CustomTaskMetric} array and per-metric instances are allocated once per tracker + * (not per call): each {@code value()} invocation reads the current long field directly, so {@code + * currentMetricsValues()} is allocation-free on the hot path. + */ +public class LanceReadMetricsTracker { + private long numFragmentsScanned; + private long numBatchesLoaded; + private long numRowsScanned; + private long datasetOpenTimeNs; + private long scannerCreateTimeNs; + private long batchLoadTimeNs; + + private final CustomTaskMetric[] taskMetrics = + new CustomTaskMetric[] { + new CustomTaskMetric() { + @Override + public String name() { + return LanceCustomMetrics.NUM_FRAGMENTS_SCANNED; + } + + @Override + public long value() { + return numFragmentsScanned; + } + }, + new CustomTaskMetric() { + @Override + public String name() { + return LanceCustomMetrics.NUM_BATCHES_LOADED; + } + + @Override + public long value() { + return numBatchesLoaded; + } + }, + new CustomTaskMetric() { + @Override + public String name() { + return LanceCustomMetrics.NUM_ROWS_SCANNED; + } + + @Override + public long value() { + return numRowsScanned; + } + }, + new CustomTaskMetric() { + @Override + public String name() { + return LanceCustomMetrics.DATASET_OPEN_TIME_NS; + } + + @Override + public long value() { + return datasetOpenTimeNs; + } + }, + new CustomTaskMetric() { + @Override + public String name() { + return LanceCustomMetrics.SCANNER_CREATE_TIME_NS; + } + + @Override + public long value() { + return scannerCreateTimeNs; + } + }, + new CustomTaskMetric() { + @Override + public String name() { + return LanceCustomMetrics.BATCH_LOAD_TIME_NS; + } + + @Override + public long value() { + return batchLoadTimeNs; + } + }, + }; + + public void addNumFragmentsScanned(long n) { + numFragmentsScanned += n; + } + + public void addNumBatchesLoaded(long n) { + numBatchesLoaded += n; + } + + public void addNumRowsScanned(long n) { + numRowsScanned += n; + } + + public void addDatasetOpenTimeNs(long ns) { + datasetOpenTimeNs += ns; + } + + public void addScannerCreateTimeNs(long ns) { + scannerCreateTimeNs += ns; + } + + public void addBatchLoadTimeNs(long ns) { + batchLoadTimeNs += ns; + } + + /** Returns current snapshot of all metrics. Called by PartitionReader.currentMetricsValues(). */ + public CustomTaskMetric[] currentMetricsValues() { + return taskMetrics; + } + + // Accessors for testing + public long getNumFragmentsScanned() { + return numFragmentsScanned; + } + + public long getNumBatchesLoaded() { + return numBatchesLoaded; + } + + public long getNumRowsScanned() { + return numRowsScanned; + } + + public long getDatasetOpenTimeNs() { + return datasetOpenTimeNs; + } + + public long getScannerCreateTimeNs() { + return scannerCreateTimeNs; + } + + public long getBatchLoadTimeNs() { + return batchLoadTimeNs; + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/metric/BaseLanceCustomMetricsTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/metric/BaseLanceCustomMetricsTest.java new file mode 100644 index 000000000..0ec7088f7 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/metric/BaseLanceCustomMetricsTest.java @@ -0,0 +1,207 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +import org.apache.spark.sql.connector.metric.CustomMetric; +import org.apache.spark.sql.connector.metric.CustomTaskMetric; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public abstract class BaseLanceCustomMetricsTest { + + @Test + void testAllMetricsReturnsSixMetrics() { + CustomMetric[] metrics = LanceCustomMetrics.allMetrics(); + assertEquals(6, metrics.length); + } + + @Test + void testMetricNamesAreUnique() { + CustomMetric[] metrics = LanceCustomMetrics.allMetrics(); + Set names = new HashSet<>(); + for (CustomMetric m : metrics) { + assertTrue(names.add(m.name()), "Duplicate metric name: " + m.name()); + } + } + + @Test + void testMetricNamesMatchConstants() { + Set expected = + new HashSet<>( + Arrays.asList( + LanceCustomMetrics.NUM_FRAGMENTS_SCANNED, + LanceCustomMetrics.NUM_BATCHES_LOADED, + LanceCustomMetrics.NUM_ROWS_SCANNED, + LanceCustomMetrics.DATASET_OPEN_TIME_NS, + LanceCustomMetrics.SCANNER_CREATE_TIME_NS, + LanceCustomMetrics.BATCH_LOAD_TIME_NS)); + CustomMetric[] metrics = LanceCustomMetrics.allMetrics(); + Set actual = new HashSet<>(); + for (CustomMetric m : metrics) { + actual.add(m.name()); + } + assertEquals(expected, actual); + } + + @Test + void testAllMetricsHaveDescriptions() { + for (CustomMetric m : LanceCustomMetrics.allMetrics()) { + assertNotNull(m.description(), "Missing description for " + m.name()); + assertTrue(m.description().length() > 0, "Empty description for " + m.name()); + } + } + + @Test + void testMetricDescriptionsAreUnique() { + CustomMetric[] metrics = LanceCustomMetrics.allMetrics(); + Set descriptions = new HashSet<>(); + for (CustomMetric m : metrics) { + assertTrue( + descriptions.add(m.description()), "Duplicate metric description: " + m.description()); + } + } + + @Test + void testCountAggregationReturnsRawLong() { + LanceCustomMetrics.NumFragmentsScannedMetric metric = + new LanceCustomMetrics.NumFragmentsScannedMetric(); + String result = metric.aggregateTaskMetrics(new long[] {3, 5, 7}); + assertEquals("15", result); + } + + @Test + void testCountAggregationEmpty() { + LanceCustomMetrics.NumFragmentsScannedMetric metric = + new LanceCustomMetrics.NumFragmentsScannedMetric(); + String result = metric.aggregateTaskMetrics(new long[] {}); + assertEquals("0", result); + } + + @Test + void testTimeAggregationFormatsDuration() { + LanceCustomMetrics.BatchLoadTimeNsMetric metric = + new LanceCustomMetrics.BatchLoadTimeNsMetric(); + assertEquals("1.5 s", metric.aggregateTaskMetrics(new long[] {500_000_000L, 1_000_000_000L})); + assertEquals("250 ms", metric.aggregateTaskMetrics(new long[] {100_000_000L, 150_000_000L})); + } + + @Test + void testFormatDurationNsThresholds() { + assertEquals("0 ns", CustomNsTimeMetric.formatDurationNs(0)); + assertEquals("0 ns", CustomNsTimeMetric.formatDurationNs(-5)); + assertEquals("999 ns", CustomNsTimeMetric.formatDurationNs(999)); + assertEquals("1 us", CustomNsTimeMetric.formatDurationNs(1_000)); + assertEquals("47 us", CustomNsTimeMetric.formatDurationNs(47_500)); + assertEquals("999 us", CustomNsTimeMetric.formatDurationNs(999_999)); + assertEquals("1 ms", CustomNsTimeMetric.formatDurationNs(1_000_000)); + assertEquals("350 ms", CustomNsTimeMetric.formatDurationNs(350_000_000)); + assertEquals("999 ms", CustomNsTimeMetric.formatDurationNs(999_999_999)); + assertEquals("1.0 s", CustomNsTimeMetric.formatDurationNs(1_000_000_000L)); + assertEquals("1.2 s", CustomNsTimeMetric.formatDurationNs(1_234_567_890L)); + assertEquals("47.3 s", CustomNsTimeMetric.formatDurationNs(47_382_910_283L)); + } + + @Test + void testNoArgConstructors() throws Exception { + // Spark instantiates metric classes via reflection — verify no-arg constructors work + Class[] metricClasses = + new Class[] { + LanceCustomMetrics.NumFragmentsScannedMetric.class, + LanceCustomMetrics.NumBatchesLoadedMetric.class, + LanceCustomMetrics.NumRowsScannedMetric.class, + LanceCustomMetrics.DatasetOpenTimeNsMetric.class, + LanceCustomMetrics.ScannerCreateTimeNsMetric.class, + LanceCustomMetrics.BatchLoadTimeNsMetric.class, + }; + for (Class clazz : metricClasses) { + Object instance = clazz.getDeclaredConstructor().newInstance(); + assertNotNull(instance, "Failed to instantiate " + clazz.getSimpleName()); + } + } + + @Test + void testTrackerInitialValues() { + LanceReadMetricsTracker tracker = new LanceReadMetricsTracker(); + CustomTaskMetric[] values = tracker.currentMetricsValues(); + assertEquals(6, values.length); + for (CustomTaskMetric m : values) { + assertEquals(0L, m.value(), "Initial value should be 0 for " + m.name()); + } + } + + @Test + void testTrackerAccumulation() { + LanceReadMetricsTracker tracker = new LanceReadMetricsTracker(); + tracker.addNumFragmentsScanned(1); + tracker.addNumFragmentsScanned(1); + tracker.addNumBatchesLoaded(5); + tracker.addNumRowsScanned(1024); + tracker.addDatasetOpenTimeNs(100_000); + tracker.addScannerCreateTimeNs(50_000); + tracker.addBatchLoadTimeNs(200_000); + + assertEquals(2, tracker.getNumFragmentsScanned()); + assertEquals(5, tracker.getNumBatchesLoaded()); + assertEquals(1024, tracker.getNumRowsScanned()); + assertEquals(100_000, tracker.getDatasetOpenTimeNs()); + assertEquals(50_000, tracker.getScannerCreateTimeNs()); + assertEquals(200_000, tracker.getBatchLoadTimeNs()); + } + + @Test + void testTrackerSnapshotReflectsLiveFields() { + // The tracker caches a single CustomTaskMetric[] at construction; each value() call + // must read the current field state, not a captured snapshot. + LanceReadMetricsTracker tracker = new LanceReadMetricsTracker(); + CustomTaskMetric[] before = tracker.currentMetricsValues(); + tracker.addNumFragmentsScanned(7); + CustomTaskMetric[] after = tracker.currentMetricsValues(); + assertSame(before, after, "currentMetricsValues() should return the same cached array"); + long observed = -1; + for (CustomTaskMetric m : after) { + if (m.name().equals(LanceCustomMetrics.NUM_FRAGMENTS_SCANNED)) { + observed = m.value(); + } + } + assertEquals(7L, observed, "value() must read live field, not capture-time value"); + } + + @Test + void testTrackerMetricNamesMatchDefinitions() { + LanceReadMetricsTracker tracker = new LanceReadMetricsTracker(); + CustomTaskMetric[] taskMetrics = tracker.currentMetricsValues(); + CustomMetric[] definitions = LanceCustomMetrics.allMetrics(); + + Set taskNames = new HashSet<>(); + for (CustomTaskMetric m : taskMetrics) { + taskNames.add(m.name()); + } + Set defNames = new HashSet<>(); + for (CustomMetric m : definitions) { + defNames.add(m.name()); + } + assertEquals( + defNames, + taskNames, + "Task metric names must match definition names for Spark to aggregate"); + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/metric/BaseLanceReadMetricsIntegrationTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/metric/BaseLanceReadMetricsIntegrationTest.java new file mode 100644 index 000000000..f76064094 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/metric/BaseLanceReadMetricsIntegrationTest.java @@ -0,0 +1,207 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.read.metric; + +import org.lance.spark.LanceDataSource; +import org.lance.spark.LanceSparkReadOptions; +import org.lance.spark.TestUtils; + +import org.apache.spark.scheduler.SparkListener; +import org.apache.spark.scheduler.SparkListenerTaskEnd; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.connector.metric.CustomMetric; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import scala.collection.JavaConverters; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public abstract class BaseLanceReadMetricsIntegrationTest { + private static SparkSession spark; + private static Dataset data; + private static MetricsCapturingListener metricsListener; + + // Spark registers custom metric accumulators using description() as the name, + // so we build a reverse lookup: description -> metric name. + private static final Map DESC_TO_NAME = new HashMap<>(); + + static { + for (CustomMetric m : LanceCustomMetrics.allMetrics()) { + DESC_TO_NAME.put(m.description(), m.name()); + } + } + + /** + * SparkListener that captures custom metric accumulator values from task-end events. Keyed by + * metric name (e.g. "numFragmentsScanned"), translated from the description that Spark uses as + * the accumulator name. + */ + static class MetricsCapturingListener extends SparkListener { + private final Map metricValues = new ConcurrentHashMap<>(); + + void reset() { + metricValues.clear(); + } + + Map getMetricValues() { + return metricValues; + } + + @Override + public void onTaskEnd(SparkListenerTaskEnd taskEnd) { + if (taskEnd.taskInfo() != null) { + JavaConverters.seqAsJavaList(taskEnd.taskInfo().accumulables()) + .forEach( + info -> { + if (info.name().isDefined() && info.update().isDefined()) { + String desc = info.name().get(); + String metricName = DESC_TO_NAME.get(desc); + if (metricName != null) { + Object value = info.update().get(); + if (value instanceof Long) { + metricValues.merge(metricName, (Long) value, Long::sum); + } + } + } + }); + } + } + } + + @BeforeAll + static void setup() { + spark = + SparkSession.builder() + .appName("lance-metrics-test") + .master("local") + .config("spark.sql.catalog.lance", "org.lance.spark.LanceNamespaceSparkCatalog") + .getOrCreate(); + metricsListener = new MetricsCapturingListener(); + spark.sparkContext().addSparkListener(metricsListener); + data = + spark + .read() + .format(LanceDataSource.name) + .option( + LanceSparkReadOptions.CONFIG_DATASET_URI, + TestUtils.getDatasetUri( + TestUtils.TestTable1Config.dbPath, TestUtils.TestTable1Config.datasetName)) + .load(); + } + + @AfterAll + static void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + @Test + void testSupportedCustomMetricsCount() { + CustomMetric[] metrics = LanceCustomMetrics.allMetrics(); + assertEquals(6, metrics.length); + } + + @Test + void testReadQueryCompletes() { + // Execute a read query and collect results to verify metrics don't break normal reading + Dataset result = data.select("x", "y"); + long count = result.count(); + assertTrue(count > 0, "Should read at least one row"); + } + + @Test + void testCountStarWithFilterCompletes() { + // COUNT(*) with filter goes through LanceCountStarPartitionReader + data.createOrReplaceTempView("metrics_test"); + Dataset result = spark.sql("SELECT COUNT(*) FROM metrics_test WHERE x > 0"); + Row row = result.collectAsList().get(0); + long count = row.getLong(0); + assertTrue(count > 0, "Filtered count should be > 0"); + } + + @Test + void testSelectAllColumnsCompletes() { + // Full row read exercises the columnar partition reader metrics path + java.util.List rows = data.collectAsList(); + assertTrue(rows.size() > 0, "Should collect at least one row"); + } + + @Test + void testReadMetricsValues() throws Exception { + metricsListener.reset(); + // Execute a full read query to trigger columnar partition reader metrics + data.select("x", "y").collect(); + // Allow listener to process events + spark.sparkContext().listenerBus().waitUntilEmpty(5000); + + Map metrics = metricsListener.getMetricValues(); + + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.NUM_FRAGMENTS_SCANNED, 0L) > 0, + "numFragmentsScanned should be > 0"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.NUM_BATCHES_LOADED, 0L) >= 1, + "numBatchesLoaded should be >= 1"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.NUM_ROWS_SCANNED, 0L) > 0, + "numRowsScanned should be > 0"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.DATASET_OPEN_TIME_NS, 0L) > 0, + "datasetOpenTimeNs should be > 0"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.SCANNER_CREATE_TIME_NS, 0L) > 0, + "scannerCreateTimeNs should be > 0"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.BATCH_LOAD_TIME_NS, 0L) > 0, + "batchLoadTimeNs should be > 0"); + } + + @Test + void testCountStarMetricsValues() throws Exception { + metricsListener.reset(); + data.createOrReplaceTempView("metrics_count_test"); + spark.sql("SELECT COUNT(*) FROM metrics_count_test WHERE x > 0").collect(); + spark.sparkContext().listenerBus().waitUntilEmpty(5000); + + Map metrics = metricsListener.getMetricValues(); + + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.NUM_FRAGMENTS_SCANNED, 0L) > 0, + "numFragmentsScanned should be > 0 for COUNT(*)"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.NUM_BATCHES_LOADED, 0L) >= 1, + "numBatchesLoaded should be >= 1 for COUNT(*)"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.NUM_ROWS_SCANNED, 0L) > 0, + "numRowsScanned should be > 0 for COUNT(*)"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.DATASET_OPEN_TIME_NS, 0L) > 0, + "datasetOpenTimeNs should be > 0 for COUNT(*)"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.SCANNER_CREATE_TIME_NS, 0L) > 0, + "scannerCreateTimeNs should be > 0 for COUNT(*)"); + assertTrue( + metrics.getOrDefault(LanceCustomMetrics.BATCH_LOAD_TIME_NS, 0L) > 0, + "batchLoadTimeNs should be > 0 for COUNT(*)"); + } +}