Add vortex rewrite: re-encode a file with the current writer

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- `vortex rewrite` re-encodes a file with the current writer. ([#749](https://github.com/dfa1/vortex-java/pull/749))
 - `Chunk.prefetch` and `VortexHandle.prefetch` start the range requests of several columns together over HTTP, merged into few requests, and an HTTP handle keeps the segments it has read. ([#744](https://github.com/dfa1/vortex-java/issues/744))
 
 ### Fixed

cli/src/main/java/io/github/dfa1/vortex/cli/RewriteCommand.java

@@ -0,0 +1,106 @@
+package io.github.dfa1.vortex.cli;
+
+import io.github.dfa1.vortex.core.error.VortexException;
+import io.github.dfa1.vortex.core.model.Edition;
+import io.github.dfa1.vortex.core.model.Editions;
+import io.github.dfa1.vortex.inspect.ByteSize;
+import io.github.dfa1.vortex.reader.VortexHandle;
+import io.github.dfa1.vortex.writer.WriteOptions;
+
+import java.io.IOException;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.ForkJoinPool;
+import java.util.stream.Collectors;
+
+@SuppressWarnings("java:S106") // CLI command: stdout is the intended output channel
+final class RewriteCommand {
+
+    private static final String USAGE = "usage: rewrite [--compact] [--no-zone-maps] [--no-global-dict]"
+            + " [--edition <id>] <file.vortex | http(s)://url> <out.vortex>";
+
+    private RewriteCommand() {
+    }
+
+    static int run(String[] args) {
+        WriteOptions options = WriteOptions.defaults().withExecutor(ForkJoinPool.commonPool());
+        List<String> positional = new ArrayList<>();
+        for (int i = 1; i < args.length; i++) {
+            switch (args[i]) {
+                case "--compact" -> options = options.withCompact(true);
+                case "--no-zone-maps" -> options = options.withZoneMaps(false);
+                case "--no-global-dict" -> options = options.withGlobalDict(false);
+                case "--edition" -> {
+                    Optional<Edition> edition = i + 1 < args.length ? edition(args[i + 1]) : Optional.empty();
+                    if (edition.isEmpty()) {
+                        return unknownEdition();
+                    }
+                    options = options.withEdition(edition.get());
+                    i++;
+                }
+                case String flag when flag.startsWith("--edition=") -> {
+                    Optional<Edition> edition = edition(flag.substring("--edition=".length()));
+                    if (edition.isEmpty()) {
+                        return unknownEdition();
+                    }
+                    options = options.withEdition(edition.get());
+                }
+                default -> positional.add(args[i]);
+            }
+        }
+        if (positional.size() != 2) {
+            System.err.println(USAGE);
+            return ExitStatus.USAGE_ERROR;
+        }
+        Path output = Path.of(positional.get(1));
+        if (Files.exists(output)) {
+            System.err.println("refusing to overwrite: " + output);
+            return ExitStatus.USAGE_ERROR;
+        }
+        try (VortexHandle handle = CliHandles.openTarget(positional.get(0))) {
+            if (handle == null) {
+                return ExitStatus.FILE_NOT_FOUND;
+            }
+            long rows = Rewriter.rewrite(handle, output, options, ProgressBar::render);
+            ProgressBar.clear();
+            System.out.printf("written: %s  (%,d rows, %s)%n", output, rows, ByteSize.format(Files.size(output)));
+            return ExitStatus.OK;
+        } catch (FileAlreadyExistsException e) {
+            // another process created the output after the check above: it is theirs, so leave it
+            ProgressBar.clear();
+            System.err.println("refusing to overwrite: " + output);
+            return ExitStatus.ERROR;
+        } catch (IOException | RuntimeException e) {
+            ProgressBar.clear();
+            System.err.println("error: " + CliHandles.describe(e));
+            if (!(e instanceof VortexException)) {
+                e.printStackTrace(System.err);
+            }
+            deleteQuietly(output);
+            return ExitStatus.ERROR;
+        }
+    }
+
+    private static int unknownEdition() {
+        System.err.println("--edition needs one of: " + Editions.ALL.stream()
+                .map(e -> e.id().toString()).collect(Collectors.joining(", ")));
+        return ExitStatus.USAGE_ERROR;
+    }
+
+    private static Optional<Edition> edition(String id) {
+        return Editions.ALL.stream().filter(e -> e.id().toString().equals(id)).findFirst();
+    }
+
+    /// A failed rewrite leaves a partial file, which a retry would then refuse to overwrite.
+    private static void deleteQuietly(Path output) {
+        try {
+            Files.deleteIfExists(output);
+        } catch (IOException _) {
+            // best effort: the original error is the one worth reporting
+        }
+    }
+}

cli/src/main/java/io/github/dfa1/vortex/cli/Rewriter.java

@@ -0,0 +1,190 @@
+package io.github.dfa1.vortex.cli;
+
+import io.github.dfa1.vortex.core.error.VortexException;
+import io.github.dfa1.vortex.core.model.ColumnName;
+import io.github.dfa1.vortex.core.model.DType;
+import io.github.dfa1.vortex.core.model.PType;
+import io.github.dfa1.vortex.csv.ProgressListener;
+import io.github.dfa1.vortex.reader.Chunk;
+import io.github.dfa1.vortex.reader.ScanIterator;
+import io.github.dfa1.vortex.reader.ScanOptions;
+import io.github.dfa1.vortex.reader.VortexHandle;
+import io.github.dfa1.vortex.reader.array.Array;
+import io.github.dfa1.vortex.reader.array.BoolArray;
+import io.github.dfa1.vortex.reader.array.ByteArray;
+import io.github.dfa1.vortex.reader.array.DoubleArray;
+import io.github.dfa1.vortex.reader.array.FloatArray;
+import io.github.dfa1.vortex.reader.array.IntArray;
+import io.github.dfa1.vortex.reader.array.LongArray;
+import io.github.dfa1.vortex.reader.array.MaskedArray;
+import io.github.dfa1.vortex.reader.array.ShortArray;
+import io.github.dfa1.vortex.reader.array.VarBinArray;
+import io.github.dfa1.vortex.writer.VortexWriter;
+import io.github.dfa1.vortex.writer.WriteOptions;
+import io.github.dfa1.vortex.writer.encode.NullableData;
+
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.function.LongConsumer;
+
+/// Re-encodes a file: scans every chunk of a source and writes it again through the
+/// [VortexWriter], so the output is whatever the current writer and `options` pick, not what the
+/// source's writer picked.
+///
+/// The first version covers a flat struct of primitive, bool and utf8 columns, nullable or not.
+/// Anything else (nested, decimal, binary, extension, f16) is rejected up front, before the
+/// output file is created, rather than written wrongly.
+final class Rewriter {
+
+    private Rewriter() {
+    }
+
+    /// Rewrites `source` into `output`.
+    ///
+    /// @param source   the file to read
+    /// @param output   the file to create; must not exist
+    /// @param options  how to write it
+    /// @param progress receives the rows written so far, and the total
+    /// @return the number of rows written
+    /// @throws IOException      if the output cannot be written
+    /// @throws VortexException  if the source has a column type this version cannot rewrite
+    static long rewrite(VortexHandle source, Path output, WriteOptions options, ProgressListener progress)
+            throws IOException {
+        if (!(source.dtype() instanceof DType.Struct schema)) {
+            throw new VortexException("rewrite needs a struct file, got " + source.dtype());
+        }
+        for (int i = 0; i < schema.fieldNames().size(); i++) {
+            requireSupported(schema.fieldNames().get(i), schema.fieldTypes().get(i));
+        }
+        long total = source.layout().rowCount();
+        long written = 0;
+        try (FileChannel channel = FileChannel.open(output, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE);
+             VortexWriter writer = VortexWriter.create(channel, schema, options);
+             ScanIterator scan = source.scan(ScanOptions.all())) {
+            while (scan.hasNext()) {
+                try (Chunk chunk = scan.next()) {
+                    // over HTTP a column is fetched when read: start all of them together, merged
+                    chunk.prefetch(schema.fieldNames());
+                    Map<ColumnName, Object> columns = new LinkedHashMap<>();
+                    for (Map.Entry<ColumnName, Chunk.Column> column : chunk.columns().entrySet()) {
+                        columns.put(column.getKey(), carrier(column.getKey(), column.getValue()));
+                    }
+                    writer.writeChunk(columns);
+                    written += chunk.rowCount();
+                    progress.onProgress(written, total);
+                }
+            }
+        }
+        if (written != total) {
+            throw new VortexException("rewrote " + written + " rows of " + total + ": the source's layout and chunks disagree");
+        }
+        return written;
+    }
+
+    private static void requireSupported(ColumnName name, DType dtype) {
+        boolean supported = switch (dtype) {
+            case DType.Primitive p -> p.ptype() != PType.F16;
+            case DType.Bool _, DType.Utf8 _ -> true;
+            default -> false;
+        };
+        if (!supported) {
+            throw new VortexException("rewrite does not support column '" + name.value() + "' of type " + dtype);
+        }
+    }
+
+    /// The writer's input for one column: a primitive array when the column is not nullable,
+    /// otherwise a [NullableData] of the values (zero at each null, as the writer requires) and a
+    /// validity per row. Utf8 is the raw bytes, `null` per null row: going through `String` would
+    /// replace any byte sequence that is not valid UTF-8, which the format does not forbid.
+    private static Object carrier(ColumnName name, Chunk.Column column) {
+        try {
+            return carrier(column.array(), column.dtype());
+        } catch (ClassCastException e) {
+            throw new VortexException("rewrite cannot read column '" + name.value() + "': unexpected array "
+                    + column.array().getClass().getSimpleName() + " for " + column.dtype(), e);
+        }
+    }
+
+    private static Object carrier(Array array, DType dtype) {
+        if (array.length() > Integer.MAX_VALUE) {
+            throw new VortexException("rewrite cannot hold a chunk of " + array.length() + " rows");
+        }
+        MaskedArray.Unwrapped unwrapped = MaskedArray.unwrap(array);
+        BoolArray validity = unwrapped.validity();
+        int rows = (int) array.length();
+        Object values = switch (dtype) {
+            case DType.Primitive p -> primitive(unwrapped.inner(), p.ptype(), validity, rows);
+            case DType.Bool _ -> bools(unwrapped.inner(), validity, rows);
+            default -> utf8(unwrapped.inner(), validity, rows);
+        };
+        if (!dtype.nullable() || dtype instanceof DType.Utf8) {
+            return values;
+        }
+        boolean[] present = new boolean[rows];
+        fill(rows, i -> present[(int) i] = valid(validity, i));
+        return new NullableData(values, present);
+    }
+
+    private static Object primitive(Array values, PType ptype, BoolArray validity, int rows) {
+        switch (ptype) {
+            case I8, U8 -> {
+                byte[] out = new byte[rows];
+                fill(rows, i -> out[(int) i] = valid(validity, i) ? ((ByteArray) values).getByte(i) : 0);
+                return out;
+            }
+            case I16, U16 -> {
+                short[] out = new short[rows];
+                fill(rows, i -> out[(int) i] = valid(validity, i) ? ((ShortArray) values).getShort(i) : 0);
+                return out;
+            }
+            case I32, U32 -> {
+                int[] out = new int[rows];
+                fill(rows, i -> out[(int) i] = valid(validity, i) ? ((IntArray) values).getInt(i) : 0);
+                return out;
+            }
+            case I64, U64 -> {
+                long[] out = new long[rows];
+                fill(rows, i -> out[(int) i] = valid(validity, i) ? ((LongArray) values).getLong(i) : 0);
+                return out;
+            }
+            case F32 -> {
+                float[] out = new float[rows];
+                fill(rows, i -> out[(int) i] = valid(validity, i) ? ((FloatArray) values).getFloat(i) : 0);
+                return out;
+            }
+            case F64 -> {
+                double[] out = new double[rows];
+                fill(rows, i -> out[(int) i] = valid(validity, i) ? ((DoubleArray) values).getDouble(i) : 0);
+                return out;
+            }
+            default -> throw new VortexException("rewrite does not support " + ptype);
+        }
+    }
+
+    private static boolean[] bools(Array values, BoolArray validity, int rows) {
+        boolean[] out = new boolean[rows];
+        fill(rows, i -> out[(int) i] = valid(validity, i) && ((BoolArray) values).getBoolean(i));
+        return out;
+    }
+
+    private static byte[][] utf8(Array values, BoolArray validity, int rows) {
+        byte[][] out = new byte[rows][];
+        VarBinArray bytes = (VarBinArray) values;
+        fill(rows, i -> out[(int) i] = valid(validity, i) ? bytes.getBytes(i) : null);
+        return out;
+    }
+
+    private static boolean valid(BoolArray validity, long i) {
+        return validity == null || validity.getBoolean(i);
+    }
+
+    private static void fill(int rows, LongConsumer body) {
+        for (long i = 0; i < rows; i++) {
+            body.accept(i);
+        }
+    }
+}

cli/src/main/java/io/github/dfa1/vortex/cli/VortexCli.java

@@ -68,6 +68,7 @@ private static int dispatch(String[] args) {
             case "select" -> SelectCommand.run(args);
             case "stats" -> StatsCommand.run(args);
             case "filter" -> FilterCommand.run(args);
+            case "rewrite" -> RewriteCommand.run(args);
             default -> {
                 System.err.println("unknown subcommand: " + args[0]);
                 printUsage(System.err);
@@ -92,6 +93,7 @@ static void printUsage(PrintStream out) {
         out.println("  select  <file.vortex> <col> [...] [--where <expr>]  project columns to CSV on stdout, optionally filtered");
         out.println("  stats   <file.vortex>               print per-column min/max statistics");
         out.println("  filter  <file.vortex> <expr>        filter rows to CSV (e.g. \"price >= 100\")");
+        out.println("  rewrite [options] <file.vortex|url> <out.vortex>  re-encode a file (--compact --no-zone-maps --no-global-dict --edition <id>)");
         out.println("Any subcommand also takes --timing: print the elapsed time to stderr, JVM start-up excluded.");
         out.println("Any subcommand also takes --allow-unknown: open a file whose extension types have no decoder,");
         out.println("  reading those columns as their storage.");