/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.apache.datafusion;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.nio.file.Files;
import java.nio.file.Path;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowReader;
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.Test;

class DataFrameTransformationsTest {
  @Test
  void countReturnsRowCount() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT * FROM (VALUES (1), (2), (3)) AS t(x)")) {
      assertEquals(3L, df.count());
    }
  }

  @Test
  void showDoesNotThrow() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT * FROM (VALUES (1), (2)) AS t(x)")) {
      df.show();
    }
  }

  @Test
  void showWithLimitDoesNotThrow() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT * FROM (VALUES (1), (2), (3)) AS t(x)")) {
      df.show(0);
      df.show(1);
      df.show(1_000_000);
    }
  }

  @Test
  void selectProjectsAndReordersColumns() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b, 3 AS c");
        DataFrame projected = source.select("b", "a");
        ArrowReader reader = projected.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      assertEquals(1, root.getRowCount());
      assertArrayEquals(
          new String[] {"b", "a"},
          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
    }
  }

  @Test
  void selectIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b")) {
      try (DataFrame first = source.select("a")) {
        assertEquals(1L, first.count());
      }
      try (DataFrame second = source.select("b")) {
        assertEquals(1L, second.count());
      }
      assertEquals(1L, source.count());
    }
  }

  @Test
  void filterRemovesRows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3), (4)) AS t(x)");
        DataFrame filtered = source.filter("x > 2")) {
      assertEquals(2L, filtered.count());
    }
  }

  @Test
  void filterIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3), (4)) AS t(x)")) {
      try (DataFrame filtered = source.filter("x > 2")) {
        assertEquals(2L, filtered.count());
      }
      assertEquals(4L, source.count());
    }
  }

  @Test
  void chainFilterSelectCount() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b UNION ALL SELECT 10 AS a, 20 AS b");
        DataFrame chained = source.filter("a > 5").select("b")) {
      assertEquals(1L, chained.count());
    }
  }

  @Test
  void methodsThrowAfterClose() {
    try (SessionContext ctx = new SessionContext()) {
      DataFrame df = ctx.sql("SELECT 1 AS x");
      df.close();
      assertThrows(IllegalStateException.class, () -> df.select("x"));
      assertThrows(IllegalStateException.class, () -> df.filter("x > 0"));
      assertThrows(IllegalStateException.class, () -> df.limit(1));
      assertThrows(IllegalStateException.class, () -> df.limit(0, 1));
      assertThrows(IllegalStateException.class, df::distinct);
      assertThrows(IllegalStateException.class, () -> df.dropColumns("x"));
      assertThrows(IllegalStateException.class, () -> df.withColumnRenamed("x", "y"));
      assertThrows(IllegalStateException.class, () -> df.withColumn("y", "x + 1"));
      assertThrows(IllegalStateException.class, () -> df.unnestColumns("x"));
      assertThrows(IllegalStateException.class, () -> df.sort(SortExpr.asc("x")));
      assertThrows(IllegalStateException.class, () -> df.repartitionRoundRobin(1));
      assertThrows(IllegalStateException.class, () -> df.repartitionHash(1, "x"));
      assertThrows(IllegalStateException.class, df::count);
      assertThrows(IllegalStateException.class, df::show);
      assertThrows(IllegalStateException.class, () -> df.show(5));
    }
  }

  @Test
  void methodsThrowAfterCollect() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      try (ArrowReader reader = df.collect(allocator)) {
        assertTrue(reader.loadNextBatch());
      }
      assertThrows(IllegalStateException.class, () -> df.select("x"));
      assertThrows(IllegalStateException.class, () -> df.filter("x > 0"));
      assertThrows(IllegalStateException.class, () -> df.limit(1));
      assertThrows(IllegalStateException.class, () -> df.limit(0, 1));
      assertThrows(IllegalStateException.class, df::distinct);
      assertThrows(IllegalStateException.class, () -> df.dropColumns("x"));
      assertThrows(IllegalStateException.class, () -> df.withColumnRenamed("x", "y"));
      assertThrows(IllegalStateException.class, () -> df.withColumn("y", "x + 1"));
      assertThrows(IllegalStateException.class, () -> df.unnestColumns("x"));
      assertThrows(IllegalStateException.class, () -> df.sort(SortExpr.asc("x")));
      assertThrows(IllegalStateException.class, () -> df.repartitionRoundRobin(1));
      assertThrows(IllegalStateException.class, () -> df.repartitionHash(1, "x"));
      assertThrows(IllegalStateException.class, df::count);
      assertThrows(IllegalStateException.class, df::show);
      assertThrows(IllegalStateException.class, () -> df.show(5));
    }
  }

  @Test
  void selectInvalidColumnThrows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(RuntimeException.class, () -> df.select("not_a_column"));
    }
  }

  @Test
  void filterMalformedPredicateThrows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(RuntimeException.class, () -> df.filter("this is not sql"));
    }
  }

  @Test
  void lineitemFilterCountAgainstSqlBaseline() throws Exception {
    Path lineitem = Path.of("tpch-data/sf1/lineitem.parquet");
    Assumptions.assumeTrue(
        Files.exists(lineitem), "TPC-H SF1 data not found; run `make tpch-data` first");

    try (SessionContext ctx = new SessionContext()) {
      ctx.registerParquet("lineitem", lineitem.toAbsolutePath().toString());

      long viaDataFrame;
      try (DataFrame df = ctx.sql("SELECT * FROM lineitem");
          DataFrame filtered = df.filter("l_orderkey < 100")) {
        viaDataFrame = filtered.count();
      }

      long viaSql;
      try (BufferAllocator allocator = new RootAllocator();
          DataFrame df = ctx.sql("SELECT COUNT(*) FROM lineitem WHERE l_orderkey < 100");
          ArrowReader reader = df.collect(allocator)) {
        assertTrue(reader.loadNextBatch());
        VectorSchemaRoot root = reader.getVectorSchemaRoot();
        viaSql = ((BigIntVector) root.getVector(0)).get(0);
      }

      assertEquals(viaSql, viaDataFrame);
    }
  }

  @Test
  void limitTakesFirstNRows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3), (4), (5)) AS t(x)");
        DataFrame limited = source.limit(2)) {
      assertEquals(2L, limited.count());
    }
  }

  @Test
  void limitWithSkipDropsLeadingRows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3), (4), (5)) AS t(x)");
        DataFrame limited = source.limit(2, 2)) {
      assertEquals(2L, limited.count());
    }
  }

  @Test
  void limitIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3)) AS t(x)")) {
      try (DataFrame limited = source.limit(1)) {
        assertEquals(1L, limited.count());
      }
      assertEquals(3L, source.count());
    }
  }

  @Test
  void limitRejectsNegativeArgs() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(IllegalArgumentException.class, () -> df.limit(-1));
      assertThrows(IllegalArgumentException.class, () -> df.limit(-1, 0));
      assertThrows(IllegalArgumentException.class, () -> df.limit(0, -1));
    }
  }

  @Test
  void distinctRemovesDuplicates() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (1), (2), (2), (3)) AS t(x)");
        DataFrame deduped = source.distinct()) {
      assertEquals(3L, deduped.count());
    }
  }

  @Test
  void distinctIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (1), (2)) AS t(x)")) {
      try (DataFrame deduped = source.distinct()) {
        assertEquals(2L, deduped.count());
      }
      assertEquals(3L, source.count());
    }
  }

  @Test
  void dropColumnsRemovesNamedColumns() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b, 3 AS c");
        DataFrame dropped = source.dropColumns("b");
        ArrowReader reader = dropped.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      assertArrayEquals(
          new String[] {"a", "c"},
          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
    }
  }

  @Test
  void dropColumnsIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b")) {
      try (DataFrame dropped = source.dropColumns("a")) {
        assertEquals(1L, dropped.count());
      }
      assertEquals(1L, source.count());
    }
  }

  @Test
  void dropColumnsSilentlyIgnoresUnknownNames() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x");
        DataFrame dropped = df.dropColumns("not_a_column");
        ArrowReader reader = dropped.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      assertArrayEquals(
          new String[] {"x"},
          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
    }
  }

  @Test
  void withColumnRenamedChangesColumnName() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b");
        DataFrame renamed = source.withColumnRenamed("a", "alpha");
        ArrowReader reader = renamed.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      assertArrayEquals(
          new String[] {"alpha", "b"},
          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
    }
  }

  @Test
  void withColumnRenamedIsNonDestructive() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b")) {
      try (DataFrame renamed = source.withColumnRenamed("a", "alpha")) {
        assertEquals(1L, renamed.count());
      }
      try (DataFrame again = source.select("a");
          ArrowReader reader = again.collect(allocator)) {
        assertTrue(reader.loadNextBatch());
        VectorSchemaRoot root = reader.getVectorSchemaRoot();
        assertArrayEquals(
            new String[] {"a"},
            root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
      }
    }
  }

  @Test
  void withColumnRenamedUnknownColumnIsNoOp() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x");
        DataFrame renamed = df.withColumnRenamed("not_a_column", "y");
        ArrowReader reader = renamed.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      assertArrayEquals(
          new String[] {"x"},
          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
    }
  }

  @Test
  void withColumnAppendsNewColumn() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b");
        DataFrame extended = source.withColumn("c", "a + b");
        ArrowReader reader = extended.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      assertEquals(1, root.getRowCount());
      assertArrayEquals(
          new String[] {"a", "b", "c"},
          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
      assertEquals(3L, ((BigIntVector) root.getVector("c")).get(0));
    }
  }

  @Test
  void withColumnReplacesExistingColumn() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b");
        DataFrame replaced = source.withColumn("b", "a * 10");
        ArrowReader reader = replaced.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      assertArrayEquals(
          new String[] {"a", "b"},
          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
      assertEquals(10L, ((BigIntVector) root.getVector("b")).get(0));
    }
  }

  @Test
  void withColumnIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b")) {
      try (DataFrame extended = source.withColumn("c", "a + b")) {
        assertEquals(1L, extended.count());
      }
      assertEquals(1L, source.count());
    }
  }

  @Test
  void withColumnUnknownColumnRefThrows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(RuntimeException.class, () -> df.withColumn("y", "not_a_column + 1"));
    }
  }

  @Test
  void withColumnRejectsNullArgs() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(IllegalArgumentException.class, () -> df.withColumn(null, "x + 1"));
      assertThrows(IllegalArgumentException.class, () -> df.withColumn("y", null));
    }
  }

  @Test
  void unnestColumnsExpandsList() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS id, [10, 20, 30] AS vals");
        DataFrame unnested = source.unnestColumns("vals");
        ArrowReader reader = unnested.collect(allocator)) {
      int rows = 0;
      while (reader.loadNextBatch()) {
        rows += reader.getVectorSchemaRoot().getRowCount();
      }
      assertEquals(3, rows);
    }
  }

  @Test
  void unnestColumnsPreserveNullsTrueKeepsNullRow() throws Exception {
    String sql =
        "SELECT id, vals FROM (VALUES "
            + "(1, [10, 20]), "
            + "(2, CAST(NULL AS BIGINT[])), "
            + "(3, [30])) AS t(id, vals)";
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql(sql);
        DataFrame unnested = source.unnestColumns(new UnnestOptions().preserveNulls(true), "vals");
        ArrowReader reader = unnested.collect(allocator)) {
      int rows = 0;
      while (reader.loadNextBatch()) {
        rows += reader.getVectorSchemaRoot().getRowCount();
      }
      assertEquals(4, rows);
    }
  }

  @Test
  void unnestColumnsPreserveNullsFalseDropsNullRow() throws Exception {
    String sql =
        "SELECT id, vals FROM (VALUES "
            + "(1, [10, 20]), "
            + "(2, CAST(NULL AS BIGINT[])), "
            + "(3, [30])) AS t(id, vals)";
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql(sql);
        DataFrame unnested =
            source.unnestColumns(new UnnestOptions().preserveNulls(false), "vals");
        ArrowReader reader = unnested.collect(allocator)) {
      int rows = 0;
      while (reader.loadNextBatch()) {
        rows += reader.getVectorSchemaRoot().getRowCount();
      }
      assertEquals(3, rows);
    }
  }

  @Test
  void unnestColumnsIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT 1 AS id, [10, 20] AS vals")) {
      try (DataFrame unnested = source.unnestColumns("vals")) {
        assertEquals(2L, unnested.count());
      }
      assertEquals(1L, source.count());
    }
  }

  @Test
  void unnestColumnsUnknownColumnThrows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(RuntimeException.class, () -> df.unnestColumns("not_a_column"));
    }
  }

  @Test
  void unnestColumnsRejectsNullArgs() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x, [1, 2] AS vals")) {
      assertThrows(
          IllegalArgumentException.class, () -> df.unnestColumns((UnnestOptions) null, "vals"));
      assertThrows(
          IllegalArgumentException.class,
          () -> df.unnestColumns(new UnnestOptions(), (String[]) null));
    }
  }

  // -- sort -----------------------------------------------------------------

  @Test
  void sortAscOrdersAscending() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT * FROM (VALUES (3), (1), (2)) AS t(x)");
        DataFrame sorted = df.sort(SortExpr.asc("x"));
        ArrowReader reader = sorted.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      BigIntVector x = (BigIntVector) root.getVector("x");
      assertEquals(3, x.getValueCount());
      assertEquals(1L, x.get(0));
      assertEquals(2L, x.get(1));
      assertEquals(3L, x.get(2));
    }
  }

  @Test
  void sortDescOrdersDescending() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT * FROM (VALUES (3), (1), (2)) AS t(x)");
        DataFrame sorted = df.sort(SortExpr.desc("x"));
        ArrowReader reader = sorted.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      BigIntVector x = (BigIntVector) root.getVector("x");
      assertEquals(3, x.getValueCount());
      assertEquals(3L, x.get(0));
      assertEquals(2L, x.get(1));
      assertEquals(1L, x.get(2));
    }
  }

  @Test
  void sortMultipleKeysMixedDirection() throws Exception {
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT * FROM (VALUES (1, 1), (1, 2), (2, 1)) AS t(x, y)");
        DataFrame sorted = df.sort(SortExpr.asc("x"), SortExpr.desc("y"));
        ArrowReader reader = sorted.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      BigIntVector x = (BigIntVector) root.getVector("x");
      BigIntVector y = (BigIntVector) root.getVector("y");
      // (1,2), (1,1), (2,1) -- x asc, y desc within each x group.
      assertEquals(1L, x.get(0));
      assertEquals(2L, y.get(0));
      assertEquals(1L, x.get(1));
      assertEquals(1L, y.get(1));
      assertEquals(2L, x.get(2));
      assertEquals(1L, y.get(2));
    }
  }

  @Test
  void sortNullsFirstPlacesNullsFirst() throws Exception {
    String sql =
        "SELECT * FROM (VALUES (CAST(1 AS BIGINT)), (CAST(NULL AS BIGINT)),"
            + " (CAST(2 AS BIGINT))) AS t(x)";
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql(sql);
        DataFrame sorted = df.sort(SortExpr.asc("x").nullsFirst(true));
        ArrowReader reader = sorted.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      BigIntVector x = (BigIntVector) root.getVector("x");
      assertEquals(3, x.getValueCount());
      assertTrue(x.isNull(0), "expected NULL first");
      assertEquals(1L, x.get(1));
      assertEquals(2L, x.get(2));
    }
  }

  @Test
  void sortNullsLastPlacesNullsLast() throws Exception {
    String sql =
        "SELECT * FROM (VALUES (CAST(1 AS BIGINT)), (CAST(NULL AS BIGINT)),"
            + " (CAST(2 AS BIGINT))) AS t(x)";
    try (BufferAllocator allocator = new RootAllocator();
        SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql(sql);
        DataFrame sorted = df.sort(SortExpr.asc("x").nullsFirst(false));
        ArrowReader reader = sorted.collect(allocator)) {
      assertTrue(reader.loadNextBatch());
      VectorSchemaRoot root = reader.getVectorSchemaRoot();
      BigIntVector x = (BigIntVector) root.getVector("x");
      assertEquals(3, x.getValueCount());
      assertEquals(1L, x.get(0));
      assertEquals(2L, x.get(1));
      assertTrue(x.isNull(2), "expected NULL last");
    }
  }

  @Test
  void sortIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (3), (1), (2)) AS t(x)")) {
      try (DataFrame sorted = source.sort(SortExpr.asc("x"))) {
        assertEquals(3L, sorted.count());
      }
      assertEquals(3L, source.count());
    }
  }

  @Test
  void sortRejectsNullArgs() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(IllegalArgumentException.class, () -> df.sort((SortExpr[]) null));
      assertThrows(IllegalArgumentException.class, () -> df.sort(new SortExpr[] {null}));
    }
  }

  @Test
  void sortUnknownColumnThrows() {
    // Sort plan-builds lazily; DataFusion may not validate column names until the plan is
    // actually executed. Drive a count() to force planning + optimization.
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(
          RuntimeException.class,
          () -> {
            try (DataFrame sorted = df.sort(SortExpr.asc("not_a_column"))) {
              sorted.count();
            }
          });
    }
  }

  // -- repartitionRoundRobin -------------------------------------------------

  @Test
  void repartitionRoundRobinPreservesRows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3), (4), (5)) AS t(x)");
        DataFrame repartitioned = source.repartitionRoundRobin(3)) {
      assertEquals(5L, repartitioned.count());
    }
  }

  @Test
  void repartitionRoundRobinIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3)) AS t(x)")) {
      try (DataFrame rp = source.repartitionRoundRobin(2)) {
        assertEquals(3L, rp.count());
      }
      assertEquals(3L, source.count());
    }
  }

  @Test
  void repartitionRoundRobinRejectsNonPositive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(IllegalArgumentException.class, () -> df.repartitionRoundRobin(0));
      assertThrows(IllegalArgumentException.class, () -> df.repartitionRoundRobin(-1));
    }
  }

  // -- repartitionHash -------------------------------------------------------

  @Test
  void repartitionHashPreservesRows() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3), (4), (5)) AS t(x)");
        DataFrame repartitioned = source.repartitionHash(3, "x")) {
      assertEquals(5L, repartitioned.count());
    }
  }

  @Test
  void repartitionHashIsNonDestructive() {
    try (SessionContext ctx = new SessionContext();
        DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3)) AS t(x)")) {
      try (DataFrame rp = source.repartitionHash(2, "x")) {
        assertEquals(3L, rp.count());
      }
      assertEquals(3L, source.count());
    }
  }

  @Test
  void repartitionHashRejectsBadArgs() {
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(IllegalArgumentException.class, () -> df.repartitionHash(0, "x"));
      assertThrows(IllegalArgumentException.class, () -> df.repartitionHash(-1, "x"));
      assertThrows(IllegalArgumentException.class, () -> df.repartitionHash(2, (String[]) null));
      assertThrows(IllegalArgumentException.class, () -> df.repartitionHash(2));
      assertThrows(
          IllegalArgumentException.class, () -> df.repartitionHash(2, new String[] {null}));
    }
  }

  @Test
  void repartitionHashUnknownColumnThrows() {
    // Repartition plan-builds lazily; DataFusion may not validate column names until the plan
    // is actually executed. Drive a count() to force planning + optimization.
    try (SessionContext ctx = new SessionContext();
        DataFrame df = ctx.sql("SELECT 1 AS x")) {
      assertThrows(
          RuntimeException.class,
          () -> {
            try (DataFrame rp = df.repartitionHash(2, "not_a_column")) {
              rp.count();
            }
          });
    }
  }
}
