test: update python expectations for lance 9.1

Cargo.lock

@@ -3421,8 +3421,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
 
 [[package]]
 name = "fsst"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "rand 0.9.5",
@@ -4777,8 +4777,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
 
 [[package]]
 name = "lance"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arc-swap",
  "arrow",
@@ -4852,8 +4852,8 @@ dependencies = [
 
 [[package]]
 name = "lance-arrow"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -4875,7 +4875,7 @@ dependencies = [
 [[package]]
 name = "lance-arrow-scalar"
 version = "58.0.0"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -4889,7 +4889,7 @@ dependencies = [
 [[package]]
 name = "lance-arrow-stats"
 version = "58.0.0"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "arrow-schema",
@@ -4898,8 +4898,8 @@ dependencies = [
 
 [[package]]
 name = "lance-bitpacking"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrayref",
  "crunchy",
@@ -4909,8 +4909,8 @@ dependencies = [
 
 [[package]]
 name = "lance-core"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -4948,8 +4948,8 @@ dependencies = [
 
 [[package]]
 name = "lance-datafusion"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow",
  "arrow-array",
@@ -4979,8 +4979,8 @@ dependencies = [
 
 [[package]]
 name = "lance-datagen"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow",
  "arrow-array",
@@ -4997,8 +4997,8 @@ dependencies = [
 
 [[package]]
 name = "lance-derive"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -5007,8 +5007,8 @@ dependencies = [
 
 [[package]]
 name = "lance-encoding"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-arith",
  "arrow-array",
@@ -5043,8 +5043,8 @@ dependencies = [
 
 [[package]]
 name = "lance-file"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-arith",
  "arrow-array",
@@ -5074,8 +5074,8 @@ dependencies = [
 
 [[package]]
 name = "lance-index"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arc-swap",
  "arrow",
@@ -5141,8 +5141,8 @@ dependencies = [
 
 [[package]]
 name = "lance-index-core"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "arrow-schema",
@@ -5164,8 +5164,8 @@ dependencies = [
 
 [[package]]
 name = "lance-io"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow",
  "arrow-arith",
@@ -5208,8 +5208,8 @@ dependencies = [
 
 [[package]]
 name = "lance-linalg"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -5225,8 +5225,8 @@ dependencies = [
 
 [[package]]
 name = "lance-namespace"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow",
  "async-trait",
@@ -5238,8 +5238,8 @@ dependencies = [
 
 [[package]]
 name = "lance-namespace-impls"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow",
  "arrow-ipc",
@@ -5293,8 +5293,8 @@ dependencies = [
 
 [[package]]
 name = "lance-select"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -5309,8 +5309,8 @@ dependencies = [
 
 [[package]]
 name = "lance-table"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow",
  "arrow-array",
@@ -5349,8 +5349,8 @@ dependencies = [
 
 [[package]]
 name = "lance-testing"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "arrow-array",
  "arrow-schema",
@@ -5363,8 +5363,8 @@ dependencies = [
 
 [[package]]
 name = "lance-tokenizer"
-version = "9.0.0-rc.1"
+version = "9.1.0-beta.2"
-source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-rc.1#cec0b7dffe2d85c7e66dbe9d1f3891c297903a1d"
+source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
 dependencies = [
  "icu_segmenter",
  "jieba-rs",

Cargo.toml

@@ -13,20 +13,20 @@ categories = ["database-implementations"]
 rust-version = "1.91.0"
 
 [workspace.dependencies]
-lance = { "version" = "=9.0.0-rc.1", default-features = false, "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-core = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-core = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-datagen = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-datagen = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-file = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-file = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-io = { "version" = "=9.0.0-rc.1", default-features = false, "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-io = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-index = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-index = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-linalg = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-linalg = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-namespace = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-namespace = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-namespace-impls = { "version" = "=9.0.0-rc.1", default-features = false, "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-namespace-impls = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-table = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-table = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-testing = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-testing = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-datafusion = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-datafusion = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-encoding = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-encoding = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
-lance-arrow = { "version" = "=9.0.0-rc.1", "tag" = "v9.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
+lance-arrow = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
 ahash = "0.8"
 # Note that this one does not include pyarrow
 arrow = { version = "58.0.0", optional = false }

java/pom.xml

@@ -28,7 +28,7 @@
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <arrow.version>15.0.0</arrow.version>
-        <lance-core.version>9.0.0-rc.1</lance-core.version>
+        <lance-core.version>9.1.0-beta.2</lance-core.version>
         <spotless.skip>false</spotless.skip>
         <spotless.version>2.30.0</spotless.version>
         <spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>

python/python/lancedb/query.py

@@ -3875,18 +3875,18 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
         >>> asyncio.run(doctest_example()) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
         RRFReranker(K=60)
             ProjectionExec: expr=[vector@0 as vector, text@3 as text, _distance@2 as _distance]
-              Take: columns="vector, _rowid, _distance, (text)"
+            Take: columns="vector, _rowid, _distance, (text)"
-                CoalesceBatchesExec: target_batch_size=1024
+              CoalesceBatchesExec: target_batch_size=1024
-                  GlobalLimitExec: skip=0, fetch=10
+                GlobalLimitExec: skip=0, fetch=10
-                    FilterExec: _distance@2 IS NOT NULL
+                  FilterExec: _distance@2 IS NOT NULL
-                      SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
+                    SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
-                        KNNVectorDistance: metric=l2
+                      KNNVectorDistance: metric=l2
-                          LanceRead: uri=..., projection=[vector], ...
+                        LanceRead: uri=..., projection=[vector], ...
             ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score]
-              Take: columns="_rowid, _score, (vector), (text)"
+            Take: columns="_rowid, _score, (vector), (text)"
-                CoalesceBatchesExec: target_batch_size=1024
+              CoalesceBatchesExec: target_batch_size=1024
-                  GlobalLimitExec: skip=0, fetch=10
+                GlobalLimitExec: skip=0, fetch=10
-                    MatchQuery: column=text, query=hello
+                  MatchQuery: column=text, query=[hello]
 
         Parameters
         ----------

python/python/tests/test_permutation.py

@@ -128,9 +128,16 @@ def test_split_hash(mem_db):
 
 def test_split_hash_with_discard(mem_db):
     """Test hash-based splitting with discard weight."""
+    total_rows = 1000
     tbl = mem_db.create_table(
         "test_table",
-        pa.table({"id": range(100), "category": ["A", "B"] * 50, "value": range(100)}),
+        pa.table(
+            {
+                "id": range(total_rows),
+                "category": [f"category-{i}" for i in range(total_rows)],
+                "value": range(total_rows),
+            }
+        ),
     )
 
     permutation_tbl = (
@@ -142,10 +149,12 @@ def test_split_hash_with_discard(mem_db):
         .execute()
     )
 
-    # Should have fewer than 100 rows due to discard
+    # Should have fewer rows due to discard, but should not be empty.
     row_count = permutation_tbl.count_rows()
-    assert row_count < 100
+    assert 0 < row_count < total_rows
-    assert row_count > 0  # But not empty
+
+    data = permutation_tbl.search(None).to_arrow().to_pydict()
+    assert set(data["split_id"]) == {0, 1}
 
 
 def test_split_sequential(mem_db):

python/python/tests/test_query.py

@@ -1273,7 +1273,7 @@ async def test_explain_plan_fts(table_async: AsyncTable):
     query = await table_async.search("dog", query_type="fts", fts_columns="text")
     plan = await query.explain_plan()
     # Should show FTS details (issue #2465 is now fixed)
-    assert "MatchQuery: column=text, query=dog" in plan
+    assert "MatchQuery: column=text, query=[dog]" in plan
     assert "GlobalLimitExec" in plan  # Default limit
 
     # Test FTS query with limit
@@ -1281,7 +1281,7 @@ async def test_explain_plan_fts(table_async: AsyncTable):
         "dog", query_type="fts", fts_columns="text"
     )
     plan_with_limit = await query_with_limit.limit(1).explain_plan()
-    assert "MatchQuery: column=text, query=dog" in plan_with_limit
+    assert "MatchQuery: column=text, query=[dog]" in plan_with_limit
     assert "GlobalLimitExec: skip=0, fetch=1" in plan_with_limit
 
     # Test FTS query with offset and limit
@@ -1289,7 +1289,7 @@ async def test_explain_plan_fts(table_async: AsyncTable):
         "dog", query_type="fts", fts_columns="text"
     )
     plan_with_offset = await query_with_offset.offset(1).limit(1).explain_plan()
-    assert "MatchQuery: column=text, query=dog" in plan_with_offset
+    assert "MatchQuery: column=text, query=[dog]" in plan_with_offset
     assert "GlobalLimitExec: skip=1, fetch=1" in plan_with_offset
 
 
@@ -1333,7 +1333,7 @@ async def test_explain_plan_with_filters(table_async: AsyncTable):
         "dog", query_type="fts", fts_columns="text"
     )
     plan_fts_filter = await query_fts_filter.where("id = 1").explain_plan()
-    assert "MatchQuery: column=text, query=dog" in plan_fts_filter
+    assert "MatchQuery: column=text, query=[dog]" in plan_fts_filter
     assert "LanceRead" in plan_fts_filter
     assert "full_filter=id = Int64(1)" in plan_fts_filter  # Should show filter details
 

rust/lancedb/src/dataloader/permutation/split.rs

@@ -761,8 +761,8 @@ mod tests {
         verify_splitter(splitter, test_data(), 50, &[11, 8, 9], false).await;
     }
 
-    #[tokio::test]
+    async fn collect_hash_split() -> RecordBatch {
-    async fn test_hash_split() {
+        let total_rows = 50;
         let data = lance_datagen::gen_batch()
             .with_seed(Seed::from(42))
             .col(
@@ -783,7 +783,7 @@ mod tests {
         );
 
         let split_batches = splitter
-            .apply(data, 10)
+            .apply(data, total_rows)
             .await
             .unwrap()
             .try_collect::<Vec<_>>()
@@ -791,20 +791,35 @@ mod tests {
             .unwrap();
 
         let schema = split_batches[0].schema();
-        let split_batch = concat_batches(&schema, &split_batches).unwrap();
+        concat_batches(&schema, &split_batches).unwrap()
+    }
 
-        // These assertions are all based on fixed seed in data generation but they match
+    #[tokio::test]
-        // up roughly to what we expect (25% discarded, 25% in split 0, 50% in split 1)
+    async fn test_hash_split() {
+        let total_rows = 50;
+        let split_batch = collect_hash_split().await;
+        let split_batch_again = collect_hash_split().await;
 
-        // 8 rows (16%) are discarded because discard_weight is 1
+        assert_eq!(split_batch.num_rows(), split_batch_again.num_rows());
-        assert_eq!(split_batch.num_rows(), 42);
+        assert_eq!(split_batch.num_columns(), split_batch_again.num_columns());
+        for (left, right) in split_batch
+            .columns()
+            .iter()
+            .zip(split_batch_again.columns().iter())
+        {
+            assert_eq!(left, right);
+        }
+
+        assert!(split_batch.num_rows() > 0);
+        assert!(split_batch.num_rows() < total_rows);
         assert_eq!(split_batch.num_columns(), 2);
 
         let split_ids = split_batch.column(1).as_primitive::<UInt64Type>().values();
         let num_in_split_0 = split_ids.iter().filter(|v| **v == 0).count();
         let num_in_split_1 = split_ids.iter().filter(|v| **v == 1).count();
 
-        assert_eq!(num_in_split_0, 12); // 24%
+        assert_eq!(num_in_split_0 + num_in_split_1, split_batch.num_rows());
-        assert_eq!(num_in_split_1, 30); // 60%
+        assert!(num_in_split_0 > 0);
+        assert!(num_in_split_1 > num_in_split_0);
     }
 }