| # 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. |
| |
| import io |
| import struct |
| |
| import pyarrow as pa |
| import pytest |
| |
| from mosaic import ( |
| ColumnStatistics, |
| MosaicReader, |
| MosaicWriter, |
| WriterOptions, |
| read_table, |
| write_table, |
| ) |
| |
| |
| def _write_to_bytes(pa_schema, data, options=None): |
| buf = io.BytesIO() |
| with MosaicWriter(buf, pa_schema, options) as writer: |
| writer.write(data) |
| return buf.getvalue() |
| |
| |
| def _reader_from_bytes(data): |
| return MosaicReader.from_input_file( |
| lambda offset, length: data[offset : offset + length], len(data) |
| ) |
| |
| |
| class TestRoundtrip: |
| def test_basic_roundtrip(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32(), nullable=False), |
| pa.field("name", pa.utf8()), |
| pa.field("score", pa.float64()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(50)), type=pa.int32()), |
| pa.array([f"user_{i}" for i in range(50)]), |
| pa.array([i * 1.5 for i in range(50)]), |
| ], |
| names=["id", "name", "score"], |
| ) |
| |
| data = _write_to_bytes(pa_schema, batch) |
| assert len(data) > 32 |
| assert data[-4:] == b"MOSA" |
| |
| with _reader_from_bytes(data) as reader: |
| assert reader.num_row_groups >= 1 |
| |
| total_rows = 0 |
| for rg in range(reader.num_row_groups): |
| rb = reader.read_row_group(rg) |
| total_rows += rb.num_rows |
| |
| ids = rb.column("id").to_pylist() |
| names = rb.column("name").to_pylist() |
| scores = rb.column("score").to_pylist() |
| |
| for j in range(rb.num_rows): |
| idx = ids[j] |
| assert names[j] == f"user_{idx}" |
| assert abs(scores[j] - idx * 1.5) < 1e-9 |
| |
| assert total_rows == 50 |
| |
| def test_null_values(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("name", pa.utf8()), |
| pa.field("value", pa.float64()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array([1, 2, 3, 4], type=pa.int32()), |
| pa.array(["hello", None, "world", None]), |
| pa.array([1.0, 2.0, None, None]), |
| ], |
| names=["id", "name", "value"], |
| ) |
| |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| rb = reader.read_row_group(0) |
| assert rb.num_rows == 4 |
| |
| names = rb.column("name") |
| values = rb.column("value") |
| |
| assert names[0].as_py() == "hello" |
| assert names[1].as_py() is None |
| assert names[2].as_py() == "world" |
| assert names[3].as_py() is None |
| |
| assert values[0].as_py() == 1.0 |
| assert values[1].as_py() == 2.0 |
| assert values[2].as_py() is None |
| assert values[3].as_py() is None |
| |
| def test_compression_none(self): |
| pa_schema = pa.schema( |
| [pa.field("x", pa.int32()), pa.field("y", pa.utf8())] |
| ) |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(20)), type=pa.int32()), |
| pa.array([f"v_{i}" for i in range(20)]), |
| ], |
| names=["x", "y"], |
| ) |
| opts = WriterOptions(compression=WriterOptions.COMPRESSION_NONE) |
| data = _write_to_bytes(pa_schema, batch, opts) |
| |
| with _reader_from_bytes(data) as reader: |
| rb = reader.read_row_group(0) |
| assert rb.num_rows == 20 |
| assert rb.column("x").to_pylist() == list(range(20)) |
| |
| def test_compression_zstd(self): |
| pa_schema = pa.schema( |
| [pa.field("x", pa.int32()), pa.field("y", pa.utf8())] |
| ) |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(100)), type=pa.int32()), |
| pa.array([f"v_{i}" for i in range(100)]), |
| ], |
| names=["x", "y"], |
| ) |
| opts = WriterOptions(compression=WriterOptions.COMPRESSION_ZSTD, zstd_level=3) |
| data = _write_to_bytes(pa_schema, batch, opts) |
| |
| with _reader_from_bytes(data) as reader: |
| rb = reader.read_row_group(0) |
| assert rb.num_rows == 100 |
| assert rb.column("x").to_pylist() == list(range(100)) |
| |
| def test_all_types(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("f_bool", pa.bool_()), |
| pa.field("f_int8", pa.int8()), |
| pa.field("f_int16", pa.int16()), |
| pa.field("f_int32", pa.int32()), |
| pa.field("f_int64", pa.int64()), |
| pa.field("f_float32", pa.float32()), |
| pa.field("f_float64", pa.float64()), |
| pa.field("f_utf8", pa.utf8()), |
| pa.field("f_binary", pa.binary()), |
| pa.field("f_decimal", pa.decimal128(10, 2)), |
| pa.field("f_date", pa.date32()), |
| pa.field("f_timestamp", pa.timestamp("ms")), |
| pa.field("f_timestamp_ns", pa.timestamp("ns")), |
| pa.field("f_timestamp_ns_tz", pa.timestamp("ns", tz="Asia/Shanghai")), |
| ] |
| ) |
| |
| ts_ns_values = [1700000000000000123, -1] |
| batch = pa.record_batch( |
| [ |
| pa.array([True, False]), |
| pa.array([42, -1], type=pa.int8()), |
| pa.array([1234, -5678], type=pa.int16()), |
| pa.array([100000, -200000], type=pa.int32()), |
| pa.array([9999999999, -9999999999], type=pa.int64()), |
| pa.array([3.14, -2.71], type=pa.float32()), |
| pa.array([2.718281828, -3.141592653]), |
| pa.array(["hello", "world"]), |
| pa.array([b"\x01\x02\x03", b"\xff\x00"], type=pa.binary()), |
| pa.array([1234567, -9876543], type=pa.decimal128(10, 2)), |
| pa.array([19000, 0], type=pa.date32()), |
| pa.array([1700000000000, 0], type=pa.timestamp("ms")), |
| pa.array(ts_ns_values, type=pa.timestamp("ns")), |
| pa.array(ts_ns_values, type=pa.timestamp("ns", tz="Asia/Shanghai")), |
| ], |
| names=[ |
| "f_bool", "f_int8", "f_int16", "f_int32", "f_int64", |
| "f_float32", "f_float64", "f_utf8", "f_binary", |
| "f_decimal", "f_date", "f_timestamp", |
| "f_timestamp_ns", "f_timestamp_ns_tz", |
| ], |
| ) |
| |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| rb = reader.read_row_group(0) |
| assert rb.num_rows == 2 |
| |
| assert rb.column("f_bool").to_pylist() == [True, False] |
| assert rb.column("f_int8").to_pylist() == [42, -1] |
| assert rb.column("f_int16").to_pylist() == [1234, -5678] |
| assert rb.column("f_int32").to_pylist() == [100000, -200000] |
| assert rb.column("f_int64").to_pylist() == [9999999999, -9999999999] |
| assert rb.column("f_utf8").to_pylist() == ["hello", "world"] |
| assert rb.column("f_binary").to_pylist() == [b"\x01\x02\x03", b"\xff\x00"] |
| |
| f32 = rb.column("f_float32").to_pylist() |
| assert abs(f32[0] - 3.14) < 1e-5 |
| assert abs(f32[1] - (-2.71)) < 1e-5 |
| |
| f64 = rb.column("f_float64").to_pylist() |
| assert abs(f64[0] - 2.718281828) < 1e-9 |
| assert abs(f64[1] - (-3.141592653)) < 1e-9 |
| |
| assert rb.schema.field("f_timestamp_ns").type == pa.timestamp("ns") |
| assert rb.column("f_timestamp_ns").cast(pa.int64()).to_pylist() == ts_ns_values |
| assert rb.schema.field("f_timestamp_ns_tz").type == pa.timestamp( |
| "ns", tz="Asia/Shanghai" |
| ) |
| assert rb.column("f_timestamp_ns_tz").cast(pa.int64()).to_pylist() == ts_ns_values |
| |
| def test_timestamp_nanos_mixed_batches_updates_dictionary(self): |
| pa_schema = pa.schema([pa.field("ts", pa.timestamp("ns"))]) |
| first_values = [1, None, 2] |
| second_values = [3 + (i % 3) for i in range(120)] |
| |
| buf = io.BytesIO() |
| with MosaicWriter(buf, pa_schema) as writer: |
| writer.write( |
| pa.record_batch( |
| [pa.array(first_values, type=pa.timestamp("ns"))], |
| names=["ts"], |
| ) |
| ) |
| writer.write( |
| pa.record_batch( |
| [pa.array(second_values, type=pa.timestamp("ns"))], |
| names=["ts"], |
| ) |
| ) |
| |
| with _reader_from_bytes(buf.getvalue()) as reader: |
| rb = reader.read_row_group(0) |
| assert rb.column("ts").cast(pa.int64()).to_pylist() == ( |
| first_values + second_values |
| ) |
| |
| def test_multiple_row_groups(self): |
| pa_schema = pa.schema( |
| [pa.field("id", pa.int32()), pa.field("data", pa.int64())] |
| ) |
| |
| opts = WriterOptions( |
| compression=WriterOptions.COMPRESSION_NONE, |
| num_buckets=1, |
| row_group_max_size=200, |
| ) |
| buf = io.BytesIO() |
| with MosaicWriter(buf, pa_schema, opts) as writer: |
| for start in range(0, 500, 50): |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(start, start + 50)), type=pa.int32()), |
| pa.array( |
| [i * 3 for i in range(start, start + 50)], |
| type=pa.int64(), |
| ), |
| ], |
| names=["id", "data"], |
| ) |
| writer.write(batch) |
| data = buf.getvalue() |
| |
| with _reader_from_bytes(data) as reader: |
| assert reader.num_row_groups > 1 |
| |
| offset = 0 |
| for rg in range(reader.num_row_groups): |
| rb = reader.read_row_group(rg) |
| ids = rb.column("id").to_pylist() |
| datas = rb.column("data").to_pylist() |
| for j in range(rb.num_rows): |
| assert ids[j] == offset + j |
| assert datas[j] == (offset + j) * 3 |
| offset += rb.num_rows |
| |
| assert offset == 500 |
| |
| def test_multiple_writes(self): |
| pa_schema = pa.schema( |
| [pa.field("x", pa.int32()), pa.field("y", pa.utf8())] |
| ) |
| |
| buf = io.BytesIO() |
| with MosaicWriter(buf, pa_schema) as writer: |
| for start in [0, 10, 20]: |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(start, start + 10)), type=pa.int32()), |
| pa.array([f"r_{i}" for i in range(start, start + 10)]), |
| ], |
| names=["x", "y"], |
| ) |
| writer.write(batch) |
| |
| data = buf.getvalue() |
| with _reader_from_bytes(data) as reader: |
| table = reader.read_all() |
| assert table.num_rows == 30 |
| xs = table.column("x").to_pylist() |
| assert xs == list(range(30)) |
| |
| def test_single_row(self): |
| pa_schema = pa.schema([pa.field("v", pa.int32())]) |
| batch = pa.record_batch( |
| [pa.array([42], type=pa.int32())], names=["v"] |
| ) |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| rb = reader.read_row_group(0) |
| assert rb.num_rows == 1 |
| assert rb.column("v")[0].as_py() == 42 |
| |
| def test_zero_rows(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("v", pa.int32(), nullable=False), |
| pa.field("s", pa.utf8(), nullable=True), |
| ] |
| ) |
| batch = pa.record_batch( |
| [pa.array([], type=pa.int32()), pa.array([], type=pa.utf8())], |
| names=["v", "s"], |
| ) |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| table = reader.read_all() |
| assert table.num_rows == 0 |
| assert table.schema.names == ["v", "s"] |
| assert table.schema.field("v").type == pa.int32() |
| assert table.schema.field("s").type == pa.utf8() |
| assert not table.schema.field("v").nullable |
| assert table.schema.field("s").nullable |
| |
| |
| class TestProjection: |
| def test_projection_subset(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("a", pa.int32()), |
| pa.field("b", pa.utf8()), |
| pa.field("c", pa.float64()), |
| pa.field("d", pa.utf8()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(20)), type=pa.int32()), |
| pa.array([f"val_{i}" for i in range(20)]), |
| pa.array([float(i) for i in range(20)]), |
| pa.array([f"extra_{i}" for i in range(20)]), |
| ], |
| names=["a", "b", "c", "d"], |
| ) |
| |
| opts = WriterOptions(num_buckets=2) |
| data = _write_to_bytes(pa_schema, batch, opts) |
| |
| with _reader_from_bytes(data) as reader: |
| reader.project(["a", "b"]) |
| |
| total_rows = 0 |
| for rg in range(reader.num_row_groups): |
| rb = reader.read_row_group(rg) |
| assert rb.num_columns == 2 |
| total_rows += rb.num_rows |
| |
| assert total_rows == 20 |
| |
| def test_projection_single_column(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("a", pa.int32()), |
| pa.field("b", pa.utf8()), |
| pa.field("c", pa.float64()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(10)), type=pa.int32()), |
| pa.array([f"v_{i}" for i in range(10)]), |
| pa.array([float(i) for i in range(10)]), |
| ], |
| names=["a", "b", "c"], |
| ) |
| |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| reader.project(["b"]) |
| rb = reader.read_row_group(0) |
| assert rb.num_columns == 1 |
| assert rb.num_rows == 10 |
| |
| def test_projection_preserves_order(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("a", pa.int32()), |
| pa.field("b", pa.utf8()), |
| pa.field("c", pa.float64()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(10)), type=pa.int32()), |
| pa.array([f"s{i}" for i in range(10)]), |
| pa.array([float(i) * 0.5 for i in range(10)]), |
| ], |
| names=["a", "b", "c"], |
| ) |
| |
| opts = WriterOptions(num_buckets=2) |
| data = _write_to_bytes(pa_schema, batch, opts) |
| |
| with _reader_from_bytes(data) as reader: |
| reader.project(["c", "a", "b"]) |
| rb = reader.read_row_group(0) |
| assert rb.num_columns == 3 |
| assert rb.schema.names == ["c", "a", "b"] |
| assert rb.column("c").to_pylist() == [i * 0.5 for i in range(10)] |
| assert rb.column("a").to_pylist() == list(range(10)) |
| assert rb.column("b").to_pylist() == [f"s{i}" for i in range(10)] |
| |
| def test_projection_empty(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("a", pa.int32()), |
| pa.field("b", pa.utf8()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(5)), type=pa.int32()), |
| pa.array([f"v{i}" for i in range(5)]), |
| ], |
| names=["a", "b"], |
| ) |
| |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| reader.project([]) |
| rb = reader.read_row_group(0) |
| assert rb.num_columns == 0 |
| assert rb.num_rows == 5 |
| |
| |
| class TestSchema: |
| def test_schema_roundtrip(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("name", pa.utf8(), nullable=True), |
| pa.field("id", pa.int32(), nullable=False), |
| pa.field("score", pa.float64(), nullable=True), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(["x"]), |
| pa.array([1], type=pa.int32()), |
| pa.array([1.0]), |
| ], |
| names=["name", "id", "score"], |
| ) |
| |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| s = reader.schema |
| assert len(s) == 3 |
| assert s.names == ["name", "id", "score"] |
| assert s.field("id").type == pa.int32() |
| assert s.field("name").type == pa.utf8() |
| assert s.field("score").type == pa.float64() |
| assert not s.field("id").nullable |
| assert s.field("name").nullable |
| |
| |
| class TestStatistics: |
| def test_stats_basic(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("name", pa.utf8()), |
| pa.field("score", pa.float64()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array([i * 10 for i in range(10)], type=pa.int32()), |
| pa.array([f"item_{i}" for i in range(10)]), |
| pa.array([i * 1.1 for i in range(10)]), |
| ], |
| names=["id", "name", "score"], |
| ) |
| |
| opts = WriterOptions(stats_columns=["id", "score"]) |
| data = _write_to_bytes(pa_schema, batch, opts) |
| |
| with _reader_from_bytes(data) as reader: |
| for rg in range(reader.num_row_groups): |
| stats = reader.get_row_group_statistics(rg) |
| assert len(stats) > 0 |
| assert isinstance(stats, dict) |
| for name, stat in stats.items(): |
| assert isinstance(stat, ColumnStatistics) |
| assert name in ("id", "score") |
| assert stat.null_count == 0 |
| assert stat.has_min_max |
| assert stat.min is not None |
| assert stat.max is not None |
| assert len(stat.min) > 0 |
| assert len(stat.max) > 0 |
| |
| id_stat = stats["id"] |
| min_id = struct.unpack(">i", id_stat.min)[0] |
| max_id = struct.unpack(">i", id_stat.max)[0] |
| assert min_id == 0 |
| assert max_id == 90 |
| |
| def test_stats_with_nulls(self): |
| pa_schema = pa.schema( |
| [pa.field("a", pa.int32()), pa.field("b", pa.int64())] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array([10, None, 5, 20], type=pa.int32()), |
| pa.array([None, None, 100, 50], type=pa.int64()), |
| ], |
| names=["a", "b"], |
| ) |
| |
| opts = WriterOptions(stats_columns=["a", "b"], num_buckets=1) |
| data = _write_to_bytes(pa_schema, batch, opts) |
| |
| with _reader_from_bytes(data) as reader: |
| stats = reader.get_row_group_statistics(0) |
| assert len(stats) == 2 |
| |
| a_stat = stats["a"] |
| assert a_stat.null_count == 1 |
| assert a_stat.has_min_max |
| min_a = struct.unpack(">i", a_stat.min)[0] |
| max_a = struct.unpack(">i", a_stat.max)[0] |
| assert min_a == 5 |
| assert max_a == 20 |
| |
| b_stat = stats["b"] |
| assert b_stat.null_count == 2 |
| assert b_stat.has_min_max |
| min_b = struct.unpack(">q", b_stat.min)[0] |
| max_b = struct.unpack(">q", b_stat.max)[0] |
| assert min_b == 50 |
| assert max_b == 100 |
| |
| def test_stats_all_null(self): |
| pa_schema = pa.schema([pa.field("x", pa.int32())]) |
| |
| batch = pa.record_batch( |
| [pa.array([None, None, None], type=pa.int32())], names=["x"] |
| ) |
| |
| opts = WriterOptions(stats_columns=["x"], num_buckets=1) |
| data = _write_to_bytes(pa_schema, batch, opts) |
| |
| with _reader_from_bytes(data) as reader: |
| stats = reader.get_row_group_statistics(0) |
| assert len(stats) == 1 |
| assert stats["x"].null_count == 3 |
| assert not stats["x"].has_min_max |
| assert stats["x"].min is None |
| assert stats["x"].max is None |
| |
| |
| def test_stats_empty_string_min(self): |
| pa_schema = pa.schema([pa.field("s", pa.utf8())]) |
| |
| batch = pa.record_batch( |
| [pa.array(["", "b"])], names=["s"] |
| ) |
| |
| opts = WriterOptions(stats_columns=["s"], num_buckets=1) |
| buf = io.BytesIO() |
| with MosaicWriter(buf, pa_schema, opts) as writer: |
| writer.write(batch) |
| |
| writer_stats = writer.get_row_group_statistics(0) |
| assert "s" in writer_stats |
| s_stat = writer_stats["s"] |
| assert s_stat.has_min_max |
| assert s_stat.min == b"" |
| assert s_stat.max == b"b" |
| assert s_stat.null_count == 0 |
| |
| data = buf.getvalue() |
| with _reader_from_bytes(data) as reader: |
| stats = reader.get_row_group_statistics(0) |
| assert stats["s"].has_min_max |
| assert stats["s"].min == b"" |
| assert stats["s"].max == b"b" |
| |
| |
| class TestConvenience: |
| def test_write_table_read_table(self): |
| table = pa.table( |
| { |
| "id": pa.array(list(range(30)), type=pa.int32()), |
| "name": pa.array([f"user_{i}" for i in range(30)]), |
| } |
| ) |
| |
| buf = io.BytesIO() |
| write_table(table, buf) |
| |
| data = buf.getvalue() |
| result = read_table( |
| lambda offset, length: data[offset : offset + length], len(data) |
| ) |
| |
| assert result.num_rows == 30 |
| assert result.column("id").to_pylist() == list(range(30)) |
| assert result.column("name").to_pylist() == [f"user_{i}" for i in range(30)] |
| |
| def test_read_table_empty_projection(self): |
| table = pa.table( |
| { |
| "id": pa.array(list(range(30)), type=pa.int32()), |
| "name": pa.array([f"user_{i}" for i in range(30)]), |
| } |
| ) |
| |
| buf = io.BytesIO() |
| write_table(table, buf) |
| |
| data = buf.getvalue() |
| result = read_table( |
| lambda offset, length: data[offset : offset + length], |
| len(data), |
| columns=[], |
| ) |
| |
| assert result.num_columns == 0 |
| assert result.num_rows == 30 |
| assert result.schema.names == [] |
| |
| def test_read_table_projection_with_zero_row_groups(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("name", pa.utf8()), |
| ] |
| ) |
| |
| buf = io.BytesIO() |
| with MosaicWriter(buf, pa_schema): |
| pass |
| |
| data = buf.getvalue() |
| |
| empty_projection = read_table( |
| lambda offset, length: data[offset : offset + length], |
| len(data), |
| columns=[], |
| ) |
| assert empty_projection.num_rows == 0 |
| assert empty_projection.num_columns == 0 |
| assert empty_projection.schema.names == [] |
| |
| name_projection = read_table( |
| lambda offset, length: data[offset : offset + length], |
| len(data), |
| columns=["name"], |
| ) |
| assert name_projection.num_rows == 0 |
| assert name_projection.num_columns == 1 |
| assert name_projection.schema.names == ["name"] |
| |
| duplicate_projection = read_table( |
| lambda offset, length: data[offset : offset + length], |
| len(data), |
| columns=["name", "name"], |
| ) |
| assert duplicate_projection.num_rows == 0 |
| assert duplicate_projection.num_columns == 1 |
| assert duplicate_projection.schema.names == ["name"] |
| |
| with _reader_from_bytes(data) as reader: |
| assert reader.num_row_groups == 0 |
| reader.project(["name"]) |
| assert reader.schema.names == ["id", "name"] |
| result = reader.read_all() |
| assert result.schema.names == ["name"] |
| |
| def test_read_all(self): |
| pa_schema = pa.schema( |
| [pa.field("x", pa.int32()), pa.field("y", pa.utf8())] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(25)), type=pa.int32()), |
| pa.array([f"row_{i}" for i in range(25)]), |
| ], |
| names=["x", "y"], |
| ) |
| |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| table = reader.read_all() |
| assert isinstance(table, pa.Table) |
| assert table.num_rows == 25 |
| assert table.column("x").to_pylist() == list(range(25)) |
| |
| |
| class TestWriter: |
| def test_estimated_file_size(self): |
| pa_schema = pa.schema( |
| [pa.field("x", pa.int32()), pa.field("y", pa.utf8())] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(100)), type=pa.int32()), |
| pa.array([f"value_{i}" for i in range(100)]), |
| ], |
| names=["x", "y"], |
| ) |
| |
| buf = io.BytesIO() |
| with MosaicWriter(buf, pa_schema) as writer: |
| writer.write(batch) |
| est = writer.estimated_file_size() |
| assert est > 0 |
| |
| def test_write_after_close_fails_before_export(self): |
| pa_schema = pa.schema([pa.field("x", pa.int32())]) |
| batch = pa.record_batch( |
| [pa.array([1, 2, 3], type=pa.int32())], names=["x"] |
| ) |
| |
| writer = MosaicWriter(io.BytesIO(), pa_schema) |
| writer.close() |
| |
| with pytest.raises(RuntimeError, match="writer is closed"): |
| writer.write(batch) |
| |
| def test_writer_stats_basic(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("name", pa.utf8()), |
| pa.field("score", pa.float64()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array([i * 10 for i in range(10)], type=pa.int32()), |
| pa.array([f"item_{i}" for i in range(10)]), |
| pa.array([i * 1.1 for i in range(10)]), |
| ], |
| names=["id", "name", "score"], |
| ) |
| |
| opts = WriterOptions(stats_columns=["id", "score"]) |
| buf = io.BytesIO() |
| writer = MosaicWriter(buf, pa_schema, opts) |
| writer.write(batch) |
| writer.close() |
| |
| assert writer.num_row_groups >= 1 |
| stats = writer.get_row_group_statistics(0) |
| assert len(stats) > 0 |
| assert isinstance(stats, dict) |
| for name, stat in stats.items(): |
| assert isinstance(stat, ColumnStatistics) |
| assert name in ("id", "score") |
| assert stat.null_count == 0 |
| assert stat.has_min_max |
| assert stat.min is not None |
| assert stat.max is not None |
| |
| id_stat = stats["id"] |
| min_id = struct.unpack(">i", id_stat.min)[0] |
| max_id = struct.unpack(">i", id_stat.max)[0] |
| assert min_id == 0 |
| assert max_id == 90 |
| |
| def test_writer_stats_with_nulls(self): |
| pa_schema = pa.schema( |
| [pa.field("a", pa.int32()), pa.field("b", pa.int64())] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array([10, None, 5, 20], type=pa.int32()), |
| pa.array([None, None, 100, 50], type=pa.int64()), |
| ], |
| names=["a", "b"], |
| ) |
| |
| opts = WriterOptions(stats_columns=["a", "b"], num_buckets=1) |
| buf = io.BytesIO() |
| writer = MosaicWriter(buf, pa_schema, opts) |
| writer.write(batch) |
| writer.close() |
| |
| assert writer.num_row_groups == 1 |
| stats = writer.get_row_group_statistics(0) |
| assert len(stats) == 2 |
| |
| a_stat = stats["a"] |
| assert a_stat.null_count == 1 |
| assert a_stat.has_min_max |
| min_a = struct.unpack(">i", a_stat.min)[0] |
| max_a = struct.unpack(">i", a_stat.max)[0] |
| assert min_a == 5 |
| assert max_a == 20 |
| |
| b_stat = stats["b"] |
| assert b_stat.null_count == 2 |
| assert b_stat.has_min_max |
| min_b = struct.unpack(">q", b_stat.min)[0] |
| max_b = struct.unpack(">q", b_stat.max)[0] |
| assert min_b == 50 |
| assert max_b == 100 |
| |
| def test_writer_stats_all_null(self): |
| pa_schema = pa.schema([pa.field("x", pa.int32())]) |
| |
| batch = pa.record_batch( |
| [pa.array([None, None, None], type=pa.int32())], names=["x"] |
| ) |
| |
| opts = WriterOptions(stats_columns=["x"], num_buckets=1) |
| buf = io.BytesIO() |
| writer = MosaicWriter(buf, pa_schema, opts) |
| writer.write(batch) |
| writer.close() |
| |
| assert writer.num_row_groups == 1 |
| stats = writer.get_row_group_statistics(0) |
| assert len(stats) == 1 |
| assert stats["x"].null_count == 3 |
| assert not stats["x"].has_min_max |
| assert stats["x"].min is None |
| assert stats["x"].max is None |
| |
| def test_writer_stats_matches_reader_stats(self): |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("score", pa.float64()), |
| ] |
| ) |
| |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(50)), type=pa.int32()), |
| pa.array([i * 2.0 for i in range(50)]), |
| ], |
| names=["id", "score"], |
| ) |
| |
| opts = WriterOptions(stats_columns=["id", "score"], num_buckets=1) |
| buf = io.BytesIO() |
| writer = MosaicWriter(buf, pa_schema, opts) |
| writer.write(batch) |
| writer.close() |
| |
| data = buf.getvalue() |
| with _reader_from_bytes(data) as reader: |
| for rg in range(writer.num_row_groups): |
| w_stats = writer.get_row_group_statistics(rg) |
| r_stats = reader.get_row_group_statistics(rg) |
| assert len(w_stats) == len(r_stats) |
| assert set(w_stats.keys()) == set(r_stats.keys()) |
| for name in w_stats: |
| assert w_stats[name].null_count == r_stats[name].null_count |
| assert w_stats[name].has_min_max == r_stats[name].has_min_max |
| assert w_stats[name].min == r_stats[name].min |
| assert w_stats[name].max == r_stats[name].max |
| |
| def test_row_group_num_rows(self): |
| pa_schema = pa.schema( |
| [pa.field("id", pa.int32()), pa.field("data", pa.int64())] |
| ) |
| |
| opts = WriterOptions( |
| compression=WriterOptions.COMPRESSION_NONE, |
| num_buckets=1, |
| row_group_max_size=200, |
| ) |
| buf = io.BytesIO() |
| with MosaicWriter(buf, pa_schema, opts) as writer: |
| for start in range(0, 500, 50): |
| batch = pa.record_batch( |
| [ |
| pa.array(list(range(start, start + 50)), type=pa.int32()), |
| pa.array( |
| [i * 2 for i in range(start, start + 50)], |
| type=pa.int64(), |
| ), |
| ], |
| names=["id", "data"], |
| ) |
| writer.write(batch) |
| data = buf.getvalue() |
| |
| with _reader_from_bytes(data) as reader: |
| assert reader.num_row_groups > 1 |
| total = 0 |
| for rg in range(reader.num_row_groups): |
| num_rows = reader.row_group_num_rows(rg) |
| assert num_rows > 0 |
| rb = reader.read_row_group(rg) |
| assert num_rows == rb.num_rows |
| total += num_rows |
| assert total == 500 |
| |
| def test_row_group_num_rows_single(self): |
| pa_schema = pa.schema([pa.field("x", pa.int32())]) |
| batch = pa.record_batch( |
| [pa.array(list(range(10)), type=pa.int32())], names=["x"] |
| ) |
| data = _write_to_bytes(pa_schema, batch) |
| |
| with _reader_from_bytes(data) as reader: |
| assert reader.num_row_groups == 1 |
| assert reader.row_group_num_rows(0) == 10 |
